perf: reduce split W4A16 routing and output overhead

[zianglih]

📌 Description

@HumansAnd

Reduce split CuTe DSL W4A16 routing and bounds overhead, and use complete output subtiles for ordinary scratch stores.

• Routing: mark top-k cutlass.Constexpr, include it in the private kernel cache key, and omit the runtime argument. This removes runtime division from routing-scale lookup and output scatter.
• Output bounds: derive whether M covers complete physical CTA clusters, cache that Boolean, and remove feature predicates/clipping only for those shapes. Route guards and the partial/padded fallback remain. Checking the full cluster span is necessary for a padded peer CTA.
• Scratch stores: for non-fused output, use the existing plain TMA store for each fully valid subtile before a route boundary; retain predicated stores for the tail. Expert output ownership and FP32 combine remain unchanged.
• Contract: public APIs, weight layouts, numerical operations, two weight-dequantization warpgroups and autotune tactics are unchanged. Matrix dimensions and token counts remain dynamic. Non-fused finalization stores BF16 expert outputs, then combines in FP32 in fixed top-k order before one BF16 output cast; it uses no atomic or TMA reduction.
• Tests: extend existing route-boundary and padding test products with top-k 2/3, both finalize modes, full-prefix-plus-partial-tail cases and aligned→padded→aligned cache reuse. Preserve existing W4A4 cases, reference calculations and tolerances.

Environment and source

• Hardware/image: one B300 SXM6 AC GPU on an 8-GPU c2 node (hu-pdx-43); GPU0 exclusively used. Driver590.48.01; explicit container nvcr.io/nvidia/pytorch:26.05-py3.
• Software: Python3.12.3; Torch2.12.0a0+5aff3928d8.nv26.5.50603568 (runtime reports2.12.0a0+5aff3928d8.nv26.05), CUDA runtime13.2; CuTe DSL4.7.1; CUPTI Python13.4.0/library13.4.58; pytest9.1.1. System nvcc/ptxas is13.2.78. The DSL extension contains bundled compiler13.3.27 and emitted PTX identifies that version; these measurements do not establish system-ptxas13.2 code generation.
• Source: before 488ffbd9fe4455c3b4c3030cf668b5603c6698e4; current 88ff7db6a558e8acbc0cd964275410cf94e5c22c. GPU runs used the baseline checkout plus the exact candidate patch; measured implementation/test hashes match this commit. No source changes followed validation.
• Execution: CUDA_VISIBLE_DEVICES=0, TRTLLM_ENABLE_PDL=1, FLASHINFER_DIST_AWARE_AUTOTUNE=0; project-local JIT caches, isolated venv preserving the image's Torch installation. Both sides use the same GPU/compiler stack.

Validation

• GPU tests: 66 passed: 48 route-boundary cases, 4 extended padding/cache cases, 2 simulated EP8 functional cases and 12 CUDA-graph cases. Full repository suite was not run. The grouped-weight fixture requires feature widths divisible by 128; W4A16 numerical coverage uses widths 256→384→256 to exercise aligned and padded clusters.
• Native: fixed GEMM2 tactic (128,64,256), cluster (1,1), M-major, N256/K512, 8 experts, PDL/CLC disabled. Top-k specialization removes runtime div.s32/reciprocal sequences for top-k 2/3 (1640→1584 SASS instructions). At M256, bounds specialization further reduces the checked fused top-k2 case to 1568 and ordinary output from 1664→1624. At M288, both fallback SASS files are byte-identical to their immediate controls; these partial-width exports were compiled but not launched. The scratch-store change keeps the checked ordinary N64 instruction count at 1624; its N16 ordinary and N64 fused control SASS are byte-identical. Registers 128, stack 0 and local memory 0 remain unchanged.
• Determinism: six forced N64/N128/N192 × top-k2/3 tail cases, each repeated three times with PDL off and on, match the immediate control bit-for-bit (36 launches per source). Input hashes, unchanged references and zero absolute differences were checked. Ordinary native code uses plain cp.async.bulk.tensor / UTMASTG.2D, with no TMA reduction; the fixed-order FP32 combiner is unchanged.
• Profiling: Nsight Systems2026.2.1 and NCU2026.1.1 --set full captured current non-fused T2048 with its saved tactic, without retuning. FC2 uses grid2×1×74, cluster2, 512 threads, 128 registers/thread and222208 bytes dynamic shared memory (plus1024 driver bytes), with no local-memory accesses. NCU used40 kernel-replay passes; six CTC metrics were unavailable and separate native exports provide source line information. Instrumented durations are diagnostic, not benchmark evidence.
• Local checks: scoped pre-commit hooks, including mypy/Ruff, syntax and git diff --check passed.

66 passed, 28 deselected, 1420 warnings in 46.07s

Warnings are existing CuTe DSL deprecations. Reproduce with:

CUDA_VISIBLE_DEVICES=0 TRTLLM_ENABLE_PDL=1 FLASHINFER_DIST_AWARE_AUTOTUNE=0 \
python -m pytest \
  tests/moe/test_cute_dsl_fused_moe.py::TestCuteDslMoeW4A16::test_route_tile_boundary_accuracy \
  tests/moe/test_cute_dsl_fused_moe.py::TestCuteDslFusedMoeFunctional::test_finalize_handles_cluster_padding_and_partial_tiles \
  tests/moe/test_cute_dsl_fused_moe.py::TestExpertParallelism::test_functional_with_ep \
  tests/moe/test_cute_dsl_fused_moe.py::TestCuteDslMoEWrapper::test_wrapper_cuda_graph \
  -k 'w4a16 or test_finalize_handles_cluster_padding_and_partial_tiles' -v --tb=short

Final performance

• Workload: upstream DeepSeek-like benchmark, H7168/I2048, 256 global experts, top-k 8, 32 local experts (single-GPU simulated EP8, TP1), SwiGLU, BF16 activations/output and NVFP4 weights. Routing is timed; weight preparation/autotuning and inter-GPU communication are excluded.
• Timing: normal autotuning for every complete sweep; ten warmups and 100 cold-L2 CUDA-graph CUPTI iterations per row, reporting medians. Fused and non-fused modes are evaluated separately; non-fused is the primary optimization target for its deterministic FP32 combine. Tokens 1/2/4 are reported and excluded from retention.

Non-fused finalize (primary)

One saved upstream sweep versus the latest complete current-source sweep on the same unchanged node/stack. Speedup is Before/After. Geometric-mean speedup over 8–8192 is 1.0056× (0.55% lower latency); this is a small measured difference, not a statistical-significance claim. The scratch-store change repeatedly improves 1024–4096 against its immediate control; 8192 varied across repeats, so its final row does not establish a stable gain.

Tokens	Before (µs)	After (µs)	Speedup
1	38.0165	38.0800	0.9983×
2	44.1925	44.0485	1.0033×
4	58.4010	58.3690	1.0005×
8	70.9120	70.7520	1.0023×
16	81.5210	81.7130	0.9977×
32	113.2335	113.1375	1.0008×
64	150.2895	149.6980	1.0040×
128	152.2260	151.8095	1.0027×
256	159.8730	159.6970	1.0011×
512	165.0575	164.4815	1.0035×
1024	182.4180	181.3775	1.0057×
2048	223.9215	222.0660	1.0084×
4096	328.1150	325.9070	1.0068×
8192	643.3660	625.4940	1.0286×

Fused finalize

Two saved upstream controls and two sweeps at b90b34744491c36d0039f44cba664191f7bb0eb0 on the same unchanged GPU/compiler stack. The later scratch-store condition is excluded from the fused SwiGLU path; source review and native control identity support reusing these measurements. These are not fresh timings of the new commit. Speedup is sqrt((Before R1 * Before R2) / (After R1 * After R2)); no row minima or fixed-tactic substitutions.

• Result: geometric-mean speedup over 8–8192 is 1.0141× (1.39% lower latency); all 11 aggregate rows improve. The top-k change accounts for the demonstrated gain. Bounds specialization reduces instruction count; no additional end-to-end gain was established. Smaller differences and autotune choices remain subject to run variation; two sweeps do not establish statistical significance. Tokens 1/2/4 are reported and excluded from retention.

Tokens	Before R1 (µs)	Before R2 (µs)	After R1 (µs)	After R2 (µs)	Speedup
1	36.6720	37.4885	36.3840	36.3845	1.0191×
2	42.7365	44.1600	42.3845	42.3680	1.0252×
4	61.1360	59.2320	59.7120	59.6800	1.0080×
8	74.8965	76.0010	74.4650	74.6240	1.0121×
16	84.4805	84.4325	84.3360	81.5850	1.0182×
32	116.4490	115.7290	115.2010	116.3695	1.0026×
64	150.1130	150.1770	149.6980	149.6970	1.0030×
128	154.2255	154.2095	153.6650	153.6010	1.0038×
256	161.7455	161.5210	160.8975	160.8975	1.0046×
512	166.4020	166.7220	165.7615	165.5540	1.0055×
1024	186.6250	186.6580	184.1945	184.2890	1.0130×
2048	232.5460	232.8340	229.3140	228.9625	1.0155×
4096	362.3400	363.1075	353.6680	353.7470	1.0255×
8192	711.8145	700.0385	675.6870	666.5500	1.0519×

All four runs select the same GEMM1/GEMM2 tactic at 8192 tokens. Other rows can select different tactics under normal autotuning.

Reproduction

In the pinned container, run from the selected revision's checkout:

git submodule update --init --recursive 3rdparty/cutlass 3rdparty/cccl 3rdparty/spdlog
python -m venv --system-site-packages ../split-w4a16-venv
source ../split-w4a16-venv/bin/activate
python - <<'PYSETUP'
import importlib.metadata as metadata
from pathlib import Path
import sys
Path(sys.prefix, "constraints.txt").write_text(
    "torch==" + metadata.version("torch") + "\n"
)
PYSETUP
PIP_CONSTRAINT="$VIRTUAL_ENV/constraints.txt" BUILD_NVEP=0 \
python -m pip install --no-build-isolation -e '.[cu13]' \
  'nvidia-cutlass-dsl[cu13]==4.7.1' 'cupti-python==13.4.0' \
  'nvidia-cuda-cupti==13.4.58' 'pytest==9.1.1'
export TORCH_EXTENSIONS_DIR="$PWD/../split-perf-cache/torch"
export FLASHINFER_WORKSPACE_BASE="$PWD/../split-perf-cache/flashinfer"

Run this from each revision's repository root for each finalize mode. The published non-fused table uses one sweep per source; the fused table uses two. It calls the existing benchmark's W4A16 arm directly, preserving its input generation, routing, wrapper, tuner and timer while avoiding other backends. Release inputs between token counts; both benchmarked revisions use identical code below.

CUDA_VISIBLE_DEVICES=0 TRTLLM_ENABLE_PDL=1 FLASHINFER_DIST_AWARE_AUTOTUNE=0 \
PYTHONPATH=.:benchmarks python - <<'PYBENCH'
import gc
import torch
from cupti import cupti
import bench_moe_deepseek as bench

use_fused_finalize = False  # Repeat with True to measure fused finalize separately.
for tokens in (2**i for i in range(14)):
    rows, _ = bench._benchmark_single(
        n=tokens, warmup=10, iters=100,
        num_local=32, local_offset=0,
        use_cuda_graph=True, use_cupti=True,
        use_wrapper=True, do_autotune=True,
        use_fused_finalize=use_fused_finalize,
        profile_cuda=False, profile_backend="cute-dsl-w4a16",
        backends=("cutedsl",), autotune_cache=None,
    )
    assert len(rows) == 1 and rows[0].backend == "CuteDSL W4A16"
    print(use_fused_finalize, tokens, rows[0].latency_ms * 1000, flush=True)
    gc.collect()
    torch.cuda.synchronize()
    torch.cuda.empty_cache()
PYBENCH

🔍 Related Issues

Follow-up to #4048. The compile-time/division audit in the MegaMoE work (#5019) motivated checking the split kernel boundary.

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• 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.).

🔬 Experimental Track

• This PR is experimental: it adds or changes code under flashinfer/experimental/ and/or an @flashinfer_experimental_api. Tracking issue: #
  • The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release.
  • Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff).
  • Tests live in tests/experimental/ and were validated on the intended hardware; a runnable example is included.
  • Nothing is registered in flashinfer/aot.py, and no experimental backend is reachable from backend="auto" without FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an @flashinfer_experimental_api or naming a backend explicitly is itself the opt-in and needs no environment variable.)
  • Test scope declared below. The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.

# One target per line: a directory or a file. (A pytest ::selector is not
# supported -- the sharding runner cannot consume one.) Must be under
# tests/experimental/ and must exist. Delete these comment lines and add yours, e.g.
#
#   tests/experimental/test_my_backend.py
#   tests/experimental/my_backend/
#
# Declaring the whole tree (tests/experimental/) is allowed but means every
# experimental PR pays for every other feature's tests, in every matrix cell.

Reviewer Notes

Summary by CodeRabbit

• Bug Fixes

  • Improved accuracy for fused mixture-of-experts processing when token batches end at tile or cluster boundaries.
  • Corrected handling of partial tiles and cluster padding across supported hidden sizes.
  • Improved reliability for configurations using different top-k routing values.
  • Fixed boundary behavior for fused-finalize processing, including cases with and without the optimized finalize path.

• Tests

  • Expanded coverage for tile tails, cluster padding, hidden sizes, top-k settings, and fused-finalize configurations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 6e90fd83-d8f8-4898-b479-d77ef987352a

📥 Commits

Reviewing files that changed from the base of the PR and between 488ffbd and 88ff7db.

📒 Files selected for processing (3)

• flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py
• flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py
• tests/moe/test_cute_dsl_fused_moe.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

The W4A16 grouped GEMM path now caches cluster-aligned and top_k variants separately. The kernel uses cluster alignment during fused-finalize stores, and tests cover partial tiles, route tails, finalize modes, and hidden sizes.

Changes

W4A16 cluster alignment

Layer / File(s)	Summary
Compile-time kernel variants flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16.py, flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py	The code computes m_cluster_aligned, adds it and top_k to the kernel cache key, passes the flag to Sm100W4A16GroupedGemmKernel, and treats top_k as compile-time data.
Fused-finalize boundary handling flashinfer/fused_moe/cute_dsl/blackwell/moe_w4a16_kernel.py	Cluster-aligned paths skip M-boundary clamping. TMA storage checks each subtile boundary. SIMT stores apply the M-boundary predicate only when needed.
Boundary accuracy validation tests/moe/test_cute_dsl_fused_moe.py	Tests cover full and half route-tile tails, top_k values 2 and 3, both finalize modes, and multiple hidden sizes.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Refactor

Merge Risk: ⚪ Minimal · up to 88ff7

No confirmed correctness or stability issue remains in the changed kernel paths. The targeted boundary tests and validation support merging.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 10 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: reducing routing and output overhead in split W4A16 kernels.
Description check	✅ Passed	The description is detailed and covers the change, related issues, tests, validation, performance results, environment, and reproduction steps. The checklist accurately leaves full-suite testing unche…
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.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 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

• X
• Mastodon
• Reddit
• LinkedIn

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