[#11932][fix] Filter CUTLASS MoE GEMM tile configs by device shared memory on SM121

[mihai-chiorean]

Summary

Consumer Blackwell (SM120 RTX PRO 6000, SM121 GB10 / DGX Spark) has a smaller opt-in shared-memory budget than SM100/B200. Some CUTLASS MoE grouped GEMM TMA warp-specialized tactics selected for the SM120 path can exceed that device limit and fail at gemm.initialize() with an opaque internal error. This was independently observed on RTX PRO 6000 in #11932 and blocks good Qwen3.6 NVFP4 performance on GB10.

This PR keeps the existing SM120 mixed-MoE/NVFP4 path, but makes tactic selection shared-memory-aware:

1. Adds a runtime shared-memory guard in moe_gemm_tma_ws_launcher.inl, checking kernel SharedStorage against cudaDevAttrMaxSharedMemoryPerBlockOptin before launch. Oversized tactics now fail with a clear diagnostic that autotuning can skip.
2. Filters SM120 candidate configs against the runtime device shared-memory limit before returning them from get_candidate_configs_sm120().
3. Preserves valid low-SMEM candidates for both NVFP4 and mixed FP8xFP4 grouped GEMM. In particular, mixed FP8xFP4 keeps its existing large tile where it fits and also receives a smaller fallback tile for 99 KiB devices.

The implementation uses runtime device attributes rather than hard-coding GB10-specific behavior.

Impact

• Without this fix: oversized tactics can fail during CUTLASS grouped-GEMM initialization on SM120/SM121, causing noisy autotuning failures and poor fallback behavior.
• With this fix: oversized tactics are skipped cleanly; valid tactics remain available, including the mixed FP8xFP4 low-SMEM fallback needed by Qwen3.6 NVFP4 workloads.

Relationship to related work

• Complements [None][fix] Fix int4 awq for sm120/121 #11561, which narrowed the outer SM120 candidate-config gate to FP4 and mixed FP8-FP4 workloads. This PR filters the inner SM120 candidate list by actual shared-memory availability.
• Required for best performance of [#14287][feat] Add Qwen3.5/3.6 MoE NVFP4 + MTP support for SM120/SM121 #15680 (Qwen3.5/3.6 MoE NVFP4 + MTP support) on SM120/SM121.
• Addresses the CUTLASS shared-memory portion of 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 #14575 for consumer Blackwell MTP workloads.

[coderabbitai]

📝 Walkthrough

Walkthrough

Two runtime shared memory validation checks were added: one filters CUTLASS GEMM candidate configurations when max shared memory is below 120KB, and another validates kernel shared memory requirements against device limits before kernel launch.

Changes

Cohort / File(s)	Summary
CUTLASS Config Filtering cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp	Added <algorithm> include and runtime device query for cudaDevAttrMaxSharedMemoryPerBlockOptin in get_candidate_configs_sm120. When max shared memory is below 120KB, candidate configs with CtaShape128x128x128B tile are filtered out, retaining only 64B K-tile variants.
MOE GEMM Launcher Validation cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl	Added pre-launch validation check verifying kernel's required shared memory (sizeof(GemmKernel::SharedStorage)) does not exceed device's maximum shared memory per block with opt-in capability. Triggered before gemm.can_implement()/gemm.initialize() calls.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: filtering CUTLASS MoE GEMM tile configs based on device shared memory constraints for SM121, which directly addresses the core issue described in the PR.
Description check	✅ Passed	The PR description includes a clear summary, impact, related work, and test plan, covering most required template content.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp`:
- Around line 561-568: The SM<120 KiB pruning only removes CtaShape128x128x128B
but leaves 64B K-tile configs that the FP4 dispatcher
(dispatchNVFP4xNVFP4GemmCTAShapeSm120) doesn't handle, causing an invalid-config
throw; update the pruning logic around candidate_configs and
CutlassGemmConfig::tile_config_sm120 to either remove the additional 64B shapes
(CtaShape128x128x64B, CtaShape128x256x64B, CtaShape256x128x64B) or, better,
whitelist only the shapes the dispatcher supports (CtaShape128x128x128B,
CtaShape128x128x256B, CtaShape256x128x128B) so candidate_configs contains only
dispatchable CutlassTileConfigSM120 values and the autotuner no longer hits the
default invalid-config path.
- Around line 554-559: The static kMaxSmem initializer currently queries device
0; change it to call cudaGetDevice() to obtain the current device and pass that
device ID into cudaDeviceGetAttribute(), adding explicit error checks for both
cudaGetDevice and cudaDeviceGetAttribute and falling back or logging on error.
Replace the literal 120 * 1024 with a named constant (e.g., kSmemThresholdBytes)
and use that constant where the literal appears. Declare the iterator variable
named it as const (e.g., const auto it) where it is defined. Also update the
file copyright year to 2026.

In
`@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl`:
- Around line 673-682: Replace the current function-static kMaxSmem that calls
cudaDeviceGetAttribute(..., 0) with a runtime query that calls cudaGetDevice()
to obtain the current device id and then calls cudaDeviceGetAttribute(&val,
cudaDevAttrMaxSharedMemoryPerBlockOptin, device) so the shared-memory check for
GemmGrouped/GemmKernel_ uses the active GPU; remove the lambda-cached static and
compute smem limit per invocation and keep the TLLM_CHECK_WITH_INFO(smem_size <=
kMaxSmem, ...) guard unchanged except that kMaxSmem is now the per-call value.
Also add the same per-device SMEM preflight guard (compute current device via
cudaGetDevice and compare sizeof(typename GemmKernel_::SharedStorage) against
that device's limit) into moe_gemm_tma_ws_mixed_input_launcher.inl before
calling gemm.can_implement().

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ee34cc95-0326-4ec6-978d-b97646526a47

📥 Commits

Reviewing files that changed from the base of the PR and between 11c40bb and 93b9c60.

📒 Files selected for processing (2)

• cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
• cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl

[karljang]

Hi @pamelap-nvidia, could you please review this PR?
My Claude suggests you as the appropriate person to review it, but please feel free to reassign someone else on the team if you prefer.

[mihai-chiorean]

@pamelap-nvidia — gentle bump on this one. PR is mergeable, all checks green, and the fix has been independently confirmed on SM120 RTX PRO 6000 by @waynehacking8 in #11932. If you're not the right reviewer for the CUTLASS heuristic / MoE launcher path, please feel free to redirect — happy to address feedback from whoever's the best fit.

[pamelap-nvidia]

/bot run --extra-stage "GB10-PyTorch-Post-Merge-1,RTXPro6000D-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2"

[tensorrt-cicd]

PR_Github #55295 [ run ] triggered by Bot. Commit: c15c02c Link to invocation

[tensorrt-cicd]

PR_Github #55295 [ run ] completed with state FAILURE. Commit: c15c02c
/LLM/main/L0_MergeRequest_PR pipeline #44248 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[mihai-chiorean]

@pamelap-nvidia I'm not sure what failed in CI and why. I updated the branch. hopefully it will pass this time. mind starting a CI run again? when this merges, it actually opens the gate for MTP for qwen3.6 35b nvfp4 with mtp on trt-llm.

[mihai-chiorean]

PR_Github #60091 [ run ] completed with state SUCCESS. Commit: 06b6a13 /LLM/main/L0_MergeRequest_PR pipeline #48478 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@karljang really appreciate your patience. I'm having a hard time figuring out the next failure. it would be really nice to know what tests blossom is running and how so I can run them myself and make sure I am confident before pinging you guys. Just an idea. Other than that... I need more logs to see what's failing this time.

Again, really appreciate your help

I would like to merge this first before #15851

[karljang]

/bot run --extra-stage "GB10-PyTorch-Post-Merge-1,RTXPro6000D-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2"

[karljang]

/bot kill

[karljang]

/bot run --extra-stage "GB10-PyTorch-Post-Merge-1,RTXPro6000D-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" --disable-fail-fast

[tensorrt-cicd]

PR_Github #60787 [ kill ] triggered by Bot. Commit: 0026d2a Link to invocation

[tensorrt-cicd]

PR_Github #60787 [ kill ] completed with state SUCCESS. Commit: 0026d2a
Successfully killed previous jobs for commit 0026d2a

Link to invocation

[tensorrt-cicd]

PR_Github #60788 [ run ] triggered by Bot. Commit: 0026d2a Link to invocation

[sunnyqgg]

LGTM

[tensorrt-cicd]

PR_Github #60788 [ run ] completed with state SUCCESS. Commit: 0026d2a
/LLM/main/L0_MergeRequest_PR pipeline #49067 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[karljang]

/bot run --extra-stage "GB10-PyTorch-Post-Merge-1,RTXPro6000D-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-1,RTXPro6000D-4_GPUs-PyTorch-Post-Merge-2" --disable-fail-fast

[tensorrt-cicd]

PR_Github #60891 [ run ] triggered by Bot. Commit: 0026d2a Link to invocation

[tensorrt-cicd]

PR_Github #60891 [ run ] completed with state SUCCESS. Commit: 0026d2a
/LLM/main/L0_MergeRequest_PR pipeline #49160 completed with status: 'SUCCESS'

CI Report

Link to invocation

[mihai-chiorean]

@rosong11 @schetlur-nv mind taking a look? CI passes now and there are 2 approvals already

[rosong11]

LGTM

[karljang]

@mihai-chiorean ,
We've got all approvals, but there are open conversations above. Could you please resolve them if possible? Thank you!