ESSAY 11 · 2026.04.20 · L4 · FINAL

How a 300-line Triton FA closes in on 80–90% of cuDNN FA-2

LLaMA-7B shapes, L4 (sm_89), under torch.compile(fullgraph=True). Three tricks combined — 4-D grid, IS_CAUSAL: tl.constexpr, torch.library.custom_op.

The question

Lesson 8's Triton FA was 2-D (N, d) and non-causal. Real LLMs don't look like that. The real shapes are (B, H, N, d) 4-D, mostly causal, and the graph must not break under torch.compile.

Meeting all three, how far can a 100-line kernel push?

Results first

78–90% of cuDNN FA-2 on LLaMA-like d=128 shapes. Tied or 13% ahead on GPT-2 d=64. 29–33× over naïve. kernel body + op wrapper = under 300 lines.

Trick 1 · 2-D → 4-D = 4 strides + 3-D grid

Every (batch, head) pair is a fully independent attention problem — don't loop over them; express the fan-out on the grid axes.

grid = (triton.cdiv(N, BLOCK_M), H, B)

# inside the kernel
pid_m = tl.program_id(0)   # BLOCK_M of queries
pid_h = tl.program_id(1)   # which head
pid_b = tl.program_id(2)   # which batch

q_base = Q_ptr + pid_b * stride_qb + pid_h * stride_qh
k_base = K_ptr + pid_b * stride_kb + pid_h * stride_kh
# everything below these four lines is identical to Lesson 8

Triton's launch starts the whole grid "as soon as resources allow." L4's 58 SMs fill up. With a loop, SMs would sit idle.

Trick 2 · Causal speed is loop skip, not the mask

Masking alone still loads the upper triangle from HBM and flips it to -inf. Real speed comes from:

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):
    ...

Average drops to N/2, so both the K/V loads and the two tl.dot calls are halved.

Verified. non-causal (1,32,2048,128) = 2.643 ms. Same shape causal = 0.784 ms. Half the FLOPs, 3.3× faster. The reason: the load pipeline amortizes proportionally to the number of tiles. FA-v2's causal optimization is this one line.

Trick 3 · tl.constexpr — fold a runtime if into compile time

Pass is_causal as a normal argument and you get a branch inside the tight K loop, tangled with the warp scheduler. Mark it tl.constexpr and Triton compiles two specialized kernels (True and False). At runtime, only one is dispatched — the branch simply isn't there.

Extra win: @triton.autotune(key=[..., "IS_CAUSAL"]) separates autotune per case. Empirically causal prefers BM=64, BN=128, non-causal prefers bigger tiles.

Trick 4 · torch.library.custom_op — dress it as an 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:
    return triton_flash_attention_mha(q, k, v, is_causal=is_causal)

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

This one block gives you three things:

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

bit-exact + zero graph breaks

[1] raw_wrapper vs torch.ops      err = 0.00e+00  (causal True/False)
[2] eager vs compiled function    err = 0.00e+00  — fullgraph=True passes
[3] eager vs compiled block       err = 0.00e+00  — AttentionBlock is a single graph
[4] schema: triton_training::flash_attention_mha(
        Tensor q, Tensor k, Tensor v, bool is_causal=False) -> Tensor

The last 20% is not Triton's gap

In the 78–90% of FA-2 range, the remaining 10–22% is what cuDNN FA-2 has:

• async copy + double/triple buffering — full overlap of smem loads with the prior tile's mma
• persistent kernel — keep blocks alive and re-dispatch tiles
• warp specialization — some warps compute, others do load-store

Triton is getting persistent / warp specialization but they're tutorial-grade. Closing the gap means going one layer deeper — CUTLASS 3.x. Next round's topic.

Trap log (digest)

• L4 stockout — us-west4-a → us-east4-c zone rotation.
• Python.h not found — install libpython3.10-dev.
• No git history — discovered mid-Lesson 9. git init + first commit should come before make vector_add.
• naive OOM — the (B,H,N,N) score tensor was 6 GB. Guard with naive_mem_bytes < 4 GB and tag "(skipped)" in the bench.
• "Not enough SMs" warning — L4's 58 SMs fall short of max_autotune_gemm's requirement. No functional impact.

What's next

Backward + autograd (register_autograd) → GQA (MQA/grouped) → closing the last 20% with persistent kernel + async copy → porting vLLM's PagedAttention. The material for Phase 2.