Skip to content

feat(attention): enable SM121 (GB10/DGX Spark) for NVFP4 attention #3897

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
saltyminty merged 4 commits into from
Jul 13, 2026
Merged

feat(attention): enable SM121 (GB10/DGX Spark) for NVFP4 attention #3897

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a5922cd
Add SM121 support
bkryu Jul 9, 2026
871222a
Put back aot compilation for CUDA 12.8 and SM120
bkryu Jul 13, 2026
c3baa30
Merge branch 'main' into sm120_nvfp4_prefill_on_sm121
bkryu Jul 13, 2026
30040e8
If else to try except
bkryu Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions benchmarks/bench_nvfp4_attention_sm120.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,10 +159,10 @@ def skip_unless_sm120() -> None:
sys.exit(0)

capability = torch.cuda.get_device_capability()
if capability != (12, 0):
if capability not in ((12, 0), (12, 1)):
print(f"Current device capability: {capability}.")
print(
"Skipping: NVFP4 attention SM120 benchmark requires compute capability (12, 0)."
"Skipping: NVFP4 attention SM120 benchmark requires compute capability (12, 0) or (12, 1)."
)
sys.exit(0)

Expand Down
4 changes: 2 additions & 2 deletions csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,8 +231,8 @@ void fwd(TensorView q_fp4, TensorView k_fp4, TensorView v_fp4_t, TensorView q_sc
cudaError_t status = cudaGetDeviceProperties(&props, q_fp4.device().device_id);
TVM_FFI_ICHECK(status == cudaSuccess)
<< "cudaGetDeviceProperties failed: " << cudaGetErrorString(status);
TVM_FFI_ICHECK(props.major == 12 && props.minor == 0)
<< "NVFP4 attention SM120 kernel requires compute capability 12.0";
TVM_FFI_ICHECK(props.major == 12 && (props.minor == 0 || props.minor == 1))
<< "NVFP4 attention SM120 kernel requires compute capability 12.0 or 12.1";

const int64_t batch = q_fp4.size(0);
const int64_t num_heads = q_fp4.size(1);
Expand Down
2 changes: 1 addition & 1 deletion flashinfer/aot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -518,7 +518,7 @@ def gen_all_modules(
add_oai_oss,
)
)
if has_sm120:
if has_sm120 or has_sm121:
jit_specs.append(gen_nvfp4_attention_sm120_module())

if add_act:
Expand Down
7 changes: 5 additions & 2 deletions flashinfer/jit/nvfp4_attention_sm120.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from pathlib import Path

from . import env as jit_env
from .core import JitSpec, gen_jit_spec, sm120a_nvcc_flags
from .core import JitSpec, current_compilation_context, gen_jit_spec


_NVFP4_ATTENTION_SM120_MODULE_NAME = "nvfp4_attention_sm120"
Expand Down Expand Up @@ -89,9 +89,12 @@ def gen_nvfp4_attention_sm120_module() -> JitSpec:
]
include_paths: list[str | Path] = []
include_paths.extend(_nvfp4_attention_sm120_include_paths())
nvcc_flags = current_compilation_context.get_nvcc_flags_list(

@saltyminty saltyminty Jul 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change causes sm120 to no longer be supported on CUDA 12.8, as _normalize_cuda_arch enforces >= 12.9. Is this intended?

@bkryu bkryu Jul 13, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right thanks @saltyminty, this was unintended. Fixed in the latest commit with a gated compilation-context path

supported_major_versions=[12]
)
return gen_jit_spec(
_NVFP4_ATTENTION_SM120_MODULE_NAME,
source_paths,
extra_cuda_cflags=sm120a_nvcc_flags + _NVFP4_ATTENTION_SM120_CUDA_FLAGS,
extra_cuda_cflags=nvcc_flags + _NVFP4_ATTENTION_SM120_CUDA_FLAGS,
extra_include_paths=include_paths,
)
4 changes: 2 additions & 2 deletions flashinfer/nvfp4_attention_sm120.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ def _preprocess_qkv(
return q.contiguous(), k.contiguous(), v.contiguous(), qk_correction


@supported_compute_capability([120])
@supported_compute_capability([120, 121])
@flashinfer_api(trace=nvfp4_attention_sm120_quantize_qkv_trace)
def nvfp4_attention_sm120_quantize_qkv(
q: torch.Tensor,
Expand Down Expand Up @@ -300,7 +300,7 @@ def _check_inputs(
return batch, num_heads, seq_len, head_dim


@supported_compute_capability([120])
@supported_compute_capability([120, 121])
@flashinfer_api(trace=nvfp4_attention_sm120_fwd_trace)
def nvfp4_attention_sm120_fwd(
q_fp4: torch.Tensor,
Expand Down
7 changes: 4 additions & 3 deletions tests/attention/test_nvfp4_attention_sm120.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,15 @@ def _patch_cutlass_dsl_operand_major_mode():
_patch_cutlass_dsl_operand_major_mode()

import flashinfer
from flashinfer.utils import is_sm120a_supported
from flashinfer.utils import is_sm120a_supported, is_sm121a_supported


def _require_sm120():
if not torch.cuda.is_available():
pytest.skip("CUDA is required")
if not is_sm120a_supported(torch.device("cuda")):
pytest.skip("SM120 GPU is required")
device = torch.device("cuda")
if not (is_sm120a_supported(device) or is_sm121a_supported(device)):
pytest.skip("SM120 or SM121 GPU is required")
Comment on lines +37 to +45

@gemini-code-assist gemini-code-assist Bot Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Instead of importing and checking individual SM120 and SM121 support functions, we can use the existing is_sm12x_supported utility function. This simplifies the check and automatically supports any future SM12x variants without requiring code changes.

Suggested change
from flashinfer.utils import is_sm120a_supported, is_sm121a_supported
def _require_sm120():
if not torch.cuda.is_available():
pytest.skip("CUDA is required")
if not is_sm120a_supported(torch.device("cuda")):
pytest.skip("SM120 GPU is required")
device = torch.device("cuda")
if not (is_sm120a_supported(device) or is_sm121a_supported(device)):
pytest.skip("SM120 or SM121 GPU is required")
from flashinfer.utils import is_sm12x_supported
def _require_sm120():
if not torch.cuda.is_available():
pytest.skip("CUDA is required")
device = torch.device("cuda")
if not is_sm12x_supported(device):
pytest.skip("SM120 or SM121 GPU is required")

@bkryu bkryu Jul 9, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Current change is not too messy so this should be fine.



def _pad_seq_len_to_128(x):
Expand Down
Loading