NVIDIA / NVIDIA/TensorRT-LLM

`1.3.0rc15` on sm_121a / GB10: MTP path fails — CUTLASS sm_120 grouped-GEMM init → SMEM shortfall → Triton fallback emits `.rs` PTX rejected by PTXAS

Open
#14,575 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Customized kernels Speculative Decoding Triton backend
Dominant language
Python
Stars
14.7k
Forks
2.8k
Avg merge
2d 23h
Merged PRs (30d)
489

Description

Summary

On nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc15 running Nemotron-3-Super-120B-A12B-NVFP4 on NVIDIA GB10 (sm_121a), enabling MTP (speculative_config.decoding_type: MTP, num_nextn_predict_layers: 3) fails at engine startup with three stacked errors:

  1. CUTLASS sm_120 grouped-GEMM kernel fails to initialize for the MTP-shaped MoE GEMM
  2. CUTLASS GroupedGEMM kernel reports insufficient shared memory at the MTP-shaped tensor sizes
  3. Triton fallback emits PTX containing the .rs (stochastic-rounding) feature, which PTXAS rejects for .target 'sm_121a'

Without MTP, the same image + same hardware + same yaml otherwise runs cleanly through a 20-min, 10-concurrent, 15%-cancel soak (validated in issue #14500 after applying a separate two-line patch suggested there).

So this issue isolates the MTP-specific path on sm_121a.

Environment

  • GPU: NVIDIA GB10 (Grace-Blackwell, integrated GPU, unified memory, sm_121a)
  • OS: DGX OS 7 (Ubuntu 24.04 aarch64)
  • Model: Nemotron-3-Super-120B-A12B-NVFP4 (MoE; hybrid Mamba/SSM + attention; NVFP4)
  • Image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc15
  • Backend: pytorch (AutoDeploy via trtllm-serve)
  • (also applied the fused_moe/quantization.py two-line patch from #14500 — needed to get past cold-start on this hardware; not relevant to this MTP failure path)

Reproducer

Launch the standard trtllm-serve command (same as #14500) and add this to the mounted extra_llm_api_options yaml:

speculative_config:
  decoding_type: MTP
  num_nextn_predict_layers: 3

Everything else identical to the baseline yaml (CUTLASS MoE backend, enable_block_reuse: false, mamba_ssm_cache_dtype: float16, mamba_ssm_stochastic_rounding: true, enable_chunked_prefill: true, max_batch_size: 8, etc.). MTP-off vs MTP-on is the only knob that flips between "soak passes clean" and the failure chain below.

Failure chain

After MTPSampler initializes successfully, the MoE autotuner begins profiling tactics. Three distinct errors stack:

1) CUTLASS sm_120 grouped-GEMM init
[Autotuner] Failed when profiling runner=<…MoERunner…>, tactic=6,
shapes=[torch.Size([1, 512]), torch.Size([512, 2688, 64]), torch.Size([0]),
        torch.Size([512, 1024, 168]), torch.Size([0])].
Error: [TensorRT-LLM][ERROR] Assertion failed: Failed to initialize cutlass TMA WS grouped gemm.
       Error: Error Internal
       (tensorrt_llm/kernels/cutlass_kernels/cutlass_instantiations/gemm_grouped/120/
        cutlass_kernel_file_gemm_grouped_sm120_M128_BS_group2.generated.cu:39)
[Autotuner] New tuning error occurs:Total failed profiling tactics occurs: 15
for custom_op=trtllm::fused_moe::gemm2.
2) CUTLASS shared-memory shortfall on the MTP-shaped GEMM
[Autotuner] Failed when profiling runner=<…MoERunner…>, tactic=8,
shapes=[torch.Size([1, 1024]), torch.Size([512, 2688, 1024]), torch.Size([0]),
        torch.Size([512, 1024, 2688]), torch.Size([0])].
Error: [TensorRT-LLM][ERROR] Assertion failed: GPU lacks the shared memory resources to run GroupedGEMM kernel
       (../tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h:174)
[Autotuner] New tuning error occurs:Total failed profiling tactics occurs: 15
for custom_op=trtllm::fused_moe::gemm1.

Note: the same model + same hardware runs CUTLASS MoE GEMM fine without MTP. The MTP-shape change is what blows the SMEM budget.

3) Triton fallback — PTXAS rejects .rs feature on sm_121a
Internal Triton PTX codegen error
ptxas /tmp/tmpacy3oa91.ptx, line 1200; error : Feature '.rs' not supported on .target 'sm_121a'
ptxas /tmp/tmpacy3oa91.ptx, line 1205; error : Feature '.rs' not supported on .target 'sm_121a'
ptxas /tmp/tmpacy3oa91.ptx, line 1210; error : Feature '.rs' not supported on .target 'sm_121a'
ptxas /tmp/tmpacy3oa91.ptx, line 1215; error : Feature '.rs' not supported on .target 'sm_121a'
ptxas /tmp/tmpacy3oa91.ptx, line 1220; error : Feature '.rs' not supported on .target 'sm_121a'
…(continues — many lines per failing kernel)

After ~9 minutes of compile failures the container exits and crashloops. Reverting MTP off (no other changes) brings the engine back up cleanly.

Test matrix (rc15 on GB10/sm_121a, Nemotron-3-Super-120B-A12B-NVFP4)

Config Result
moe_config.backend: CUTLASS, MTP off, with the #14500 patch applied ✅ 20-min / 10-concurrent / 15%-cancel soak clean; 0 engine deaths, 0 cancel hangs
moe_config.backend: CUTLASS, MTP on, with the #14500 patch applied ❌ this issue — three stacked errors above
moe_config.backend: TRTLLM, MTP either NotImplementedError: TRTLLMGenFusedMoE does not support SM120 and above (separate gap)

Why this matters

MTP is the largest documented throughput knob for this model. The combination of "MTP off is the only viable config on GB10" + "the documented MTP+mamba slot-accounting fix [PR #13151] needs MTP on to validate" leaves sm_121a/GB10 deployments unable to take the throughput recovery from #13151 even when running the latest release.

Likely fix surfaces

  1. The Triton MTP MoE kernel emitting .rs — if .rs (stochastic rounding) is not actually required for this particular kernel on sm_120+, gate the emission so it's omitted for sm_121a. If it IS required, sm_121a (the integrated GB10 SASS variant) may need a different SASS path or an explicit non-.rs fallback codegen
  2. The CUTLASS gemm_grouped_sm120 MTP-shape SMEM blowup — if MTP-shaped MoE GEMM legitimately needs more SMEM than baseline, a kernel variant that fits in the integrated-GPU SMEM budget would prevent the Triton fallback from being needed at all
  3. Documentation — clarifying that decoding_type: MTP is currently unsupported on sm_121a (until either of the above lands) so users on GB10 don't keep rediscovering this

Cross-reference

  • Issue #14500 — separate cold-start hang on rc15+GB10 (now resolved by a two-line patch from contributor ssam18, who also flagged that the .rs/PTXAS path was likely separate)
  • The yaml-comment chain that originally documented the rc14 PTXAS issue (which now matches what we see in rc15) — this is the same wall, just hit on a different trigger path

Happy to capture full failing PTX dumps, TLLM_LOG_LEVEL=DEBUG logs, py-spy at any of the three failure points, or test a candidate patch on this hardware.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the MTP configuration with trtllm-serve and capture the three failure points. Read the referenced generated CUTLASS file, cutlass_kernel_file_gemm_grouped_sm120_M128_BS_group2.generated.cu, and moe_gemm_template_dispatch.h; compare the MTP-on and MTP-off autotuner shapes. Done means MTP starts successfully on sm_121a without the shared-memory or PTXAS failures.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.