B200 Flash Attention4 Planner

E — Executable steps (Execution)

Once the skill is activated, the agent must execute the following process:

1. Fix the math semantics
   • Q/K/V shapes, layout, head_dim, causal, GQA ratio, scale, the output, and whether LSE is needed.
2. Define the streaming state
   • The initial values and per-block update formulas of row_max, row_sum, and O for each row.
   • Be explicit about using natural exp or the equivalent exp2 scaling; the reference must match.
3. Define the two MMAs
   • score MMA: Q×Kᵀ→S.
   • value MMA: P×V→O.
   • Write out both tile shapes, dtypes, SMEM operands, and TMEM outputs.
4. Plan the S/P/O TMEM
   • After S is ready, read it into registers for mask/softmax; write P back to TMEM; accumulate O in TMEM and rescale when necessary.
   • Give the regions, strides, stages, and safe-reuse conditions.
5. Assign warp roles
   • TMA load, MMA issue, softmax stage 0/1, correction/epilogue, TMA store.
   • Check that the collective scopes are complete.
6. Build the barrier graph
   • Q/K/V ready→score/value MMA.
   • S ready→softmax.
   • P ready + O safe→value MMA.
   • final O ready→epilogue→store.
   • Draw the scalar mailbox's full/empty protocol separately.
7. Implement mask/GQA
   • causal tiles: handle fully-skipped blocks, fully-valid blocks, and diagonal-boundary masking separately.
   • GQA: make the Q-head-to-KV-head mapping, reuse, and scheduler coordinates explicit.
8. Implement rescale and writeback
   • The O correction is a full TMEM→register→TMEM tile operation; it cannot be removed.
   • After the loop ends: O / row_sum, cast, SMEM staging, TMA store drain.
9. Validate
   • Compare against a high-precision PyTorch reference.
   • Cover causal/noncausal, different sequence lengths, head_dim, GQA ratios, tail tiles, and extreme logits.
   • If extending to the training forward, verify the LSE definition and scale are consistent; backward requires a separate design.

Required outputs

1. Conclusion: the current choice/diagnosis; do not use a vague "it could be any of them".
2. Evidence or assumptions: which items come from user data, and which are hypotheses awaiting verification.
3. Contract/table/timeline: the auditable intermediate artifacts corresponding to this skill.
4. Minimal validation: correctness tests, boundary tests, and one falsifiable experiment.
5. Risks and fallback: alternative paths when hardware, version, or resource requirements are not met.

B — Boundaries (Boundary) ★

Do not use when

• Only a stable implementation already provided by PyTorch/FlashAttention needs to be called.
• The user asks for the full backward, but the saved intermediates and gradient algorithm are not yet defined.

Failure modes

• Not rescaling the old row_sum/O.
• Normalizing P too early, breaking the streaming accumulation.
• Reusing the S/P/O TMEM regions before their consumers finish.
• Using the same slow path for the causal mask on both full blocks and boundary blocks.
• Wrong GQA head mapping.

Limitations

• This skill centers on the forward structure from the book; training backward, dropout, variable-length packed sequences, and distributed attention require additional design.

Related skills

• depends-on: b200-tcgen05-mma-contract-builder, b200-tmem-lifecycle-planner, b200-mbarrier-protocol-auditor
• contrasts-with: b200-gemm-optimization-ladder
• composes-with: b200-tma-pipeline-designer, b200-warp-specialized-debugger, b200-kernel-roofline-triage

Audit info

• Validation passed: V1 ✓ / V2 ✓ / V3 ✓
• Test definitions: 6 (3 should_trigger / 2 should_not_trigger / 1 edge_case)
• Hardware validation: not performed; must be verified on a target B200
• Distilled: 2026-06-25