Skip to content

feat(moe): support input-side router weights in CuTe DSL W4A4/W4A16 MoE - #4366

Draft
zianglih wants to merge 4 commits into
flashinfer-ai:mainfrom
zianglih:agent/sm100-router-weight-on-input
Draft

zianglih wants to merge 4 commits into
flashinfer-ai:mainfrom
zianglih:agent/sm100-router-weight-on-input

Conversation

@zianglih

@zianglih zianglih commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

📌 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, let a_r = activation(W1_r x). The two placements are algebraically identical because FC2 is linear:

  • Default: 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:

  • W4A4 uses one shared epilogue to multiply the activated FP32 value by the FP32 route weight before its quantization paths diverge.
  • W4A4 per-tensor scaling directly quantizes that routed FP32 value to FP4. W4A4 per-token scaling first stores it as BF16/FP16, then invokes the existing standalone per-token FP4 quantizer; the route multiplication itself remains FP32.
  • W4A16 multiplies the activated FP32 value by the route weight before the BF16 FC2-input store.
  • The default path continues applying each route weight once after FC2.

The implementation covers:

  • SM100-family W4A4 and W4A16 kernels.
  • SwiGLU and ReLU2.
  • Atomic fused finalize and deterministic two-stage finalize, without applying route weights twice.
  • The direct cute_dsl_fused_moe_nvfp4 functional API and CuteDslMoEWrapper.
  • Unified ExecutionConfig.apply_router_weight_on_input, forwarded through CuteDslNvfp4Runner to W4A4 and W4A16.
  • Unified backends that do not consume this option now reject it instead of silently retaining output-side placement.
  • Tracing, reference computation, autotuner signatures, and cache keys.
  • --apply-router-weight-on-input in both bench_moe_deepseek.py and bench_cute_dsl_moe_distributed.py, including distributed profiler-worker propagation.

apply_router_weight_on_input=False remains the backward-compatible default.

Existing SM120 status

Upstream main already has an internal apply_router_weight_on_input option in the SM120/SM12x W4A16 run_w4a16_moe path. It is not exposed through the unified ExecutionConfig, CuteDslMoEWrapper, or SM100/SM103 cute_dsl_fused_moe_nvfp4 APIs 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 computes ReLU(p_r * z_r)^2 = p_r^2 * ReLU(z_r)^2 for nonnegative router weights, rather than the algebraically identical placement p_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

🚀 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

  • I have installed pre-commit by running pip install pre-commit (or used your preferred method).
  • I have installed the hooks with pre-commit install.
  • I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

  • Tests have been added or updated as needed.
  • All tests are passing (unittest, etc.).

Validated head: d684167f69d124bfcaf725ad69d485f39ae9605b

Local 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/infra devbox flashinfer-pr4366-cu132 with 8x NVIDIA B200, driver 580.126.09, CUDA 13.2, PyTorch 2.13.0+cu132, CuTe DSL 4.7.0, FlashInfer 0.6.18, and cupti-python==13.2.0.

  • test_wrapper_accuracy collects 288 cases with independent axes for token count, activation, router-weight placement, finalize mode, W4A4/W4A16 mode, top-k, and expert count.
  • A representative 24-case final-head numerical slice passed in 10.96 seconds. It covers W4A4 per-tensor, W4A4 per-token, and W4A16; SwiGLU and ReLU2; input- and output-side placement; and fused and deterministic finalize at 128 tokens, top-k 2, and 256 experts.
  • The full 288-case matrix was collected but not executed.
  • Review-fix validation passed 12 targeted B200 tests: ten direct functional cases spanning W4A4 per-tensor, W4A4 per-token, W4A16, SwiGLU/ReLU2, and both finalize paths; one deterministic W4A16 wrapper case covering the corrected route-scale shared-memory plan; and one direct W4A16 grouped-GEMM caller case.
  • Explicit probes also confirmed unsupported unified backends reject the option and cross-GPU token_final_scales are rejected before the low-level gather launch.
  • test_finalize_handles_cluster_padding_and_partial_tiles now passes a 12-case Cartesian matrix: W4A4 per-tensor/per-token and W4A16 × hidden sizes 256/384 × output/input-side router placement.
  • W4A16 test_route_tile_boundary_accuracy now 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-cache
  • CUTE_DSL_CACHE_DIR=/root/flashinfer-pr4366-cache/cute-dsl
  • FLASHINFER_NVCC_THREADS=16

The logs contain no compiler/JIT-build or cubin-download messages, and process monitoring observed no actual nvcc or cicc process 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-finalize selects deterministic two-stage finalize. --ep 8 is 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-finalize

Input-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-input

Raw 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' Speedup columns compare CuTe DSL against TRTLLM, not input-side against output-side placement.

Deterministic two-stage finalize — output-side
DeepSeek-V3 MoE Performance Benchmark
GPU: NVIDIA B200
CuteDSL API: Wrapper
Per-token activation: True
Initial activation quantization: False
CuteDSL modes: W4A4 and W4A16
Tensor parallelism simulation: TP=1
CUDA profiler capture: False
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: output

========================================================================================================================
DeepSeek-V3 MoE Benchmark: CuteDSL W4A4/W4A16 vs TRTLLM (EP=8, TP=1)
========================================================================================================================
Model: hidden=7168, intermediate=2048, experts=256, top_k=8
EP Config: 32 local experts (simulating 8-way parallelism)
TP Config: intermediate size 2048 (simulating 1-way parallelism)
CUDA Graph: enabled, CUPTI: enabled
Routing bias scale: 0.01 (larger values tend to create expert imbalance)
Timed initial activation quantization for FP4-activation backends: excluded; W4A16 consumes BF16 directly
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: output
CUTLASS omitted: it does not consume the per-token activation scale.
------------------------------------------------------------------------------------------------------------------------
Tokens |  CuteDSL W4A4   |  CuteDSL W4A16  |     TRTLLM      | Speedup vs TRTLLM  |  Winner  | Active  |     Stats
       |      ms  TFLOPS |      ms  TFLOPS |      ms  TFLOPS |     W4A4     W4A16 |          | experts | min/max/median
------------------------------------------------------------------------------------------------------------------------
     1 |   0.038     2.3 |   0.036     2.4 |   0.031     2.8 |     0.81x    0.85x |  TRTLLM  |       2 |   0/  1/   0.00
     2 |   0.051     3.5 |   0.044     4.0 |   0.043     4.1 |     0.84x    0.98x |  TRTLLM  |       4 |   0/  1/   0.00
     4 |   0.070     5.0 |   0.059     6.0 |   0.056     6.3 |     0.80x    0.96x |  TRTLLM  |       7 |   0/  1/   0.00
     8 |   0.087     8.1 |   0.073     9.7 |   0.071     9.9 |     0.82x    0.98x |  TRTLLM  |      10 |   0/  2/   0.00
    16 |   0.096    14.7 |   0.082    17.2 |   0.083    17.0 |     0.87x    1.01x |  W4A16   |      13 |   0/  2/   0.00
    32 |   0.131    21.5 |   0.116    24.3 |   0.111    25.4 |     0.85x    0.96x |  TRTLLM  |      21 |   0/  3/   1.00
    64 |   0.167    33.8 |   0.151    37.4 |   0.146    38.7 |     0.87x    0.97x |  TRTLLM  |      29 |   0/  5/   2.00
   128 |   0.170    66.1 |   0.155    72.9 |   0.150    75.1 |     0.88x    0.97x |  TRTLLM  |      30 |   0/  8/   3.00
   256 |   0.181   124.9 |   0.163   138.6 |   0.160   140.8 |     0.89x    0.98x |  TRTLLM  |      32 |   1/ 11/   7.00
   512 |   0.185   244.3 |   0.166   272.5 |   0.170   265.4 |     0.92x    1.03x |  W4A16   |      32 |   7/ 23/  14.00
  1024 |   0.195   463.6 |   0.187   482.5 |   0.320   281.7 |     1.65x    1.71x |  W4A16   |      32 |  17/ 42/  28.50
  2048 |   0.216   833.6 |   0.232   778.9 |   0.323   558.0 |     1.49x    1.40x |   W4A4   |      32 |  38/ 74/  57.50
  4096 |   0.280  1286.4 |   0.351  1028.9 |   0.334  1080.4 |     1.19x    0.95x |   W4A4   |      32 |  81/147/ 117.00
------------------------------------------------------------------------------------------------------------------------
Speedup > 1.0 means that CuTe DSL mode is faster than the comparison backend
Deterministic two-stage finalize — input-side
DeepSeek-V3 MoE Performance Benchmark
GPU: NVIDIA B200
CuteDSL API: Wrapper
Per-token activation: True
Initial activation quantization: False
CuteDSL modes: W4A4 and W4A16
Tensor parallelism simulation: TP=1
CUDA profiler capture: False
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: input

========================================================================================================================
DeepSeek-V3 MoE Benchmark: CuteDSL W4A4/W4A16 vs TRTLLM (EP=8, TP=1)
========================================================================================================================
Model: hidden=7168, intermediate=2048, experts=256, top_k=8
EP Config: 32 local experts (simulating 8-way parallelism)
TP Config: intermediate size 2048 (simulating 1-way parallelism)
CUDA Graph: enabled, CUPTI: enabled
Routing bias scale: 0.01 (larger values tend to create expert imbalance)
Timed initial activation quantization for FP4-activation backends: excluded; W4A16 consumes BF16 directly
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: input
CUTLASS omitted: it does not consume the per-token activation scale.
------------------------------------------------------------------------------------------------------------------------
Tokens |  CuteDSL W4A4   |  CuteDSL W4A16  |     TRTLLM      | Speedup vs TRTLLM  |  Winner  | Active  |     Stats
       |      ms  TFLOPS |      ms  TFLOPS |      ms  TFLOPS |     W4A4     W4A16 |          | experts | min/max/median
------------------------------------------------------------------------------------------------------------------------
     1 |   0.038     2.3 |   0.037     2.4 |   0.031     2.9 |     0.81x    0.84x |  TRTLLM  |       2 |   0/  1/   0.00
     2 |   0.051     3.5 |   0.044     4.0 |   0.045     4.0 |     0.88x    1.02x |  W4A16   |       4 |   0/  1/   0.00
     4 |   0.070     5.0 |   0.059     6.0 |   0.056     6.3 |     0.80x    0.95x |  TRTLLM  |       7 |   0/  1/   0.00
     8 |   0.087     8.1 |   0.073     9.7 |   0.071     9.9 |     0.82x    0.97x |  TRTLLM  |      10 |   0/  2/   0.00
    16 |   0.095    14.8 |   0.082    17.1 |   0.083    17.1 |     0.86x    1.00x |  W4A16   |      13 |   0/  2/   0.00
    32 |   0.131    21.5 |   0.116    24.4 |   0.111    25.4 |     0.85x    0.96x |  TRTLLM  |      21 |   0/  3/   1.00
    64 |   0.167    33.7 |   0.150    37.5 |   0.144    39.0 |     0.86x    0.96x |  TRTLLM  |      29 |   0/  5/   2.00
   128 |   0.171    66.0 |   0.155    72.8 |   0.150    75.1 |     0.88x    0.97x |  TRTLLM  |      30 |   0/  8/   3.00
   256 |   0.181   124.8 |   0.163   138.2 |   0.160   141.1 |     0.88x    0.98x |  TRTLLM  |      32 |   1/ 11/   7.00
   512 |   0.185   243.8 |   0.166   271.9 |   0.170   264.9 |     0.92x    1.03x |  W4A16   |      32 |   7/ 23/  14.00
  1024 |   0.194   464.8 |   0.187   482.0 |   0.321   281.0 |     1.65x    1.72x |  W4A16   |      32 |  17/ 42/  28.50
  2048 |   0.215   840.7 |   0.231   780.3 |   0.325   554.3 |     1.52x    1.41x |   W4A4   |      32 |  38/ 74/  57.50
  4096 |   0.278  1297.6 |   0.366   986.8 |   0.334  1081.4 |     1.20x    0.91x |   W4A4   |      32 |  81/147/ 117.00
------------------------------------------------------------------------------------------------------------------------
Speedup > 1.0 means that CuTe DSL mode is faster than the comparison backend

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.

Tokens W4A4 output (ms) W4A4 input (ms) W4A4 delta W4A16 output (ms) W4A16 input (ms) W4A16 delta
1 0.038 0.038 +0.00% 0.036 0.037 +2.78%
2 0.051 0.051 +0.00% 0.044 0.044 +0.00%
4 0.070 0.070 +0.00% 0.059 0.059 +0.00%
8 0.087 0.087 +0.00% 0.073 0.073 +0.00%
16 0.096 0.095 -1.04% 0.082 0.082 +0.00%
32 0.131 0.131 +0.00% 0.116 0.116 +0.00%
64 0.167 0.167 +0.00% 0.151 0.150 -0.66%
128 0.170 0.171 +0.59% 0.155 0.155 +0.00%
256 0.181 0.181 +0.00% 0.163 0.163 +0.00%
512 0.185 0.185 +0.00% 0.166 0.166 +0.00%
1024 0.195 0.194 -0.51% 0.187 0.187 +0.00%
2048 0.216 0.215 -0.46% 0.232 0.231 -0.43%
4096 0.280 0.278 -0.71% 0.351 0.366 +4.27%

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_100 CUTLASS 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 has nvcc/JIT-build/cubin-fetch = 0/0/0.

Output-side command

unset FLASHINFER_NVFP4_4OVER6 \
  FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 \
  FLASHINFER_NVFP4_4OVER6_ERR_MODE \
  FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH \
  FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH

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 \
python3 benchmarks/bench_moe_deepseek.py \
  --ep 8 \
  --num-tokens 1,2,4,8,16,32,64,128,256,512,1024,2048,4096 \
  --no-fused-finalize

Input-side command

unset FLASHINFER_NVFP4_4OVER6 \
  FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 \
  FLASHINFER_NVFP4_4OVER6_ERR_MODE \
  FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH \
  FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH

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 \
python3 benchmarks/bench_moe_deepseek.py \
  --ep 8 \
  --num-tokens 1,2,4,8,16,32,64,128,256,512,1024,2048,4096 \
  --no-fused-finalize \
  --apply-router-weight-on-input

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
DeepSeek-V3 MoE Performance Benchmark
GPU: NVIDIA B200
CuteDSL API: Wrapper
Per-token activation: False
Initial activation quantization: False
CuteDSL modes: W4A4 and W4A16
Tensor parallelism simulation: TP=1
CUDA profiler capture: False
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: output

===============================================================================================================================================================
DeepSeek-V3 MoE Benchmark: CuteDSL W4A4/W4A16 vs CUTLASS vs TRTLLM (EP=8, TP=1)
===============================================================================================================================================================
Model: hidden=7168, intermediate=2048, experts=256, top_k=8
EP Config: 32 local experts (simulating 8-way parallelism)
TP Config: intermediate size 2048 (simulating 1-way parallelism)
CUDA Graph: enabled, CUPTI: enabled
Routing bias scale: 0.01 (larger values tend to create expert imbalance)
Timed initial activation quantization for FP4-activation backends: excluded; W4A16 consumes BF16 directly
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: output
---------------------------------------------------------------------------------------------------------------------------------------------------------------
Tokens |  CuteDSL W4A4   |  CuteDSL W4A16  |     CUTLASS     |     TRTLLM      | Speedup vs CUTLASS | Speedup vs TRTLLM  |  Winner  | Active  |     Stats
       |      ms  TFLOPS |      ms  TFLOPS |      ms  TFLOPS |      ms  TFLOPS |     W4A4     W4A16 |     W4A4     W4A16 |          | experts | min/max/median
---------------------------------------------------------------------------------------------------------------------------------------------------------------
     1 |   0.031     2.9 |   0.036     2.4 |   0.045     2.0 |   0.025     3.6 |     1.46x    1.24x |     0.80x    0.68x |  TRTLLM  |       2 |   0/  1/   0.00
     2 |   0.042     4.2 |   0.043     4.1 |   0.055     3.2 |   0.035     5.0 |     1.32x    1.27x |     0.83x    0.80x |  TRTLLM  |       4 |   0/  1/   0.00
     4 |   0.058     6.1 |   0.059     6.0 |   0.073     4.8 |   0.052     6.8 |     1.26x    1.25x |     0.89x    0.89x |  TRTLLM  |       7 |   0/  1/   0.00
     8 |   0.074     9.6 |   0.072     9.7 |   0.088     8.0 |   0.064    11.0 |     1.19x    1.21x |     0.87x    0.88x |  TRTLLM  |      10 |   0/  2/   0.00
    16 |   0.083    16.9 |   0.081    17.4 |   0.099    14.3 |   0.074    18.9 |     1.18x    1.22x |     0.89x    0.92x |  TRTLLM  |      13 |   0/  2/   0.00
    32 |   0.114    24.6 |   0.116    24.2 |   0.127    22.1 |   0.104    27.2 |     1.11x    1.09x |     0.91x    0.89x |  TRTLLM  |      21 |   0/  3/   1.00
    64 |   0.150    37.7 |   0.151    37.4 |   0.159    35.5 |   0.139    40.6 |     1.06x    1.06x |     0.93x    0.92x |  TRTLLM  |      29 |   0/  5/   2.00
   128 |   0.153    73.9 |   0.154    73.3 |   0.166    67.9 |   0.142    79.2 |     1.09x    1.08x |     0.93x    0.93x |  TRTLLM  |      30 |   0/  8/   3.00
   256 |   0.161   139.8 |   0.162   139.0 |   0.172   130.9 |   0.151   148.9 |     1.07x    1.06x |     0.94x    0.93x |  TRTLLM  |      32 |   1/ 11/   7.00
   512 |   0.163   276.4 |   0.168   269.1 |   0.177   254.8 |   0.161   280.3 |     1.08x    1.06x |     0.99x    0.96x |  TRTLLM  |      32 |   7/ 23/  14.00
  1024 |   0.166   543.3 |   0.187   482.0 |   0.190   475.0 |   0.306   295.1 |     1.14x    1.01x |     1.84x    1.63x |   W4A4   |      32 |  17/ 42/  28.50
  2048 |   0.176  1023.6 |   0.232   776.0 |   0.215   837.6 |   0.300   600.7 |     1.22x    0.93x |     1.70x    1.29x |   W4A4   |      32 |  38/ 74/  57.50
  4096 |   0.219  1650.7 |   0.360  1001.9 |   0.275  1311.3 |   0.301  1196.7 |     1.26x    0.76x |     1.38x    0.84x |   W4A4   |      32 |  81/147/ 117.00
---------------------------------------------------------------------------------------------------------------------------------------------------------------
Speedup > 1.0 means that CuTe DSL mode is faster than the comparison backend
Default activation, deterministic two-stage finalize — input-side
DeepSeek-V3 MoE Performance Benchmark
GPU: NVIDIA B200
CuteDSL API: Wrapper
Per-token activation: False
Initial activation quantization: False
CuteDSL modes: W4A4 and W4A16
Tensor parallelism simulation: TP=1
CUDA profiler capture: False
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: input

===============================================================================================================================================================
DeepSeek-V3 MoE Benchmark: CuteDSL W4A4/W4A16 vs CUTLASS vs TRTLLM (EP=8, TP=1)
===============================================================================================================================================================
Model: hidden=7168, intermediate=2048, experts=256, top_k=8
EP Config: 32 local experts (simulating 8-way parallelism)
TP Config: intermediate size 2048 (simulating 1-way parallelism)
CUDA Graph: enabled, CUPTI: enabled
Routing bias scale: 0.01 (larger values tend to create expert imbalance)
Timed initial activation quantization for FP4-activation backends: excluded; W4A16 consumes BF16 directly
CuteDSL finalize: deterministic two-stage
CuteDSL router weight placement: input
---------------------------------------------------------------------------------------------------------------------------------------------------------------
Tokens |  CuteDSL W4A4   |  CuteDSL W4A16  |     CUTLASS     |     TRTLLM      | Speedup vs CUTLASS | Speedup vs TRTLLM  |  Winner  | Active  |     Stats
       |      ms  TFLOPS |      ms  TFLOPS |      ms  TFLOPS |      ms  TFLOPS |     W4A4     W4A16 |     W4A4     W4A16 |          | experts | min/max/median
---------------------------------------------------------------------------------------------------------------------------------------------------------------
     1 |   0.031     2.9 |   0.037     2.4 |   0.045     2.0 |   0.026     3.4 |     1.47x    1.23x |     0.83x    0.70x |  TRTLLM  |       2 |   0/  1/   0.00
     2 |   0.042     4.2 |   0.044     4.0 |   0.055     3.2 |   0.035     5.0 |     1.32x    1.26x |     0.83x    0.80x |  TRTLLM  |       4 |   0/  1/   0.00
     4 |   0.058     6.1 |   0.059     6.0 |   0.073     4.8 |   0.052     6.8 |     1.26x    1.25x |     0.90x    0.88x |  TRTLLM  |       7 |   0/  1/   0.00
     8 |   0.073     9.6 |   0.073     9.7 |   0.087     8.1 |   0.064    10.9 |     1.19x    1.20x |     0.88x    0.89x |  TRTLLM  |      10 |   0/  2/   0.00
    16 |   0.083    17.0 |   0.081    17.3 |   0.099    14.3 |   0.074    19.0 |     1.19x    1.21x |     0.90x    0.91x |  TRTLLM  |      13 |   0/  2/   0.00
    32 |   0.114    24.7 |   0.116    24.4 |   0.128    22.1 |   0.104    27.1 |     1.12x    1.10x |     0.91x    0.90x |  TRTLLM  |      21 |   0/  3/   1.00
    64 |   0.149    37.8 |   0.150    37.7 |   0.159    35.5 |   0.138    40.9 |     1.06x    1.06x |     0.92x    0.92x |  TRTLLM  |      29 |   0/  5/   2.00
   128 |   0.152    73.9 |   0.154    73.1 |   0.166    68.1 |   0.142    79.1 |     1.09x    1.07x |     0.93x    0.92x |  TRTLLM  |      30 |   0/  8/   3.00
   256 |   0.161   140.1 |   0.162   138.8 |   0.175   129.1 |   0.152   148.3 |     1.08x    1.08x |     0.94x    0.94x |  TRTLLM  |      32 |   1/ 11/   7.00
   512 |   0.163   277.0 |   0.169   267.3 |   0.180   250.0 |   0.161   280.7 |     1.11x    1.07x |     0.99x    0.95x |  TRTLLM  |      32 |   7/ 23/  14.00
  1024 |   0.165   545.4 |   0.187   481.1 |   0.188   478.9 |   0.307   294.0 |     1.14x    1.00x |     1.85x    1.64x |   W4A4   |      32 |  17/ 42/  28.50
  2048 |   0.174  1037.2 |   0.232   778.0 |   0.215   837.2 |   0.303   594.5 |     1.24x    0.93x |     1.74x    1.31x |   W4A4   |      32 |  38/ 74/  57.50
  4096 |   0.210  1714.8 |   0.351  1028.6 |   0.278  1297.0 |   0.302  1192.8 |     1.32x    0.79x |     1.44x    0.86x |   W4A4   |      32 |  81/147/ 117.00
---------------------------------------------------------------------------------------------------------------------------------------------------------------
Speedup > 1.0 means that CuTe DSL mode is faster than the comparison backend

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.

Tokens W4A4 output (ms) W4A4 input (ms) W4A4 delta W4A16 output (ms) W4A16 input (ms) W4A16 delta
1 0.031 0.031 +0.00% 0.036 0.037 +2.78%
2 0.042 0.042 +0.00% 0.043 0.044 +2.33%
4 0.058 0.058 +0.00% 0.059 0.059 +0.00%
8 0.074 0.073 -1.35% 0.072 0.073 +1.39%
16 0.083 0.083 +0.00% 0.081 0.081 +0.00%
32 0.114 0.114 +0.00% 0.116 0.116 +0.00%
64 0.150 0.149 -0.67% 0.151 0.150 -0.66%
128 0.153 0.152 -0.65% 0.154 0.154 +0.00%
256 0.161 0.161 +0.00% 0.162 0.162 +0.00%
512 0.163 0.163 +0.00% 0.168 0.169 +0.60%
1024 0.166 0.165 -0.60% 0.187 0.187 +0.00%
2048 0.176 0.174 -1.14% 0.232 0.232 +0.00%
4096 0.219 0.210 -4.11% 0.360 0.351 -2.50%

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

    • Added an option to apply MoE router weights to the FC2 input instead of the output.
    • Exposed the setting through fused MoE APIs, wrappers, tracing, and benchmark command-line tools.
    • Added support across relevant BF16, NVFP4, and W4A16 execution paths, with output-side weighting remaining the default.
  • Tests

    • Expanded accuracy coverage for router-weight placement and multiple activation types.

@coderabbitai

coderabbitai Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Changes

The 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

Layer / File(s) Summary
Public configuration and benchmark wiring
flashinfer/fused_moe/api.py, flashinfer/fused_moe/cute_dsl/fused_moe.py, flashinfer/trace/templates/moe.py, benchmarks/*
Adds the option to execution APIs, trace inputs, benchmark CLIs, and startup reporting.
Runner and cache propagation
flashinfer/fused_moe/cute_dsl/tuner.py, flashinfer/fused_moe/runners.py
Stores the option, includes it in cache keys, validates backend support, and forwards it to NVFP4 and W4A16 runners.
NVFP4 activation weighting
flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py, flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
Validates and forwards token scales, then applies route weights to activated FC2 inputs before conversion.
W4A16 grouped GEMM and finalization
flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py, flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py, flashinfer/fused_moe/cute_dsl/blackwell/*finalize*
Loads route metadata for fused activation paths and avoids applying router weights twice.
Conditional unpermutation scaling
flashinfer/fused_moe/cute_dsl/moe_utils.py, csrc/moe_utils_binding.cu, csrc/nv_internal/tensorrt_llm/kernels/cuteDslKernels/moeUtils.*
Adds scaled and unscaled unpermutation kernel variants selected at runtime.
Reference and accuracy validation
tests/moe/utils.py, tests/moe/test_cute_dsl_fused_moe.py
Updates reference calculations and covers both placement modes with multiple activation and tile configurations.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

Suggested labels: op: gemm

Suggested reviewers: aleozlx, jiahanc, yzh119

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 64.10% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly and concisely summarizes the main change: input-side router-weight support for CuTe DSL W4A4 and W4A16 MoE.
Description check ✅ Passed The description explains the motivation, scope, implementation details, validation, benchmarks, related issue, and reviewer focus areas.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@ziang-and
ziang-and force-pushed the agent/sm100-router-weight-on-input branch 4 times, most recently from 12c1796 to 1258c31 Compare August 8, 2026 06:51
@ziang-and
ziang-and force-pushed the agent/sm100-router-weight-on-input branch from 1258c31 to 99336aa Compare August 8, 2026 06:55
@zianglih zianglih changed the title feat(moe): support input-side router weights in CuTe DSL NVFP4 MoE feat(moe): support input-side router weights in CuTe DSL W4A4/W4A16 MoE Aug 8, 2026
@zianglih
zianglih marked this pull request as ready for review August 8, 2026 09:19
@zianglih

zianglih commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

@coderabbitai review

@coderabbitai

coderabbitai Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor
⚠️ Action not completed

Already reviewed.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Cover input-side routing weights through the direct functional API.

The helper defaults apply_router_weight_on_input to False. The shown calls from test_numerical_accuracy and test_deterministic_finalize_numerical_accuracy do 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 win

Document the packed and scalar route-scaling paths.

This hot path selects packed f32x2 operations or a scalar loop with vectorized_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 win

Document 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

📥 Commits

Reviewing files that changed from the base of the PR and between b1d9585 and 99336aa.

📒 Files selected for processing (19)
  • benchmarks/bench_cute_dsl_moe_distributed.py
  • benchmarks/bench_moe_deepseek.py
  • csrc/moe_utils_binding.cu
  • csrc/nv_internal/tensorrt_llm/kernels/cuteDslKernels/moeUtils.cu
  • csrc/nv_internal/tensorrt_llm/kernels/cuteDslKernels/moeUtils.h
  • flashinfer/fused_moe/api.py
  • flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
  • flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py
  • flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py
  • flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.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/fused_moe.py
  • flashinfer/fused_moe/cute_dsl/moe_utils.py
  • flashinfer/fused_moe/cute_dsl/tuner.py
  • flashinfer/fused_moe/runners.py
  • flashinfer/trace/templates/moe.py
  • tests/moe/test_cute_dsl_fused_moe.py
  • tests/moe/utils.py

Comment thread flashinfer/fused_moe/api.py
Comment thread flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py
Comment thread flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py
Comment thread flashinfer/fused_moe/cute_dsl/fused_moe.py
@zianglih

zianglih commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

Review-summary follow-up for aadf875:

  • The existing direct functional fused and deterministic accuracy matrices now explicitly exercise apply_router_weight_on_input=True. Together they cover W4A4 per-tensor, W4A4 per-token, W4A16, SwiGLU/ReLU2, and both finalize paths without adding a duplicate test family or increasing collection size; output-side direct coverage remains in the existing padding/EP paths.
  • A targeted direct-functional B200 slice passed all 10 selected cases.
  • Placement rationale is documented once at the orchestration boundary and once at the complementary W4A16 finalize branch. I did not duplicate it in the generic unpermute implementation because apply_topk_scales=False is also a valid generic unweighted-reduction contract.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
tests/moe/test_cute_dsl_fused_moe.py (1)

1475-1475: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add 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

📥 Commits

Reviewing files that changed from the base of the PR and between 99336aa and aadf875.

📒 Files selected for processing (5)
  • flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py
  • flashinfer/fused_moe/cute_dsl/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py
  • flashinfer/fused_moe/cute_dsl/fused_moe.py
  • flashinfer/fused_moe/runners.py
  • tests/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

@zianglih

Copy link
Copy Markdown
Contributor Author

need to rebase

@zianglih
zianglih marked this pull request as draft August 17, 2026 21:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants