NVIDIA / NVIDIA/TensorRT-LLM

[Bug]: B12x NVFP4 MoE JIT fails on SM120/SM121 with CUDA 12 CuTe DSL payload

Open
#15,853 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

System Info
  • GPU: Blackwell SM121 / DGX Spark observed for full-model repro; SM120 target reproduced with the same FlashInfer B12x micro-kernel target override
  • Model: Qwen3 MoE NVFP4
  • Backend: PyTorch, MoE backend CUTEDSL, FlashInfer B12x NVFP4 MoE decode path
  • Relevant package state: nvidia-cutlass-dsl-libs-cu13==4.5.0 installed, but active cutlass.base_dsl.version_info.CUDA_VERSION reported CUDA 12.9 because the active CuTe DSL native payload came from the base CUDA 12 payload
Who can help?

@Tracin @Fridah-nv

Information
  • The official example scripts
  • My own modified scripts
Tasks
  • An officially supported task in the examples folder
  • My own task or dataset
Reproduction

Use an NVFP4 Qwen3 MoE model on Blackwell SM121 and force the PyTorch MoE backend to CUTEDSL so TRT-LLM selects the FlashInfer/CuTe DSL B12x fused MoE path. The selector promotes to CuteDslB12xFusedMoE, then FlashInfer JIT enters the Blackwell SM12x MoE micro-kernel path and instantiates the NVFP4/MXF4 MMA atom.

The failing environment has nvidia-cutlass-dsl-libs-cu13==4.5.0 installed, but the active CuTe DSL runtime reports CUDA 12.9. In that state, the generated PTX contains internal _mma.block_scale... instructions rather than the CUDA 13 public mma.sync.aligned...kind::mxf4nvf4 spelling. ptxas rejects the generated PTX during JIT compilation.

A minimal FlashInfer B12x micro-kernel repro shows the same ptxas failure for both SM121 and an SM120 target override when the active CuTe DSL payload is CUDA 12.x. Reinstalling the CUDA 13 CuTe DSL native payload so CUDA_VERSION reports 13.1 makes the SM121 B12x path compile and run.

Expected behavior

TRT-LLM should not select the B12x CuTe DSL/FlashInfer NVFP4 MoE path when the active CuTe DSL native payload is CUDA 12.x or unavailable, because that combination JIT-compiles invalid PTX for the B12x NVFP4 MMA atom.

When the active CuTe DSL native payload reports CUDA 13 or newer, TRT-LLM should allow B12x on supported SM12x targets.

Actual behavior

With CUDA 12.x active in CuTe DSL, initialization fails during FlashInfer/CuTe DSL JIT with ptxas errors similar to:

ptxas application ptx input, line ...; error : Unexpected instruction types specified for `_mma`
DSLRuntimeError: ICE
RuntimeError: Executor worker died during initialization

This happens after TRT-LLM has already selected CuteDslB12xFusedMoE, so the process fails during model initialization instead of falling back to a safe MoE backend.

Root cause hypothesis

CUDA 12.x CuTe DSL lowers the FlashInfer B12x NVFP4/MXF4 MMA atom to an internal _mma.block_scale... PTX form that the active ptxas rejects. CUDA 13.x CuTe DSL lowers the same operation to the public kind::mxf4nvf4 opcode form, which compiles successfully.

The immediate TRT-LLM fix is a runtime capability gate: for B12x on supported SM12x targets, require the active CuTe DSL native payload to report CUDA 13 or newer; otherwise fall back to CUTLASS with an actionable warning.

Validation
  • Reinstalling the CUDA 13 CuTe DSL native payload made CUDA_VERSION report CUDA 13.1.
  • Full-model Qwen3 NVFP4 on SM121 with CuteDslB12xFusedMoE selected initialized successfully, generated coherent output, and improved throughput versus the CUTLASS MoE path.
  • MTP was separately verified active on the CUDA 13 B12x path via iteration stats and instantiated model modules; no MTP code change is required for the B12x gate itself.
Related PR
  • Fix/guard PR: #15851
Before submitting a new issue...
  • I searched existing issues for B12x, SM121, DGX Spark, CuTe DSL, FlashInfer, _mma, and block_scale and did not find this exact failure mode.

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 at the MoE backend selector that promotes to CuteDslB12xFusedMoE, then trace how the active CuTe DSL CUDA_VERSION is exposed and used. Reproduce with the minimal FlashInfer B12x micro-kernel on SM120 or SM121; done means CUDA 12.x or unavailable payloads avoid B12x with an actionable warning and CUDA 13+ permits it.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.