[Cute,Fwd,Sm120] Fix FlashAttentionForwardSm120 runtime errors on SM120

[blake-snc]

FlashAttentionForwardSm120 fails at dispatch on real SM120 hardware (validated on SM121a / DGX Spark GB10) with three distinct symptoms that all stem from one root cause. FlashAttentionForwardBase.__init__ assigns self.arch = BaseDSL._get_dsl().get_arch_enum(), which on SM120 returns sm_121a and overwrites the class level arch = 80 declaration on FlashAttentionForwardSm120. Three symptoms follow:

1. AttributeError: 'FlashAttentionForwardSm120' object has no attribute 'is_split_kv' on every call. The base __init__ does not accept is_split_kv (only Sm90 and Sm100 do), but the shared Sm80.__call__ reads self.is_split_kv directly.
2. On the non varlen path, self.use_tma_O = self.arch >= Arch.sm_90 resolves True on sm_121a, so the TMA O branch runs and fails with AttributeError: 'NoneType' object has no attribute '_trait' at cpasync/helpers.py:209 (TMA atom is None because SM120 has no TMA O support).
3. On the varlen path, ragged = self.use_tma_O and has_cu_seqlens is True for the same reason, hitting an incompatible layout path and failing at flash_fwd.py:398 with expects coord and shape of view are weakly congruent, but got '!cute.layout<"(?,?):(?{i64 div=8},1)">', '!cute.coord<"(_,_,?)">'.

Fix: add an __init__ override on FlashAttentionForwardSm120 that runs after super().__init__() and restores self.arch = Arch.sm_80, sets self.is_split_kv = False, and sets self.pack_gqa = False.

The pack_gqa override addresses a separate preexisting issue. Sm80.__call__'s pack GQA path does not call pack_gqa_layout before handing tensors to PackGQA methods, while Sm90.__call__ (L273) and Sm100.__call__ (L504) do. Porting those three lines into Sm80.__call__ verbatim is not sufficient, because the Sm80 mainloop's tile sizing assumes tile_m divides the seqlen dimension cleanly, which fails for the packed (qhead_per_kvhead, seqlen_q) layout when qhead_per_kvhead does not divide tile_m (for example, cute.local_tile raises on an unaligned division). A proper Sm80 pack GQA fix needs a follow up that also adjusts the tile scheduler. Forcing pack_gqa = False on the Sm120 subclass sidesteps both issues without touching the Sm80 class, keeping blast radius off SM8.

All changes are contained to FlashAttentionForwardSm120. Nothing outside that class is touched, so SM8 / SM9 / SM10 / SM11 dispatch paths are unaffected.

Scope

This is a correctness fix that unblocks FlashAttentionForwardSm120 end to end. It is not a performance improvement. FA4's perf advantage over FA2 comes from tcgen05 and TMEM on SM100. SM120 has neither, so FlashAttentionForwardSm120 compiles down to the same SM80 era mma.sync path as FA2 with additional dispatch complexity on top. On the realistic Qwen 3 prefill shapes I measured, patched FA4 on SM120 runs within roughly 15 percent of FA2 at longer sequences and is slower at short sequences. Users on SM120 are better served by FA2 today for pure attention throughput. The motivations to land this fix anyway are:

• Correctness. The class currently crashes at the first dispatch; any code path that constructs it is dead.
• Feature reach on SM120. FA4 only features such as score_mod, block sparse masking (Add SM80/SM120 block-sparse forward attention support #2389), and split KV FlashDecoding (Add SM120 split-KV (FlashDecoding) with FP32 partial outputs #2336) become available on SM120 once the crash is resolved.
• A working baseline for the SM120 FA4 kernel to evolve from.

Relationship to in flight SM120 feature PRs

The self.arch = Arch.sm_80 portion of this fix also appears in three open SM120 feature PRs (#2336 split KV, #2348 paged KV, #2389 block sparse), each of which adds its own __init__ carrying that one line as part of its feature scope. I am filing this as a standalone bug fix because the arch issue alone blocks every SM120 dispatch path regardless of which feature is enabled, and a small focused review surface is easier to land quickly. If this PR merges first, those three feature PRs rebase to a one line diff. If one of them merges first, this PR rebases to drop the self.arch line. I own all four branches and am happy to handle the rebase either way.

Validation

SM121a (DGX Spark GB10), bfloat16 and float16, causal and non causal, dense and varlen, six shape configs: 48 / 48 pass, max abs diff <= 0.0156 against a PyTorch fp32 reference. GQA configs pass through the non packed path (Qwen style H_q / H_kv = 16 / 2 at D = 128, MQA 4 / 1 at D = 128, MHA D = 96).

Known limitations

• FA4 on SM120 is slower than FA2 on SM120 today. Users should continue to prefer FA2 for pure attention throughput on SM120 until the SM120 mainloop closes the gap. vLLM's existing _is_fa4_supported() gate already excludes SM120, so end users are not impacted by default.
• PackGQA disabled on SM120 pending the follow up Sm80 fix. Functionally correct through the non packed path.
• Split KV not supported on SM120 in this kernel variant. Consistent with the dispatch assertion at interface.py:753. Add SM120 split-KV (FlashDecoding) with FP32 partial outputs #2336 adds full SM120 split KV support as a feature.
• head_size > 128 not supported on SM120 for this kernel variant due to the 99 KB SMEM budget. Matches vLLM's existing fa_utils.py gate that routes head_size > 128 to FA2 on Blackwell.

AI assistance

AI assistance was used to author this patch and its test plan. I validated each change on SM121a hardware and reviewed every line before pushing.

Contributed by Second Nature Computing (https://joinsecondnature.com)

[vient]

Bump, master still breaks on sm_120. @tridao can you review this please?

[vanbukin]

Any plans/updates?

[blake-snc]

Closing this PR. The three failure modes it addressed have since been fixed on main by other means: the SM120 shim now forces the SM80 arch path after __init__, and use_tma_O is guarded below SM 12.0. Our SM120 hardware is currently allocated to other work, so I could not re-verify at head; closing on code inspection of current main. Reachable at github.com/bledden if anything here resurfaces.