SM120 attention kernels exist but are blocked by wiring issues (fmha_v2, backend selector, MLA)

[blake-snc]

Summary

We've done a deep investigation of FlashInfer's SM120 support on an NVIDIA DGX Spark (GB10, SM121a, CUDA 13.0) and found that SM120 attention kernels already exist in the codebase but are not wired into the runtime pipeline. Three specific wiring issues prevent SM120 from using its optimized paths:

Issue 1: FMHA_V2 SM120 kernels gated behind ENABLE_SM120

Location: flashinfer/jit/attention/fmha_v2/generator_utils.py:6772-6786

SM120-specific FMHA_V2 kernels are only generated if the ENABLE_SM120 environment variable is set:

if "ENABLE_SM120" in os.environ:
    enumerate_hmma_flash_kernels(specs, sm=120, dtype="fp16")
    enumerate_hmma_flash_kernels(specs, sm=120, dtype="bf16")
    # ...

Three SM120 kernel files are referenced in jit/attention/modules.py:1918-1920 (BF16 and FP8 variants) but never get generated in standard builds. This means gen_trtllm_fmha_v2_module() can't find them at runtime.

Expected: SM120 kernels should be generated by default when building for SM120 targets.

Issue 2: determine_attention_backend() routes SM120 to generic FA2

Location: flashinfer/utils.py:455-497

def determine_attention_backend(device, ...):
    if is_sm90a_supported(device) and is_fa3_backend_supported(...):
        return "fa3"
    else:
        return "fa2"

This only checks for SM90. SM120 falls through to generic FA2 (Ampere-level), even though SM120-specific FMHA_V2 kernels exist. The function should have an SM120 path that routes to the fmha_v2 SM120 backend.

Issue 3: determine_mla_backend() routes SM120 to FA2 instead of XQA

Location: flashinfer/utils.py:561-562

def determine_mla_backend(device):
    return "fa3" if is_sm90a_supported(device) else "fa2"

SM120 gets fa2 for MLA, but FlashInfer already has a functional XQA MLA backend specifically for SM120 (data/csrc/xqa/mla_sm120.cu, compiled with supported_major_versions=[12] in jit/xqa.py:137). This function should return "xqa" for SM120.

What Currently Works on SM120

Component	Status
XQA decode (MHA)	✅ Works (xqa.py checks [9, 10, 12])
XQA decode (MLA, FP8)	✅ Code exists (mla_sm120.cu) but Issue #2166 reports failures
FMHA_V2 SM120 prefill	❌ Kernels exist but never generated or routed to
SM120 GEMM (FP8, FP4)	✅ Works via dedicated SM120 CUTLASS kernels

Proposed Fixes

1. Remove ENABLE_SM120 gate in generator_utils.py (or enable by default for SM120 targets)
2. Add SM120 branch in determine_attention_backend() to route to fmha_v2 SM120 kernels
3. Add SM120 branch in determine_mla_backend() to route to XQA

Environment

• NVIDIA DGX Spark (GB10, SM121a)
• CUDA 13.0, aarch64
• flashinfer-python 0.6.2+cu130
• PyTorch 2.9.1+cu130

Related issues: #1147, #2166, #2294

From Second Nature Computing — testing on DGX Spark.

[blake-snc]

Validation Update from DGX Spark (GB10, SM121a, CUDA 13.0)

We've completed hands-on validation and have corrections to the original report:

Issue 3 Correction: determine_mla_backend() returning "fa2" is CORRECT

We validated that FA2 MLA (BatchMLAPagedAttentionWrapper with backend="fa2") works perfectly on SM120 with both BF16 and FP16:

# Both pass on SM120:
wrapper = BatchMLAPagedAttentionWrapper(float_workspace_buffer=ws, backend="fa2")
wrapper.plan(qo_indptr, kv_indptr, kv_indices, kv_len_arr, num_heads=128,
             head_dim_ckv=512, head_dim_kpe=64, page_size=16, ...)
out = wrapper.run(q_nope, q_pe, ckv_cache, kpe_cache)
# → shape [4, 128, 512], no NaN ✅

So determine_mla_backend() returning "fa2" for SM120 is actually correct — the JIT-compiled FA2 MLA kernels work. The XQA MLA FP8 path (trtllm_batch_decode_with_kv_cache_mla) is a separate API that also works on SM120 but is FP8-only.

Full Validated Results on SM120

Operation	Status	Notes
FA2 single prefill (BF16/FP16)	✅ PASS	single_prefill_with_kv_cache
FA2 paged decode (FP16)	✅ PASS	BatchDecodeWithPagedKVCacheWrapper
FA2 batch prefill paged (FP16)	✅ PASS	BatchPrefillWithPagedKVCacheWrapper
FA2 MLA decode (BF16/FP16)	✅ PASS	BatchMLAPagedAttentionWrapper
XQA decode auto/forced (FP16)	✅ PASS	trtllm_batch_decode_with_kv_cache correctly selects XQA
XQA MLA FP8 decode	✅ PASS	trtllm_batch_decode_with_kv_cache_mla with float8_e4m3fn
trtllm-gen prefill	❌ FAIL	"Unsupported architecture" at fmhaRunner.cuh:30
XQA MLA BF16 decode	❌ FAIL	"XQA MLA only supports fp8 operation on SM120 GPUs"

Remaining Valid Issues

Issue 1 (fmha_v2 ENABLE_SM120 gating) remains valid — the SM120-optimized fmha_v2 kernels exist but are only generated with ENABLE_SM120=1. With ENABLE_SM120=1, determine_attention_backend() still returns "fa2" so the fmha_v2 kernels are never selected.

Issue 2 (attention backend) remains valid as a performance optimization — FA2 works but fmha_v2 SM120 kernels could potentially be faster.

trtllm-gen prefill failure is the biggest gap — fmhaRunner.cuh:30 hard-fails on SM120. This blocks any trtllm-gen prefill usage on SM120.

[blake-snc]

Validated by Second Nature Computing on NVIDIA DGX Spark (GB10, SM121a).

[blake-snc]

Status update after reading the current code more carefully:

Issue 1 (FMHA_V2 SM120 gated behind ENABLE_SM120): The ENABLE_SM120 gate in generator_utils.py:6874 only affects AOT builds (pip wheel generation). The JIT runtime path uses generate_kernels.py's enumerate_kernels(), which generates SM120 kernels directly without the gate — and there's a dedicated gen_trtllm_fmha_v2_sm120_module() in modules.py that calls it. So this isn't blocking at JIT time.

Issue 2 (determine_attention_backend() routes SM120 to FA2): Fixed by #2598 (open, pending review) — adds SM12x routing to a CuTe DSL backend, which works end-to-end on SM121a.

Issue 3 (determine_mla_backend() routes SM120 to FA2): FA2 MLA works correctly on SM121a (validated). #2689 (open, pending review) adds BF16 XQA MLA kernel for SM120 as a follow-on improvement.

The main remaining gap is that if someone builds a pip wheel without ENABLE_SM120=1, the SM120 fmha_v2 kernels are excluded from the wheel. But that's only relevant for the fmha_v2/trtllm backend path — once #2598 merges, SM120 gets the CuTe DSL path instead.