Conversation
📝 WalkthroughWalkthroughChangesThe CuteDSL MoE pipeline now supports applying router weights before FC2 instead of during output reduction. The option propagates through APIs, runners, kernels, unpermutation, tracing, benchmarks, reference calculations, and accuracy tests. Router-weight placement
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 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 |
12c1796 to
1258c31
Compare
1258c31 to
99336aa
Compare
|
@coderabbitai review |
|
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/moe/test_cute_dsl_fused_moe.py (1)
1538-1592: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winCover input-side routing weights through the direct functional API.
The helper defaults
apply_router_weight_on_inputtoFalse. The shown calls fromtest_numerical_accuracyandtest_deterministic_finalize_numerical_accuracydo not override it. The direct API therefore never tests the new mode for either finalize path.Add direct functional cases that pass
True, including W4A4 per-tensor and per-token quantization, W4A16, SwiGLU, and ReLU2. The wrapper matrix does not validate the direct API path.🤖 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_fused_moe.py` around lines 1538 - 1592, The direct functional tests around the helper invoking cute_dsl_fused_moe_nvfp4 currently always use the default apply_router_weight_on_input=False. Add explicit direct API test cases with apply_router_weight_on_input=True for both finalize paths, covering W4A4 per-tensor and per-token quantization, W4A16, SwiGLU, and ReLU2; ensure these cases exercise the functional call rather than only wrapper tests.
🧹 Nitpick comments (2)
flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py (1)
3000-3014: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the packed and scalar route-scaling paths.
This hot path selects packed
f32x2operations or a scalar loop withvectorized_f32. State that the packed path reduces epilogue instruction cost and that the scalar loop supports non-vectorized variants.Proposed comment
if cutlass.const_expr(self.apply_router_weight_on_input): # Apply the route in FP32 after the expert activation # and before the FC2-input store or quantization. + # Use packed f32x2 operations to reduce epilogue cost. + # Use the scalar loop for non-vectorized kernel variants.As per coding guidelines, “For performance-critical hot paths, document the rationale for special algorithmic choices and potential alternatives in a code comment.”
🤖 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 3000 - 3014, Add concise comments in the route-scaling branch around vectorized_f32 explaining that the packed mul_packed_f32x2 path reduces epilogue instruction cost, while the scalar loop supports non-vectorized variants. Preserve the existing scaling logic and placement.Source: Coding guidelines
flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py (1)
2036-2045: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the router-weight placement decision in each hot path.
State why input-side scaling occurs after activation and before the FC2-input store. State that output-side scaling is the alternative and changes finite-precision ordering. State that later finalization or unpermutation must then skip duplicate scaling.
flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py#L2036-L2045: document the gated input-side scaling rationale.flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py#L2137-L2152: document the non-gated input-side scaling rationale.flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py#L2173-L2198: document the complementary fused-finalize behavior.flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py#L1969-L1974: document why input-side mode uses neutral finalization scaling.csrc/nv_internal/tensorrt_llm/kernels/cuteDslKernels/moeUtils.cu#L195-L205: document why unscaled reduction is valid only after input-side scaling.As per coding guidelines, “For performance-critical hot paths, document the rationale for special algorithmic choices and potential alternatives in a code comment.”
🤖 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/moe_w4a16_kernel.py` around lines 2036 - 2045, Document the router-weight placement rationale at all five sites: in moe_w4a16_kernel.py lines 2036-2045 and 2137-2152, explain gated and non-gated input-side scaling after activation and before the FC2-input store; at lines 2173-2198, state that later fused finalization or unpermutation must skip duplicate scaling and describe output-side scaling as the alternative with different finite-precision ordering; in blockscaled_contiguous_grouped_gemm_finalize_fusion.py lines 1969-1974, explain why input-side mode uses neutral finalization scaling; and in moeUtils.cu lines 195-205, explain that unscaled reduction is valid only after input-side scaling.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/fused_moe/api.py`:
- Around line 207-216: Update the backend validation in check_support to reject
apply_router_weight_on_input=True for TRTLLM and B12x runners, which do not
implement input-side routing-weight application. Raise NotImplementedError for
those unsupported combinations, while preserving the existing supported CuteDSL
path.
In `@flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py`:
- Around line 1037-1044: Update _compute_stages_and_tmem_cols() so its
finalize_metadata_bytes shared-memory carveout is reserved when either
use_fused_finalize or load_router_weight is enabled, matching the sRouteScale
allocation condition and preserving deterministic input-side mode.
In `@flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py`:
- Line 221: Update every direct caller of _run_grouped_gemm, including the
invocation in test_cute_dsl_fused_moe.py, to provide
apply_router_weight_on_input=False. Preserve existing call behavior while
satisfying the newly required argument.
In
`@flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`:
- Around line 534-545: Update the token_final_scales validation in the
apply_router_weight_on_input path to require token_final_scales.device to match
a.device, not merely any CUDA device. Preserve the existing dtype, contiguity,
shape, and missing-tensor checks, and reject tensors from different CUDA devices
before passing their pointer to the kernel.
In `@flashinfer/fused_moe/cute_dsl/fused_moe.py`:
- Around line 332-333: At the `token_final_scales` and
`apply_router_weight_on_input` placement boundary in the fused MoE flow, add a
concise comment documenting that enabling input-side router weighting applies
scales after activation and before FC2, so GEMM2 and deterministic unpermutation
must skip post-GEMM2 scaling; mention output-side scaling as the alternative and
guard against double scaling.
---
Outside diff comments:
In `@tests/moe/test_cute_dsl_fused_moe.py`:
- Around line 1538-1592: The direct functional tests around the helper invoking
cute_dsl_fused_moe_nvfp4 currently always use the default
apply_router_weight_on_input=False. Add explicit direct API test cases with
apply_router_weight_on_input=True for both finalize paths, covering W4A4
per-tensor and per-token quantization, W4A16, SwiGLU, and ReLU2; ensure these
cases exercise the functional call rather than only wrapper tests.
---
Nitpick comments:
In
`@flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`:
- Around line 3000-3014: Add concise comments in the route-scaling branch around
vectorized_f32 explaining that the packed mul_packed_f32x2 path reduces epilogue
instruction cost, while the scalar loop supports non-vectorized variants.
Preserve the existing scaling logic and placement.
In `@flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py`:
- Around line 2036-2045: Document the router-weight placement rationale at all
five sites: in moe_w4a16_kernel.py lines 2036-2045 and 2137-2152, explain gated
and non-gated input-side scaling after activation and before the FC2-input
store; at lines 2173-2198, state that later fused finalization or unpermutation
must skip duplicate scaling and describe output-side scaling as the alternative
with different finite-precision ordering; in
blockscaled_contiguous_grouped_gemm_finalize_fusion.py lines 1969-1974, explain
why input-side mode uses neutral finalization scaling; and in moeUtils.cu lines
195-205, explain that unscaled reduction is valid only after input-side scaling.
🪄 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: 0c6beb14-6b7b-4c67-a97f-899cf3dc698e
📒 Files selected for processing (19)
benchmarks/bench_cute_dsl_moe_distributed.pybenchmarks/bench_moe_deepseek.pycsrc/moe_utils_binding.cucsrc/nv_internal/tensorrt_llm/kernels/cuteDslKernels/moeUtils.cucsrc/nv_internal/tensorrt_llm/kernels/cuteDslKernels/moeUtils.hflashinfer/fused_moe/api.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/blackwell/moe_w4a16.pyflashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.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.pyflashinfer/fused_moe/cute_dsl/moe_utils.pyflashinfer/fused_moe/cute_dsl/tuner.pyflashinfer/fused_moe/runners.pyflashinfer/trace/templates/moe.pytests/moe/test_cute_dsl_fused_moe.pytests/moe/utils.py
|
Review-summary follow-up for aadf875:
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/moe/test_cute_dsl_fused_moe.py (1)
1475-1475: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd input-side coverage to boundary-focused tests.
The route-tile and cluster-padding tests still use the default
apply_router_weight_on_input=False. Add input-side cases to those tests so routing-weight placement is also validated at partial-tile and padding boundaries.🤖 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_fused_moe.py` at line 1475, Extend the boundary-focused route-tile and cluster-padding tests around the existing apply_router_weight_on_input configuration to cover both False and True. Ensure the new input-side cases validate routing-weight placement at partial-tile and padding boundaries while preserving the existing default-side coverage.
🤖 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.
Nitpick comments:
In `@tests/moe/test_cute_dsl_fused_moe.py`:
- Line 1475: Extend the boundary-focused route-tile and cluster-padding tests
around the existing apply_router_weight_on_input configuration to cover both
False and True. Ensure the new input-side cases validate routing-weight
placement at partial-tile and padding boundaries while preserving the existing
default-side coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 83ac2d1e-52d9-4d95-aadd-d331213ff7f4
📒 Files selected for processing (5)
flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.pyflashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pyflashinfer/fused_moe/cute_dsl/fused_moe.pyflashinfer/fused_moe/runners.pytests/moe/test_cute_dsl_fused_moe.py
🚧 Files skipped from review as they are similar to previous changes (3)
- flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
- flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py
- flashinfer/fused_moe/cute_dsl/fused_moe.py
|
need to rebase |
📌 Description
@HumansAnd
This PR adds opt-in input-side router-weight placement to the SM100-family CuTe DSL MoE implementation for both W4A4 and W4A16. The branch is rebased onto current
main, including merged W4A16 support from #4048.Motivation
RL stacks may use different MoE implementations across execution paths—for example, FlashInfer for rollout or inference and native Megatron for training or recompute. Native Megatron can pass routing probabilities into the expert MLP, while FlashInfer historically applies them to each route after FC2 during finalization.
Moving the scalar across linear FC2 is algebraically equivalent in exact arithmetic, so this option does not change the model-level MoE function. It is nevertheless numerically different when the multiplication moves across the W4A4 FC2-input quantization boundary or the W4A16 BF16 FC2-input store. The opt-in placement lets FlashInfer match another implementation's finite-precision ordering for closer cross-stack numerical alignment and easier parity debugging. The existing output-side behavior remains the default.
For route
r, leta_r = activation(W1_r x). The two placements are algebraically identical because FC2 is linear:sum_r p_r * W2_r(a_r)apply_router_weight_on_input=True:sum_r W2_r(p_r * a_r)The flag scales the post-activation FC2 input; it never scales the FC1 preactivation.
At the finite-precision boundary:
The implementation covers:
cute_dsl_fused_moe_nvfp4functional API andCuteDslMoEWrapper.ExecutionConfig.apply_router_weight_on_input, forwarded throughCuteDslNvfp4Runnerto W4A4 and W4A16.--apply-router-weight-on-inputin bothbench_moe_deepseek.pyandbench_cute_dsl_moe_distributed.py, including distributed profiler-worker propagation.apply_router_weight_on_input=Falseremains the backward-compatible default.Existing SM120 status
Upstream
mainalready has an internalapply_router_weight_on_inputoption in the SM120/SM12x W4A16run_w4a16_moepath. It is not exposed through the unifiedExecutionConfig,CuteDslMoEWrapper, or SM100/SM103cute_dsl_fused_moe_nvfp4APIs extended by this PR.The current SM120 implementation routes the weight through FC1 (
mul_topk_weights=True) and disables it on FC2. In its ReLU2 epilogue, however, the multiplication occurs before the activation. That computesReLU(p_r * z_r)^2 = p_r^2 * ReLU(z_r)^2for nonnegative router weights, rather than the algebraically identical placementp_r * ReLU(z_r)^2. SwiGLU is nonlinear as well. This is a potential correctness and cross-backend API-consistency bug if callers expect this flag to mean post-activation FC2-input placement.SM120 is outside this PR's SM100/SM103 implementation scope. A follow-up should move its multiplication after activation, or explicitly rename and document the existing preactivation semantics.
🔍 Related Issues
main.🚀 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.).Validated head:
d684167f69d124bfcaf725ad69d485f39ae9605bLocal static validation:
pre-commit run --files <review-fix files>: passed, including mypy, Ruff check, and Ruff format.git diff --check origin/main...HEAD: passed.CUDA 13.2 validation used the bare
c1/infradevboxflashinfer-pr4366-cu132with 8x NVIDIA B200, driver 580.126.09, CUDA 13.2, PyTorch 2.13.0+cu132, CuTe DSL 4.7.0, FlashInfer 0.6.18, andcupti-python==13.2.0.test_wrapper_accuracycollects 288 cases with independent axes for token count, activation, router-weight placement, finalize mode, W4A4/W4A16 mode, top-k, and expert count.token_final_scalesare rejected before the low-level gather launch.test_finalize_handles_cluster_padding_and_partial_tilesnow passes a 12-case Cartesian matrix: W4A4 per-tensor/per-token and W4A16 × hidden sizes 256/384 × output/input-side router placement.test_route_tile_boundary_accuracynow passes 12 cases: route tiles 8/16/32/64/128/192 × output/input-side router placement in deterministic SwiGLU mode.Performance benchmark
All reportable arms use the same warmed compiled cache:
FLASHINFER_WORKSPACE_BASE=/root/flashinfer-pr4366-cacheCUTE_DSL_CACHE_DIR=/root/flashinfer-pr4366-cache/cute-dslFLASHINFER_NVCC_THREADS=16The logs contain no compiler/JIT-build or cubin-download messages, and process monitoring observed no actual
nvccorciccprocess during the reportable runs. The benchmark's normal in-process tactic autotuning still occurs before its timed table.All arms use CUPTI 13.2.0 hardware-activity timing with CUDA graphs enabled.
--no-fused-finalizeselects deterministic two-stage finalize.--ep 8is a single-process compute-shard simulation with 32 local experts; it does not launch eight distributed ranks or measure inter-rank communication.4-over-6 per-token activation contract
Output-side command
cd /root/flashinfer-src CUDA_VISIBLE_DEVICES=0 \ PYTHONUNBUFFERED=1 \ PYTHONPATH=/root/flashinfer-src \ FLASHINFER_WORKSPACE_BASE=/root/flashinfer-pr4366-cache \ CUTE_DSL_CACHE_DIR=/root/flashinfer-pr4366-cache/cute-dsl \ FLASHINFER_NVCC_THREADS=16 \ FLASHINFER_NVFP4_4OVER6=1 \ FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=1 \ FLASHINFER_NVFP4_4OVER6_ERR_MODE=MSE \ FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 \ FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1 \ python3 benchmarks/bench_moe_deepseek.py \ --ep 8 \ --num-tokens 1,2,4,8,16,32,64,128,256,512,1024,2048,4096 \ --use-per-token-activation \ --no-fused-finalizeInput-side command
cd /root/flashinfer-src CUDA_VISIBLE_DEVICES=0 \ PYTHONUNBUFFERED=1 \ PYTHONPATH=/root/flashinfer-src \ FLASHINFER_WORKSPACE_BASE=/root/flashinfer-pr4366-cache \ CUTE_DSL_CACHE_DIR=/root/flashinfer-pr4366-cache/cute-dsl \ FLASHINFER_NVCC_THREADS=16 \ FLASHINFER_NVFP4_4OVER6=1 \ FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=1 \ FLASHINFER_NVFP4_4OVER6_ERR_MODE=MSE \ FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 \ FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1 \ python3 benchmarks/bench_moe_deepseek.py \ --ep 8 \ --num-tokens 1,2,4,8,16,32,64,128,256,512,1024,2048,4096 \ --use-per-token-activation \ --no-fused-finalize \ --apply-router-weight-on-inputRaw output
Autotuner progress is omitted below; the benchmark preamble and final tables are copied from the raw logs with trailing whitespace normalized. The raw tables'
Speedupcolumns compare CuTe DSL against TRTLLM, not input-side against output-side placement.Deterministic two-stage finalize — output-side
Deterministic two-stage finalize — input-side
Input-side versus output-side comparison
Delta = (input-side - output-side) / output-side; negative is faster and positive is slower. Values use the printed 0.001 ms resolution.W4A4's arithmetic-mean delta is -0.16% across the full sweep and -0.18% from 128 through 4096 tokens. W4A16's corresponding deltas are +0.46% and +0.64%. Most nonzero changes are one microsecond and should be treated as timing resolution or run-to-run noise. W4A16 at 4096 tokens is 0.015 ms slower (+4.27%) and should be sampled repeatedly before drawing a performance conclusion.
Default activation contract (no 4-over-6, no per-token activation)
This additional pair explicitly unsets the previous 4-over-6/fast-math environment and omits
--use-per-token-activation. Without per-token activation, the benchmark also exercises and reports the CUTLASS comparison backend.The first output-side warmup compiled the previously absent
fused_moe_100CUTLASS extension and fetched 118 TRTLLM cubins into the same persistent cache; that warmup was discarded. The input-side warmup and both reportable arms observed no compiler processes. Each reportable log also hasnvcc/JIT-build/cubin-fetch = 0/0/0.Output-side command
Input-side command
Raw output
Autotuner progress is omitted below. CUTLASS and TRTLLM are independently rerun comparison baselines and are not affected by the CuTe DSL placement flag, so their cross-arm timing drift is not part of the placement comparison.
Default activation, deterministic two-stage finalize — output-side
Default activation, deterministic two-stage finalize — input-side
Input-side versus output-side comparison
Delta = (input-side - output-side) / output-side; negative is faster and positive is slower. Values use the printed 0.001 ms resolution.W4A4's arithmetic-mean delta is -0.66% across the full sweep and -1.08% from 128 through 4096 tokens. W4A16's corresponding deltas are +0.30% and -0.32%. Most nonzero changes are only one microsecond and should be treated as timing resolution or run-to-run noise. The 4096-token changes are nine microseconds faster for both modes; they need repeated sampling before being treated as a placement performance effect.
Reviewer Notes
Please focus on post-activation placement, the W4A4 FP4 and W4A16 BF16 boundaries, avoiding duplicate scaling in both finalize modes, and propagation through functional, wrapper, and unified execution APIs.
Summary by CodeRabbit
New Features
Tests