[Bugfix] Fix two SM80/SM120 forward kernel bugs: missing is_split_kv default, mDynamicCausal NameError

[tgmerritt]

Found while enabling FA4 for SM120/SM121 (consumer Blackwell — RTX PRO 6000, GB10/DGX Spark) in vLLM v0.24.0. Both bugs are unconditional and architecture-agnostic — they'd break literal SM80 (Ampere) hardware too, since FlashAttentionForwardSm120 subclasses FlashAttentionForwardSm80 and shares this exact code path. Confirmed both still present on main.

Bug 1: self.is_split_kv never set on Sm80/Sm120

FlashAttentionForwardBase.__init__ doesn't accept or set is_split_kv — only the Sm90/Sm100 family of constructors do (they receive it explicitly from interface.py's dispatch). But the shared __call__/kernel() code in this file references self.is_split_kv unconditionally, raising AttributeError on any Sm80 or Sm120 instantiation.

interface.py's SM120 dispatch branch already asserts not is_split_kv before constructing FlashAttentionForwardSm120 (split_kv isn't supported on SM 12.0), so defaulting self.is_split_kv = False in Sm80.__init__ is safe.

Bug 2: mDynamicCausal NameError

__call__'s signature accepts mDynamicCausal, but never forwards it into the paired @cute.kernel-decorated kernel() function's own parameter list or the .launch() call's args. kernel()'s body references the bare name anyway:

psc = mDynamicCausal[batch_size] if const_expr(mDynamicCausal is not None) else None

Since kernel() is traced in isolation by the CuTe-DSL JIT compiler, this raises NameError: name 'mDynamicCausal' is not defined on every call, regardless of causal=True/False — verified both crash identically.

A complete fix would thread mDynamicCausal through kernel()'s parameter list and the .launch() call. This patch instead defaults it to None right before use, restoring prior (pre-dynamic-causal) behavior for callers not using that feature — i.e. unblocks the overwhelmingly common case (static causal/non-causal) without the larger plumbing change.

Not fixed by #154 "Fix dynamic_causal" (merged) — that PR fixes a different dynamic_causal bug in mask.py's masking logic, not this one.

Verification

Verified end-to-end on real GB10 (SM121) hardware via vLLM's actual production call path (flash_attn_varlen_func), bf16, head_dim=128, two varlen sequences (lengths 64/96), causal=True. Output matches torch.nn.functional.scaled_dot_product_attention reference: max abs diff 0.0078, mean abs diff 0.00025 (normal bf16 rounding). Without these two fixes the same call crashes (NameError, then AttributeError if only one is fixed).

Note: tested via vLLM's vendored copy of this file (renames flash_attn.cute imports to vllm.vllm_flash_attn.cute for its package namespace) since that's the integration path I was working in — same two line-level changes, not re-run against a from-scratch build of this exact tree (would require a fresh ~70min CUDA compile). The fix is identical logic in both cases.

AI assistance disclosure

Found and fixed with Claude Code (Anthropic) while building a custom vLLM v0.24.0 image for GB10. I reviewed and understand the change.