feat: MNNVL Allreduce quant fusion and performance optimization - #3385
Merged
aleozlx merged 15 commits intoMay 26, 2026
Merged
Conversation
Brings non-quant latency-communication optimizations from the latcomm reference impl into FlashInfer's trtllm_mnnvl_allreduce kernels: - add typed kernel params and raw-bit Lamport dirty polling - keep OOB threads in CTA arrival and dirty-buffer clearing - skip local-rank polling and sanitize polled Lamport payloads - fuse twoshot allreduce with residual add and RMSNorm - stage fused RMSNorm inputs with SM90+ TMA/mbarrier helpers
Keep the local-rank Lamport poll skip, but store the local shard and remote payloads in rank-indexed slots so oneshot and twoshot reductions accumulate ranks in the same 0..world_size-1 order everywhere.
timlee0212
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
May 21, 2026 19:39
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds FP8/NVFP4 quantized-output support to MNNVL AllReduce fusion: new quant enums/params, unified kernel-params struct, Lamport synchronization hardening, FP8/FP4 quant helpers and emission, expanded dispatch/validation, Python/C++ API extensions, JIT flag, docs entry, and distributed tests. ChangesMNNVL AllReduce FP8/FP4 Quantization Support
Sequence Diagram(s)sequenceDiagram
participant User
participant PythonAPI as flashinfer.comm.allreduce
participant PyWrapper as trtllm_mnnvl_allreduce_fusion (py)
participant CppEntry as trtllm_mnnvl_allreduce_fusion (c++)
participant Kernel
User->>PythonAPI: call allreduce_fusion(..., pattern, quant_type, output_scale, layout_code)
PythonAPI->>PyWrapper: prepare args, allocate quant_out/scale_out (if needed)
PyWrapper->>CppEntry: forward extended args (quant_type, quant_out, sf_out, output_scale, layout_code)
CppEntry->>Kernel: launch with AllReduceKernelParams (quant/scaling/layout)
Kernel-->>CppEntry: writes quant_out, scaling_out, residual_out, output
CppEntry-->>PyWrapper: returns kernel outputs
PyWrapper-->>PythonAPI: returns (quant_out, scale_out, residual_out, output)
PythonAPI-->>User: deliver tensors
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related issues
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 |
5 tasks
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request implements FP8 and NVFP4 quantization fusion for the MNNVL all-reduce backend, extending the unified allreduce_fusion API and adding a specialized Python wrapper. The CUDA kernels were refactored to handle quantization outputs and scaling factors, alongside JIT compilation updates for SM90/SM100 targets. Review feedback highlights potential race conditions due to premature Programmatic Dependent Launch (PDL) triggers, suggests enforcing scalar constraints on quantization scales in Python, and recommends utilizing the sanitizeLamportPayload utility for better code consistency.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh (1)
845-895:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winReject partial tails in the oneshot path.
oneshotAllreduceFusionKerneldoes full-widthPackedTypeloads/stores for input, residual, gamma, and output, but this dispatch path never checks thattokenDimis divisible bysizeof(float4) / sizeof(T). On a tail row, the last live thread will read and write past the token boundary.twoshotAllreduceFusionDispatchalready has this guard; oneshot needs the same protection.🛡️ Proposed guard
template <typename T> cudaError_t oneshotAllreduceFusionDispatch(AllReduceFusionParams const& params) { int const numTokens = params.numTokens; int const tokenDim = params.tokenDim; int const eltsPerThread = sizeof(float4) / sizeof(T); + + FLASHINFER_CHECK( + tokenDim % eltsPerThread == 0, + "[MNNVL AllReduceOneShot] token_dim must be divisible by %d", + eltsPerThread); auto [blockSize, clusterSize, loadsPerThread] = adjustGridConfig(numTokens, tokenDim, eltsPerThread, true);🤖 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 `@include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh` around lines 845 - 895, oneshotAllreduceFusionDispatch currently launches oneshotAllreduceFusionKernel which performs full-width PackedType (float4) loads/stores, but it never rejects tails so a non-divisible tokenDim will read/write out of bounds; add the same divisibility guard used in twoshotAllreduceFusionDispatch: compute eltsPerThread (sizeof(float4)/sizeof(T)) and check tokenDim % eltsPerThread == 0 (or equivalent ceil_div check) before launching the kernel and return cudaErrorInvalidValue with a clear FLASHINFER_ERROR/FLASHINFER_CHECK message referencing oneshotAllreduceFusionKernel, PackedType/float4 and trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE if relevant so that non-multiple tokenDim values are rejected rather than causing out-of-bounds accesses.
🤖 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 `@csrc/trtllm_mnnvl_allreduce.cu`:
- Around line 49-51: The code currently assigns sf_layout from layout_code
allowing any integer except SWIZZLED_8x4 to propagate; instead whitelist
supported QuantizationSFLayout values by checking layout_code.value() and only
accepting QuantizationSFLayout::LINEAR or QuantizationSFLayout::SWIZZLED_128x4
(for both places where sf_layout is computed, e.g., the first occurrence
producing sf_layout and the later occurrence at 79-81), and handle any other
value by setting a safe default or returning/logging an error before assigning
to params.sfLayout so invalid layout codes cannot slip through.
In `@flashinfer/comm/allreduce.py`:
- Around line 696-698: In the TRTLLM MoE reduction branch in allreduce.py the
scale_factor argument is being replaced with 1.0 for any non-float/int, which
drops valid 0-D or 1-element torch.Tensor scalars; update the logic that sets
scale_factor so that if scale_factor is an int/float you keep it, if it's a
torch.Tensor with numel() == 1 you preserve the tensor (do not coerce to 1.0),
otherwise fall back to 1.0; reference the scale_factor variable used in the
TRTLLM MoE reduction path and ensure torch is used to detect tensor types and
numel when implementing the fix.
In `@flashinfer/comm/trtllm_mnnvl_ar.py`:
- Around line 657-780: The helper is missing device checks for caller-provided
tensors and may pass CPU/other-GPU pointers to
module.trtllm_mnnvl_allreduce_fusion; before the native call, validate that
residual_in, gamma, output, residual_out, quant_out, and scale_out (if not None)
are on input.device and either move them with .to(device=input.device,
dtype=their_expected_dtype) or raise a clear ValueError; implement these
checks/implicit moves just before the module.trtllm_mnnvl_allreduce_fusion
invocation so the native kernel always receives tensors on input.device.
In `@flashinfer/jit/comm.py`:
- Around line 40-42: The call to current_compilation_context.get_nvcc_flags_list
in comm.py limits supported_major_versions to [9, 10]; update this invocation to
include 11 and 12 so it becomes [9, 10, 11, 12] (i.e., change the
supported_major_versions argument passed to get_nvcc_flags_list where nvcc_flags
is assigned) to align with the JIT architecture guideline.
In `@include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh`:
- Around line 782-788: The kernel always writes the pre-norm tensor via
params.prenormedPtr (e.g., inside the RMSNormFusion path where
PackedVec<PackedType,T> residualIn, packedAccum and the write
"*reinterpret_cast<PackedType*>(¶ms.prenormedPtr[threadOffset]) =
packedAccum.packed" occur), which can dereference nullptr when callers omit
residual_out; guard these writes by checking prenormedPtr for non-null before
performing the reinterpret_cast write (same pattern used for optional
outputPtr), and apply the same null-check fix to the other occurrence around
lines 1188-1195 so both fused paths only write prenormedPtr when
params.prenormedPtr != nullptr.
In `@tests/comm/test_trtllm_mnnvl_allreduce.py`:
- Around line 31-40: The functions fp8_quant and dequant are shadowing the
built-in name `input`; rename that parameter (e.g., to `x` or `tensor`) in both
fp8_quant and dequant signatures and all internal references, update associated
type hints and any callers within the test file to use the new parameter name,
and ensure the behavior remains identical (fp8_quant still computes qinput using
the new param and dequant still multiplies by scale and casts to dtype).
- Around line 826-827: Replace the local per-rank check in
test_mnnvl_allreduce_quant_unified that uses torch.cuda.get_device_capability()
with a centralized helper that uses flashinfer.utils.get_compute_capability(...)
to determine major>=10 and then performs a cross-rank agreement (e.g., gather a
boolean support flag via torch.distributed.all_gather_object or similar) so all
ranks decide the same thing; call that helper at the start of the test and skip
the entire test on all ranks if NVFP4/FP4 is unsupported, ensuring the decision
happens before entering any dist.barrier() or loop to avoid hangs (update the
test to reference the helper instead of the inline torch.cuda check and ensure
the skip is executed consistently across ranks).
---
Outside diff comments:
In `@include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh`:
- Around line 845-895: oneshotAllreduceFusionDispatch currently launches
oneshotAllreduceFusionKernel which performs full-width PackedType (float4)
loads/stores, but it never rejects tails so a non-divisible tokenDim will
read/write out of bounds; add the same divisibility guard used in
twoshotAllreduceFusionDispatch: compute eltsPerThread (sizeof(float4)/sizeof(T))
and check tokenDim % eltsPerThread == 0 (or equivalent ceil_div check) before
launching the kernel and return cudaErrorInvalidValue with a clear
FLASHINFER_ERROR/FLASHINFER_CHECK message referencing
oneshotAllreduceFusionKernel, PackedType/float4 and
trtllm_allreduce_fusion::details::CVT_FP4_SF_VEC_SIZE if relevant so that
non-multiple tokenDim values are rejected rather than causing out-of-bounds
accesses.
🪄 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: 2ba39c8e-d9cc-441f-b511-7e453710de3f
📒 Files selected for processing (7)
csrc/trtllm_mnnvl_allreduce.cudocs/api/comm.rstflashinfer/comm/allreduce.pyflashinfer/comm/trtllm_mnnvl_ar.pyflashinfer/jit/comm.pyinclude/flashinfer/comm/trtllm_mnnvl_allreduce.cuhtests/comm/test_trtllm_mnnvl_allreduce.py
aleozlx
approved these changes
May 21, 2026
Member
|
/bot run |
Collaborator
Collaborator
|
[FAILED] Pipeline #52146554: 10/20 passed |
Merged
1 task
Member
|
re-running failed T4 job |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
Performance Change Dashboard: report.html
For M <= 8, fused oneshot is 6.24% faster latency-weighted and ar_only oneshot is 4.26% faster. The main benefit is from avoiding loading the local buffer in lamport polling.
Fused two-shot gets benefit for relatively larger batch size. The main benefit is from reducing CTA sync overhead by using named barrier instead of syncthread.
🔍 Related Issues
None
🚀 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
Behavior / Validation
Documentation
Tests