You write two small Python functions instead. One adjusts an attention score (score_mod), adding a positional bias, capping it, whatever your variant needs. The other answers a single question (mask_mod): should this position be computed at all. torch.compile turns both into one fused Triton kernel, and autograd generates the backward pass for free.
The interesting part is why it takes two functions and not one. The scoring function can already express any mask, since returning negative infinity drops the position out of the softmax. But that computes everything and then throws it away. Declaring the mask separately is what lets PyTorch skip whole blocks before doing the work, which is worth around 2× on a causal mask, and the compiler cannot recover that intent from arbitrary Python on its own.
In ten minutes we'll write one, use create_block_mask to see the block structure it produces, and look at what it costs, which is roughly 90% of FlashAttention-2 on the forward pass and 85% on the backward, on PyTorch's own benchmarks.
By the end, you'll be able to implement an attention variant with no existing kernel in pure Python — and know exactly what it costs.