[TRTLLM-6744][feat] Remove input_sf swizzle for module WideEPMoE #6231
Conversation
📝 WalkthroughWalkthroughThis change introduces a new boolean parameter, Changes
Sequence Diagram(s)sequenceDiagram
participant PythonUser as Python User
participant TorchModule as Torch MoE Module
participant TorchOp as torch.ops.trtllm.fused_moe
participant FusedMoeRunner as FusedMoeRunner (C++)
participant CutlassMoeFCRunner as CutlassMoeFCRunner (C++)
participant Kernel as expandInputRowsKernel
PythonUser->>TorchModule: forward_chunk(..., input_sf, ...)
TorchModule->>TorchOp: fused_moe(..., input_sf, swizzled_input_sf, ...)
TorchOp->>FusedMoeRunner: runMoe(..., input_sf, swizzled_input_sf, ...)
FusedMoeRunner->>CutlassMoeFCRunner: runMoe(..., input_sf, swizzled_input_sf, ...)
CutlassMoeFCRunner->>Kernel: expandInputRowsKernelLauncher(..., input_sf, swizzled_input_sf, ...)
Kernel->>Kernel: Branch on swizzled_input_sf for layout
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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ 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: 0
🔭 Outside diff range comments (1)
tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
217-241: Missing parameter in fake registration function.The fake registration function for
fused_moe(starting at line 217) is missing the newswizzled_input_sfparameter in its signature, while the actual function has it. This inconsistency could cause issues with torch.compile compatibility.Add the missing parameter to maintain signature consistency:
@torch.library.register_fake("trtllm::fused_moe") def _( input: torch.Tensor, token_selected_experts: torch.Tensor, token_final_scales: torch.Tensor, fc1_expert_weights: torch.Tensor, fc1_expert_biases: Optional[torch.Tensor], fc2_expert_weights: torch.Tensor, fc2_expert_biases: Optional[torch.Tensor], output_dtype: torch.dtype, quant_scales: List[torch.Tensor], input_sf: Optional[torch.Tensor] = None, + swizzled_input_sf: bool = True, tp_size: int = 1, tp_rank: int = 0, ep_size: int = 1, ep_rank: int = 0, cluster_size: int = 1, cluster_rank: int = 0, enable_alltoall: bool = False, use_deepseek_fp8_block_scale: bool = False, use_w4a8_group_scaling: bool = False, use_mxfp8_act_scaling: bool = False, min_latency_mode: bool = False, tune_max_num_tokens: int = 8192, ):
🧹 Nitpick comments (1)
cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu (1)
1066-1084: Consider refactoring to reduce code duplication.The logic correctly implements the conditional behavior based on
swizzled_input_sf, but there's significant code duplication between the two branches. Consider refactoring:if (input_sf) { - if (swizzled_input_sf) - { - auto const sf_in - = cvt_quant_to_fp4_get_sf_out_offset<TmaWarpSpecializedGroupedGemmInput::ElementSF, NumThreadsPerSF, - VecSize>(std::nullopt /* batchIdx */, source_token_id, elem_idx, std::nullopt /* numRows */, - num_cols, const_cast<TmaWarpSpecializedGroupedGemmInput::ElementSF*>(input_sf), - FP4QuantizationSFLayout::SWIZZLED); - *sf_out = *sf_in; - } - else - { - auto const sf_in - = cvt_quant_to_fp4_get_sf_out_offset<TmaWarpSpecializedGroupedGemmInput::ElementSF, NumThreadsPerSF, - VecSize>(std::nullopt /* batchIdx */, source_token_id, elem_idx, std::nullopt /* numRows */, - num_cols, const_cast<TmaWarpSpecializedGroupedGemmInput::ElementSF*>(input_sf), - FP4QuantizationSFLayout::LINEAR); - *sf_out = *sf_in; - } + auto const layout = swizzled_input_sf ? FP4QuantizationSFLayout::SWIZZLED : FP4QuantizationSFLayout::LINEAR; + auto const sf_in + = cvt_quant_to_fp4_get_sf_out_offset<TmaWarpSpecializedGroupedGemmInput::ElementSF, NumThreadsPerSF, + VecSize>(std::nullopt /* batchIdx */, source_token_id, elem_idx, std::nullopt /* numRows */, + num_cols, const_cast<TmaWarpSpecializedGroupedGemmInput::ElementSF*>(input_sf), layout); + *sf_out = *sf_in; }
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h(2 hunks)cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h(2 hunks)cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h(1 hunks)cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu(9 hunks)cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp(2 hunks)cpp/tensorrt_llm/thop/moeOp.cpp(6 hunks)cpp/tensorrt_llm/thop/moeUtilOp.cpp(1 hunks)cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu(1 hunks)tensorrt_llm/_torch/custom_ops/torch_custom_ops.py(2 hunks)tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py(1 hunks)tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py(1 hunks)
🔇 Additional comments (23)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)
293-293: LGTM! Addition of swizzled_input_sf parameter is consistent with the broader MoE interface update.The addition of
swizzled_input_sf=Trueto thefused_moeoperator call is correctly aligned with the codebase-wide introduction of this parameter. TheTruevalue indicates that input scaling factors are swizzled in the CutlassFusedMoE path, which is consistent with the existing behavior and the PR's objective to differentiate swizzling behavior between CutlassFusedMoE and WideEPMoE implementations.cpp/tensorrt_llm/thop/moeUtilOp.cpp (1)
86-86: LGTM! Parameter addition aligns with the coordinated MoE interface update.The addition of the
trueargument for theswizzled_input_sfparameter toexpandInputRowsKernelLauncheris consistent with the broader codebase changes. The hardcodedtruevalue indicates that input scaling factors are always swizzled in this utility operation path, which appears appropriate given that this is a low-level utility function that likely expects pre-processed (swizzled) data.cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp (1)
960-960: LGTM! Consistent parameter addition across both compilation paths.The addition of the
trueargument forswizzled_input_sfto bothmMOERunner->runMoecalls maintains consistency between theUSING_OSS_CUTLASS_MOE_GEMMand non-OSS compilation paths. Thetruevalue correctly indicates that input scaling factors are swizzled in the plugin execution path, which aligns with the expected behavior and the broader interface update across the MoE codebase.Also applies to: 971-971
cpp/micro_benchmarks/mixtureOfExpertsBackendBenchmarkFixture.h (1)
983-983: LGTM: Consistent with the newswizzled_input_sfparameter interface.The addition of
trueas the third parameter aligns with the updatedrunMoesignature that now includes theswizzled_input_sfboolean parameter. This indicates that the benchmark assumes input scaling factors are swizzled, which is consistent with the pattern seen in other parts of the codebase.Both conditional compilation paths (#ifdef USING_OSS_CUTLASS_MOE_GEMM and #else) are updated consistently.
Also applies to: 995-995
cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu (1)
1176-1186: LGTM! Test updated to match new runMoe interface.The addition of the
swizzled_input_sfparameter (set totrue) correctly updates the test to match the new method signature. Setting this totruemaintains the existing swizzling behavior for standard MoE operations, which is appropriate for this test code.tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (1)
681-681: Parameter addition looks correct.The addition of
swizzled_input_sf=Falseparameter is consistent with the PR objective to remove input_sf swizzle for the WideEPMoE module.However, I notice the AI summary mentions that conditional swizzling code was removed (specifically lines checking
if self.has_nvfp4:and applyingswizzle_sf), but I don't see any removed lines in the provided code. This suggests a potential inconsistency between the summary and the actual changes shown.Likely an incorrect or invalid review comment.
cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h (1)
61-62: Clean kernel interface update.The addition of the
swizzled_input_sfparameter to theexpandInputRowsKernelLauncherfunction template is well-positioned and properly typed. The parameter placement afterinput_sfis logical and maintains interface consistency.tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (2)
132-132: Good parameter addition with sensible default.The addition of
swizzled_input_sf: bool = Trueparameter is well-placed and maintains backward compatibility with the default value.
202-202: Parameter correctly forwarded to underlying implementation.The parameter is properly passed through to the underlying runner call.
cpp/tensorrt_llm/thop/moeOp.cpp (4)
221-224: LGTM: Method signature updated correctly.The new
swizzled_input_sfparameter is properly positioned after the relatedinput_sfparameter and follows correct C++ const conventions.
322-322: LGTM: Parameter correctly passed to kernel calls.The
swizzled_input_sfparameter is consistently passed to both OSS and non-OSS kernel runner calls in the correct position.Also applies to: 336-336
358-361: LGTM: Consistent method signature update.The
runMoeMinLantencymethod signature is updated consistently withrunMoe, maintaining proper parameter positioning and const conventions.
453-453: LGTM: Consistent parameter passing in runMoeMinLantency.The kernel calls in
runMoeMinLantencyconsistently pass the new parameter to both OSS and non-OSS code paths, matching the pattern established inrunMoe.Also applies to: 467-467
cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h (2)
395-402: LGTM: Interface signature updated correctly.The
runMoemethod signature inCutlassMoeFCRunnerInterfaceproperly adds theswizzled_input_sfparameter with correct positioning, typing, and const conventions.
542-549: LGTM: Implementation signature matches interface.The
CutlassMoeFCRunner::runMoemethod signature correctly implements the interface with matching parameter positioning, typing, and the proper override specifier.cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu (8)
1046-1046: LGTM: Well-designed backward compatibility.The addition of the
swizzled_input_sfparameter with a default value oftruemaintains backward compatibility while enabling the new functionality. The parameter name clearly communicates its purpose.
1467-1468: LGTM: Consistent kernel interface extension.The addition of
input_sfandswizzled_input_sfparameters to the kernel template correctly propagates the new functionality. The parameter ordering is logical and maintains consistency with the overall design pattern.
1569-1571: LGTM: Correct parameter propagation.The call to
writeSFcorrectly passes the newswizzled_input_sfparameter, maintaining the intended behavior flow from kernel to device function.
1671-1672: LGTM: Launcher interface extension.The launcher function signature correctly extends to include the new parameters, maintaining consistency with the kernel interface changes.
1748-1749: LGTM: Complete parameter propagation.The kernel launch correctly passes the new
input_sfandswizzled_input_sfparameters, completing the propagation chain from the public interface through to the kernel execution.
1759-1760: LGTM: Template instantiation consistency.The macro correctly includes the new parameters to ensure all template instantiations have consistent signatures with the updated function template.
1641-1642: LGTM: Correct call site usage.The call to
expandInputRowsKernelLaunchercorrectly passes the newinput_sfandswizzled_input_sfparameters in the proper order, maintaining the intended functionality.
3422-3423: LGTM: Public API extension.The
runMoemethod signature correctly extends the public API with the newinput_sf_voidandswizzled_input_sfparameters. The parameter naming is clear and the ordering is logical, providing the entry point for controlling input scaling factor layout interpretation.
StudyingShao
force-pushed
the
jiangs/1.0.0rc4/WideEPMoE_rm_swizzle
branch
from
d2608f4 to
92792a5
Compare
|
/bot run |
|
PR_Github #12470 [ run ] triggered by Bot |
|
PR_Github #12470 [ run ] completed with state |
StudyingShao
force-pushed
the
jiangs/1.0.0rc4/WideEPMoE_rm_swizzle
branch
from
92792a5 to
9658755
Compare
|
/bot run |
|
PR_Github #13980 [ run ] triggered by Bot |
StudyingShao
changed the title
StudyingShao
changed the title
|
PR_Github #13980 [ run ] completed with state |
|
/bot run |
|
PR_Github #13992 [ run ] triggered by Bot |
|
PR_Github #13992 [ run ] completed with state |
StudyingShao
force-pushed
the
jiangs/1.0.0rc4/WideEPMoE_rm_swizzle
branch
from
942de6e to
04d0c08
Compare
|
/bot run |
|
PR_Github #14184 [ run ] triggered by Bot |
|
PR_Github #14184 [ run ] completed with state |
|
/bot run |
|
PR_Github #14245 [ run ] triggered by Bot |
|
/bot run |
|
@coderabbitai resolve |
✅ Actions performedComments resolved. |
|
PR_Github #14435 [ run ] triggered by Bot |
|
PR_Github #14435 [ run ] completed with state |
Signed-off-by: Jiang Shao <[email protected]>
|
/bot run |
|
PR_Github #14459 [ run ] triggered by Bot |
|
PR_Github #14459 [ run ] completed with state |
|
/bot run |
|
PR_Github #14486 [ run ] triggered by Bot |
|
PR_Github #14486 [ run ] completed with state |
|
/bot run |
|
PR_Github #14504 [ run ] triggered by Bot |
|
PR_Github #14504 [ run ] completed with state |
juney-nvidia
changed the title
|
Added "SW Architecture" label since follow-up actions are needed to make this PR and the related foundation code better for future maintenance. |
…DIA#6231) Signed-off-by: Jiang Shao <[email protected]>
4 participants
Summary by CodeRabbit
New Features
Refactor
Style
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.