feat(comm): add Blackwell MNNVL CuTe DSL all-reduce fusion backend - #4358
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 865937f44e7c4266eb4b8712d4394f9e445fbc08 and fa1d17c. 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe pull request adds an MNNVL CuTe DSL AllReduce fusion backend with LL, BT, and HT protocols, static routing presets, symmetric-memory support, CuTe primitives, standard and MoE dispatch, and distributed numerical-contract tests. ChangesMNNVL CuTe DSL backend
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant AllReduceFusion
participant MNNVLWorkspace
participant Protocol
participant SymmetricMemory
participant CuTeKernels
AllReduceFusion->>MNNVLWorkspace: dispatch standard or MoE fusion
MNNVLWorkspace->>Protocol: select route and invoke operation
Protocol->>SymmetricMemory: access mailbox state
Protocol->>CuTeKernels: launch publication, reduction, and RMSNorm kernels
CuTeKernels-->>AllReduceFusion: return normalized and optional residual outputs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (9)
flashinfer/comm/mnnvl_cutedsl/runtime.py (1)
23-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the
stream=-1DLPack contract.
stream=-1requests no producer-consumer synchronization. That choice is what makes the wrapper CUDA-graph safe. Add a short comment with this rationale, because the repository style asks for documented intentional departures.♻️ Proposed fix
def __dlpack__(self, stream=None): + # stream=-1 requests no sync from the producer, which keeps the export + # legal during CUDA graph capture. The caller stream is ignored on + # purpose; kernels launch on the same stream as the source tensor. return self.tensor.__dlpack__(stream=-1)🤖 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 `@flashinfer/comm/mnnvl_cutedsl/runtime.py` around lines 23 - 33, Add a concise comment in _GraphSafeDLPack.__dlpack__ documenting that stream=-1 intentionally disables producer-consumer synchronization and preserves CUDA-graph safety; leave the existing DLPack behavior unchanged.Source: Coding guidelines
flashinfer/comm/mnnvl_cutedsl/__init__.py (1)
42-52: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnsorted
__all__in both new package initializers. Ruff reports RUF022 at both sites. The shared root cause is that the exported names are grouped by meaning instead of the isort-style order the rule enforces.
flashinfer/comm/mnnvl_cutedsl/__init__.py#L42-L52: move"KernelTarget"after the four*_CONFIGentries.flashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.py#L33-L47: move"LLAllReduceTuning","LLCollectiveTuning", and"LLFinalizeTuning"before theLL_*constants.🤖 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 `@flashinfer/comm/mnnvl_cutedsl/__init__.py` around lines 42 - 52, Sort the exported names in both __all__ lists according to Ruff RUF022. In flashinfer/comm/mnnvl_cutedsl/__init__.py:42-52, place KernelTarget after the four *_CONFIG entries; in flashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.py:33-47, place LLAllReduceTuning, LLCollectiveTuning, and LLFinalizeTuning before the LL_* constants.Source: Linters/SAST tools
flashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py (2)
688-707: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert the tuning invariants that this kernel relies on.
_LamportResidualRMSNormDeviceKernelderives its whole index mapping from integer division with no validation. Several tuning combinations would compile and run, then produce silently wrong reductions:
self.rank_waves = tp // rank_lanesat Line 704 truncates. Iftp % rank_lanes != 0, the wave loop never visits the highest source ranks, so those contributions are dropped.- The lane reduction at Lines 850-851 only emits butterfly offsets from
(1, 2, 4). That coversrank_lanes <= 8. Forrank_lanes == 16the offset8is missing, so half of each group is never summed.self.groups_per_cta = threads // rank_lanesandself.fragments = hidden // VEC_BF16also truncate. Ahiddenthat is not a multiple ofVEC_BF16silently drops the tail.The shipped presets satisfy all of these. The failure mode for a new preset is a wrong numerical result rather than an error, so add explicit checks in
__init__.🛡️ Proposed fix
) -> None: + if rank_lanes not in (1, 2, 4, 8): + raise ValueError("rank_lanes must be a power of two in [1, 8]") + if tp % rank_lanes: + raise ValueError("tp must be divisible by rank_lanes") + if threads % rank_lanes or threads % WARP_SIZE: + raise ValueError("threads must be a multiple of rank_lanes and WARP_SIZE") + if hidden % VEC_BF16: + raise ValueError("hidden must be a multiple of VEC_BF16") self.hidden = hidden self.tp = tp🤖 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 `@flashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py` around lines 688 - 707, In _LamportResidualRMSNormDeviceKernel.__init__, add explicit validation that tp is divisible by rank_lanes, rank_lanes is no greater than 8, threads is divisible by rank_lanes, and hidden is divisible by VEC_BF16 before deriving rank_waves, groups_per_cta, and fragments. Raise a clear exception identifying the violated tuning invariant, while preserving the existing derived values for valid presets.
264-286: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument why the scalar path publishes
ACTIVE_STAGEbefore the mailbox store.
_ScalarFinalizePublishDeviceKernelstoresACTIVE_STAGEat Line 264, before thestmc_bf16x2at Line 283._QuadFinalizePublishDeviceKernelstores it after the mailbox store at Line 502._SharedOnlyPublishDeviceKernelexposes the same choice as therelease_before_storetuning knob.The order is safe, because the consumer spins on the Lamport sentinel rather than on the stage flag. The asymmetry between the two finalize kernels is still not obvious to a later reader. Add a comment that states the invariant, or expose the same
release_before_storeknob here.🤖 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 `@flashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py` around lines 264 - 286, Add a concise comment in _ScalarFinalizePublishDeviceKernel immediately before the ACTIVE_STAGE store explaining that publishing it before stmc_bf16x2 is intentional and safe because consumers synchronize by spinning on the Lamport sentinel, not the stage flag; clarify that this ordering differs from the quad finalize path.Source: Coding guidelines
flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py (1)
17-29: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the
L2_EVICT_FIRSTcache-hint constant.
0x12F0000000000000is an encodedcreatepolicydescriptor. Add a short comment that states the policy it encodes and where the encoding comes from. That keeps the constant reviewable.🤖 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 `@flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py` around lines 17 - 29, Document the L2_EVICT_FIRST constant with a concise comment stating that 0x12F0000000000000 encodes the createpolicy L2-evict-first cache policy and identifying the source of that encoding.tests/comm/test_mnnvl_cutedsl_config.py (1)
20-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGuard collection when the CuTe DSL is absent.
flashinfer.comm.mnnvl_cutedsllazy-loadspresets, but this test explicitly imports presets/kernel modules that importcutlass.cuteandcuda.bindings.driver. Add a pytest import guard for CuTe DSL availability so these CPU routing tests skip instead of failing at collection whennvidia-cutlass-dslis not installed.♻️ Proposed fix
import pytest import torch +pytest.importorskip("cutlass.cute", reason="requires nvidia-cutlass-dsl") + from flashinfer.comm.mnnvl_cutedsl import (🤖 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_mnnvl_cutedsl_config.py` around lines 20 - 42, Add a module-level pytest skip guard before the CuTe DSL-dependent imports in tests/comm/test_mnnvl_cutedsl_config.py, checking whether the required nvidia-cutlass-dsl dependency is available. Ensure the guard skips collection when unavailable while allowing the existing MNNVLCuteDSLConfig and kernel preset tests to run normally when it is installed.flashinfer/comm/mnnvl_cutedsl_ar.py (1)
374-390: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRun the PDL comparison after shape and capacity validation.
workspace._uses_pdl(pattern, m)callsroutes.select(m)with an unvalidatedm. If a caller passes a token count aboveworkspace.max_token_num, route selection can fail here and hide the precise capacity error raised later at Line 434 or Line 487. Move this warning block after the per-pattern validation, wheremis already bounded. The block also recomputesmthat the pattern branches derive again.♻️ Suggested reordering
- m = None - if pattern == AllReduceFusionPattern.kARResidualRMSNorm and input.ndim: - m = input.shape[0] - elif ( - pattern == AllReduceFusionPattern.kMoEFinalizeARResidualRMSNorm - and expanded_idx_to_permuted_idx is not None - and expanded_idx_to_permuted_idx.ndim - ): - m = expanded_idx_to_permuted_idx.shape[0] - if m is not None: - preset_pdl = workspace._uses_pdl(pattern, m) - if launch_with_pdl != preset_pdl: - logger.warning( - "launch_with_pdl does not match the selected MNNVL CuTe DSL " - "preset; using enable_pdl=%s", - preset_pdl, - ) if rms_eps != workspace.rms_eps:Then call a small helper that emits the warning inside each pattern branch, after the capacity check validates
m.🤖 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 `@flashinfer/comm/mnnvl_cutedsl_ar.py` around lines 374 - 390, Move the launch_with_pdl comparison out of the shared pre-validation block and emit it within each relevant pattern branch after shape and capacity validation completes. Reuse the branch-local validated m when calling workspace._uses_pdl(pattern, m), avoiding recomputation and ensuring oversized token counts raise the existing capacity error before route selection.flashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.py (1)
28-37: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSort
__all__to satisfy Ruff RUF022.Ruff reports
__all__as unsorted. Apply isort-style ordering so lint passes.♻️ Proposed ordering
__all__ = [ - "HT_FINALIZE_GB300_TP8_H8192_K10", "HT_ALL_REDUCE_GB300_TP16_H8192", "HT_ALL_REDUCE_GB300_TP8_H8192", "HT_FINALIZE_GB300_TP16_H8192_K10", + "HT_FINALIZE_GB300_TP8_H8192_K10", "HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049", "HT_FINALIZE_GB300_TP8_H8192_K10_M_LE_2048", "HTAllReduceTuning", "HTFinalizeTuning", ]🤖 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 `@flashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.py` around lines 28 - 37, Sort the entries in __all__ using isort-style ordering so the exported kernel symbols and tuning classes satisfy Ruff RUF022, without changing the exported names.Source: Linters/SAST tools
flashinfer/comm/mnnvl_cutedsl/kernel_ht/protocol.py (1)
212-224: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the aliased placeholder arguments with explicit empty tensors.
AllReduceRMSNormHTKernelpasseslocal_contributionas both the routed-output and the expert-weights argument, and it passesindex_argas the permuted-indices argument. The kernel is compiled withtop_k=0, sometadata_chunksis 0 and these arguments are never dereferenced. The aliasing is only safe because of that compile-time constant. A future change that readsexpert_weightsorpermuted_indiceswhentop_k == 0would silently read the contribution buffer or the counter buffer.Add a short comment that records the
top_k == 0invariant, or pass dedicated zero-size placeholder tensors.🤖 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 `@flashinfer/comm/mnnvl_cutedsl/kernel_ht/protocol.py` around lines 212 - 224, In AllReduceRMSNormHTKernel’s self._compiled invocation, stop reusing local_contribution and index_arg for the expert-weights and permuted-indices placeholders. Pass dedicated empty tensors for those arguments, or add a concise comment documenting that top_k == 0 makes metadata_chunks zero and these arguments must remain unused.
🤖 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/mnnvl_cutedsl/kernel_bt/__init__.py`:
- Around line 31-43: Reorder the entries in __all__ according to numeric-aware
natural sorting required by RUF022, placing each TP8 preset before the
corresponding TP16 preset while preserving the existing symbol names and
exports.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_bt/protocol.py`:
- Around line 169-172: Update _PathKwargs and _BTPath to carry add_residual and
include_shared_expert, then validate optional operands before kernel argument
selection. At flashinfer/comm/mnnvl_cutedsl/kernel_bt/protocol.py lines 169-172,
raise ValueError when add_residual is true without residual_source; at lines
216-219, raise ValueError when include_shared_expert is true without
shared_output, rather than substituting norm_output.
- Around line 490-507: Update _create_state after both sentinel fill_ operations
to synchronize the device, then execute the group barrier before returning
BTProtocolState. Ensure the synchronization and barrier occur after
prenorm.tensor is initialized and before any rank can begin using the returned
state.
- Around line 369-376: The _compile_finalize configuration path must validate
supported dimensions before kernel dispatch. Require elements_per_thread to be
one of 1, 2, 4, or VEC_BF16, and require hidden_size to be divisible by
VEC_BF16; reject invalid values before selecting
_ScalarFinalizeUnicastDeviceKernel, _VectorFinalizeUnicastDeviceKernel, or
_NarrowVectorFinalizeUnicastDeviceKernel.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py`:
- Around line 273-286: Update the constructor for the class containing the shown
scalar publish logic to validate that hidden is even, raising a clear error for
odd values before kernel execution. Keep the existing lane predicate and mailbox
addressing unchanged, and allow the shipped even hidden sizes to initialize
normally.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_ll/protocol.py`:
- Around line 152-179: The fallback in _launch_collective is unsafe when the
compiled collective was built with add_residual=True, because it substitutes
norm_output for a missing residual_source and may read uninitialized memory.
Track the add_residual setting on _LLPath (and include it in _path_kwargs if
needed), then validate in _launch_collective or the public __call__ paths for
FinalizeAllReduceRMSNormLLKernel and AllReduceRMSNormLLKernel that
residual_source is provided whenever add_residual is enabled. Keep the existing
fallback only for the add_residual=False case.
- Around line 452-456: Update the LLProtocolState initialization in the protocol
setup to avoid relying on unsupported torch versions: either pin
requirements.txt to a PyTorch minimum that supports torch.uint32 and
to_cute(state.stage_state, 4), or change stage_state to a supported dtype such
as torch.int32 and update dependent conversions consistently.
In `@flashinfer/comm/mnnvl_cutedsl/symmetric_buffer.py`:
- Around line 50-89: Add a supported minimum torch version to requirements.txt
and document that SymmetricBuffer.create/rendezvous depend on the corresponding
private torch.distributed._symmetric_memory APIs, including get_backend, empty,
rendezvous, multicast_ptr, and get_remote_tensor. Keep the dependency
documentation scoped to these Symmetric Memory paths.
In `@tests/comm/test_mnnvl_cutedsl_numerical_contract.py`:
- Around line 61-65: Move the SM100 capability check before the conditional
process-group initialization, replacing the raw torch capability comparison with
flashinfer.utils.is_sm100a_supported(device). Preserve the existing pytest.skip
behavior, and leave dist.init_process_group to run only after the device-support
check passes.
---
Nitpick comments:
In `@flashinfer/comm/mnnvl_cutedsl_ar.py`:
- Around line 374-390: Move the launch_with_pdl comparison out of the shared
pre-validation block and emit it within each relevant pattern branch after shape
and capacity validation completes. Reuse the branch-local validated m when
calling workspace._uses_pdl(pattern, m), avoiding recomputation and ensuring
oversized token counts raise the existing capacity error before route selection.
In `@flashinfer/comm/mnnvl_cutedsl/__init__.py`:
- Around line 42-52: Sort the exported names in both __all__ lists according to
Ruff RUF022. In flashinfer/comm/mnnvl_cutedsl/__init__.py:42-52, place
KernelTarget after the four *_CONFIG entries; in
flashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.py:33-47, place
LLAllReduceTuning, LLCollectiveTuning, and LLFinalizeTuning before the LL_*
constants.
In `@flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py`:
- Around line 17-29: Document the L2_EVICT_FIRST constant with a concise comment
stating that 0x12F0000000000000 encodes the createpolicy L2-evict-first cache
policy and identifying the source of that encoding.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.py`:
- Around line 28-37: Sort the entries in __all__ using isort-style ordering so
the exported kernel symbols and tuning classes satisfy Ruff RUF022, without
changing the exported names.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_ht/protocol.py`:
- Around line 212-224: In AllReduceRMSNormHTKernel’s self._compiled invocation,
stop reusing local_contribution and index_arg for the expert-weights and
permuted-indices placeholders. Pass dedicated empty tensors for those arguments,
or add a concise comment documenting that top_k == 0 makes metadata_chunks zero
and these arguments must remain unused.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py`:
- Around line 688-707: In _LamportResidualRMSNormDeviceKernel.__init__, add
explicit validation that tp is divisible by rank_lanes, rank_lanes is no greater
than 8, threads is divisible by rank_lanes, and hidden is divisible by VEC_BF16
before deriving rank_waves, groups_per_cta, and fragments. Raise a clear
exception identifying the violated tuning invariant, while preserving the
existing derived values for valid presets.
- Around line 264-286: Add a concise comment in
_ScalarFinalizePublishDeviceKernel immediately before the ACTIVE_STAGE store
explaining that publishing it before stmc_bf16x2 is intentional and safe because
consumers synchronize by spinning on the Lamport sentinel, not the stage flag;
clarify that this ordering differs from the quad finalize path.
In `@flashinfer/comm/mnnvl_cutedsl/runtime.py`:
- Around line 23-33: Add a concise comment in _GraphSafeDLPack.__dlpack__
documenting that stream=-1 intentionally disables producer-consumer
synchronization and preserves CUDA-graph safety; leave the existing DLPack
behavior unchanged.
In `@tests/comm/test_mnnvl_cutedsl_config.py`:
- Around line 20-42: Add a module-level pytest skip guard before the CuTe
DSL-dependent imports in tests/comm/test_mnnvl_cutedsl_config.py, checking
whether the required nvidia-cutlass-dsl dependency is available. Ensure the
guard skips collection when unavailable while allowing the existing
MNNVLCuteDSLConfig and kernel preset tests to run normally when it is installed.
🪄 Autofix
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 Plus
Run ID: f9142e1d-9e18-487a-b486-0fcd08665cfd
📥 Commits
Reviewing files that changed from the base of the PR and between 6c85301 and e57acd7da3be950e2dfe3dc4a47d357e7d981920.
📒 Files selected for processing (20)
flashinfer/comm/allreduce.pyflashinfer/comm/mnnvl_cutedsl/__init__.pyflashinfer/comm/mnnvl_cutedsl/config.pyflashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.pyflashinfer/comm/mnnvl_cutedsl/kernel_bt/__init__.pyflashinfer/comm/mnnvl_cutedsl/kernel_bt/device_kernels.pyflashinfer/comm/mnnvl_cutedsl/kernel_bt/protocol.pyflashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.pyflashinfer/comm/mnnvl_cutedsl/kernel_ht/device_kernel.pyflashinfer/comm/mnnvl_cutedsl/kernel_ht/protocol.pyflashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.pyflashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.pyflashinfer/comm/mnnvl_cutedsl/kernel_ll/protocol.pyflashinfer/comm/mnnvl_cutedsl/presets.pyflashinfer/comm/mnnvl_cutedsl/runtime.pyflashinfer/comm/mnnvl_cutedsl/symmetric_buffer.pyflashinfer/comm/mnnvl_cutedsl_ar.pyflashinfer/trace/templates/comm.pytests/comm/test_mnnvl_cutedsl_config.pytests/comm/test_mnnvl_cutedsl_numerical_contract.py
e57acd7 to
2d41951
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (4)
flashinfer/comm/mnnvl_cutedsl/__init__.py (1)
42-52: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSort the new
__all__declarations.Ruff reports RUF022 for these declarations. Use the configured natural ordering.
flashinfer/comm/mnnvl_cutedsl/__init__.py#L42-L52: Sort all exported names.flashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.py#L28-L37: Sort HT preset exports and tuning types.flashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.py#L33-L47: Sort LL preset exports and tuning types.Confidence: high. As per coding guidelines, match established project style.
🤖 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 `@flashinfer/comm/mnnvl_cutedsl/__init__.py` around lines 42 - 52, Sort every exported name in the __all__ declarations using Ruff’s configured natural ordering: reorder all entries in flashinfer/comm/mnnvl_cutedsl/__init__.py (lines 42-52), the HT preset exports and tuning types in flashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.py (lines 28-37), and the LL preset exports and tuning types in flashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.py (lines 33-47); preserve the exported names and only change their ordering.Sources: Coding guidelines, Linters/SAST tools
flashinfer/comm/mnnvl_cutedsl/kernel_bt/device_kernels.py (2)
670-693: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument why the vector variant clamps
routed_indexto 0.
_ScalarFinalizeUnicastDeviceKernel(Line 174) and_NarrowVectorFinalizeUnicastDeviceKernel(Line 383) keeprouted_index == -1and suppress the load withload_*_predicated. This kernel instead rewritesrouted_indexto0at Line 683 and then loads row 0 unconditionally at Line 765, becauseload_global_u32x4has no predicated form. Both designs are safe, but the difference is not obvious from the code.Add a short comment stating that the clamp keeps the unpredicated 16-byte load in bounds and that
weight == 0.0discards the value.The three metadata staging loops (Lines 162-185, 371-393, 670-693) and the three destination-offset computations (Lines 268-283, 561-584, 825-840) are otherwise identical. A shared
@cute.jithelper would keep the masking rule in one place.As per coding guidelines: "Match established project style, efficiency, complexity, verbosity, and defensiveness; document intentional departures with rationale."
🤖 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 `@flashinfer/comm/mnnvl_cutedsl/kernel_bt/device_kernels.py` around lines 670 - 693, Add a concise comment in the metadata staging loop of _VectorFinalizeUnicastDeviceKernel immediately before the routed_index == -1 clamp, documenting that clamping to zero keeps the later unpredicated 16-byte load in bounds and that the zero weight discards the loaded value. Do not alter the existing masking behavior or refactor the related loops.Source: Coding guidelines
1054-1071: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider tracking per-rank arrival to reduce spin traffic.
The retry loop re-reads all
tp_sizemailbox fragments on every pass. One late peer therefore forces repeated 16-byte loads for every already-arrived rank. Attp_size=16each retry costs 256 bytes of NVLink read traffic per thread.A per-rank arrival mask would keep the already-arrived words in
rank_valuesand re-read only the outstanding ranks.♻️ Sketch of a per-rank arrival mask
rank_values.fill(Uint32(0)) - dirty = active - while dirty: - dirty = False - for source_rank in cutlass.range_constexpr(self.tp_size): - source_offset = ( - (Int64(stage) * self.tp_size + source_rank) * self.local_capacity - + Int64(local_token) - ) * self.hidden_size + Int64(fragment) * VEC_BF16 - source_pointer = cute.make_ptr( - BFloat16, - (contribution_mailbox.iterator + source_offset).llvm_ptr, - cute.AddressSpace.gmem, - assumed_align=16, - ) - packed = load_global_u32x4(source_pointer, volatile=True) - dirty = dirty | fragment_has_negative_zero(packed) - for word in cutlass.range_constexpr(4): - rank_values[source_rank, word] = packed[word] + pending = Uint32(0) + if active: + pending = Uint32((1 << self.tp_size) - 1) + while pending != Uint32(0): + for source_rank in cutlass.range_constexpr(self.tp_size): + if (pending & Uint32(1 << source_rank)) != Uint32(0): + source_offset = ( + (Int64(stage) * self.tp_size + source_rank) + * self.local_capacity + + Int64(local_token) + ) * self.hidden_size + Int64(fragment) * VEC_BF16 + source_pointer = cute.make_ptr( + BFloat16, + (contribution_mailbox.iterator + source_offset).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source_pointer, volatile=True) + if not fragment_has_negative_zero(packed): + for word in cutlass.range_constexpr(4): + rank_values[source_rank, word] = packed[word] + pending = pending & ~Uint32(1 << source_rank)
self.tp_sizemust stay at or below 32 for aUint32mask.🤖 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 `@flashinfer/comm/mnnvl_cutedsl/kernel_bt/device_kernels.py` around lines 1054 - 1071, Update the retry loop around rank_values and fragment_has_negative_zero to maintain a per-rank arrival mask, with a Uint32 mask and an explicit self.tp_size <= 32 constraint. Skip mailbox loads for ranks already marked arrived, retain their values in rank_values, and continue polling only outstanding ranks until all have arrived.flashinfer/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py (1)
570-575: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the
rowsshared-buffer reuse contract.The RMS path reinterprets
rowsasrms_stage_slotsslices ofhiddenelements, while the finalize pipeline used the same allocation asstagesslices ofshard_elements. The reuse is safe, but the reasoning is not local. It rests on three separate facts: the capacity guard at lines 148-152, theproducer_tailcall at line 373, and the fact that consumers release every issued stage before they passfinalize_joinat line 484.Add a short comment at the reuse site that states the invariant and points to the capacity guard. This prevents a future change to
stagesorrms_pipeline_stagesfrom silently overflowing the staging region.📝 Proposed comment
+ # `rows` is re-used here as `rms_stage_slots` + # slices of `hidden`. Capacity is guaranteed by + # the check in `__init__` (rms_stage_slots * + # hidden <= shard_elements * stages), and the + # finalize pipeline has fully drained via + # `producer_tail` plus `finalize_join`. store_shared_u32x4( rows.iterator + stage_slot * self.hidden + pack * VEC_BF16, packed_prenorm, )🤖 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 `@flashinfer/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py` around lines 570 - 575, At the rows reuse site in the RMS path, add a short comment documenting that rows is reinterpreted from finalize stages of shard_elements into rms_stage_slots of hidden elements, and that this is safe because the earlier capacity guard covers both layouts. Reference the existing capacity guard symbol or condition without changing the storage or synchronization logic.Source: Coding guidelines
🤖 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/mnnvl_cutedsl/kernel_bt/device_kernels.py`:
- Around line 620-623: Update the StaticProfile construction or validation path
to reject hidden_size values that are not divisible by VEC_BF16 before creating
kernels. Ensure invalid values fail validation rather than reaching
_VectorFinalizeUnicastDeviceKernel, whose fragments calculation truncates the
unaligned tail.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py`:
- Around line 245-251: Clamp the CTA count produced by _resolve_ctas() to the
maximum resident CTAs supported by the compiled kernel’s register/shared-memory
occupancy before assigning or using self.active_ctas. Apply this limit to both
user-provided persistent_ctas and the default SM-scaled value, preserving the
existing launch grid while ensuring grid=(self.active_ctas, ...) never exceeds
resident capacity.
- Around line 609-621: The RMS group-to-barrier dispatch ladder is duplicated
across four RMS phase sites. In
flashinfer/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py at lines 609-621,
653-665, 799-811, and 833-845, add one helper after the RMS barriers are
constructed to select and wait on the correct barrier, then replace each
duplicated ladder with calls to that helper while preserving the existing
rms_token_groups behavior.
---
Nitpick comments:
In `@flashinfer/comm/mnnvl_cutedsl/__init__.py`:
- Around line 42-52: Sort every exported name in the __all__ declarations using
Ruff’s configured natural ordering: reorder all entries in
flashinfer/comm/mnnvl_cutedsl/__init__.py (lines 42-52), the HT preset exports
and tuning types in flashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.py (lines
28-37), and the LL preset exports and tuning types in
flashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.py (lines 33-47); preserve the
exported names and only change their ordering.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_bt/device_kernels.py`:
- Around line 670-693: Add a concise comment in the metadata staging loop of
_VectorFinalizeUnicastDeviceKernel immediately before the routed_index == -1
clamp, documenting that clamping to zero keeps the later unpredicated 16-byte
load in bounds and that the zero weight discards the loaded value. Do not alter
the existing masking behavior or refactor the related loops.
- Around line 1054-1071: Update the retry loop around rank_values and
fragment_has_negative_zero to maintain a per-rank arrival mask, with a Uint32
mask and an explicit self.tp_size <= 32 constraint. Skip mailbox loads for ranks
already marked arrived, retain their values in rank_values, and continue polling
only outstanding ranks until all have arrived.
In `@flashinfer/comm/mnnvl_cutedsl/kernel_ht/device_kernel.py`:
- Around line 570-575: At the rows reuse site in the RMS path, add a short
comment documenting that rows is reinterpreted from finalize stages of
shard_elements into rms_stage_slots of hidden elements, and that this is safe
because the earlier capacity guard covers both layouts. Reference the existing
capacity guard symbol or condition without changing the storage or
synchronization logic.
🪄 Autofix
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 Plus
Run ID: ad138622-41d3-43e4-aaaf-c654a999a51a
📥 Commits
Reviewing files that changed from the base of the PR and between 3c57ef1 and 2d41951c30d87a497818a4c525f9db0d7f6cf58d.
📒 Files selected for processing (20)
flashinfer/comm/allreduce.pyflashinfer/comm/mnnvl_cutedsl/__init__.pyflashinfer/comm/mnnvl_cutedsl/config.pyflashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.pyflashinfer/comm/mnnvl_cutedsl/kernel_bt/__init__.pyflashinfer/comm/mnnvl_cutedsl/kernel_bt/device_kernels.pyflashinfer/comm/mnnvl_cutedsl/kernel_bt/protocol.pyflashinfer/comm/mnnvl_cutedsl/kernel_ht/__init__.pyflashinfer/comm/mnnvl_cutedsl/kernel_ht/device_kernel.pyflashinfer/comm/mnnvl_cutedsl/kernel_ht/protocol.pyflashinfer/comm/mnnvl_cutedsl/kernel_ll/__init__.pyflashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.pyflashinfer/comm/mnnvl_cutedsl/kernel_ll/protocol.pyflashinfer/comm/mnnvl_cutedsl/presets.pyflashinfer/comm/mnnvl_cutedsl/runtime.pyflashinfer/comm/mnnvl_cutedsl/symmetric_buffer.pyflashinfer/comm/mnnvl_cutedsl_ar.pyflashinfer/trace/templates/comm.pytests/comm/test_mnnvl_cutedsl_config.pytests/comm/test_mnnvl_cutedsl_numerical_contract.py
🚧 Files skipped from review as they are similar to previous changes (11)
- flashinfer/trace/templates/comm.py
- flashinfer/comm/allreduce.py
- flashinfer/comm/mnnvl_cutedsl/presets.py
- flashinfer/comm/mnnvl_cutedsl/runtime.py
- flashinfer/comm/mnnvl_cutedsl/symmetric_buffer.py
- flashinfer/comm/mnnvl_cutedsl/kernel_ll/device_kernels.py
- flashinfer/comm/mnnvl_cutedsl/kernel_bt/protocol.py
- tests/comm/test_mnnvl_cutedsl_numerical_contract.py
- flashinfer/comm/mnnvl_cutedsl/config.py
- flashinfer/comm/mnnvl_cutedsl/kernel_ht/protocol.py
- flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py
|
/bot run tests/comm |
|
[FAILED] Pipeline #61459490 — 16/18 executed test jobs passed Compared with nightly #61182354. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
No individual test or infrastructure failures could be extracted. |
|
/bot run tests/comm |
|
[FAILED] Pipeline #61519875 — 15/18 executed test jobs passed Compared with nightly #61367193. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
Failure detailsNew relative to nightly (attribution uncertain)
|
Add a BF16 Blackwell backend with low-latency, balanced, and high-throughput protocols for fused MoE finalize and all-reduce tails. Integrate workspace and configuration routing with the unified all-reduce API, and add routing and numerical-contract coverage.
Head branch was pushed to by a user without write access
23922f9 to
fb5ec27
Compare
|
/bot run tests/comm |
5 participants
📌 Description
This PR adds a BF16 Blackwell backend for the latency-critical tail of tensor-parallel MoE layers. It supports two patterns through the unified
allreduce_fusionAPI:The initial profiles target GB300,
H=8192,Top-K=10, and TP8/TP16, covering decode through prefill. The implementation uses CuTe DSL, PyTorch Symmetric Memory, and MNNVL/NVLink multicast.Design goal
The goal is one backend that remains efficient from small-
Mdecode to large-Mprefill. For small payloads, LL and BT aggressively use Programmatic Dependent Launch (PDL): stages are enqueued together but keep independent grid, CTA, cluster, and warp configurations. At largeM, HT switches to persistent warp-specialized execution.M/ decodeMM/ prefillLL protocol
Kernel 1 performs local finalize and publishes BF16 contributions with STMC into a triple-buffered symmetric mailbox. Kernel 2 performs the ordered Lamport reduction, residual add, optional prenorm output, and RMSNorm.
BT protocol
Kernel 1 finalizes and unicasts shards to their owners. Kernel 2 reduces owner-local contributions, adds the residual, and broadcasts prenorm through STMC. Kernel 3 materializes the output and runs RMSNorm with its own launch configuration.
HT protocol
HT uses persistent CTAs with loader, finalize/RMS, publisher, and reduction warp roles. Contributions are reduced with LDMC over
H / TPshards and prenorm values are published with STMC, overlapping communication and computation at largeM.Dispatch and workspace integration
MNNVLCuteDSLAllReduceFusionWorkspaceowns the compiled variants, symmetric state, and static routing profile. It compiles only the protocols required by the requested capacity and completes symmetric-memory rendezvous before use.The default GB300 profile uses the following protocol boundaries:
M <= 2324–48(P0),49–703(P1)M >= 704M <= 1516–256(P0),257–1024(P1)M >= 1025M <= 78–52(P0),53–703(P1)M >= 704M <= 56–512(P0),513–959(P1)M >= 960Profiles are keyed by TP size, hidden size, Top-K, and dtype, and map contiguous
Mranges to a protocol and tuning preset. Runtime arguments are validated against the compiled workspace contract.Performance
The figure reports maximum-rank latency on GB300 for BF16,
H=8192,Top-K=10, TP8/TP16, andM=1..8192. The upper row includes MoE finalize; the lower row starts from a materialized rank-local contribution. Baselines use either FlashInfer MNNVL AR/residual/RMSNorm or NCCL AllReduce + fused add/RMSNorm. Every plotted point passed its reference check.At the common power-of-two points covered by the default routed protocol, the current measurements provide the following conservative speedup ranges:
1.74–3.20xfaster2.32–8.29xfaster1.74–3.19xfaster2.48–11.36xfaster1.03–1.68xfaster1.62–7.93xfaster1.05–1.80xfaster1.81–11.30xfasterLL leads in decode, BT covers the middle range, and persistent HT scales best into prefill.
Correctness strategy and numerical contract
Correctness uses two layers. The common end-to-end sweep checks the new backend and existing FlashInfer/NCCL paths against an independent unfused reference. Every one of the 70 graph invocations at each reported point is validated on every rank. Dedicated distributed tests then force LL, BT, and HT individually and apply a reference matching each protocol's reduction order.
multimem_one_shot_all_reduce_out, matching HT's multicast reduction semantics.The dedicated tests also:
2^24,+1,-2^24inputs to expose incorrect association order.residual_outandnorm_outseparately, using raw BF16 bits or explicit ULP distance.weight_bias, and every default routing boundary.🚀 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 my preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Please focus on the LL/BT/HT split and crossover points, PDL stage boundaries, symmetric-memory lifetime, and the protocol-specific bitwise/ULP contracts. The initial built-in profile is scoped to Blackwell BF16, GB300,
H=8192,Top-K=10, and TP8/TP16.The workspace must be fully constructed before its first invocation, and calls sharing one workspace must not overlap.
Summary by CodeRabbit
New Features
Bug Fixes