FlashInfer — DGX Spark (SM121) Current Support Audit
Date: 2026-04-24
Scope: Audit of SM121 (compute capability 12.1, DGX Spark) support across FlashInfer ops, APIs, backends, and testing.
Goal: Improve Spark usability on FlashInfer by identifying existing gaps.
Note: skip to "15. Consolidated Action Items" for TLDR.
Architecture Context
| Gencode |
Target |
CUDA Requirement |
compute_120a / sm_120a |
RTX 5090 (SM120) |
CUDA 12.8 |
compute_120f / sm_120f |
SM12x family (120 + 121) |
CUDA 12.9 |
compute_121a / sm_121a |
DGX Spark (SM121) |
CUDA 12.9 |
SM120 vs SM121: Implications for FlashInfer
Per the CUDA 13.2 Programming Guide — Compute Capabilities (Tables 29–33), SM120 and SM121 share a single "12.x" spec column: identical shared memory (99 KB/block, 100 KB/SM), warps (48/SM), registers (64K/SM), tensor core data types (TF32, BF16, FP16, FP8, FP6, FP4, INT8), TMA, and cluster support. Table 28 confirms compute_120f is compatible with both CC 12.0 and 12.1. The chip-level differences (SM121 = aarch64 host, unified LPDDR5x, 84 SMs vs SM120's x86_64, discrete GDDR7, up to 170 SMs) do not affect kernel correctness — only SM-count-based heuristics and host toolchain.
f vs a compilation targets: performance implications
Per the PTX ISA 9.3 (§9.7.15.5.14 for dense mma, §9.7.15.6.3 for sparse mma.sp), 120f (family) is a subset of 120a/121a (arch-specific). Most SM12x features — including dense NVFP4/MXFP4 block-scaled MMA — are family-specific and available via 120f from PTX ISA 8.8. The arch-specific restriction (requiring sm_120a/sm_121a) applies only to sparse MMA (mma.sp) with .kind::mxf4nvf4/.kind::mxf4:
| PTX Feature |
compute_120 (baseline) |
compute_120f (family) |
compute_120a / compute_121a (arch) |
FP4/FP6/FP8 conversions (cvt) |
No |
Yes |
Yes |
FP4/FP6 MMA, block-scaled MMA (.kind::mxf8f6f4) |
No |
Yes |
Yes |
Dense NVFP4 MMA (.kind::mxf4nvf4) |
No |
Yes (PTX ISA ≥ 8.8) |
Yes |
Dense MXFP4 MMA (.kind::mxf4) |
No |
Yes (PTX ISA ≥ 8.8) |
Yes |
Sparse NVFP4 MMA (mma.sp .kind::mxf4nvf4) |
No |
No |
Yes |
Sparse MXFP4 MMA (mma.sp .kind::mxf4) |
No |
No |
Yes |
FlashInfer's b12x GEMM kernel (dense_blockscaled_gemm_sm120_b12x.py) uses cute.nvgpu.warp.MmaMXF4NVF4Op — a dense block-scaled MMA atom. This means sm_120f wheels already emit native NVFP4/MXFP4 dense MMA instructions. Only sparse NVFP4/MXFP4 MMA (not currently used in FlashInfer) requires sm_120a/sm_121a.
Correction (2026-06-12): An earlier version of this table stated that ALL .kind::mxf4nvf4/.kind::mxf4 MMA was arch-specific-only. Per Harry-Chen's feedback and PTX ISA §9.7.15.5.14 vs §9.7.15.6.3, the restriction is limited to the sparse (mma.sp) variant.
Wheel build arch targets
The release wheel CI (.github/workflows/release.yml) builds the following SM12x targets:
| CUDA Version |
SM12x target in FLASHINFER_CUDA_ARCH_LIST |
| < 12.9 |
12.0a (SM120 arch-specific; no SM121 support) |
| >= 12.9, < 13.0 |
12.0f (family; covers SM120 + SM121; includes dense NVFP4/MXFP4 MMA) |
| >= 13.0 |
12.0f (same) |
No 121a target is included in any wheel build. However, since FlashInfer's FP4 kernels use dense block-scaled MMA (MmaMXF4NVF4Op, not sparse), the 120f wheels already emit native NVFP4/MXFP4 MMA instructions. The 121a target would only be needed for future sparse MMA kernels. JIT compilation on SM121 still produces 121a cubins for forward-compatibility.
What this means for FlashInfer
- Any kernel that runs on SM120 is architecturally compatible with SM121. Hardcoded
sm_version == "sm_120" or minor == 0 checks in backend dispatch are over-restricting without justification.
- Separate cubins matter for JIT but not for wheels.
CompilationContext correctly emits sm_121a for JIT on SM121. Since FlashInfer uses only dense block-scaled MMA (available on 120f), the wheels do not lose FP4 GEMM performance. The 121a JIT path provides forward-compatibility for future sparse MMA use.
- SMEM capacity maps, tile configs, and MMA atoms are identical. Code using
get_smem_capacity_in_bytes("sm_120") returns the correct value for SM121 too — the only risk is KeyError if a map expects an "sm_121" key and doesn't have one.
Key utility functions
Function (utils.py) |
Behavior for SM121 |
is_sm12x_supported |
True (CUDA ≥ 12.9) |
is_sm121a_supported |
True (CUDA ≥ 12.9) |
is_sm120a_supported |
False (12.0 only) |
device_support_pdl |
True (major >= 9) |
determine_attention_backend |
Returns fa2 (fa3 is SM90-only) |
determine_gemm_backend |
Returns sm80 (no SM12x branch) |
1. Attention
1.1 Decode (batch paged KV)
| Backend |
CC Gate |
SM121 Status |
Notes |
| fa2 |
Generic |
Supported |
Default via determine_attention_backend for non-FP8 tensor-core decode |
| fa3 |
SM90 only |
Not available |
determine_attention_backend returns fa3 only on SM90 |
| xqa |
major in {9, 10, 12} |
Supported |
trtllm_batch_decode auto → xqa when major != 10 |
| trtllm-gen |
major == 10 only |
Not available |
Auto selects this only for SM100 |
| cudnn |
Runtime-dependent |
Unknown |
No explicit CC exclusion; depends on cuDNN + driver |
| MLA tensor path |
device_arch == 80 |
Not available |
SM80-only legacy path; use XQA MLA instead |
Dispatch: decode.py:~2429 (trtllm batch), ~1057 (tensor-core), ~1934 (MLA legacy).
1.2 Prefill (batch paged KV)
| Backend |
CC Gate |
SM121 Status |
Notes |
| fa2 |
Generic |
Supported |
Default auto path |
| fa3 |
SM90 only |
Not available |
|
| cudnn |
Runtime-dependent |
Unknown |
No explicit CC block |
| cutlass FMHA |
is_sm100a_supported || is_sm110a_supported |
Not available |
Explicitly excluded for SM12x (no tcgen05); error message names RTX 5090 / DGX Spark |
| FMHAv2 (TRT-LLM) |
is_sm12x_supported |
Supported |
Intended Blackwell 12x path; FP8 FMHAv2 explicitly blocked |
| cute-dsl |
SM100+ code paths |
Uncertain |
Kernels use SM100-oriented helpers (get_smem_capacity_in_bytes("sm_100")); needs on-device validation |
Dispatch: prefill.py:~132–163 (CUTLASS FMHA gate), ~1973 (auto), ~4416 (FMHAv2).
1.3 MLA (Multi-Latent Attention)
| Backend |
CC Gate |
SM121 Status |
Notes |
| xqa |
major == 12 |
Supported |
Auto for decode MLA on SM12x |
| xqa (FP8/NVFP4) |
major in {12} |
Supported |
Docstring says "SM120" but gate is major == 12 (includes 121) |
| trtllm-gen |
major == 10 |
Not available |
|
| cute-dsl |
cc[0] >= 10 |
Supported (gate) |
SM121 passes the check; SM100-oriented internals need validation |
| fa2 |
Generic |
Supported |
Fallback path |
Dispatch: mla/_core.py:~677 (trtllm MLA decode), ~686 (XQA MLA).
1.4 XQA
| Backend |
CC Gate |
SM121 Status |
Notes |
| xqa |
major in {9, 10, 12} |
Supported |
|
| xqa_mla |
major in {12} |
Supported |
JIT gen_xqa_module_mla uses supported_major_versions=[12] |
| xqa NVFP4 |
major in {12} |
Supported |
Error message says "SM120 GPUs" — misleading for SM121 |
Dispatch: xqa.py:~316 (general), ~536 (MLA). JIT: jit/xqa.py:~97 (xqa), ~172 (xqa_mla).
1.5 POD / Cascade / Sparse / Attention Sink
| API |
SM121 Status |
Notes |
| POD |
Inherits prefill + decode |
Same as §1.1 / §1.2 |
Cascade (merge_state) |
Supported |
No arch filter |
| Sparse |
Supported (fa2) |
Auto → fa2 on SM121 |
| Attention sink |
Supported (fa2) |
Auto → fa2 on SM121 |
1.6 cuDNN Attention
No explicit CC exclusion in cudnn/prefill.py or cudnn/decode.py. Runtime-dependent on cuDNN version and driver support for Spark hardware.
2. GEMM
2.1 mm_fp4 / FP4 Dense GEMM
| Backend |
CC Gate |
SM121 Status |
Notes |
| cutlass |
[100, 103, 110, 120, 121] |
Supported |
get_gemm_sm120_module_cutlass_fp4() for sm_major == 12 |
| cudnn |
[100, 103, 110, 120, 121] |
Supported |
MXFP4 on 12x needs cuDNN ≥ 9.14.0 |
| b12x |
[120, 121] (decorator) |
Supported (gate) but excluded from auto |
Heuristic checks major == 12 and minor == 0 — SM121 never gets b12x in auto |
| trtllm |
[100, 103] |
Not available |
|
| cute-dsl |
[100, 103] |
Not available |
Uses SM100-style kernel; not enabled for SM12x |
Key issue: _heuristic_func_mm_fp4 (gemm_base.py:5065–5071) sets is_sm120 = major == 12 and minor == 0; only then prepends b12x. SM121 falls through to cutlass/cudnn only. A prior internal FP4 audit also identified this as the likely root cause of DGX Spark NVFP4 performance complaints, along with related issues: AOT building fp4_quantization_121 while runtime redirects to 120f, the sm120a FP4 quant module being dead code, and an upstream nvidia-cutlass-dsl blocker (missing "sm_121" in SMEM_CAPACITY_MAP, pending cutlass-dsl 4.5).
2.2 mm_bf16 / bmm_bf16
| Backend |
CC Gate |
SM121 Status |
Notes |
| cudnn |
[80, 86, 89, 90, 100, 103] |
Not available |
121 not in list |
| cutlass |
[100, 103] |
Not available |
|
| tgv |
[100, 103] |
Not available |
|
SM121 has no BF16 mm/bmm backend through the standard bf16_gemm_sm100 path. Falls back to generic segment GEMM or PyTorch.
2.3 mm_mxfp8 / bmm_mxfp8
| Backend |
CC Gate |
SM121 Status |
Notes |
| cutlass |
[120, 121] for mm; [120, 121] for bmm |
Supported |
Enforces 1D swizzled scales, 128×4 block, K/N multiples of 32 on SM12x |
| cute-dsl |
[100, 103] |
Not available |
|
| trtllm |
[100, 103] |
Not available |
|
| cudnn (bmm) |
[100, 103] |
Not available |
|
Dispatch: gemm_base.py:~3984+ (mm_mxfp8), ~7503+ (bmm_mxfp8).
2.4 bmm_fp8
| Backend |
CC Gate |
SM121 Status |
Notes |
| cutlass_sm12x |
_match_sm_version(["120", "121"]) |
Supported |
Uses get_gemm_sm120_module_cutlass_fp8 |
| cudnn |
CC-dependent |
Unknown |
|
| cublas |
Generic |
Supported |
|
2.5 Groupwise / Blockscaled GEMM
| API |
CC Gate |
SM121 Status |
Notes |
gemm_fp8_nt_groupwise |
[100, 103, 120, 121] |
Supported |
Uses get_gemm_sm120_module() on SM12x |
gemm_fp8_nt_blockscaled |
[100, 103, 120, 121] |
Supported |
|
group_gemm_fp8_nt_groupwise |
[100, 103, 120, 121] |
Supported |
Correctness issues for num_groups > 1 on SM120/121 (documented) |
group_gemm_mxfp4_nt_groupwise |
[120, 121] |
Supported |
|
group_gemm_nvfp4_nt_groupwise |
[120, 121] |
Supported |
|
2.6 DeepGEMM
| API |
CC Gate |
SM121 Status |
group_deepgemm_fp8_nt_groupwise |
[100, 103] |
Not available |
batch_deepgemm_fp8_nt_groupwise |
[100, 103] |
Not available |
2.7 Other GEMM
| API |
SM121 Status |
Notes |
SegmentGEMMWrapper |
Not SM12x-optimized |
Uses determine_gemm_backend → "sm80"; no SM12x branch |
routergemm (tinygemm_bf16, etc.) |
Not available |
[100, 103] only |
tgv_gemm_sm100 |
Not available |
SM10x stack |
mm_fp8 (TRT-LLM low-latency) |
Unknown |
No explicit CC gate in Python |
2.8 GEMM JIT
All gen_gemm_sm120_* generators use get_nvcc_flags_list(supported_major_versions=[12]) — CompilationContext supplies sm_121a when device is SM121. CUTLASS kernel generation (generate_kernels.py:1049) explicitly checks has_arch(120) or has_arch(121).
3. MoE (Mixture of Experts)
3.1 cutlass_fused_moe (FP8)
| Backend |
CC Gate |
SM121 Status |
Notes |
| cutlass SM120 |
major*10+minor → "120"/"121" both → gen_cutlass_fused_moe_sm120_module |
Supported |
JIT uses supported_major_versions=[12]; FP8 MoE may fall back to SM89 tactics when native occupancy is zero |
| trtllm MoE |
arch[0] >= 10 |
Supported |
No 120/121 split |
3.2 b12x_fused_moe (NVFP4 MoE)
| Backend |
CC Gate |
SM121 Status |
Notes |
| b12x |
@supported_compute_capability([120, 121]) |
Supported (gate) |
Docstrings say SM120/SM121; uses dense_blockscaled_gemm_sm120 internals |
3.3 cute_dsl_fused_moe_nvfp4
| Backend |
CC Gate |
SM121 Status |
| cute-dsl |
[100, 103] |
Not available |
4. Norm
4.1 RMSNorm / LayerNorm / Gemma Norm
| API |
Backend |
SM121 Status |
Notes |
rmsnorm |
CUDA JIT or CuTe DSL |
Supported |
No CC gate; PDL enabled (__CUDA_ARCH__ >= 900) |
fused_add_rmsnorm |
CUDA JIT or CuTe DSL |
Supported |
|
gemma_rmsnorm |
CUDA JIT or CuTe DSL |
Supported |
|
gemma_fused_add_rmsnorm |
CUDA JIT or CuTe DSL |
Supported |
|
layernorm |
CUDA JIT or CuTe DSL |
Supported |
|
fused_add_layernorm |
CUDA JIT or CuTe DSL |
Supported |
|
All *_quant variants |
CUDA JIT or CuTe DSL |
Supported |
|
4.2 fused_rmsnorm_silu
| Dtype |
CC Gate |
SM121 Status |
Notes |
| BF16/FP16 |
sm_version >= 80 |
Supported |
|
| FP8 output |
sm_version >= 89 |
Supported |
|
| NVFP4 output |
sm_version >= 100 |
Supported |
SM121 = 121 ≥ 100; uses SM100+ sweep LUT |
4.3 rmsnorm_fp4quant / add_rmsnorm_fp4quant
CuTe DSL only (requires nvidia-cutlass-dsl). Uses get_sm_version = major*10+minor → 121. No explicit SM121 exclusion. Supported assuming CuTe DSL + CUDA support FP4 PTX on SM121.
5. RoPE
| API |
Backend |
SM121 Status |
Notes |
apply_rope / *_inplace |
CUDA JIT |
Supported |
No CC gate |
apply_rope_pos_ids / *_inplace |
CUDA JIT |
Supported |
|
apply_llama31_rope / *_inplace |
CUDA JIT |
Supported |
|
apply_rope_with_cos_sin_cache / *_inplace |
CUDA JIT |
Supported |
|
rope_quantize_fp8 |
CUDA JIT |
Supported |
|
mla_rope_quantize_fp8 |
CUDA JIT |
Supported |
|
All RoPE kernels use PDL via __CUDA_ARCH__ >= 900 (SM121 qualifies). No Triton RoPE in-tree.
6. Activation
| API |
Backend |
SM121 Status |
Notes |
silu_and_mul |
CUDA JIT |
Supported |
PDL enabled on SM121 |
gelu_and_mul |
CUDA JIT |
Supported |
|
gelu_tanh_and_mul |
CUDA JIT |
Supported |
|
silu_and_mul_scaled_nvfp4_experts_quantize |
TRT-LLM FP4 JIT |
Supported |
Uses get_fp4_quantization_module("121") → remaps to "120f" on CUDA ≥ 12.9 |
Triton silu_and_mul also available with no SM-specific guards.
7. Sampling
| API |
Backend |
SM121 Status |
Notes |
sampling_from_logits / *_probs |
CUDA JIT |
Supported |
Arch-agnostic |
top_k_* / top_p_* / min_p_* |
CUDA JIT |
Supported |
|
chain_speculative_sampling |
CUDA JIT |
Supported |
|
top_k (clusters path) |
CUDA JIT |
Not available |
can_use_clusters_topk requires cap[0] == 10 (SM100 only) |
8. Quantization
| API |
Backend |
SM121 Status |
Notes |
fp4_quantize (TRT-LLM style) |
CUDA JIT (120f) |
Supported |
Runtime remaps "121" → "120f" on CUDA ≥ 12.9 |
nvfp4_* / mxfp4_* |
CUDA JIT |
Supported |
@supported_compute_capability includes 121 |
Generic quantization module |
CUDA JIT |
Supported |
No per-SM override |
mxfp8_quantization |
CUDA JIT |
Supported |
gen_mxfp8_quantization_sm100_module uses all arches from context |
9. Mamba
| API |
Backend |
SM121 Status |
Notes |
selective_state_update |
CUDA JIT (SM100 generator) |
Supported |
major >= 10 → uses gen_selective_state_update_sm100_module; not in AOT list — JIT on first use |
SSDCombined |
CuTe DSL |
Supported |
major >= 10 |
AOT gap: gen_selective_state_update_sm100_module is not in the default AOT build list. SM121 users with FLASHINFER_DISABLE_JIT will hit a missing-module error.
10. Communication
| API |
Backend |
SM121 Status |
Notes |
alltoall |
CUDA JIT (generic) |
Supported |
All arches from compilation context |
vllm_comm |
CUDA JIT (generic) |
Supported |
|
trtllm_allreduce |
CUDA JIT |
Not available |
supported_major_versions=[9, 10] — SM12 excluded |
trtllm_mnnvl |
AOT only |
Not available |
Only built when has_sm100 |
moe_alltoall |
AOT only |
Not available |
Only built when has_sm100 |
11. GDN (Gated Delta Net)
| API |
Backend |
SM121 Status |
Notes |
gdn_prefill |
CUDA JIT (SM90 path) |
Broken |
Blackwell CuTe path requires is_sm100a_supported — not SM121. Falls through to SM90 path, but that JIT module is compiled with sm90a_nvcc_flags (compute_90a,code=sm_90a), producing SM90-only SASS with no PTX fallback. Will fail to load on SM121. |
gdn_decode |
CUDA JIT (SM90 path) |
Broken (same issue) |
Also compiled with sm90a flags; incompatible with SM121 |
12. Page (KV Cache Ops)
KV cache append/management via generic CUDA JIT. No CC gates. Supported.
Cross-Cutting Issues Summary
| # |
Severity |
Area |
Issue |
| 1 |
HIGH |
GEMM |
b12x FP4 backend excluded from auto on SM121 (heuristic checks minor == 0) |
| 2 |
HIGH |
GEMM |
mm_bf16 / bmm_bf16 have no backend available for SM121 |
| 3 |
MEDIUM |
Quantization |
AOT builds fp4_quantization_121 module but runtime always uses 120f — wasted artifact + JIT penalty if 120f not prebuilt |
| 4 |
MEDIUM |
Mamba |
gen_selective_state_update_sm100_module missing from AOT — JIT penalty on first SSU call; breaks FLASHINFER_DISABLE_JIT |
| 5 |
MEDIUM |
Comm |
trtllm_allreduce JIT limited to SM 9/10 — no fused allreduce on SM121 |
| 6 |
HIGH |
GDN |
GDN prefill falls through to SM90 path on SM121, but the JIT module is compiled with sm90a_nvcc_flags (compute_90a,code=sm_90a) — incompatible with SM121; would fail to load at runtime |
| 7 |
LOW |
Sampling |
Clusters top-k optimization SM100-only; SM121 uses slower fallback |
| 8 |
LOW |
GEMM |
SegmentGEMMWrapper not SM12x-optimized (uses sm80 backend) |
| 9 |
INFO |
Quantization |
sm120a FP4 quant module is dead code (always remapped to 120f on CUDA ≥ 12.9) |
| 10 |
INFO |
GEMM |
SM100 FP4 CUTLASS compiles for SM12x needlessly (supported_major_versions=[10, 11, 12]) |
Test-related and documentation issues are covered in §13.5 and §14.3 respectively.
Key Files Reference
| File |
Relevance |
flashinfer/compilation_context.py |
Arch normalization (120a vs 120f vs 121a) |
flashinfer/utils.py |
is_sm12x_supported, is_sm121a_supported, determine_attention_backend, determine_gemm_backend |
flashinfer/jit/core.py |
NVCC flag definitions (sm120f_nvcc_flags, sm121a_nvcc_flags) |
flashinfer/gemm/gemm_base.py |
All GEMM backend dispatch, b12x runner, heuristics |
flashinfer/fused_moe/core.py |
MoE dispatch |
flashinfer/fused_moe/cute_dsl/b12x_moe.py |
b12x MoE API + CC gate |
flashinfer/quantization/fp4_quantization.py |
FP4 quant module selection + 121→120f redirect |
flashinfer/aot.py |
AOT build module list + detect_sm_capabilities() |
flashinfer/prefill.py |
Prefill backend dispatch incl. CUTLASS FMHA exclusion |
flashinfer/decode.py |
Decode backend dispatch |
flashinfer/mla/_core.py |
MLA backend dispatch |
flashinfer/xqa.py |
XQA CC gates |
flashinfer/mamba/selective_state_update.py |
Mamba SSU dispatch |
flashinfer/jit/comm.py |
Comm module JIT (trtllm_allreduce SM9/10 only) |
flashinfer/gdn_prefill.py |
GDN prefill SM100a vs SM90 branching |
Part 2: Test Coverage & Arch-Check Audit
13. Unit Test SM121 Skip / xfail / Waiver Inventory
13.1 Attention Tests
| Test File |
Condition |
SM121 Impact |
Justified? |
test_batch_attention.py:207–249 |
xfail when cap[0] == 12 |
xfail on SM121 (and SM120) — tile size/stages |
Yes — known limitation for SM12x family |
test_fmha_v2_prefill.py:486–489 |
is_sm90a_supported or is_sm120a_supported |
SM121 SKIPPED — uses is_sm120a (12.0-only) not is_sm12x |
Gap — message says "SM12x" but condition excludes SM121 |
test_fmha_v2_prefill.py:492–506 |
is_sm120_plus = is_sm120a_supported(...) |
SM121 excluded from FP8/sliding/SQV sub-tests |
Gap — should use is_sm12x_supported if FMHAv2 is intended for SM12x |
test_xqa.py:123–126 |
cap[0] not in [9, 10, 12] |
SM121 included |
OK |
test_xqa.py:468–471 |
cap[0] not in [12] |
SM121 included |
OK |
test_xqa_batch_decode.py:393–396 |
cap[0] not in [9, 10, 12] |
SM121 included |
OK |
test_xqa_batch_decode.py:582–585 |
cap[0] not in [12] |
SM121 included but reason says "SM120 GPUs" only |
Minor doc gap |
test_xqa_mla_batch_decode.py:25–27 |
cap[0] != 12 |
SM121 included |
OK |
test_trtllm_gen_mla.py:276–279 |
cap[0] != 12 for xqa backend |
SM121 included |
OK |
13.2 GEMM Tests
| Test File |
Condition |
SM121 Impact |
Justified? |
test_mm_fp4.py:40–45 |
cap[0] != 12 or cap[1] != 0 |
SM121 SKIPPED for b12x backend — reason: "b12x only supports SM120" |
Aligned with current heuristic; gap if b12x should work on SM121 |
test_mm_fp4.py:105–106 |
cuDNN xfail for FP4/MXFP4 |
SM120-specific cuDNN quirk |
OK |
test_bmm_mxfp8.py:23–24 |
cap[0] in [11, 12] → skip |
SM121 SKIPPED — "Not tested on SM110/SM120/SM121" |
Gap — MXFP8 BMM cutlass lists [120, 121] in supported_compute_capability but test skips both |
test_bmm_fp8.py:20–23 |
cap[0] not in [10, 11, 12] for cutlass |
SM121 included |
OK |
test_group_gemm_fp4.py:118–121 |
cap[0] not in [12] |
SM121 included |
OK |
test_groupwise_scaled_gemm_fp8.py:202–212 |
group_size > 1 and cap[0] in [12] → skip |
SM121 SKIPPED (same as SM120) — known correctness issue |
Justified — documented bug |
test_groupwise_scaled_gemm_mxfp4.py:257–335 |
cap[0] not in [10, 12] |
SM121 included |
OK |
test_mm_mxfp8.py:88–518 |
cap[0] == 12 for swizzled scales |
SM121 included |
OK |
test_mm_mxfp8_sm120.py:12–24 |
cc[0] == 12 |
SM121 included (despite "SM120" naming) |
OK |
13.3 MoE Tests
| Test File |
Condition |
SM121 Impact |
Justified? |
test_b12x_fused_moe.py:88–90 |
@not_sm121 (_is_sm121()) |
SM121 SKIPPED for ALL b12x MoE tests — 10+ test methods |
Major gap — API decorator lists [120, 121] but tests say "not supported on SM121" |
test_trtllm_cutlass_fused_moe.py:486–488 |
cap[0] not in [10, 11, 12] |
SM121 included for NVFP4 MoE |
OK |
test_trtllm_cutlass_fused_moe.py:1826–1828 |
not is_sm100a_supported |
SM121 skipped for unswizzled_input_sf |
Questionable — only checks SM100a |
test_trtllm_cutlass_fused_moe.py:1981–1984 |
not is_sm100a and not is_sm12x |
SM121 included |
OK |
13.4 Utils / Other Tests
| Test File |
Condition |
SM121 Impact |
Justified? |
test_jit_example.py:169–171 |
xfail when cap == (12, 1) |
SM121 xfail — "Numerical accuracy issue on SM 121 (Spark)" |
Justified if reproducible; should verify if still needed |
test_fp4_quantize.py / test_fp4_quantize_padding.py |
is_sm12x_supported |
SM121 included |
OK |
test_gen_module_symlink_race_condition.py |
is_sm100a or is_sm12x |
SM121 included |
OK |
13.5 Test Gap Summary
| # |
Severity |
Test File |
Issue |
| T1 |
HIGH |
test_fmha_v2_prefill.py |
Uses is_sm120a_supported instead of is_sm12x_supported — SM121 skips all FMHAv2 prefill tests |
| T2 |
HIGH |
test_b12x_fused_moe.py |
All b12x MoE tests have @not_sm121 — contradicts API decorator [120, 121] |
| T3 |
MEDIUM |
test_bmm_mxfp8.py |
Skips SM12x entirely despite MXFP8 BMM cutlass supporting [120, 121] |
| T4 |
LOW |
test_jit_example.py |
SM121-specific xfail for numerical accuracy — needs re-verification |
| T5 |
LOW |
test_xqa_batch_decode.py |
NVFP4 XQA reason string says "SM120" but condition includes SM121 |
| T6 |
LOW |
test_trtllm_cutlass_fused_moe.py |
unswizzled_input_sf test only checks is_sm100a — SM12x not considered |
14. Arch-Check Over-Restriction Audit
Audit of all sm120/sm121/sm120a/sm121a/sm120f checks in non-test source code for cases where logic is more restrictive than necessary.
14.1 Confirmed Over-Restrictions
| # |
File |
Lines |
Pattern |
Issue |
Suggested Fix |
| A1 |
utils.py |
574–576 |
is_sm120f_supported: requires minor == 0 |
Name implies SM12x family on CUDA 12.9, but implementation excludes SM121 |
Change to major == 12 (or rename to clarify SM120-only intent). Currently unused in the codebase. |
| A2 |
gemm/gemm_base.py |
5065–5070 |
is_sm120 = major == 12 and minor == 0 in _heuristic_func_mm_fp4 |
b12x backend never selected in auto for SM121 |
Change to major == 12 if b12x is validated on SM121 |
| A3 |
gemm/kernels/dense_blockscaled_gemm_sm120.py |
1591–1594 |
if sm_version != "sm_120": raise ValueError(...) |
Rejects any SM version string other than "sm_120" |
Allow "sm_121" (or any sm_12*); architecturally compatible (same MMA atom, same smem) |
| A4 |
cute_dsl/attention/compat.py |
29–36 |
_TMEM_MAX_ALLOC_COLUMNS_MAP has "sm_120": 512 but no "sm_121" |
Would raise KeyError if "sm_121" ever passed to fallback path |
Add "sm_121": 512 |
14.2 Confirmed Intentional (Not Over-Restricting)
| # |
File |
Lines |
Pattern |
Why Intentional |
| B1 |
compilation_context.py |
47–54 |
SM120 → 0f, SM121 → 1a (separate gencodes) |
Required for correct cubins — SM120 SASS can cause illegal instructions on SM121 |
| B2 |
quantization/fp4_quantization.py |
157–165 |
"120" / "121" both → "120f" on CUDA ≥ 12.9 |
Correct family-level consolidation |
| B3 |
jit/attention/fmha_v2/fmha_library.py |
Various |
kspec.sm == 120 as Blackwell bucket |
Design choice: 120 = logical SM12x FMHA label; NVCC gets actual arch from CompilationContext |
| B4 |
jit/gemm/cutlass/generate_kernels.py |
1049 |
has_arch(120) or has_arch(121) |
Explicitly includes both |
| B5 |
fused_moe/core.py |
206–207 |
backend in ("120", "121") → shared module |
Both mapped to same JIT module |
| B6 |
aot.py |
380–540 |
Separate has_sm120, has_sm121, has_sm120f |
AOT needs per-variant control; not over-restricting |
| B7 |
All @supported_compute_capability decorators |
Various |
Lists include both 120 and 121 |
Consistently inclusive across all audited APIs |
14.3 Cosmetic / Documentation Issues
| # |
File |
Lines |
Issue |
| C1 |
prefill.py |
~4423, 4464 |
Docstrings say "SM120 (Blackwell)" without mentioning SM121 |
| C2 |
decode.py |
~2338 |
Mentions sm_120/sm_121 but some messages say just "SM120" |
| C3 |
xqa.py |
~310 |
NVFP4 error message says "SM120 GPUs" but gate is major == 12 |
| C4 |
fused_moe/cute_dsl/blackwell_sm12x/ |
Various |
Function names launch_sm120_*, allocate_sm120_* — naming suggests SM120-only but used for SM12x family |
| C5 |
gemm/kernels/dense_blockscaled_gemm_sm120.py |
Various |
File and function names say SM120; architecturally covers SM12x |
15. Consolidated Action Items
Must Fix (blocking SM121 correctness or coverage)
Should Fix (improve robustness)
Nice to Have (documentation / cosmetic)
