feat(moe): MXFP8 x MXFP4 CuTe-DSL fused MoE for SM100, plus large-batch routing locality - #4440
Merged
jiahanc merged 7 commits intoAug 13, 2026
Conversation
Adds a mixed-precision path to the Blackwell CuTe-DSL MoE kernels: MXFP8 (E4M3 + linear block-32 E8M0) activations against packed MXFP4 (E2M1 + MMA-layout block-32 E8M0) weights, an MXFP8 FC1 intermediate, and a BF16 finalized output. - Both grouped-GEMM kernels now take separate a_dtype/b_dtype, size SMEM from smem_alloc_*_dtype, and load the narrow operand via TMA UNPACK_U8. - The FC1 epilogue emits the same exact UE8M0 codes as mxfp8_quantize and scalar-stores them as bytes, avoiding the unsupported vector<2xE8M0> conversion/store lowering. - New public APIs cute_dsl_fused_moe_mxfp8_mxfp4 and CuteDslMxfp8Mxfp4MoEWrapper, with their own tactic space and autotune cache namespace so NVFP4 results and cache entries are unaffected. - GEMM2 can_implement no longer requires N to divide the MMA N tile: the finalize epilogue already clamps its bulk-reduce to valid_columns, and a partial-tile case is covered by the new grouped-GEMM test. AI-assisted with Claude Code. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
…tches The cooperative routing kernel assigns each CTA a grid stride of expanded indices, and an expert's rows land in CTA-arrival order, so the rows of one expert are drawn from the whole batch. A grouped-GEMM tile then gathers across the entire activation tensor, and the footprint grows with the batch: measured on B200 (hidden 6144, 256 experts, top_k 8, 32 local, tile 256) a GEMM1 tile touches 94 distinct 2 MiB pages at 64K tokens, 190 at 128K and 228 at 256K, which outgrows the uTLB. Add an opt-in ordering where each CTA owns one contiguous span instead, so a tile gathers from a few narrow token windows and the footprint stays flat at ~28-38 pages. moe_sort enables it from 65536 tokens, where the gather is large enough for the locality to pay for itself: M GEMM1 full MoE pipeline 65536 1.00-1.10x 0.99-1.03x 131072 1.40-1.44x 1.17-1.18x 262144 1.43-1.47x 1.24-1.28x The span is sized to keep every CTA busy rather than fixed at the per-thread maximum, and each CTA stops at its own span end -- the batch-size check alone would let the trailing iterations run into the next CTA's span and permute those routes twice. The flag defaults to off, so every other caller of the routing kernels, including the TRT-LLM MoE backend, keeps its current ordering bit for bit. AI-assisted with Claude Code. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
vitamin-chaos
requested review from
Aneureka,
IwakuraRein,
StudyingShao,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
August 10, 2026 12:44
Contributor
📝 WalkthroughWalkthroughThis PR adds CuTeDSL MXFP8-activation/MXFP4-weight fused MoE support for Blackwell GPUs. It adds mixed-width grouped GEMM kernels, public APIs, autotuning, traces, tests, and large-batch contiguous routing windows. ChangesMXFP8/MXFP4 fused MoE
Contiguous route windows
Estimated code review effort: 5 (Critical) | ~120 minutes Mergeability Score: 🔵 Low · up to The new MXFP8×MXFP4 wrapper may use the ambient CUDA device instead of its configured device when launching work or recording events, which can fail or run on the wrong GPU in multi-device use. This is a localized, mergeable risk requiring explicit owner follow-up. Possibly related PRs
Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Caller
participant CuteDslMxfp8Mxfp4MoEWrapper
participant RoutingKernel
participant GroupedGEMM
Caller->>CuteDslMxfp8Mxfp4MoEWrapper: submit MXFP8 activations and MXFP4 weights
CuteDslMxfp8Mxfp4MoEWrapper->>RoutingKernel: sort routed tokens
RoutingKernel->>GroupedGEMM: provide routed tensors and scales
GroupedGEMM->>CuteDslMxfp8Mxfp4MoEWrapper: return BF16 output
CuteDslMxfp8Mxfp4MoEWrapper->>Caller: return MoE result
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (7)
tests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py (2)
81-110: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueShare the quantization helpers with the grouped-GEMM test.
_interleave_linear_and_gateand_quantize_mxfp4_groupedare also defined intests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.pyat Lines 44-54 and 115-150. The two copies differ: the grouped-GEMM version assertsrows % (2 * group_size) == 0and also returns a dequantized reference. Move one implementation intotests/moe/utils.pyso the interleaving contract stays defined in one place.🤖 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/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py` around lines 81 - 110, Move _interleave_linear_and_gate and _quantize_mxfp4_grouped into tests/moe/utils.py, using the grouped-GEMM implementations as the canonical versions, including the rows divisibility assertion and dequantized reference return. Remove the duplicate local definitions from both test modules and import the shared helpers instead.
37-47: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the exact NVFP4 tactic count.
Line 39 asserts
len(ALL_MOE_TACTICS) == 16. That couples this mixed-path test to the NVFP4 tactic space. An unrelated change to the NVFP4 tactics breaks this test without indicating a mixed-path regression. Lines 40 and 47 already prove namespace separation.🤖 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/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py` around lines 37 - 47, Remove the exact count assertion for ALL_MOE_TACTICS from the test, while preserving the mixed-path tactic count, uniqueness, runner separation, and tile-size assertions.flashinfer/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py (2)
785-786: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClarify that this method returns the unfiltered tactic space.
CuteDslFusedMoEMxfp8Mxfp4Runner.get_valid_tacticsfilters tactics againstcan_implementfor both kernels. This wrapper method returns every entry inALL_MXFP8_MXFP4_MOE_TACTICS, so some returned tactics fail at launch for the wrapper's configuredhidden_sizeandintermediate_size. Add a docstring that states the method lists the tactic space rather than the shape-valid subset, or rename it toget_all_tactics.🤖 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/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py` around lines 785 - 786, Clarify the contract of CuteDslFusedMoEMxfp8Mxfp4Runner.get_valid_tactics by adding a docstring stating that it returns the complete ALL_MXFP8_MXFP4_MOE_TACTICS space without filtering for can_implement or the configured shapes; keep the existing return behavior unchanged.
591-643: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider bucketing capacities when
max_num_tokensis unset.Without
max_num_tokens,_capacity_forreturns the exactnum_tokens. Each distinct token count then allocates a new output tensor, a new_ScratchBufferspair, and two sort buffers. Nothing is ever evicted. A serving workload with variable batch sizes grows VRAM until every observed token count is cached.The class docstring documents this, and
max_num_tokensavoids it. If you want a safer default, roundnum_tokensup to a power-of-two capacity so the cache stays bounded.🤖 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/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py` around lines 591 - 643, Update _capacity_for to bucket token counts when max_num_tokens is unset, rounding num_tokens up to a power-of-two capacity before _get_output calls _allocate_capacity. Preserve the existing max_num_tokens validation and return behavior, and ensure the bucket is at least num_tokens so returned output slicing remains correct.flashinfer/trace/templates/moe.py (1)
3692-3698: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCorrect the comment about
top_k.The comment states that no tensor on
run()carries anum_expertsortop_kdim.token_selected_expertsandtoken_final_scalesstay in_cute_dsl_mxfp8_mxfp4_wrapper_inputs, and both declaredim_namesof["num_tokens", "top_k"]. The claim holds fornum_expertsonly.📝 Proposed comment fix
_cute_dsl_mxfp8_mxfp4_wrapper_axes = dict(cute_dsl_fused_moe_mxfp8_mxfp4_trace.axes) -# No tensor on run() carries a num_experts or top_k dim once the scalars move -# to __init__, so both axes have to be free variables here. +# num_experts has no tensor source once its scalar moves to __init__, so it +# has to be a free variable. top_k is still carried by token_selected_experts +# and token_final_scales; it is declared Var for symmetry with num_experts. _cute_dsl_mxfp8_mxfp4_wrapper_axes["num_experts"] = Var(🤖 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/trace/templates/moe.py` around lines 3692 - 3698, Update the comment above _cute_dsl_mxfp8_mxfp4_wrapper_axes to state that no run() tensor carries a num_experts dimension, while token_selected_experts and token_final_scales still carry top_k; preserve the explanation that num_experts must remain a free variable and remove the incorrect claim about top_k.flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py (1)
2758-2773: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThese two unpack checks are unreachable for this kernel.
can_implementrejects any configuration wherea_major != "k"orb_major != "k"at Line 2864. The guards here requirea_major == "m"orb_major == "n", so they never evaluate toFalseon a path that otherwise succeeds. Keep them only if you intend them as forward-compatible guards, and add a short comment saying so.🤖 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/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py` around lines 2758 - 2773, Remove the unreachable unpack checks in can_implement around use_2cta_instrs, or explicitly retain them only as forward-compatible guards with a short explanatory comment. Preserve the existing validation that rejects configurations where a_major is not "k" or b_major is not "k".flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py (1)
3863-3892: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAlign the comment with the symmetric check.
The comment states that B is the only unpacked operand in this gather kernel, but the code applies
check_contiguous_128_elementsto A and B. The code is correct and defensive. Reword the comment so it does not contradict the applied check.📝 Proposed comment change
- # UNPACK_U8 requires the contiguous dimension of a sub-byte operand to - # be a multiple of 128 elements. In this gather kernel B is the only - # unpacked operand, but keep the check symmetric for clarity. + # UNPACK_U8 requires the contiguous dimension of a sub-byte operand to + # be a multiple of 128 elements. Only B is unpacked today; the check + # is applied to both operands so a future A-side FP4 path stays valid.🤖 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/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py` around lines 3863 - 3892, Update the comment above the symmetric check in can_implement to avoid claiming that B is the only unpacked operand; describe that the 128-element contiguous-dimension validation is applied to both A and B defensively, while preserving the existing check_contiguous_128_elements logic.
🤖 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/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py`:
- Around line 66-81: Make the resources returned by
_get_mixed_functional_resources thread-local per device so concurrent calls to
cute_dsl_fused_moe_mxfp8_mxfp4 never share main_event, memset_event, or
aux_stream. Preserve the existing resource initialization and device
association, ensuring each host thread’s _moe_core_impl_mxfp8_mxfp4 record/wait
sequence uses its own objects.
In `@flashinfer/fused_moe/cute_dsl/mixed_tuner.py`:
- Around line 369-379: The fallback at
flashinfer/fused_moe/cute_dsl/mixed_tuner.py:369-379 must raise a ValueError
naming the unsupported tokens, hidden, intermediate, and expert shape instead of
returning DEFAULT_MXFP8_MXFP4_MOE_TACTIC. At
flashinfer/fused_moe/cute_dsl/mixed_tuner.py:386-391, validate _fallback_tactic
with both BlockScaledContiguousGatherGroupedGemmKernel.can_implement and
Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel.can_implement before
calling _extract_tactic_params, and raise if either predicate rejects it.
In `@tests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py`:
- Around line 11-22: Replace the hand-rolled _is_sm100_family checks with
flashinfer.utils.is_sm100a_supported(device), preserving the existing
sm100_required skip marker. Apply this change in
tests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py lines 11-22 and
tests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.py lines 30-41; both sites
require the same capability check, including supported SM100-family variants and
CUDA version validation.
In `@tests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.py`:
- Around line 267-271: Update the comment above the n=192 assertion to describe
the E8M0 scale layout’s requirement that the N/column extent align to complete
128-column groups, rather than referring to 128-row groups. Keep the assertions
and the separate explanation of partial trailing N tiles unchanged.
In `@tests/moe/test_moe_sort_route_windows.py`:
- Around line 45-48: Update the sm100_required skip condition to use
flashinfer.utils.is_sm100a_supported(torch.device("cuda")) instead of checking
CUDA availability and compute capability directly; import the utility and torch
as needed, while preserving the existing skip reason.
---
Nitpick comments:
In
`@flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`:
- Around line 3863-3892: Update the comment above the symmetric check in
can_implement to avoid claiming that B is the only unpacked operand; describe
that the 128-element contiguous-dimension validation is applied to both A and B
defensively, while preserving the existing check_contiguous_128_elements logic.
In
`@flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py`:
- Around line 2758-2773: Remove the unreachable unpack checks in can_implement
around use_2cta_instrs, or explicitly retain them only as forward-compatible
guards with a short explanatory comment. Preserve the existing validation that
rejects configurations where a_major is not "k" or b_major is not "k".
In `@flashinfer/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py`:
- Around line 785-786: Clarify the contract of
CuteDslFusedMoEMxfp8Mxfp4Runner.get_valid_tactics by adding a docstring stating
that it returns the complete ALL_MXFP8_MXFP4_MOE_TACTICS space without filtering
for can_implement or the configured shapes; keep the existing return behavior
unchanged.
- Around line 591-643: Update _capacity_for to bucket token counts when
max_num_tokens is unset, rounding num_tokens up to a power-of-two capacity
before _get_output calls _allocate_capacity. Preserve the existing
max_num_tokens validation and return behavior, and ensure the bucket is at least
num_tokens so returned output slicing remains correct.
In `@flashinfer/trace/templates/moe.py`:
- Around line 3692-3698: Update the comment above
_cute_dsl_mxfp8_mxfp4_wrapper_axes to state that no run() tensor carries a
num_experts dimension, while token_selected_experts and token_final_scales still
carry top_k; preserve the explanation that num_experts must remain a free
variable and remove the incorrect claim about top_k.
In `@tests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py`:
- Around line 81-110: Move _interleave_linear_and_gate and
_quantize_mxfp4_grouped into tests/moe/utils.py, using the grouped-GEMM
implementations as the canonical versions, including the rows divisibility
assertion and dequantized reference return. Remove the duplicate local
definitions from both test modules and import the shared helpers instead.
- Around line 37-47: Remove the exact count assertion for ALL_MOE_TACTICS from
the test, while preserving the mixed-path tactic count, uniqueness, runner
separation, and tile-size assertions.
🪄 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: 6c235d1b-429c-45f1-940f-2f3de8e63144
📒 Files selected for processing (19)
csrc/moe_utils_binding.cudocs/api/fused_moe.rstflashinfer/__init__.pyflashinfer/fused_moe/__init__.pyflashinfer/fused_moe/cute_dsl/__init__.pyflashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pyflashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.pyflashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pyflashinfer/fused_moe/cute_dsl/blockscaled_contiguous_grouped_gemm_finalize_fusion.pyflashinfer/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.pyflashinfer/fused_moe/cute_dsl/mixed_tuner.pyflashinfer/fused_moe/cute_dsl/tuner.pyflashinfer/trace/templates/moe.pyinclude/flashinfer/trtllm/fused_moe/RoutingKernel.cuhinclude/flashinfer/trtllm/fused_moe/RoutingKernel.htests/moe/test_cute_dsl_moe_can_implement.pytests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.pytests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.pytests/moe/test_moe_sort_route_windows.py
5 tasks
- Make the functional path's async-memset event/stream resources thread-local, so one thread's memset_event.wait() can no longer be satisfied by another thread's record and start GEMM2 over a buffer that was never zeroed. - Stop returning DEFAULT_MXFP8_MXFP4_MOE_TACTIC when no tactic passes can_implement, matching the NVFP4 runner's gh flashinfer-ai#3957 behaviour of returning an empty list so the autotuner skips a bucket it cannot profile. - Gate the MXFP8 x MXFP4 and route-window tests on flashinfer.utils.is_sm100a_supported so they honour the CUDA 12.8+ requirement instead of only checking the compute-capability major version. - Clarify that the finalize kernel's n % 128 guard groups the N extent rather than rows, so it no longer reads as a different dimension from the adjacent partial-N-tile note. AI-assisted with Claude Code. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
The no-autotune fallback matched one exact problem shape and returned hand-measured tactics for it. The measured edge was under 1% on the full pipeline, no test guarded the entries against kernel changes, and the sibling NVFP4 runner has no equivalent table, so the fallback now returns DEFAULT_MXFP8_MXFP4_MOE_TACTIC like CuteDslFusedMoENvfp4Runner.forward does. Autotuning is unaffected. Also drop the concrete problem dimensions from the mUseContiguousRouteWindows comment; the speedups it cites stand on their own. AI-assisted with Claude Code. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
Follows the NVFP4 wrapper (flashinfer-ai#3404): CUDA graph capture records allocations from its private pool, so pre-sizing workspace for a maximum batch never helped capture and only cost memory. Drop the output, routing, intermediate and scale buffers and keep just the persistent stream and events, which do have to exist before capture. max_num_tokens is now accepted but ignored, matching CuteDslMoEWrapper. This also removes the per-batch-size buffer caches, which grew without bound when max_num_tokens was unset, and makes run() return a caller-owned tensor instead of a view of wrapper storage that the next call overwrites. AI-assisted with Claude Code. Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
Collaborator
|
/bot run tests/moe |
Collaborator
jiahanc
approved these changes
Aug 12, 2026
jiahanc
left a comment
jiahanc
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, thanks for the work!
Collaborator
|
[SUCCESS] Pipeline #62264435: 18/18 executed test jobs passed |
Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
Contributor
|
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. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
flashinfer/fused_moe/cute_dsl/mixed_tuner.py (3)
347-358: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate
tacticasOptional.The parameter defaults to
Nonebut is annotatedTuple[Any, ...]. PEP 484 does not allow an implicitOptional, and Ruff reports RUF013 here. Line 354 already handlesNone, so the correct annotation isOptional[Tuple[Any, ...]]. That change also removes the need for thetype: ignore[assignment]comment.
Optionalis already imported throughtypingin this module's import list? Check line 26: it importsAny, Callable, Dict, List, Tupleonly, so addOptional.♻️ Proposed change
-from typing import Any, Callable, Dict, List, Tuple +from typing import Any, Callable, Dict, List, Optional, Tupledef forward( # type: ignore[override] self, inputs: List[torch.Tensor], - tactic: Tuple[Any, ...] = None, # type: ignore[assignment] + tactic: Optional[Tuple[Any, ...]] = None, do_preparation: bool = False, **kwargs: Any, ) -> torch.Tensor:🤖 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/fused_moe/cute_dsl/mixed_tuner.py` around lines 347 - 358, Update the forward method’s tactic parameter annotation to Optional[Tuple[Any, ...]], add Optional to the typing imports, and remove the now-unnecessary type: ignore[assignment] comment while preserving the existing None handling.Source: Linters/SAST tools
269-286: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse one source for
num_local_experts.Line 274 derives
num_local_expertsfromw1_weight.shape[0]. Line 284 passesself.num_local_expertstoget_max_num_permuted_tokens. The public entry points validate that the two agree, so the current behavior is correct. Reading the local variable in both places would make that invariant explicit and remove the question for a future reader.♻️ Proposed change
- permuted_m = get_max_num_permuted_tokens( - num_tokens, - self.top_k, - self.num_local_experts, - params["tile_size"], - ) + permuted_m = get_max_num_permuted_tokens( + num_tokens, + self.top_k, + num_local_experts, + params["tile_size"], + )🤖 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/fused_moe/cute_dsl/mixed_tuner.py` around lines 269 - 286, Update the tactic loop in the relevant tuning method to pass the locally derived num_local_experts value to get_max_num_permuted_tokens instead of self.num_local_experts, using one consistent source for the expert count.
239-254: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueDrop
device.indexfrom the MXFP8/MXFP4 autotune cache key.
compute_capabilityalready separates entries by GPU architecture.device.indexrepeats the 32-tactic sweep for each GPU ordinal and makes persisted configs ordinal-specific. The sibling NVFP4 runners omitdevice.index; keep the cache-key behavior consistent.♻️ Proposed change
return ( "mxfp8_mxfp4_v2", self.enable_pdl, device.type, - device.index, compute_capability, int(self.activation_type),🤖 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/fused_moe/cute_dsl/mixed_tuner.py` around lines 239 - 254, Update get_cache_key_extras in the MXFP8/MXFP4 tuner to remove device.index from the returned cache-key tuple. Preserve compute_capability and all other existing key components so cached configurations remain separated by GPU architecture rather than device ordinal.flashinfer/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py (1)
679-680: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument that this enumeration is unfiltered.
CuteDslFusedMoEMxfp8Mxfp4Runner.get_valid_tacticsfiltersALL_MXFP8_MXFP4_MOE_TACTICSthrough both kernels'can_implementpredicates for the current shapes. This wrapper method returns the full list with no filtering, because it has noinputsorprofileargument. A caller can read the name and then pass a tactic that the kernels reject for its shapes.A short docstring that states the list is shape-independent would prevent that reading.
♻️ Proposed docstring
def get_valid_tactics(self) -> list: + """Return every MXFP8 x MXFP4 tactic, without shape filtering. + + The returned tactics are not validated against the kernels' + ``can_implement`` predicates for a specific problem shape. Pass + ``tactic=None`` to ``run`` to let the autotuner select a tactic that + both grouped GEMMs accept. + """ return list(ALL_MXFP8_MXFP4_MOE_TACTICS)🤖 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/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py` around lines 679 - 680, Add a concise docstring to CuteDslFusedMoEMxfp8Mxfp4Runner.get_valid_tactics stating that it returns the complete, unfiltered ALL_MXFP8_MXFP4_MOE_TACTICS list and is independent of input shapes because no kernel capability checks are performed.
🤖 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/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py`:
- Around line 500-540: Update the wrapper’s run method to execute its entire
body inside with torch.cuda.device(self.device), ensuring current-stream, GEMM
launches, and event operations use the wrapper’s configured CUDA device. Locate
the run method associated with the constructor above and preserve its existing
logic within the device context.
---
Nitpick comments:
In `@flashinfer/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.py`:
- Around line 679-680: Add a concise docstring to
CuteDslFusedMoEMxfp8Mxfp4Runner.get_valid_tactics stating that it returns the
complete, unfiltered ALL_MXFP8_MXFP4_MOE_TACTICS list and is independent of
input shapes because no kernel capability checks are performed.
In `@flashinfer/fused_moe/cute_dsl/mixed_tuner.py`:
- Around line 347-358: Update the forward method’s tactic parameter annotation
to Optional[Tuple[Any, ...]], add Optional to the typing imports, and remove the
now-unnecessary type: ignore[assignment] comment while preserving the existing
None handling.
- Around line 269-286: Update the tactic loop in the relevant tuning method to
pass the locally derived num_local_experts value to get_max_num_permuted_tokens
instead of self.num_local_experts, using one consistent source for the expert
count.
- Around line 239-254: Update get_cache_key_extras in the MXFP8/MXFP4 tuner to
remove device.index from the returned cache-key tuple. Preserve
compute_capability and all other existing key components so cached
configurations remain separated by GPU architecture rather than device ordinal.
🪄 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: 556ee918-deff-4547-b3d6-1150e7faee07
📒 Files selected for processing (20)
csrc/moe_utils_binding.cudocs/api/fused_moe.rstflashinfer/__init__.pyflashinfer/fused_moe/__init__.pyflashinfer/fused_moe/cute_dsl/__init__.pyflashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pyflashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.pyflashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pyflashinfer/fused_moe/cute_dsl/blockscaled_contiguous_grouped_gemm_finalize_fusion.pyflashinfer/fused_moe/cute_dsl/fused_moe_mxfp8_mxfp4.pyflashinfer/fused_moe/cute_dsl/mixed_tuner.pyflashinfer/fused_moe/cute_dsl/tuner.pyflashinfer/trace/templates/moe.pyinclude/flashinfer/trtllm/fused_moe/RoutingKernel.cuhinclude/flashinfer/trtllm/fused_moe/RoutingKernel.htests/moe/test_cute_dsl_moe_can_implement.pytests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.pytests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.pytests/moe/test_moe_sort_route_windows.pytests/trace/template_registry.py
🚧 Files skipped from review as they are similar to previous changes (17)
- flashinfer/fused_moe/cute_dsl/tuner.py
- flashinfer/fused_moe/init.py
- include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh
- csrc/moe_utils_binding.cu
- docs/api/fused_moe.rst
- flashinfer/fused_moe/cute_dsl/init.py
- flashinfer/init.py
- tests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py
- include/flashinfer/trtllm/fused_moe/RoutingKernel.h
- tests/moe/test_moe_sort_route_windows.py
- tests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.py
- flashinfer/trace/templates/moe.py
- tests/moe/test_cute_dsl_moe_can_implement.py
- flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
- flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
- flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
- flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…ch routing locality (flashinfer-ai#4440) ## 📌 Description Two related changes to the Blackwell CuTe-DSL MoE path. **1. MXFP8 x MXFP4 fused MoE (SM100/SM103)** — `e3bc88b3` A mixed-precision path: MXFP8 (E4M3 + linear block-32 E8M0) activations against packed MXFP4 (E2M1 + MMA-layout block-32 E8M0) weights, with an MXFP8 FC1 intermediate and BF16 finalized output. - Both grouped-GEMM kernels take separate `a_dtype`/`b_dtype`, size SMEM from `smem_alloc_*_dtype`, and load the narrow operand via TMA `UNPACK_U8`. - The FC1 epilogue emits the same exact UE8M0 codes as `mxfp8_quantize` and scalar-stores them as bytes, avoiding the unsupported `vector<2xE8M0>` conversion/store lowering. - New public APIs `cute_dsl_fused_moe_mxfp8_mxfp4` and `CuteDslMxfp8Mxfp4MoEWrapper`, with their own tactic space and autotune cache namespace so NVFP4 results and cache entries are unaffected. - GEMM2 `can_implement` no longer requires N to divide the MMA N tile: the finalize epilogue already clamps its bulk-reduce to `valid_columns`, which `tests/moe/test_cute_dsl_moe_can_implement.py` documents. A partial-N-tile case is added to the grouped-GEMM test to cover it. **2. Contiguous per-CTA route windows in `moe_sort` at large batches** — `5e0f9e64` The cooperative routing kernel assigns each CTA a grid stride of expanded indices, and an expert's rows land in CTA-arrival order, so one expert's rows are drawn from the whole batch. A grouped-GEMM tile then gathers across the entire activation tensor, and the footprint grows with the batch: measured on B200 (hidden 6144, 256 experts, top_k 8, 32 local, tile 256) a GEMM1 tile touches 94 distinct 2 MiB pages at 64K tokens, 190 at 128K and 228 at 256K, which outgrows the uTLB. With `mUseContiguousRouteWindows` each CTA owns one contiguous span instead, so a tile gathers from a few narrow token windows and the footprint stays flat at ~28-38 pages. `moe_sort` enables it from 65536 tokens: | tokens | GEMM1 | full MoE pipeline | 2 MiB pages / tile | |---:|:--|:--|:--| | 65536 | 1.00-1.10x | 0.99-1.03x | 94 -> 38 | | 131072 | 1.40-1.44x | **1.17-1.18x** | 190 -> 32 | | 262144 | 1.43-1.47x | **1.24-1.28x** | 228 -> 29 | The flag defaults to off, so every other caller of the routing kernels, including the TRT-LLM MoE backend, keeps its current ordering bit for bit. ## 🔍 Related Issues <!-- link related issues here --> ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). New tests: `tests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.py`, `tests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py`, `tests/moe/test_moe_sort_route_windows.py`. Run on B200 (SM100): - `tests/moe/test_cute_dsl_fused_moe.py` plus the three MXFP8 x MXFP4 / `can_implement` files — 564 passed - `tests/moe/test_moe_sort_route_windows.py` plus MXFP8 x MXFP4 end to end — 12 passed - `tests/trace/test_fi_trace_template_consistency.py`, `tests/trace/test_template_init.py` — 818 passed, 167 skipped ## Reviewer Notes - The routing change touches `RoutingKernel.cuh`, shared with the TRT-LLM MoE backend. It is gated behind a `DataBase` flag that only `moe_sort` sets, so TRT-LLM paths are unchanged; that separation is the main thing worth checking. - Scope of the new ordering: all three CuTe-DSL MoE entry points share `moe_sort` (`cute_dsl_fused_moe_nvfp4`, `cute_dsl_fused_moe_mxfp8_mxfp4` and the W4A16 path), so all three pick it up above the threshold. NVFP4 uses the same gather kernel so the mechanism applies, but its packed-FP4 activations are half the bytes per token, so its benefit is not that much. - Only the cooperative routing path honours the flag, so it applies while that path is selected (roughly 65536 to ~570K tokens at 256 experts / top_k 8 on a 148-SM B200). Above that the multi-kernel path runs and the flag is ignored. - This is a locality change, not an ordering guarantee: rows within an expert still land in atomic-arrival order and the permutation stays non-deterministic run to run, exactly as before. Only each CTA's source range becomes contiguous. - Two subtleties in the coop kernel are load-bearing: the per-CTA span is sized to keep every CTA busy rather than fixed at the per-thread maximum, and each CTA must stop at its own span end. Without the latter the trailing iterations run into the next CTA's span — still below `expandedIdxSize`, so the batch-size check alone does not stop them — and those routes get counted and permuted twice. `test_moe_sort_route_windows.py` asserts per-expert row counts against the routing histogram, which is what catches this class of bug; a bijection check alone does not, since a duplicated permutation is still self-consistent. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added MXFP8 activation × MXFP4 weight support for fused Mixture-of-Experts workloads on compatible GPUs. * Added public APIs, reusable wrappers, autotuning, tracing, and documentation for mixed-precision fused MoE operations. * Improved large-batch routing with contiguous route windows for more localized processing. * **Bug Fixes** * Improved validation, scaling, alignment, memory handling, and output quantization. * Preserved compatibility with existing NVFP4 APIs. * **Tests** * Added functional, integration, routing, and numerical correctness coverage. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Tiekai Bi <tiekaib@nvidia.com>
This was referenced Aug 19, 2026
feat(moe_ep): SM100 W4A8 (MXFP8xMXFP4) CuTeDSL split kernel backend with MXFP8 packed dispatch
#4529
Merged
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
Two related changes to the Blackwell CuTe-DSL MoE path.
1. MXFP8 x MXFP4 fused MoE (SM100/SM103) —
e3bc88b3A mixed-precision path: MXFP8 (E4M3 + linear block-32 E8M0) activations against
packed MXFP4 (E2M1 + MMA-layout block-32 E8M0) weights, with an MXFP8 FC1
intermediate and BF16 finalized output.
a_dtype/b_dtype, size SMEM fromsmem_alloc_*_dtype, and load the narrow operand via TMAUNPACK_U8.mxfp8_quantizeandscalar-stores them as bytes, avoiding the unsupported
vector<2xE8M0>conversion/store lowering.
cute_dsl_fused_moe_mxfp8_mxfp4andCuteDslMxfp8Mxfp4MoEWrapper, with their own tactic space and autotune cachenamespace so NVFP4 results and cache entries are unaffected.
can_implementno longer requires N to divide the MMA N tile: thefinalize epilogue already clamps its bulk-reduce to
valid_columns, whichtests/moe/test_cute_dsl_moe_can_implement.pydocuments. A partial-N-tilecase is added to the grouped-GEMM test to cover it.
2. Contiguous per-CTA route windows in
moe_sortat large batches —5e0f9e64The cooperative routing kernel assigns each CTA a grid stride of expanded
indices, and an expert's rows land in CTA-arrival order, so one expert's rows
are drawn from the whole batch. A grouped-GEMM tile then gathers across the
entire activation tensor, and the footprint grows with the batch: measured on
B200 (hidden 6144, 256 experts, top_k 8, 32 local, tile 256) a GEMM1 tile
touches 94 distinct 2 MiB pages at 64K tokens, 190 at 128K and 228 at 256K,
which outgrows the uTLB.
With
mUseContiguousRouteWindowseach CTA owns one contiguous span instead, soa tile gathers from a few narrow token windows and the footprint stays flat at
~28-38 pages.
moe_sortenables it from 65536 tokens:The flag defaults to off, so every other caller of the routing kernels,
including the TRT-LLM MoE backend, keeps its current ordering bit for bit.
🔍 Related Issues
🚀 Pull Request Checklist
✅ 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.).New tests:
tests/moe/test_cute_dsl_mxfp8_mxfp4_grouped_gemm.py,tests/moe/test_cute_dsl_mxfp8_mxfp4_fused_moe.py,tests/moe/test_moe_sort_route_windows.py.Run on B200 (SM100):
tests/moe/test_cute_dsl_fused_moe.pyplus the three MXFP8 x MXFP4 /can_implementfiles — 564 passedtests/moe/test_moe_sort_route_windows.pyplus MXFP8 x MXFP4 end to end — 12 passedtests/trace/test_fi_trace_template_consistency.py,tests/trace/test_template_init.py— 818 passed, 167 skippedReviewer Notes
RoutingKernel.cuh, shared with the TRT-LLM MoEbackend. It is gated behind a
DataBaseflag that onlymoe_sortsets, soTRT-LLM paths are unchanged; that separation is the main thing worth checking.
moe_sort(cute_dsl_fused_moe_nvfp4,cute_dsl_fused_moe_mxfp8_mxfp4andthe W4A16 path), so all three pick it up above the threshold. NVFP4 uses the same gather kernel so the
mechanism applies, but its packed-FP4 activations are half the bytes per
token, so its benefit is not that much.
path is selected (roughly 65536 to ~570K tokens at 256 experts / top_k 8 on a
148-SM B200). Above that the multi-kernel path runs and the flag is ignored.
still land in atomic-arrival order and the permutation stays
non-deterministic run to run, exactly as before. Only each CTA's source range
becomes contiguous.
to keep every CTA busy rather than fixed at the per-thread maximum, and each
CTA must stop at its own span end. Without the latter the trailing iterations
run into the next CTA's span — still below
expandedIdxSize, so thebatch-size check alone does not stop them — and those routes get counted and
permuted twice.
test_moe_sort_route_windows.pyasserts per-expert row countsagainst the routing histogram, which is what catches this class of bug; a
bijection check alone does not, since a duplicated permutation is still
self-consistent.
Summary by CodeRabbit
New Features
Bug Fixes
Tests