Skip to content

[None][feat] Ungate fused MoE for SM120/SM121 (GB10/DGX Spark) #11997

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

Open
scottgl9 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[None][feat] Ungate fused MoE for SM120/SM121 (GB10/DGX Spark) #11997

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
63b21c5
[None][feat] Ungate fused MoE for SM120/SM121 (GB10/DGX Spark)
scottgl9 Mar 7, 2026
dc3fa35
Merge branch 'main' into gb10-sm120-support
scottgl9 Mar 7, 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
6 changes: 3 additions & 3 deletions tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -808,13 +808,13 @@ def get_valid_tactics(self, inputs: List[torch.Tensor],
# Add CuteDSL tactics if available
if self._is_backend_allowed("cutedsl"):
if IS_CUTLASS_DSL_AVAILABLE:
# Check SM version first - CuteDSL NVFP4 only supports SM 100 (B200)
# Check SM version first - CuteDSL NVFP4 only supports Blackwell
sm_version = get_sm_version()
if sm_version not in [100, 103]:
if sm_version not in [100, 103, 120, 121]:
if self._is_only_backend("cutedsl"):
# Explicitly forced CuteDSL but SM version not supported
raise ValueError(
f"CuteDSL NVFP4 backend requires SM 100 (B200) or SM 103 (B300), but got SM {sm_version}. "
f"CuteDSL NVFP4 backend requires Blackwell (SM100/103/120/121), but got SM {sm_version}. "
f"CuteDSL NVFP4 is not supported on this GPU architecture. "
"Please add other backends to allowed_backends.")
else:
Expand Down
4 changes: 2 additions & 2 deletions tensorrt_llm/_torch/model_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,12 +250,12 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> str:
if architecture == "GptOssForCausalLM":
sm_version = get_sm_version()
# Select the best performing backend based on SM version
if 100 <= sm_version < 120: # Blackwell
if 100 <= sm_version < 120 or sm_version in (120, 121): # Blackwell
return "TRTLLM"

@pamelap-nvidia pamelap-nvidia Mar 13, 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.

Unfortunately TRTLLM MOE backend doesn't support sm120/121.

@xxi-nv xxi-nv Apr 8, 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.

@syuoni Could you help to confirm whether the cuteDSL supports sm120/121? I don't think we have supported it,

elif 90 <= sm_version < 100: # Hopper
return "TRITON"
else:
return "CUTLASS" # Fallback to CUTLASS for other SM versions (e.g., SM120)
return "CUTLASS" # Fallback for other SM versions
Comment on lines +253 to +258

@coderabbitai coderabbitai Bot Mar 7, 2026 •
edited
Loading

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.

⚠️ Potential issue | 🟡 Minor

Don't auto-route every 11x GPU to TRTLLM.

Line 253 treats any 100 <= sm_version < 120 as Blackwell, but the rest of this change only enables {100, 103, 120, 121}. That means AUTO would start picking TRTLLM for a future 11x device before the backend gates are actually updated.

♻️ Proposed fix 
-            if 100 <= sm_version < 120 or sm_version in (120, 121):  # Blackwell
+            if sm_version in (100, 103, 120, 121):  # Blackwell
                 return "TRTLLM"
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if 100 <= sm_version < 120 or sm_version in (120, 121): # Blackwell
return "TRTLLM"
elif 90 <= sm_version < 100: # Hopper
return "TRITON"
else:
return "CUTLASS" # Fallback to CUTLASS for other SM versions (e.g., SM120)
return "CUTLASS" # Fallback for other SM versions
if sm_version in (100, 103, 120, 121): # Blackwell
return "TRTLLM"
elif 90 <= sm_version < 100: # Hopper
return "TRITON"
else:
return "CUTLASS" # Fallback for other SM versions
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/model_config.py` around lines 253 - 258, The current
conditional uses a broad range (100 <= sm_version < 120) to return "TRTLLM",
which will wrongly auto-route any future 11x GPUs; narrow this to only the
explicitly supported SMs (e.g., check sm_version in the explicit set {100, 103,
120, 121}) so AUTO only selects "TRTLLM" for those specific SM versions—update
the branch that inspects sm_version (the block returning "TRTLLM") to use an
explicit membership test instead of a numeric range.

@karljang karljang Mar 12, 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.

The proposed fix looks reasonable to me, could you apply the fix?

@coderabbitai coderabbitai Bot Mar 12, 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.

@karljang, sure! Here's the plan to apply the fix:

  • Replace the broad range check (100 <= sm_version < 120 or sm_version in (120, 121)) with an explicit set membership test (sm_version in (100, 103, 120, 121)) in resolve_moe_backend.
🧠 Learnings used 
Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 11508
File: tests/unittest/_torch/sampler/test_beam_search_util.py:71-71
Timestamp: 2026-02-13T10:15:37.120Z
Learning: In TensorRT-LLM (Python requires >=3.10 and <4 as per setup.py), you can use Python 3.10+ features (e.g., PEP 585 generics like dict[str, int], list[str], etc.) throughout the codebase, and you do not need to add from __future__ import annotations. This applies to all Python files, including tests (e.g., tests/unittest/...); ensure tests and code consistently rely on Python 3.10+ features where applicable.

Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 11983
File: tensorrt_llm/_torch/pyexecutor/llm_request.py:237-240
Timestamp: 2026-03-06T11:45:02.068Z
Learning: In tensorrt_llm/_torch/pyexecutor/llm_request.py, LogProbStorage.append() relies on next(iter(prob.values())).logprob to accumulate cum_log_probs when cum_log_probs is None. This path assumes prob is a non-empty dict because TorchSampler places the sampled logprob first when num_logprobs >= 0. Therefore, no guard for empty dicts is needed here. If future changes may yield empty prob, consider adding a guard or a clearer invariant.

Learnt from: yibinl-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 12009
File: tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py:296-299
Timestamp: 2026-03-09T12:34:56.631Z
Learning: In tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py, BasicAVTransformerBlock should not be flagged for a config mismatch when config.parallel.dit_ulysses_size > 1. The function setup_sequence_parallelism() returns use_ulysses=True for dit_ulysses_size > 1, or raises a RuntimeError/ValueError/NotImplementedError; it never returns use_ulysses=False in that case. Treat this as intentional and correct; do not flag as a mismatch between raw config checks and setup_sequence_parallelism()'s result.

💡 Use @coderabbitai plan to create a plan for code changes

@karljang karljang Mar 12, 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.

Instead, using is_blackwell might be a better approach.

@coderabbitai coderabbitai Bot Mar 12, 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.

@karljang, great idea! Using is_blackwell() is cleaner and future-proof — any new Blackwell SM version only needs to be added to is_blackwell() in _utils.py once.

Here's the updated plan:

  • Add is_blackwell to the existing tensorrt_llm._utils import in model_config.py
  • Replace the broad/explicit-set SM check with is_blackwell(sm_version) in resolve_moe_backend

sm_version is passed explicitly to is_blackwell(sm_version) to avoid a second call to get_sm_version() (already needed for the Hopper branch), keeping it efficient and readable.

🧠 Learnings used 
Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 11508
File: tests/unittest/_torch/sampler/test_beam_search_util.py:71-71
Timestamp: 2026-02-13T10:15:37.120Z
Learning: In TensorRT-LLM (Python requires >=3.10 and <4 as per setup.py), you can use Python 3.10+ features (e.g., PEP 585 generics like dict[str, int], list[str], etc.) throughout the codebase, and you do not need to add from __future__ import annotations. This applies to all Python files, including tests (e.g., tests/unittest/...); ensure tests and code consistently rely on Python 3.10+ features where applicable.

Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 11983
File: tensorrt_llm/_torch/pyexecutor/llm_request.py:237-240
Timestamp: 2026-03-06T11:45:02.068Z
Learning: In tensorrt_llm/_torch/pyexecutor/llm_request.py, LogProbStorage.append() relies on next(iter(prob.values())).logprob to accumulate cum_log_probs when cum_log_probs is None. This path assumes prob is a non-empty dict because TorchSampler places the sampled logprob first when num_logprobs >= 0. Therefore, no guard for empty dicts is needed here. If future changes may yield empty prob, consider adding a guard or a clearer invariant.

Learnt from: yibinl-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 12009
File: tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py:296-299
Timestamp: 2026-03-09T12:34:56.631Z
Learning: In tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py, BasicAVTransformerBlock should not be flagged for a config mismatch when config.parallel.dit_ulysses_size > 1. The function setup_sequence_parallelism() returns use_ulysses=True for dit_ulysses_size > 1, or raises a RuntimeError/ValueError/NotImplementedError; it never returns use_ulysses=False in that case. Treat this as intentional and correct; do not flag as a mismatch between raw config checks and setup_sequence_parallelism()'s result.

💡 Use @coderabbitai plan to create a plan for code changes


return "CUTLASS"

Expand Down
15 changes: 8 additions & 7 deletions tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import torch
import torch.nn.functional as F

from tensorrt_llm._utils import get_sm_version, is_sm_100f
from tensorrt_llm._utils import get_sm_version, is_blackwell, is_sm_100f
from tensorrt_llm.models.modeling_utils import QuantAlgo

from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec,
Expand Down Expand Up @@ -62,8 +62,8 @@ def cute_dsl_fp8_group_blockwise_gemm_ref(
a_tmp = a.as_strided((m, k, 1), (k, 1, m * k))
b_tmp = b.permute(1, 2, 0)

# Note: we have different output scale shape for fp8_quantize_1x128, so we need to handle it differently for sm100 and other archs.
if is_sm_100f():
# Note: we have different output scale shape for fp8_quantize_1x128, so we need to handle it differently for Blackwell and other archs.
if is_blackwell():
input_scale_tmp = a_sf.permute(1, 0).as_strided((m, w_k, 1),
(1, m, m * w_k))
else:
Expand Down Expand Up @@ -339,7 +339,7 @@ def can_implement(
Check if CuteDslFusedMoE can implement the given quantization algorithm.

CuteDslFusedMoE supports:
- NVFP4: SM in {100, 103}
- NVFP4: SM in {100, 103, 120, 121}

@xxi-nv xxi-nv Jun 8, 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.

Have you tested it locally? @syuoni It seems that we don't have the kernel to support SM120, right?

@syuoni syuoni Jun 8, 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.

No, sm120 is not supported. This PR doesn't add new kernels, so I would expect crash on sm120 GPUs.


Does NOT support unquantized mode. Output dtype is hardcoded to bfloat16.
Does NOT support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit).
Expand Down Expand Up @@ -381,11 +381,12 @@ def can_implement(
"CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)"
)

# NVFP4 - SM in {100, 103}
# NVFP4 - SM in {100, 103, 120, 121} (Blackwell family)
if quant_algo == QuantAlgo.NVFP4:
if sm_version not in {100, 103}:
if sm_version not in {100, 103, 120, 121}:
return _warn_and_return(
f"NVFP4 requires SM100 or SM103, got SM{sm_version}")
f"NVFP4 requires Blackwell (SM100/103/120/121), got SM{sm_version}"
)
return True, None

return _warn_and_return(
Expand Down
15 changes: 5 additions & 10 deletions tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ class TRTLLMGenFusedMoE(MoE):
aux_stream_dict (Optional[Dict[AuxStreamType, torch.cuda.Stream]]): Auxiliary CUDA streams for overlapping.

MoE torch custom op:
Only support min-latency mode now (SM100 Blackwell only).
Only support min-latency mode now (Blackwell: SM100/103/120/121).
Quant: fp8 block scales quant and nvfp4 quant and w4a16_mxfp4 quant
FusedMoE Op: routing(topK, etc.) + scatter + gemm1 + swiglu + gemm2 + finalize MoeRoute

Expand Down Expand Up @@ -105,7 +105,7 @@ def can_implement(
"""
Check if TRTLLMGenFusedMoE can implement the given quantization algorithm.

TRTLLMGenFusedMoE only supports SM in {100, 103} and the following quantizations:
TRTLLMGenFusedMoE only supports SM in {100, 103, 120, 121} and the following quantizations:
- NVFP4
- FP8_BLOCK_SCALES
- W4A8_NVFP4_FP8
Expand All @@ -129,10 +129,10 @@ def can_implement(

sm_version = get_sm_version()

# TRTLLMGenFusedMoE requires SM in {100, 103}
if sm_version not in {100, 103}:
# TRTLLMGenFusedMoE requires SM in {100, 103, 120, 121} (Blackwell family)
if sm_version not in {100, 103, 120, 121}:

@xxi-nv xxi-nv Jun 8, 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.

TRTLLM-GEN can not support sm120, @nekorobov please help to confirm.

@nekorobov nekorobov Jun 8, 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.

Confirm

return _warn_and_return(
f"TRTLLMGenFusedMoE requires SM100 or SM103, got SM{sm_version}"
f"TRTLLMGenFusedMoE requires Blackwell (SM100/103/120/121), got SM{sm_version}"
)

# Check dtype_activation: only bfloat16 is supported
Expand Down Expand Up @@ -201,11 +201,6 @@ def __init__(
activation_type=activation_type,
)

sm_version = get_sm_version()
if sm_version >= 120:
raise NotImplementedError(
"TRTLLMGenFusedMoE does not support SM120 and above.")

assert not self.smart_router, "Smart router is not supported in TRTLLMGenFusedMoE."

# Note: Load balancer initialization is handled by base class _init_load_balancer()
Expand Down
12 changes: 12 additions & 0 deletions tensorrt_llm/_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -781,6 +781,18 @@ def is_sm_100f(sm_version=None):
return sm_version == 100 or sm_version == 103


def is_sm_120f(sm_version=None):
if sm_version is None:
sm_version = get_sm_version()
return sm_version == 120 or sm_version == 121


def is_blackwell(sm_version=None):
if sm_version is None:
sm_version = get_sm_version()
return is_sm_100f(sm_version) or is_sm_120f(sm_version)


def print_all_stacks():
"""Print stack traces for all threads"""
for thread_id, frame in sys._current_frames().items():
Expand Down
12 changes: 12 additions & 0 deletions tests/integration/defs/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1919,6 +1919,18 @@ def is_sm_100f(sm_version=None):
return sm_version == 100 or sm_version == 103


def is_sm_120f(sm_version=None):
if sm_version is None:
sm_version = get_sm_version()
return sm_version == 120 or sm_version == 121


def is_blackwell(sm_version=None):
if sm_version is None:
sm_version = get_sm_version()
return is_sm_100f(sm_version) or is_sm_120f(sm_version)


def get_gpu_device_list():
"get device list"
with tempfile.TemporaryDirectory() as temp_dirname:
Expand Down
Loading