[TRTLLM-6898][feat] make fused_moe_cute_dsl work on blackwell #6616
Conversation
Signed-off-by: Mindy Li <[email protected]>
📝 WalkthroughWalkthroughThis update refactors the FP8 block scales fused MoE quantization logic by splitting SM100-specific resmoothing and scale transformation from the base class into a new subclass. It introduces conditional selection of FP8 GEMM implementations via a new flag, updates test coverage for both implementations, and adjusts test decorators and configuration handling. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Model
participant DeepSeekFP8BlockScalesFusedMoEMethod
participant DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm
User->>Model: Instantiate Fused MoE with quantization
Model->>DeepSeekFP8BlockScalesFusedMoEMethod: load_weights(...)
alt SM version 100 and DeepGemm
Model->>DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm: load_weights(...)
DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm->>DeepSeekFP8BlockScalesFusedMoEMethod: super().load_weights(...)
DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm->>DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm: Resmooth weights, transform quant scales
else Other cases
DeepSeekFP8BlockScalesFusedMoEMethod->>DeepSeekFP8BlockScalesFusedMoEMethod: super().load_weights(...)
end
sequenceDiagram
participant Linear
participant FP8BlockScalesLinearMethod
Linear->>FP8BlockScalesLinearMethod: apply(...)
alt use_cute_dsl_blockscaling_mm is True
FP8BlockScalesLinearMethod->>FP8BlockScalesLinearMethod: Use Cute DSL FP8 GEMM
else
FP8BlockScalesLinearMethod->>FP8BlockScalesLinearMethod: Use Deep GEMM FP8 GEMM
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. 📜 Recent review detailsConfiguration used: .coderabbit.yaml 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
⏰ 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)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![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 (3)
tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)
1142-1142: Decorator change leaves neighbouring tests inconsistent
test_fp8_block_scales_4gpusright below still uses@skip_no_hopper.
If FP8 Block-Scales for the Cutlass backend also became Blackwell-only, consider updating that decorator for consistency; otherwise please document why the capability split differs between the plain and Cute-DSL paths.tensorrt_llm/_torch/modules/fused_moe/quantization.py (1)
433-433: Class naming could be more descriptiveThe class name
DeepSeekFP8BlockScalesFusedMoEMethodCuteDslis quite long and the "CuteDsl" suffix might not clearly indicate its purpose. Consider a more descriptive name likeDeepSeekFP8BlockScalesFusedMoEMethodBaseorDeepSeekFP8BlockScalesFusedMoEMethodSimpleto better reflect that this is the base implementation without SM-specific optimizations.tests/unittest/_torch/modules/test_fused_moe.py (1)
956-957: Document the use_trtllmgen_mm parameterConsider adding a docstring or inline comment to explain when and why this parameter should be set to True, based on the accuracy issues mentioned in line 657.
model_config: ModelConfig = ModelConfig(), - use_trtllmgen_mm: bool = False): + use_trtllmgen_mm: bool = False): # Set True to avoid accuracy issues with deepgemm mm
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py(2 hunks)tensorrt_llm/_torch/modules/fused_moe/quantization.py(2 hunks)tensorrt_llm/_torch/modules/gated_mlp.py(3 hunks)tensorrt_llm/_torch/modules/linear.py(3 hunks)tests/integration/defs/accuracy/test_llm_api_pytorch.py(4 hunks)tests/unittest/_torch/modules/test_fused_moe.py(6 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile = ...).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL = ...).
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a Python class in the constructor.
For interfaces that may be used outside a file, prefer docstrings over comments in Python.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for classes and functions in Python, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without it.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.
Files:
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.pytensorrt_llm/_torch/modules/gated_mlp.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytensorrt_llm/_torch/modules/linear.pytests/unittest/_torch/modules/test_fused_moe.py
**/*.{cpp,h,hpp,cc,cxx,cu,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
Files:
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.pytensorrt_llm/_torch/modules/gated_mlp.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytensorrt_llm/_torch/modules/linear.pytests/unittest/_torch/modules/test_fused_moe.py
🧠 Learnings (1)
📚 Learning: in tensorrt-llm testing, it's common to have both cli flow tests (test_cli_flow.py) and pytorch api ...
Learnt from: moraxu
PR: NVIDIA/TensorRT-LLM#6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.
Applied to files:
tests/integration/defs/accuracy/test_llm_api_pytorch.pytests/unittest/_torch/modules/test_fused_moe.py
⏰ 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: Pre-commit Check
🔇 Additional comments (15)
tensorrt_llm/_torch/modules/linear.py (3)
464-464: LGTM: New parameter properly added to constructor.The
use_trtllmgen_mmparameter is correctly added with appropriate type annotation and default value.
481-481: LGTM: Parameter properly stored as instance variable.The parameter is correctly stored as an instance variable for later use in the
applymethod.
575-575: Backend selection on Blackwell is correctThe
if get_sm_version() == 100 and not module.use_trtllmgen_mmbranch ensures:
- On Blackwell (SM100) with use_trtllmgen_mm=False ⇒ deep_gemm.fp8_gemm_nt is used
- In all other cases (including SM100 + use_trtllmgen_mm=True) ⇒ trtllm.fp8_block_scaling_gemm is used
This matches the PR objective to let users override deep_gemm on Blackwell via use_trtllmgen_mm.
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py (2)
12-13: LGTM: Import statement properly updated.The imports are correctly expanded to include the new quantization method classes needed for the refactored quantization system.
143-153: LGTM: Quantization method selection logic is well-structured.The
_get_quant_methodhelper method provides clear quantization method selection with:
- Proper check for quantization configuration existence
- Specific handling for FP8 block scales quantization
- Clear error message for unsupported modes
- Fallback to unquantized method when no quantization is enabled
The logic aligns with the refactored quantization classes mentioned in the AI summary.
tensorrt_llm/_torch/modules/gated_mlp.py (3)
41-41: LGTM: Parameter properly added to GatedMLP constructor.The
use_trtllmgen_mmparameter is correctly added with appropriate type annotation and default value, maintaining consistency with the Linear class.
79-79: LGTM: Flag properly propagated to gate_up_proj Linear layer.The
use_trtllmgen_mmflag is correctly passed to the gate_up_proj Linear instance, ensuring consistent matrix multiplication behavior.
97-97: LGTM: Flag properly propagated to down_proj Linear layer.The
use_trtllmgen_mmflag is correctly passed to the down_proj Linear instance, maintaining consistency across both Linear layers in the GatedMLP.tests/integration/defs/accuracy/test_llm_api_pytorch.py (3)
988-988: Verify gating:@skip_pre_blackwellmakes the test Blackwell-onlyThe test used to be skipped when no Hopper was present; with
@skip_pre_blackwellit will now run only on Blackwell parts, excluding Hopper/H200.
Double-check that the kernel path exercised byCuteDslFusedMoEreally requires Blackwell and is not supported on Hopper, otherwise we may lose coverage.
1013-1018: Switch tocuda_graph_configlooks correct – ensure legacyuse_cuda_graphis removedThe new dict drops the deprecated
use_cuda_graphflag in favour of the explicitcuda_graph_config.
Confirm that:
- No stale
use_cuda_graphkeyword is still passed further down the call chain (grep in the repo).cuda_graph_config=Nonepreserves the previous “disabled” behaviour whencuda_graphisFalse.If both are true, the change is good.
1177-1182: Same update for 4-GPU variant – keep the two call sites in syncThe 4-GPU test mirrors the single-GPU one; good consistency.
As above, run a quick sweep to ensure no other helper buildspytorch_configwith the old flag.tensorrt_llm/_torch/modules/fused_moe/quantization.py (1)
562-598: SM version 100 specific optimizations look correctThe implementation properly handles Blackwell-specific optimizations with appropriate conditionals and expert ID filtering. The resmoothing and scale transformation logic aligns with the PR objective of enabling fused_moe_cute_dsl on Blackwell.
tests/unittest/_torch/modules/test_fused_moe.py (3)
17-17: Import cleanup looks goodThe removal of
skip_non_hopper_unittestand retention ofskip_pre_blackwellaligns with the PR's focus on Blackwell platform support.
561-577: Test function properly updated for CuteDsl implementationThe test function renaming and decorator change correctly reflect the shift to the new CuteDsl implementation and Blackwell platform targeting.
672-697: Multi-GPU test properly updatedThe multi-GPU test correctly mirrors the single-GPU test changes, maintaining consistency across the test suite.
Signed-off-by: Mindy Li <[email protected]>
limin2021
changed the title
Signed-off-by: Mindy Li <[email protected]>
|
/bot run |
Signed-off-by: Mindy Li <[email protected]>
|
PR_Github #14109 [ run ] triggered by Bot |
Signed-off-by: Mindy Li <[email protected]>
|
/bot run |
|
PR_Github #14111 [ run ] triggered by Bot |
|
PR_Github #14109 [ run ] completed with state |
|
PR_Github #14111 [ run ] completed with state |
…tep-0 Signed-off-by: Mindy Li <[email protected]>
Signed-off-by: Mindy Li <[email protected]>
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 0
🔭 Outside diff range comments (2)
tensorrt_llm/_torch/modules/fused_moe/quantization.py (1)
1-1: Add required NVIDIA copyright header.The file is missing the required NVIDIA copyright header as specified in the coding guidelines for all TensorRT-LLM source files.
Add the copyright header at the beginning of the file:
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from abc import ABC, abstractmethodtests/unittest/_torch/modules/test_fused_moe.py (1)
697-828: Significant code duplication between test functions.This entire test function
test_fused_moe_fp8_blockwise_cutlassis nearly identical to the CuteDsl version above, with only minor differences:
- Uses
@skip_non_hopper_unittestdecorator- Instantiates
CutlassFusedMoEinstead ofCuteDslFusedMoE- Uses
use_cute_dsl_blockscaling_mm=Truein reference model (Line 816)Consider extracting the common test logic into a helper function to reduce duplication:
def _test_fused_moe_fp8_blockwise_common(moe_class, use_cute_dsl_mm, dtype, num_experts, seq_len, hidden_size, RoutingMethodCls, WeightLoadingMode, mapping=None): # Common test logic here # ... ref_fused_moe = RefGatedMLPFusedMoE( # ... other params use_cute_dsl_blockscaling_mm=use_cute_dsl_mm, ) # ... rest of common logic @skip_pre_blackwell @pytest.mark.parametrize(...) def test_fused_moe_fp8_blockwise_cute_dsl(...): return _test_fused_moe_fp8_blockwise_common( CuteDslFusedMoE, True, dtype, num_experts, seq_len, hidden_size, RoutingMethodCls, WeightLoadingMode, mapping ) @skip_non_hopper_unittest @pytest.mark.parametrize(...) def test_fused_moe_fp8_blockwise_cutlass(...): return _test_fused_moe_fp8_blockwise_common( CutlassFusedMoE, True, dtype, num_experts, seq_len, hidden_size, RoutingMethodCls, WeightLoadingMode, mapping )
🧹 Nitpick comments (2)
tensorrt_llm/_torch/modules/fused_moe/quantization.py (2)
569-571: Add class docstring to explain the specialized functionality.The class lacks documentation explaining its purpose and how it differs from the base
DeepSeekFP8BlockScalesFusedMoEMethod.Add a docstring to document the class:
class DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm( DeepSeekFP8BlockScalesFusedMoEMethod): + """ + Specialized DeepSeek FP8 block scales quantization method for DeepGemm backend. + + This subclass adds SM version 100-specific processing including weight resmoothing + and scale transformation for the DeepGemm implementation. + """
599-609: Consider extracting scale transformation logic to reduce duplication.The scale transformation logic for w3_w1 and w2 scales is very similar and could be extracted into a helper method to improve maintainability.
Consider extracting the transformation logic:
+ def _transform_scale_parameter(self, module, scale_tensor, weight_tensor, param_name): + """Transform scale tensor into required layout and create new parameter.""" + transformed_scale = transform_sf_into_required_layout( + scale_tensor, + mn=weight_tensor.shape[1], + k=weight_tensor.shape[2], + recipe=(1, 128, 128), + num_groups=weight_tensor.shape[0], + is_sfa=False) + setattr(module, param_name, nn.Parameter(transformed_scale, requires_grad=False)) + if get_sm_version() == 100: - transfromed_w3_w1_scale = transform_sf_into_required_layout( - module.quant_scales[0], - mn=module.w3_w1_weight.shape[1], - k=module.w3_w1_weight.shape[2], - recipe=(1, 128, 128), - num_groups=module.w3_w1_weight.shape[0], - is_sfa=False) - module.w3_w1_weight_scaling_factor = nn.Parameter( - transfromed_w3_w1_scale, requires_grad=False) - transfromed_w2_scale = transform_sf_into_required_layout( - module.quant_scales[1], - mn=module.w2_weight.shape[1], - k=module.w2_weight.shape[2], - recipe=(1, 128, 128), - num_groups=module.w3_w1_weight.shape[0], - is_sfa=False) - module.w2_weight_scaling_factor = nn.Parameter(transfromed_w2_scale, - requires_grad=False) + self._transform_scale_parameter( + module, module.quant_scales[0], module.w3_w1_weight, + 'w3_w1_weight_scaling_factor') + self._transform_scale_parameter( + module, module.quant_scales[1], module.w2_weight, + 'w2_weight_scaling_factor')
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/modules/fused_moe/quantization.py(1 hunks)tests/integration/defs/accuracy/test_llm_api_pytorch.py(4 hunks)tests/unittest/_torch/modules/test_fused_moe.py(7 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/integration/defs/accuracy/test_llm_api_pytorch.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile = ...).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL = ...).
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a Python class in the constructor.
For interfaces that may be used outside a file, prefer docstrings over comments in Python.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without it.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.
Files:
tensorrt_llm/_torch/modules/fused_moe/quantization.pytests/unittest/_torch/modules/test_fused_moe.py
**/*.{cpp,h,cu,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
Files:
tensorrt_llm/_torch/modules/fused_moe/quantization.pytests/unittest/_torch/modules/test_fused_moe.py
🔇 Additional comments (11)
tensorrt_llm/_torch/modules/fused_moe/quantization.py (2)
574-588: LGTM! Weight resmoothing logic is correctly implemented.The expert ID collection and weight resmoothing for SM version 100 is properly implemented. The logic correctly:
- Collects both initial and shared expert IDs
- Filters weights by expert membership to avoid unnecessary processing
- Applies resmoothing to matching weight/scale pairs
- Updates the weights dictionary in-place
591-610: Scale transformation and parameter reassignment looks correct.The SM100-specific scale transformation logic properly:
- Transforms both w3_w1 and w2 scales using the required layout format
- Creates new Parameter objects with the transformed scales
- Maintains the correct tensor shapes and device placement
- Calls setup_quant_scales to finalize the configuration
tests/unittest/_torch/modules/test_fused_moe.py (9)
563-563: LGTM! Decorator update aligns with Blackwell platform support.The decorator change from
@skip_non_hopper_unittestto@skip_pre_blackwellcorrectly reflects that this test is specifically for Blackwell platform capabilities.
575-581: Function renamed to reflect CuteDsl implementation testing.The function has been appropriately renamed from
test_fused_moe_fp8_blockwisetotest_fused_moe_fp8_blockwise_cute_dslto clearly indicate it tests the CuteDsl implementation.
674-695: Reference model correctly configured for CuteDsl testing.The reference model instantiation properly sets
use_cute_dsl_blockscaling_mm=Trueto match the implementation being tested, ensuring accurate comparison.
831-831: LGTM! Decorator update for multi-GPU test.The decorator change from
@skip_non_hopper_unittestto@skip_pre_blackwellis consistent with the single-GPU test changes.
836-836: Function renamed to reflect CuteDsl multi-GPU testing.The function has been appropriately renamed to
test_fused_moe_fp8_blockwise_cute_dsl_multi_gpufor clarity.
840-840: Test function call updated correctly.The call has been updated to invoke the renamed CuteDsl test function, maintaining consistency.
847-847: Weight loading mode parameter correctly specified.The explicit specification of
MoEWeightLoadingMode.VANILLAensures consistent test behavior.
1115-1116: New parameter addition supports implementation selection.The addition of
use_cute_dsl_blockscaling_mm: bool = Falseparameter allows the reference model to use different matrix multiplication implementations for comparison purposes.
1133-1133: Parameter correctly passed to GatedMLP constructor.The
use_cute_dsl_blockscaling_mmparameter is properly forwarded to theGatedMLPconstructor, enabling the correct implementation path.
|
PR_Github #14214 [ run ] triggered by Bot |
Signed-off-by: Mindy Li <[email protected]>
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/modules/test_fused_moe.py (1)
575-828: Consider refactoring to reduce code duplication.The two test functions
test_fused_moe_fp8_blockwise_cute_dslandtest_fused_moe_fp8_blockwise_cutlasshave nearly identical logic with only minor differences in implementation class and reference model configuration. This violates the DRY principle and increases maintenance burden.Consider extracting common test logic into a helper function:
def _test_fused_moe_fp8_blockwise_impl(moe_class, use_cute_dsl_mm, dtype, num_experts, seq_len, hidden_size, RoutingMethodCls, WeightLoadingMode, mapping=None): # Common test implementation # ... fused_moe = moe_class(...) ref_fused_moe = RefGatedMLPFusedMoE(..., use_cute_dsl_blockscaling_mm=use_cute_dsl_mm) # ... @skip_pre_blackwell @pytest.mark.parametrize(...) def test_fused_moe_fp8_blockwise_cute_dsl(...): return _test_fused_moe_fp8_blockwise_impl(CuteDslFusedMoE, True, ...) @skip_non_hopper_unittest @pytest.mark.parametrize(...) def test_fused_moe_fp8_blockwise_cutlass(...): return _test_fused_moe_fp8_blockwise_impl(CutlassFusedMoE, False, ...)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/unittest/_torch/modules/test_fused_moe.py(8 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile = ...).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL = ...).
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a Python class in the constructor.
For interfaces that may be used outside a file, prefer docstrings over comments in Python.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without it.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.
Files:
tests/unittest/_torch/modules/test_fused_moe.py
**/*.{cpp,h,cu,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
Files:
tests/unittest/_torch/modules/test_fused_moe.py
⏰ 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: Pre-commit Check
🔇 Additional comments (5)
tests/unittest/_torch/modules/test_fused_moe.py (5)
563-695: LGTM! Test refactoring properly separates CuteDsl implementation testing.The function rename and decorator change align with the PR objective to enable fused_moe_cute_dsl on Blackwell. The use of
CuteDslFusedMoEand theuse_cute_dsl_blockscaling_mm=Trueparameter in the reference model ensure this test specifically validates the new implementation path.
836-894: Multi-GPU test structure looks good.The separation of multi-GPU tests for both implementations follows a clean pattern, with appropriate decorator updates. The delegation to single-GPU test functions maintains consistency.
However, this inherits the same reference model configuration issue from
test_fused_moe_fp8_blockwise_cutlassthat should be addressed.
1152-1170: LGTM! Parameter addition follows good practices.The new
use_cute_dsl_blockscaling_mmparameter is properly added with a sensible default value and correctly passed through to theGatedMLPconstructor. This maintains backward compatibility while enabling the new functionality.
17-17: Import usage is appropriate for platform-specific testing.The continued import of
skip_non_hopper_unittestis correct since it's still used for the Cutlass-specific tests that don't support pre-Hopper platforms, whileskip_pre_blackwellis used for the new CuteDsl tests that require Blackwell or newer.
564-894: Comprehensive test coverage for both implementations.The test suite now provides thorough coverage for both CuteDsl and Cutlass implementations, including:
- Single-GPU and multi-GPU scenarios
- Multiple weight loading modes
- Consistent parameter ranges across implementations
This ensures both implementation paths are properly validated.
Signed-off-by: Mindy Li <[email protected]>
|
PR_Github #14217 [ run ] triggered by Bot |
|
PR_Github #14214 [ run ] completed with state |
|
/bot run |
|
PR_Github #14223 [ run ] triggered by Bot |
|
PR_Github #14217 [ run ] completed with state |
|
PR_Github #14223 [ run ] completed with state |
…tep-0 Signed-off-by: Mindy Li <[email protected]>
|
/bot run |
|
PR_Github #14348 [ run ] triggered by Bot |
…tep-0 Signed-off-by: Mindy Li <[email protected]>
|
/bot run |
|
PR_Github #14413 [ run ] triggered by Bot |
|
PR_Github #14348 [ run ] completed with state |
|
PR_Github #14413 [ run ] completed with state |
…tep-0 Signed-off-by: Mindy Li <[email protected]>
|
/bot run |
|
PR_Github #14520 [ run ] triggered by Bot |
|
PR_Github #14520 [ run ] completed with state |
Signed-off-by: Rakib Hasan <[email protected]>
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Chores
Description
Test Coverage
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.