LESSON 09 · 2026.04.20 · L4

MHA + Causal FA — 80% of cuDNN in 300 lines

Extend the 2-D Triton FA to (B, H, N, d) + a causal mask. Register with torch.library.custom_op, and with torch.compile(fullgraph=True) an entire LLaMA-style attention block lands in a single graph.

Five things that changed

Accuracy (max rel_err vs fp32 reference)

Speed — LLaMA-7B causal, d=128, fp16

On GPT-2 (d=64) short shapes: 1.13× ahead at (16,12,512), tied at (8,12,1024).

Lesson 1 · Causal speed isn't from the mask

Just masking still loads the upper triangle and flips it to -inf. The real speed comes from loop skip:

if IS_CAUSAL:
    end_n = tl.minimum(N, (pid_m + 1) * BLOCK_M)
else:
    end_n = N
for start_n in range(0, end_n, BLOCK_N):
    ...

On average it drops to N/2, so both the K/V loads and the two tl.dots are halved. Verified: non-causal (1,32,2048,128) 2.643 ms vs causal 0.784 ms — half the FLOPs, but 3.3× faster (pipeline amortization).

Lesson 2 · tl.constexpr = fold away the runtime if

Pass is_causal as a plain arg and you get a branch inside the tight K loop plus warp-scheduler tangles. With tl.constexpr, Triton compiles two specialized kernels (True/False). At runtime, only one is dispatched — the branch doesn't even exist. Using autotune(key=[...,"IS_CAUSAL"]) also separates the autotune space: empirically, causal prefers BM=64 / BN=128 while non-causal likes bigger tiles.

Lesson 3 · Three gifts from custom_op

@custom_op("triton_training::flash_attention_mha",
           mutates_args=(), device_types="cuda")
def flash_attention_mha(q, k, v, is_causal: bool = False) -> Tensor: ...

@flash_attention_mha.register_fake
def _(q, k, v, is_causal): return torch.empty_like(q)

1. torch.compile doesn't break the graph. register_fake supplies shape inference so Dynamo doesn't see an unknown op. Under fullgraph=True, the entire LLaMA AttentionBlock lives in one graph.
2. Serialization-preserving. ONNX/AOT/TorchScript record triton_training::flash_attention_mha as-is.
3. Drop-in torch.ops.<ns>.<op> path. Downstreams like vLLM call without importing Triton. The same pattern vLLM uses to expose custom kernels.

Lesson 4 · The ROI of 300 lines

• kernel 100 lines + op wrapper 76 lines + autotune/register_fake ≈ 300 lines
• 78–90% of cuDNN FA-2 on d=128 LLaMA-7B causal
• Tied or 13% ahead on GPT-2 (d=64)
• 29–33× over naive
• Drops into an attention block under torch.compile(fullgraph=True)

Doing the same in CUDA: 5–10× the code (Lesson 6's FA v1 was ~500 lines + 150 lines of host), plus manual re-tuning on every new GPU. This is the whole reason to learn Triton.

The remaining 20% gap

The 22% left to FA-2 isn't Triton's gap — it's what cuDNN FA-2 has:

• async copy + double/triple buffering (smem load ↔ prior tile mma overlap)
• persistent kernel scheduling
• warp specialization (compute warps vs load-store warps)

Triton is getting persistent / warp-spec but it's tutorial-stage. Closing this gap takes CUTLASS 3.x — next round.