[feat] Add gemma RMS AR fusion - #3322
Conversation
📝 WalkthroughWalkthroughAdds a ChangesAllReduce Fusion weight_bias Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. 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 Gemma/Qwen3.5 style RMSNorm to AllReduce fusion kernels by introducing a weight_bias parameter across the CUDA kernels and Python API. The implementation is accompanied by new benchmarks and comprehensive tests. Feedback from the reviewer correctly identifies that weight_bias should not be included in the mutates_args list of the register_custom_op decorators in flashinfer/comm/trtllm_ar.py, as it is an input scalar rather than an in-place modified tensor.
| @@ -242,6 +242,7 @@ def trtllm_custom_all_reduce( | |||
| "scale_factor", | |||
| "layout_code", | |||
| "block_quant_group_size", | |||
| "weight_bias", | |||
There was a problem hiding this comment.
The weight_bias parameter is an input scalar and should not be included in the mutates_args list of the register_custom_op decorator. This list is intended for arguments that are modified in-place by the operation.
| @@ -317,6 +320,7 @@ def trtllm_allreduce_fusion( | |||
| "norm_out", | |||
| "quant_out", | |||
| "scale_out", | |||
| "weight_bias", | |||
There was a problem hiding this comment.
The weight_bias parameter is an input scalar and should not be included in the mutates_args list. Only tensors that are modified in-place should be listed here.
| mutates_args=[ | ||
| "residual_out", | ||
| "norm_out", | ||
| "quant_out", | ||
| "scale_out", | ||
| "weight_bias", | ||
| ], |
There was a problem hiding this comment.
The weight_bias parameter is an input scalar and should not be included in mutates_args. Additionally, the reformatting of this list into multiple lines is unnecessary and deviates from the style used for other similar operations in this repository.
mutates_args=["residual_out", "norm_out", "quant_out", "scale_out"],
56cb99d to
da4259a
Compare
|
/bot run |
![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 (2)
tests/comm/test_allreduce_fusion_moe_unified_api.py (1)
267-275:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd unsupported-architecture skip gating in these test entrypoints.
These tests only gate on GPU count; they should also skip on unsupported SM targets via
flashinfer.utilscapability checks to avoid false failures on unsupported architectures.As per coding guidelines "
tests/**/*.py: Skip tests on unsupported GPU architectures usingflashinfer.utilscheck functions likeis_sm90a_supported(),is_sm100a_supported(), andget_compute_capability()."Also applies to: 404-415
🤖 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/comm/test_allreduce_fusion_moe_unified_api.py` around lines 267 - 275, The test entrypoint function test_moe_finalize_allreduce_unified_api currently only skips based on GPU count; add an architecture capability check at the start of this function (and the other similar test functions around lines 404-415) using flashinfer.utils helpers (e.g., call get_compute_capability() and/or is_sm90a_supported()/is_sm100a_supported()) and call pytest.skip(...) when the device compute capability is unsupported; ensure you import the needed helpers from flashinfer.utils and perform the skip before allocating or using torch.cuda to avoid running the test on unsupported SM targets.tests/comm/test_trtllm_allreduce_fusion.py (1)
470-481:⚠️ Potential issue | 🟠 Major | ⚡ Quick winGate this suite on supported compute capability, not just GPU count.
Please add an architecture capability skip using
flashinfer.utilschecks in the shared execution path for this suite, so unsupported SMs are skipped consistently.As per coding guidelines "
tests/**/*.py: Skip tests on unsupported GPU architectures usingflashinfer.utilscheck functions likeis_sm90a_supported(),is_sm100a_supported(), andget_compute_capability()."🤖 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/comm/test_trtllm_allreduce_fusion.py` around lines 470 - 481, The test currently only skips by GPU count; add a compute-capability gate in test_trtllm_allreduce_fusion using flashinfer.utils so unsupported SMs are skipped consistently: import and call is_sm90a_supported(), is_sm100a_supported() or get_compute_capability() at the start of the test_trtllm_allreduce_fusion function (before running GPU-dependent code) and call pytest.skip(...) when the current device compute capability is not supported; ensure you reference the flashinfer.utils helpers (is_sm90a_supported, is_sm100a_supported, get_compute_capability) and keep the existing world_size vs torch.cuda.device_count() check after the capability check so tests on unsupported architectures are skipped early.
🤖 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 `@flashinfer/comm/trtllm_mnnvl_ar.py`:
- Around line 302-303: The custom-op wrapper's parameter weight_bias currently
has type Optional[float] = None but is always forwarded to the CUDA op; change
its signature to use a numeric default (e.g., float = 0.0) so deprecated callers
that omit it continue to receive 0.0 instead of None. Update the function/method
signature that declares weight_bias (in the wrapper around the RMSNorm/custom op
in trtllm_mnnvl_ar.py) from Optional[float] = None to float = 0.0 and adjust any
type hints or callers if necessary to match the non-optional float default.
In `@tests/comm/test_gemma_ar_fusion.py`:
- Around line 164-169: In test_gemma_rmsnorm_ar_fusion, add GPU architecture
checks using flashinfer.utils to skip the test on unsupported SMs: import and
call get_compute_capability() and the helper predicates (is_sm90a_supported(),
is_sm100a_supported()) and if the current compute capability isn't supported,
call pytest.skip(...) similar to the existing world_size/GPU-count guard; ensure
checks are performed before heavy GPU usage so the test exits early when the
device SM is unsupported.
---
Outside diff comments:
In `@tests/comm/test_allreduce_fusion_moe_unified_api.py`:
- Around line 267-275: The test entrypoint function
test_moe_finalize_allreduce_unified_api currently only skips based on GPU count;
add an architecture capability check at the start of this function (and the
other similar test functions around lines 404-415) using flashinfer.utils
helpers (e.g., call get_compute_capability() and/or
is_sm90a_supported()/is_sm100a_supported()) and call pytest.skip(...) when the
device compute capability is unsupported; ensure you import the needed helpers
from flashinfer.utils and perform the skip before allocating or using torch.cuda
to avoid running the test on unsupported SM targets.
In `@tests/comm/test_trtllm_allreduce_fusion.py`:
- Around line 470-481: The test currently only skips by GPU count; add a
compute-capability gate in test_trtllm_allreduce_fusion using flashinfer.utils
so unsupported SMs are skipped consistently: import and call
is_sm90a_supported(), is_sm100a_supported() or get_compute_capability() at the
start of the test_trtllm_allreduce_fusion function (before running GPU-dependent
code) and call pytest.skip(...) when the current device compute capability is
not supported; ensure you reference the flashinfer.utils helpers
(is_sm90a_supported, is_sm100a_supported, get_compute_capability) and keep the
existing world_size vs torch.cuda.device_count() check after the capability
check so tests on unsupported architectures are skipped early.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 73515f0e-c1c5-44e0-a6e1-14c4b2be1ac7
📥 Commits
Reviewing files that changed from the base of the PR and between ed0f5f8 and 81323a86b4b50fd4525b63ff6490d478022cf0f4.
📒 Files selected for processing (14)
benchmarks/bench_gemma_ar_fusion.pycsrc/trtllm_allreduce_fusion.cucsrc/trtllm_mnnvl_allreduce.cucsrc/trtllm_moe_allreduce_fusion.cuflashinfer/comm/allreduce.pyflashinfer/comm/trtllm_ar.pyflashinfer/comm/trtllm_mnnvl_ar.pyflashinfer/trace/templates/comm.pyinclude/flashinfer/comm/trtllm_allreduce_fusion.cuhinclude/flashinfer/comm/trtllm_mnnvl_allreduce.cuhinclude/flashinfer/comm/trtllm_moe_allreduce_fusion.cuhtests/comm/test_allreduce_fusion_moe_unified_api.pytests/comm/test_gemma_ar_fusion.pytests/comm/test_trtllm_allreduce_fusion.py
| def test_gemma_rmsnorm_ar_fusion(world_size): | ||
| """End-to-end Gemma RMSNorm correctness for fused AllReduce path.""" | ||
| available_gpus = torch.cuda.device_count() | ||
| if world_size > available_gpus: | ||
| pytest.skip(f"world_size {world_size} > available GPUs {available_gpus}") | ||
|
|
There was a problem hiding this comment.
Add compute-capability skip for unsupported GPUs.
This test should include flashinfer.utils architecture checks (in addition to GPU-count checks) to prevent unsupported-SM failures.
As per coding guidelines "tests/**/*.py: Skip tests on unsupported GPU architectures using flashinfer.utils check functions like is_sm90a_supported(), is_sm100a_supported(), and get_compute_capability()."
🤖 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/comm/test_gemma_ar_fusion.py` around lines 164 - 169, In
test_gemma_rmsnorm_ar_fusion, add GPU architecture checks using flashinfer.utils
to skip the test on unsupported SMs: import and call get_compute_capability()
and the helper predicates (is_sm90a_supported(), is_sm100a_supported()) and if
the current compute capability isn't supported, call pytest.skip(...) similar to
the existing world_size/GPU-count guard; ensure checks are performed before
heavy GPU usage so the test exits early when the device SM is unsupported.
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
81323a8 to
e7b9cdf
Compare
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
4 participants
📌 Description
Verified from framework side. See perf at vllm-project/vllm#42646
Add weight bias to RMS norm AR fusion to support gemma and qwen3.5 RMS
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
Tests
Documentation