Feature: Add support for L40 FusedMoE in cutlass path #1973
Conversation
Signed-off-by: Amir Klein <[email protected]>
amirkl94
requested review from
aleozlx,
bkryu,
cyx-6,
joker-eph,
wenscarl,
yongwww and
yzh119
as code owners
Summary of ChangesHello @amirkl94, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces comprehensive support for L40 GPUs (SM89 architecture) within the CUTLASS FusedMoE kernels. It resolves critical compilation issues and a runtime crash related to shared memory allocation for a specific GEMM tactic on SM89, ensuring stable and efficient operation. The changes also include updates to the Python JIT compilation infrastructure to enable proper module generation and loading for L40. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdd explicit SM89 support and FP4/FP8 gating adjustments across CUTLASS heuristics, MOE GEMM dispatch, and Flashinfer JIT: new SM89 NVCC flags and JIT module generator, tightened ENABLE_FP4 guards, and split SM89 vs SM>=120 handling for grouped FP8 heuristics. Changes
Sequence Diagram(s)sequenceDiagram
actor User
participant PyCore as flashinfer/fused_moe/core.py
participant JITGen as fused_moe.gen_cutlass_fused_moe_sm89_module
participant JITCore as flashinfer/jit/core.py
participant CppDisp as moe_gemm_template_dispatch.h
participant CUTLASS as cutlass_kernels
User->>PyCore: get_cutlass_fused_moe_module(backend="89")
PyCore->>JITGen: build_and_load()
JITGen->>JITCore: use sm89_nvcc_flags + BF16/FP8 (+FP8_BLOCK_SCALE?)
JITGen->>PyCore: return compiled module
PyCore-->>User: loaded module
Note over CppDisp,CUTLASS: SM dispatch for FP formats
alt FP8 or grouped FP8
CppDisp->>CppDisp: if sm == 89 (new explicit branch)
CppDisp->>CUTLASS: dispatch SM89-specific tile configs
else if sm >= 120
CppDisp->>CUTLASS: dispatch SM>=120 tile configs
else
CppDisp->>CUTLASS: default/fallback path
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Potential attention points:
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for L40 GPUs (sm_89) in the CUTLASS FusedMoE path. The changes include fixing compilation issues, removing a problematic GEMM tactic for sm_89 that was causing crashes, and adding the necessary build configurations for this architecture. The changes are logical and well-implemented. I have one suggestion to improve code clarity when constructing compiler flags.
| nvcc_flags = sm89_nvcc_flags + [ | ||
| "-DENABLE_BF16", | ||
| "-DENABLE_FP8", | ||
| "-DENABLE_FP8_BLOCK_SCALE" if is_cuda_version_at_least("12.8") else "", | ||
| "-DUSING_OSS_CUTLASS_MOE_GEMM", | ||
| ] |
There was a problem hiding this comment.
Conditionally adding an empty string to the list of compiler flags can be slightly confusing and less clean. It's better to conditionally append the flag to the list to avoid adding empty elements. This improves code readability and maintainability.
| nvcc_flags = sm89_nvcc_flags + [ | |
| "-DENABLE_BF16", | |
| "-DENABLE_FP8", | |
| "-DENABLE_FP8_BLOCK_SCALE" if is_cuda_version_at_least("12.8") else "", | |
| "-DUSING_OSS_CUTLASS_MOE_GEMM", | |
| ] | |
| nvcc_flags = sm89_nvcc_flags + [ | |
| "-DENABLE_BF16", | |
| "-DENABLE_FP8", | |
| "-DUSING_OSS_CUTLASS_MOE_GEMM", | |
| ] | |
| if is_cuda_version_at_least("12.8"): | |
| nvcc_flags.append("-DENABLE_FP8_BLOCK_SCALE") |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h (1)
42-42: Fix misleading comment about SM120 vs SM100.The comment mentions
CUTLASS_ARCH_MMA_SM100_SUPPORTEDbut this function checks forCUTLASS_ARCH_MMA_SM120_SUPPORTED(line 35). Update the comment to accurately reflect the macro being checked.Apply this diff:
- return false; // CUTLASS_ARCH_MMA_SM100_SUPPORTED is set when Blackwell kernels are enabled + return false; // CUTLASS_ARCH_MMA_SM120_SUPPORTED is set when Blackwell kernels are enabled
🧹 Nitpick comments (2)
flashinfer/jit/fused_moe.py (1)
80-88: SM89 module generation is correctly implemented.The function appropriately:
- Uses
sm89_nvcc_flagswhich excludes FP4 support for L40- Omits Hopper-specific TMA GEMM flags (correct for Ada architecture)
- Includes conditional FP8 block scale support for CUDA ≥12.8
Optional: Consider iterable unpacking for cleaner syntax.
As suggested by Ruff, you could use iterable unpacking instead of concatenation:
- nvcc_flags = sm89_nvcc_flags + [ + nvcc_flags = [ + *sm89_nvcc_flags, "-DENABLE_BF16", "-DENABLE_FP8", "-DENABLE_FP8_BLOCK_SCALE" if is_cuda_version_at_least("12.8") else "", "-DUSING_OSS_CUTLASS_MOE_GEMM", ]This is a minor style improvement and can be deferred.
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h (1)
696-709: LGTM: Correctly reorganizes SM89 dispatch to avoid shared memory issues.The reorganized control flow properly addresses the L40 (SM89) issue by:
- Routing FP8 workloads to Sm89 kernels with runtime validation (line 703)
- Routing non-FP8 workloads to Sm80 kernels (lines 707-708)
This aligns with the kernel implementation in
moe_cutlass_kernel.hwhich shows SM89 architecture reusing Sm80 kernels for non-FP8 types, and prevents the "GPU lacks the shared memory resources" assertion mentioned in the PR objectives.Optional suggestion: Consider adding a brief comment explaining why non-FP8 on SM89 uses the Sm80 path, to help future maintainers understand the shared memory constraint that motivated this design.
Apply this diff to add a clarifying comment:
} else { + // Non-FP8 workloads on SM89 (L40) reuse Sm80 kernels to avoid + // Sm89-specific tactics that exceed shared memory limits dispatchMoeGemmToCutlass<T, WeightType, ScaleBiasType, cutlass::arch::Sm80, EpilogueTag>( inputs, multi_processor_count_); }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp(1 hunks)csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h(1 hunks)csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h(2 hunks)flashinfer/fused_moe/core.py(2 hunks)flashinfer/jit/core.py(1 hunks)flashinfer/jit/fused_moe.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
flashinfer/fused_moe/core.py (2)
flashinfer/jit/fused_moe.py (1)
gen_cutlass_fused_moe_sm89_module(80-87)flashinfer/jit/core.py (1)
build_and_load(272-284)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h (2)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h (6)
__nv_fp8_e5m2(91-93)cutlass(114-116)cutlass(120-122)cutlass(127-129)cutlass(132-134)cutlass(140-142)csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h (1)
cutlass(40-677)
flashinfer/jit/fused_moe.py (2)
flashinfer/jit/core.py (2)
JitSpec(185-284)gen_jit_spec(287-353)flashinfer/jit/cpp_ext.py (1)
is_cuda_version_at_least(86-87)
🪛 Ruff (0.14.1)
flashinfer/jit/fused_moe.py
81-86: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
🔇 Additional comments (6)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp (1)
161-168: LGTM! Targeted fix for L40 shared memory constraints.The separation of SM89 from SM >= 120 handling correctly removes
CtaShape16x256x128_WarpShape16x64x128for L40 GPUs. This config's larger K dimension (128 vs 64) would exceed L40's shared memory capacity in GROUPED_GEMM mode, causing the assertion failure mentioned in the PR. The fix is minimal, well-scoped, and maintains full functionality for newer architectures.flashinfer/jit/core.py (1)
93-96: SM89 flags correctly omit FP4 support.The implementation appropriately excludes
-DFLASHINFER_ENABLE_FP4_E2M1for SM89 (L40), unlike other architectures that usecommon_nvcc_flags. This aligns with the PR objective of fixing compilation issues for L40 by removing problematic GEMM tactics.flashinfer/fused_moe/core.py (2)
40-40: LGTM!Import added correctly to support SM89 module generation.
289-290: LGTM!Backend dispatch for SM89 follows the established pattern and correctly invokes the new module generation function.
flashinfer/jit/fused_moe.py (1)
21-27: LGTM!Import correctly updated to include
sm89_nvcc_flagsneeded for the new SM89 module generation function.csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h (1)
692-694: LGTM: Proper FP4 validation for SM89.The explicit compile-time check correctly prevents FP4 usage on SM < 90, consistent with the earlier validation for SM75-80 at lines 685-688.
.../nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h
Outdated
Show resolved
Hide resolved
.../nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h
Outdated
Show resolved
Hide resolved
Signed-off-by: Amir Klein <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Deploy Docs
🔇 Additional comments (1)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h (1)
35-36: LGTM: Guard is appropriate for FP4-only SM120 specialization.The
ENABLE_FP4guard is correct here. SinceisValidSM120MOESpecialisationonly supports__nv_fp4_e2m1types (lines 37-40), requiringENABLE_FP4at the preprocessor level prevents compilation errors when FP4 support is not available.
| #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && \ | ||
| defined(ENABLE_FP4) // TODO Is there a better choice |
There was a problem hiding this comment.
Overly restrictive guard blocks non-FP4 Blackwell specializations.
Requiring ENABLE_FP4 at the top level prevents ALL Blackwell MOE specializations from being valid when FP4 support is not compiled, including non-FP4 cases like FP8/FP8 or BF16/BF16 that should work regardless. Line 53 explicitly supports cutlass::platform::is_same<T, WeightType>::value, which matches any same-type pair.
The Hopper implementation (lines 68-89) uses a better pattern: conditional compilation guards only the FP4-specific parts inside the return statement, allowing non-FP4 specializations to work without ENABLE_FP4.
Apply this diff to follow the Hopper pattern:
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) && \
- defined(ENABLE_FP4) // TODO Is there a better choice
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
return (cutlass::platform::is_same<T, WeightType>::value ||
+#ifdef ENABLE_FP4
(cutlass::platform::is_same<T, __nv_fp8_e4m3>::value &&
- cutlass::platform::is_same<WeightType, __nv_fp4_e2m1>::value)) &&
+ cutlass::platform::is_same<WeightType, __nv_fp4_e2m1>::value) ||
+#else
+ false ||
+#endif
+ false) &&
cutlass::platform::is_same<EpilogueTag, cutlass_extensions::EpilogueOpDefault>::value &&
Fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE;Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h
around lines 51-52, the current top-level guard requires ENABLE_FP4 and blocks
all Blackwell (SM100) MOE specializations when FP4 is not enabled; change this
to only guard FP4-specific logic: keep the CUTLASS_ARCH_MMA_SM100_SUPPORTED
check on the specialization but remove ENABLE_FP4 from the #if, then wrap only
the FP4-specific condition(s) inside the specialization (the return expression)
with #if defined(ENABLE_FP4) / #endif as the Hopper implementation does (lines
~68-89), ensuring the same-type WeightType check remains unguarded so non-FP4
types (FP8/BF16/etc.) compile normally.
There was a problem hiding this comment.
The issue here is:
cutlass::platform::is_same<WeightType, __nv_fp4_e2m1>::value
Which I think isn't defined unless ENABLE_FP4 is defined.
There was a problem hiding this comment.
@amirkl94, you're absolutely right that __nv_fp4_e2m1 is only defined when ENABLE_FP4 is set! That's why my suggestion included the #ifdef ENABLE_FP4 guard specifically around the line that references that type:
#ifdef ENABLE_FP4
(cutlass::platform::is_same<T, __nv_fp8_e4m3>::value &&
cutlass::platform::is_same<WeightType, __nv_fp4_e2m1>::value) ||
#endifThe key insight is that the first condition in the return statement:
cutlass::platform::is_same<T, WeightType>::valuedoesn't reference __nv_fp4_e2m1 at all—it just checks if T and WeightType are the same type (e.g., FP8/FP8, BF16/BF16). This condition should work regardless of whether ENABLE_FP4 is defined.
By moving the top-level ENABLE_FP4 requirement into the return expression (guarding only the FP4-specific logic), non-FP4 Blackwell specializations can compile and work correctly even when FP4 support is disabled. This matches the pattern already used in isValidHopperMOESpecialisation at lines 68-89.
Let me know if you'd like me to refine the suggestion further!
| CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64, | ||
| CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64, | ||
| CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64, | ||
| CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64}; |
There was a problem hiding this comment.
Are there any SM89 GPUs that can support the CtaShape16x256x128_WarpShape16x64x128, or is this an SM120 addition?
Having a small M value like this helps with low latency cases, so I'd want to understand why its not supported before disabling it
There was a problem hiding this comment.
At the very least can you leave a comment saying what the difference between the two lists are, so people dont have to manually compare the items
There was a problem hiding this comment.
Are there any SM89 GPUs that can support the
CtaShape16x256x128_WarpShape16x64x128, or is this an SM120 addition? Having a small M value like this helps with low latency cases, so I'd want to understand why its not supported before disabling it
This is a removal from the sm89 path. When I tested it on an L40 GPU I got Assertion failed: GPU lacks the shared memory resources to run GroupedGEMM kernel .
It might be that on other sm89 GPUs it will pass, the main issue is that this was the default tactic that was chosen when trying to use FusedMoE, I believe that moving it to be the last will also fix my issue.
There was a problem hiding this comment.
I tried moving this tile config to be the last one and now the default tactic won't fail on l40. The issue is that if autotuner is on then the tactics that use this tile config will report an error with a stacktrace which looks bad.
@yzh119 do you think it'll be ok to change the errors that happen in the autotuner to be debug logs? Otherwise it means users will get spammed with error messages when they run autotuning on L40 FusedMoE.
There was a problem hiding this comment.
I think the autotuner should still output warnings, but just make them say "Skipping tactic x due to error. This tactic may not be supported by the current GPU architecture".
That said I know there is a difference of opinion on whether we should proactively filter them as you have done here, the argument being that we should be able to do the due diligence to determine what tactics are supported so that we can raise an error when a tactic fails when it shouldn't. So I can see either side.
|
There are some compilation errors on CI, @amirkl94 would you mind taking a look? |
Signed-off-by: Amir Klein <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h (1)
691-699: LGTM: FP8/W4AFP8 routing to SM89 is correct.The restructured logic correctly routes FP8 and W4AFP8 workloads to the SM89 architecture for the SM 80-89 range, which is essential for L40 support. The runtime check at line 697 appropriately prevents FP8 usage on SM80-88 hardware that lacks FP8 support. The static_assert placement within the ENABLE_FP8 block is consistent with the codebase pattern (see lines 244-247).
Minor: Consider clarifying the error message.
The error message at line 697 mentions "fp8" but the condition includes
use_w4afp8. For clarity, consider updating the message to reflect both supported types.- TLLM_CHECK_WITH_INFO(sm_ == 89, "For sm >= 80 and < 90, fp8 is only supported with sm == 89"); + TLLM_CHECK_WITH_INFO(sm_ == 89, "For sm >= 80 and < 90, fp8 and w4afp8 are only supported with sm == 89");
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h (3)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h (6)
__nv_fp8_e5m2(91-93)cutlass(114-116)cutlass(120-122)cutlass(127-129)cutlass(132-134)cutlass(140-142)csrc/nv_internal/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h (1)
cutlass(40-677)csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h (1)
multi_processor_count_(341-341)
🔇 Additional comments (1)
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h (1)
700-712: LGTM: Non-FP8 handling and FP4 guard are correct.The else branch correctly handles non-FP8 cases for SM 80-89:
- FP4 types are appropriately blocked on SM < 90 (line 703)
- Non-FP8 workloads dispatch to SM80 architecture, which is an intentional kernel reuse strategy for SM89 as documented in the kernel implementation
The duplication of the dispatch call (lines 705-706 and 709-710) is necessary for conditional compilation when ENABLE_FP4 is not defined.
Pushed a commit to fix this |
3 participants
📌 Description
Fixed a few compilation issues for L40, and removed 1 gemm tactic for
sm == 89that crashes due to:🧪 Tests
Ran
pytest tests/moe/test_trtllm_cutlass_fused_moe.pymanually on an L40 GPU and verified all tests passed.Summary by CodeRabbit
New Features
Bug Fixes / Compatibility
Performance