[#15853][fix] Gate B12x MoE on CUDA13 CuTe DSL#15851
Conversation
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
mihai-chiorean
changed the title
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
mihai-chiorean
changed the title
📝 WalkthroughWalkthroughAdds a runtime capability gate to CuteDslB12xFusedMoE that checks the active CUTLASS CuTe DSL CUDA major version via a new get_runtime_disable_reason helper, wires this check into both can_implement and backend selection in create_moe.py to fall back to CutlassFusedMoE when CUDA 13 is unavailable, and updates unit tests to mock the new CUDA major detection across gating scenarios. ChangesB12x runtime gating
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CreateMoe as create_moe.py
participant B12x as CuteDslB12xFusedMoE
participant DslRuntime as CuTe DSL Runtime
participant Cutlass as CutlassFusedMoE
CreateMoe->>B12x: get_runtime_disable_reason(sm_version)
B12x->>B12x: check sm_version in supported list
B12x->>DslRuntime: _get_cutlass_dsl_cuda_major()
DslRuntime-->>B12x: CUDA major or None
alt CUDA major < 13 or unavailable
B12x-->>CreateMoe: disable reason string
CreateMoe->>Cutlass: select as fallback (log warning)
else CUDA major >= 13
B12x-->>CreateMoe: None
CreateMoe->>CreateMoe: continue flashinfer probe / select B12x
end
Related Issues: Related PRs: None mentioned Suggested labels: bug, moe, fused_moe Suggested reviewers: Tracin, Fridah-nv 🐰 A rabbit checks the CUDA before it jumps, 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py (3)
181-191: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate test:
test_get_moe_cls_cutedsl_selects_b12x_on_supported_smvstest_get_moe_cls_cutedsl_selects_b12x_on_supported_sm_with_cuda13.These two tests are functionally identical (same SM parametrize, same CUDA-major=13 mock, same assertion
cls is CuteDslB12xFusedMoE). Suggest consolidating into a single test.Also applies to: 208-218
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py` around lines 181 - 191, The test `test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm` duplicates `test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm_with_cuda13` with the same `sm_version` parametrization, the same `CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major` mock returning 13, and the same `get_moe_cls(cfg)` assertion. Consolidate the two into a single test by keeping one representative name and removing the redundant copy, while preserving the shared setup around `ModelConfig`, `QuantConfig`, `QuantAlgo.NVFP4`, and the `get_moe_cls` assertion.Source: Path instructions
51-58: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate test:
test_can_implement_accepts_supported_sm_with_nvfp4vstest_can_implement_accepts_supported_sm_with_cuda13_cute_dsl.Both tests mock the same
sm_versionparametrize source and_get_cutlass_dsl_cuda_major=13, and assert the exact same outcome (ok,reason is None). Consider removing one (keep the more descriptive..._with_cuda13_cute_dslname) to avoid duplicate maintenance.Also applies to: 85-93
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py` around lines 51 - 58, The test `test_can_implement_accepts_supported_sm_with_nvfp4` duplicates `test_can_implement_accepts_supported_sm_with_cuda13_cute_dsl`, since both use the same `sm_version` parametrization, patch `CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major` to 13, and assert the same `can_implement` result. Remove the redundant test (prefer keeping the more descriptive `test_can_implement_accepts_supported_sm_with_cuda13_cute_dsl`) and, if needed, update any nearby assertions in `CuteDslB12xFusedMoE.can_implement` coverage to preserve distinct behavior coverage.Source: Path instructions
61-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage gap:
_get_cutlass_dsl_cuda_major's own implementation is never exercised.Every test in this file mocks
_get_cutlass_dsl_cuda_majordirectly, so the actual import/getattr/isinstancelogic inside it (ImportError handling, non-intmajor, missingCUDA_VERSIONattribute) has no dedicated test. Consider adding a focused unit test in this file that monkeypatchessys.modules["cutlass.base_dsl.version_info"](or similarly stubs thecutlasspackage) to verify_get_cutlass_dsl_cuda_major()itself returns the right value/Nonefor: (a) module not importable, (b)CUDA_VERSION.majormissing, (c)CUDA_VERSION.majornon-int, (d) a valid int major. Coverage is otherwise sufficient for the gating behavior end-to-end (can_implement + get_moe_cls); this gap is isolated to the low-level version-parsing helper and could reasonably be addressed as a small follow-up rather than blocking this PR.Also applies to: 194-218
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py` around lines 61 - 93, Add focused unit coverage for CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major itself, since current tests only mock it and never exercise its import/version parsing path. Create tests that stub the cutlass.base_dsl.version_info import (via sys.modules or equivalent) to verify the helper returns None when the module is missing, when CUDA_VERSION or CUDA_VERSION.major is absent, and when major is non-int, plus returns the int major for a valid module. Keep the existing can_implement tests as-is; this change should target the helper’s own import/getattr/isinstance logic.Source: Path instructions
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py (1)
123-144: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a docstring to the new public
get_runtime_disable_reason.This classmethod is now called cross-file from
create_moe.get_moe_cls, but it only has an inline comment explaining the CUDA 12 vs 13 PTX issue, no docstring describing its contract (input/return semantics).📝 Suggested docstring
`@classmethod` def get_runtime_disable_reason(cls, sm_version: int) -> Optional[str]: + """Return a human-readable reason this backend can't run on ``sm_version`` + right now, or ``None`` if all runtime requirements are met. + + Checked here (rather than only in ``can_implement``) so + ``create_moe.get_moe_cls`` can fall back to CutlassFusedMoE before + ever attempting to build/JIT this backend. + + Args: + sm_version: SM version as returned by ``get_sm_version()``. + """ if sm_version not in cls._SUPPORTED_SM_VERSIONS:As per coding guidelines, "For interfaces used outside a file, prefer docstrings over comments; reserve comments for code within functions or interfaces local to a file."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py` around lines 123 - 144, Add a docstring to CuteDslB12xFusedMoE.get_runtime_disable_reason describing its contract for external callers such as create_moe.get_moe_cls: document that it takes an sm_version integer, returns a human-readable disable reason string when the runtime is unsupported or unavailable, and returns None when the runtime can be used. Keep the existing CUDA 12 vs 13 logic intact, and place the docstring on the classmethod itself so the public interface is documented instead of relying on the inline comment.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py`:
- Around line 146-160: The CUTLASS runtime check is using an internal import
path in _get_cutlass_dsl_cuda_major, which is brittle and can break if CUTLASS
internals change. Update that helper to read the public
cutlass.CUDA_VERSION.major API directly, and keep
_is_cutlass_dsl_runtime_available using the returned major version for its
threshold check.
---
Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py`:
- Around line 123-144: Add a docstring to
CuteDslB12xFusedMoE.get_runtime_disable_reason describing its contract for
external callers such as create_moe.get_moe_cls: document that it takes an
sm_version integer, returns a human-readable disable reason string when the
runtime is unsupported or unavailable, and returns None when the runtime can be
used. Keep the existing CUDA 12 vs 13 logic intact, and place the docstring on
the classmethod itself so the public interface is documented instead of relying
on the inline comment.
In `@tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py`:
- Around line 181-191: The test
`test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm` duplicates
`test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm_with_cuda13` with the
same `sm_version` parametrization, the same
`CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major` mock returning 13, and the
same `get_moe_cls(cfg)` assertion. Consolidate the two into a single test by
keeping one representative name and removing the redundant copy, while
preserving the shared setup around `ModelConfig`, `QuantConfig`,
`QuantAlgo.NVFP4`, and the `get_moe_cls` assertion.
- Around line 51-58: The test
`test_can_implement_accepts_supported_sm_with_nvfp4` duplicates
`test_can_implement_accepts_supported_sm_with_cuda13_cute_dsl`, since both use
the same `sm_version` parametrization, patch
`CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major` to 13, and assert the same
`can_implement` result. Remove the redundant test (prefer keeping the more
descriptive `test_can_implement_accepts_supported_sm_with_cuda13_cute_dsl`) and,
if needed, update any nearby assertions in `CuteDslB12xFusedMoE.can_implement`
coverage to preserve distinct behavior coverage.
- Around line 61-93: Add focused unit coverage for
CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major itself, since current tests only
mock it and never exercise its import/version parsing path. Create tests that
stub the cutlass.base_dsl.version_info import (via sys.modules or equivalent) to
verify the helper returns None when the module is missing, when CUDA_VERSION or
CUDA_VERSION.major is absent, and when major is non-int, plus returns the int
major for a valid module. Keep the existing can_implement tests as-is; this
change should target the helper’s own import/getattr/isinstance logic.
🪄 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: Enterprise
Run ID: c47b1f54-6586-4d54-a7b0-6039ef5d1b05
📒 Files selected for processing (3)
tensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.pytests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
Addressed the CodeRabbit review in
Validation after the follow-up: focused pytest |
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
|
Checked the family-modes angle on e4cf2cc, since NVFP4_AWQ and W4A8_NVFP4_FP8 share the NVFP4 QuantMode bit (mode.py:427) and previously entered the gate via has_nvfp4(): can_implement already exact-matches QuantAlgo.NVFP4 (fused_moe_cute_dsl_b12x.py:108), so the tighter gate can't regress them -- it just makes the selector consistent with the backend's own check. Source-read only, on the same tree I ran the #12704 tile probes from. |
leslie-fang25
left a comment
leslie-fang25
left a comment
There was a problem hiding this comment.
LGTM. Better land after #12704
thank you!! appreciate it! @leslie-fang25 any chance you can help me with CI on #12704 ? it's approve, but I've been waiting on CI for awhile. |
3 participants
Fixes #15853
Summary
_mma.block_scalefailureValidation
36 passedNotes
The root cause was package payload ordering:
nvidia-cutlass-dsl-libs-cu13==4.5.0was installed, but the active files came fromnvidia-cutlass-dsl-libs-base==4.5.0, so CuTe DSL reported CUDA 12.9 and generated PTX withoutkind::mxf4nvf4. Reinstalling the cu13 libs after base made CuTe DSL report CUDA 13.1 and the B12x path compiled and ran on SM121.The same CUDA 12.x lowering failure was reproduced for an SM120 target, so the guard applies to the whole current B12x SM12x support set.
MTP did not require an additional code change here; it works once the B12x path is available and acceptance is read from iteration stats rather than request-level perf metrics.
Summary by CodeRabbit
Bug Fixes
Tests