feat: add SM_121 (GB10 consumer Blackwell) support for FA4

[TyGu1]

Summary

Adds SM_121 (GB10 / DGX Spark consumer Blackwell) support to TransformerEngine, enabling Flash Attention 4 on GB10 hardware. SM_121 is a minor SM_120 variant sharing the same ISA (SM80-era MMA + CpAsync, no TMA) but with a slightly different arch integer (12,1 vs 12,0).

Primary test result (GB10 hardware): test_dpa_fa4_sm121 — 3/3 PASS

Full test results with root-cause analysis for all remaining failures: see te_sm121_test_results.txt in this PR.

---

Changes (6 files)

1. build_tools/utils.py

• get_cuda_include_dirs() rewritten to be additive — always supplements the system CUDA toolkit with headers from the nvidia-cuda-* pip wheels. This fixes builds on systems where only pip-installed CUDA is available (DGX Spark default configuration).
• get_nccl_include_dirs() kept as a compatibility shim (returns empty list).

2. transformer_engine/common/CMakeLists.txt

• Added unconditional NCCL header search so the CMake build finds NCCL headers from pip wheels independently of the Python-side discovery.
• Added SM_121 arch classify block (clones the SM_120 pattern): 121f for CUDA ≥ 12.9, else 121a.

3. transformer_engine/common/fused_attn/fused_attn.cpp

• Extended sm_arch_ == 120 cuDNN special-case to (sm_arch_ == 120 || sm_arch_ == 121). Body is identical — SM_121 uses the same cuDNN path as SM_120.

4. transformer_engine/pytorch/attention/dot_product_attention/utils.py

• Updated FA4 support comment to include SM_121. No logic change — SM_121 already passes the >= (8, 0) gate.

5. tests/pytorch/attention/test_attention.py

• Added test_dpa_fa4_sm121 gated on device_compute_capability == (12, 1), using model_configs_fa4_sm121 (MHA-only, no GQA/splits — those fail due to upstream flash-attn-4 limitations noted in the test comments and in te_sm121_test_results.txt).

6. te_sm121_test_results.txt

• Full GB10 test run output with root-cause analysis for all failures (all upstream flash-attn-4 issues, not SM_121 or TE defects).

---

Flash Attention 4 SM_121 Bringup Findings

During SM_121 bringup with flash-attn-4 v4.0.0b16, 5 bugs in the flash-attn-4 package were found and patched locally (in the installed package, not in TE). These should be fixed upstream in flash-attn-4:

#	File	Symptom	Root Cause
1	flash_fwd.py:658	AttributeError: 'NoneType'.._trait	FlashAttentionForwardSm120.__init__ clobbers class arch=80 with Arch.sm_121; sm_121 >= sm_90 → tries tma_atom_O=None
2	interface.py bwd	UnboundLocalError: dQ_single_wg	SM_12x backward branch never sets dQ_single_wg, used in compile_key for arch//10 in [8,9,12]
3	flash_bwd.py:440	DSLRuntimeError: None to Float32	compute_softmax_scale_log2 returns None; was clobbering softmax_scale needed by kernel
4	utils.py:486	TypeError: atomicrmw() got 'res'	res=T.f32() kwarg removed from nvvm.atomicrmw in nvidia-cutlass-dsl 4.5.2
5	interface.py fwd+bwd	DSLRuntimeError: Int32 expected, got int	cute.compile requires Int32, not plain Python int, for Optional[Int32] params; SM_80/12x path had no implicit coercion

---

Test Results Summary (GB10 / SM_121 hardware)

test_dpa_fa4_sm121 (primary):  3/3 PASS  ✓

Isolation results by category:
  test_dpa_fa4_base           :  3/7  (gqa/splits → upstream flash-attn-4 bugs)
  test_dpa_fa4_sm121          :  3/3  PASS ✓
  test_dpa_fa4_mask           :  3/6  (padding → TE cuDNN backend, not FA4)
  test_dpa_fa4_varlen         :  3/8  (thd layout + gqa → upstream FA4/cuDNN)
  test_dpa_fa4_sliding_window :  0/8  (SM_80 bwd has no local-mask impl)

All failures are pre-existing upstream flash-attn-4 limitations affecting the SM_80 code path (which SM_121 shares). None are SM_121-specific regressions.

---

Testing

Tested on NVIDIA GB10 (DGX Spark), device_compute_capability = (12, 1):

pytest tests/pytorch/attention/test_attention.py::test_dpa_fa4_sm121 -v
# → 3 passed

GitHub CI (compile + import, CPU-only) cannot run FA4 kernels. SM_121 hardware results are attached via te_sm121_test_results.txt. Requesting /te-ci GPU CI trigger from @cyanguwa or team.

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

This PR adds SM_121 (GB10 / DGX Spark consumer Blackwell) support to TransformerEngine by extending every SM_120-specific runtime guard and build-system arch list to also cover SM_121, and by reworking get_cuda_include_dirs() to include pip-wheel headers on top of any system CUDA toolkit so that DGX Spark's pip-only CUDA environment builds correctly.

• The C++ and Python runtime guards (fused_attn.cpp, fused_attn_f16_arbitrary_seqlen.cu, utils.py, context_parallel.py) are consistently and correctly extended to cover SM_121 alongside SM_120.
• The CMake find_path(NVTE_NCCL_INCLUDE_DIR ...) block added for pip-based NCCL discovery has no PATHS or HINTS pointing to site-packages, so it will silently fail on DGX Spark (pip-only NCCL) — the exact environment the PR targets.
• get_nccl_include_dirs() is introduced as a new function that always returns [] and is immediately called as a no-op; it should be removed.

Confidence Score: 3/5

The runtime attention-backend changes are correct and consistent; the CMake NCCL discovery silently fails on pip-only DGX Spark, leaving the primary stated use case potentially broken at the cmake build level.

FA4 kernel routing, cuDNN workarounds, and Python backend-selection logic are all correctly mirrored from SM_120 to SM_121, and the dedicated test passes 3/3. The gap is find_path in CMakeLists.txt having no path hints for pip wheel directories, so nccl.h won't resolve on DGX Spark where NCCL comes only from pip.

transformer_engine/common/CMakeLists.txt (NCCL find_path needs HINTS for pip wheel paths) and build_tools/utils.py (no-op get_nccl_include_dirs + CUDA version comment vs test evidence)

Important Files Changed

Filename	Overview
build_tools/utils.py	Rewrites get_cuda_include_dirs() to be additive (system toolkit + pip wheels); adds no-op get_nccl_include_dirs() shim; adds SM_121 to arch list for CUDA >= 12.9 (but tests ran on 12.8)
transformer_engine/common/CMakeLists.txt	Adds SM_121 arch block (mirrors SM_120 pattern); adds unconditional NCCL find_path that won't locate pip-installed headers in site-packages without PATHS/HINTS
transformer_engine/common/fused_attn/fused_attn.cpp	Extends cuDNN backend special-case from sm_arch_ == 120 to (120
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu	All five SM_120 exclusion guards consistently updated to also exclude SM_121
transformer_engine/pytorch/attention/dot_product_attention/utils.py	SM_121 added to all SM_120 equality/membership guards for FP8 disabling, KV-caching disabling, MLA head-dim check, and THD layout checks
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Adds (12, 1) to the softmax_lse_in_packed_format exclusion guard; backward pass uses ctx value saved from forward
tests/pytorch/attention/test_attention.py	Adds test_dpa_fa4_sm121 gated on device_compute_capability == (12, 1) with MHA-only configs; GQA and splits excluded with clear upstream-limitation comments
build_tools/pytorch.py	Imports and calls get_nccl_include_dirs() which always returns []; the call is a no-op
te_sm121_test_results.txt	Committed test-run artifact documenting GB10 results; unconventional to merge into the source tree

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Runtime: sm_arch_ detected] --> B{sm_arch_ == 120 or 121?}
    B -->|Yes| C[fused_attn.cpp: cuDNN special-case path]
    B -->|No| D[Default cuDNN path]
    C --> E{cudnn < 9.18.1?}
    E -->|Yes| F[No Backend]
    E -->|No| G{deterministic training?}
    G -->|Yes| F
    G -->|No| H{T3HD/TH3D layout?}
    H -->|Yes| F
    H -->|No| I[NVTE_F16_arbitrary_seqlen backend]
    I --> J[fused_attn_f16_arbitrary_seqlen.cu]
    J --> K{sm_arch_ != 120 and != 121?}
    K -->|True| L[Packed ragged stats + token-count dims]
    K -->|False| M[BHSD-like strides, max_seqlen]
    N[Python: device_compute_capability] --> O{cap in 12,0 or 12,1?}
    O -->|Yes| P[Disable FP8, FlashAttn KV-cache, MLA bwd checks]
    O -->|No| Q[Standard backend selection]
    P --> R[FA4 still enabled cap >= 8,0]
    Q --> R
    S[Build: cuda_archs / CMakeLists] --> T{CUDA >= 12.9?}
    T -->|Yes| U[Include SM_121 in arch list]
    T -->|No CUDA >= 12.8| V[SM_120 only, no SM_121]

Loading

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

find_path won't locate pip-installed NCCL headers

find_path without PATHS or HINTS searches only standard system include paths (e.g., /usr/include, /usr/local/include) and any CMAKE_PREFIX_PATH entries. On DGX Spark — the primary target of this PR — NCCL headers land in a site-packages tree like ~/.local/lib/python3.x/site-packages/nvidia/nccl/include/, which is never in CMake's default search roots. The QUIET flag silences the failure, NVTE_NCCL_INCLUDE_DIR is left empty, and target_include_directories is never called. Any translation unit that includes nccl.h (e.g., comm_gemm_overlap.h) will then fail to compile in a pip-only CUDA environment — the exact scenario this PR is supposed to fix. A HINTS pointing to the Python-resolved path (passed in via -DNVTE_NCCL_INCLUDE_HINT from build_ext.py) or an execute_process call to query Python would close the gap.

[greptile-apps]

Dead shim with a misleading "compatibility" label

get_nccl_include_dirs() is a brand-new function (introduced in this PR) that immediately returns []. The docstring calls it a "compatibility shim", but there are no pre-existing callers to stay compatible with — pytorch.py is also calling it for the first time in this same PR. The only concrete effect of the include_dirs.extend(get_nccl_include_dirs()) call in pytorch.py is extend([]), which is a no-op. The function and its import can be removed without any behavioral change; NCCL headers are now properly included via get_cuda_include_dirs().

[greptile-apps]

Comment contradicts the test evidence

The inline comment says "confirmed under CUDA >= 12.9; minimum version for 12.8 needs verification", yet te_sm121_test_results.txt (committed in this same PR) explicitly shows CUDA 12.8 in its header line and reports 3/3 passing. Either SM_121 should be added to the >= (12, 8) arch list (if NVCC 12.8 accepts sm_121), or the test results file should clarify that the FA4 kernel path uses JIT compilation (bypassing NVCC's arch list) so the CUDA version constraint only affects native TE CUDA kernels.

[greptile-apps]

Test-result artifact committed to source tree

Committing generated test output files to the repository is generally discouraged — they become stale as the code evolves, are not run as part of CI, and add noise to git log. The conventional approach is to attach them to the PR description (as a GitHub comment or linked gist) rather than merging them into the source tree. If the intent is permanent documentation, consider a location under docs/ with a clear label rather than the repo root.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!