feat(pytorch): add CuTeDSL W4A16 MoE backend

[qescccczmr]

Motivation

Add a Hopper CuTe DSL W4A16 provider for compressed-tensors MoE without expanding packed INT4 weights into a persistent BF16 cache.

Modification

• Implement BF16 WGMMA RS with register-side INT4 dequantization, shared-memory BF16 activations and FP32 accumulation.
• Add two-stage TMA buffering, small-batch split-K and route-count-based tile selection.
• Keep FusedMoEW4A16BuildSpec and add provider selection through LMDEPLOY_W4A16_MOE_BACKEND=auto|triton|cute, defaulting to auto. auto prefers CuTe when the build-time format, device and optional kernel imports are compatible, and DeepEP is available when EP > 1; otherwise it selects the existing Triton builder. Explicit cute is strict; explicit triton retains the existing Triton path.
• Support local experts and EP through the existing DeepEP normal/low-latency dispatch and combine paths, including invalid routes and masked expert counts.

Production changes remain limited to three files: the CuTe kernel, CUDA compressed-tensors MoE backend and environment configuration. Three related test files now cover the legacy builder, provider compatibility and Hopper CuTe numerical/graph behavior. Communication optimizations, router changes, documentation, unrelated tests and benchmarks remain excluded.

BC-breaking

The default is auto: unsupported CuTe environments retain the existing Triton builder without requiring a new environment setting. CuTe requires Hopper SM90, BF16 activations/output, INT4 group-size 32 and optional nvidia-cutlass-dsl / cuda-python dependencies. Explicit cute remains strict. Both providers require DeepEP for EP > 1; automatic local-kernel selection does not remove that communication dependency or hide invalid EP configuration errors.

Use cases

Select the CuTe provider with LMDEPLOY_W4A16_MOE_BACKEND=cute for supported compressed-tensors W4A16 MoE models on Hopper. No TP all-reduce or DeepEP communication kernel is replaced by this PR.

Benchmark

LMDeploy native Triton W4A16 baseline

• Hardware/software: one NVIDIA H200, PyTorch 2.13.0+cu129, Triton 3.7.1, CuTe DSL 4.4.2, CUDA Python 13.3.1.
• Workload: Kimi TP8-shard shape, 384 experts, hidden size 7168, per-shard intermediate size 256, top-k 8, INT4 group-size 32. Both paths receive the same packed weights, BF16 activations/scales, FP32 routing weights and INT32 expert IDs, generated with seed 20260908; renormalize=False.
• Timing: complete local MoE, including route preparation where needed, gate/up GEMM, SiLU, down GEMM and weighted reduction. Warm CUDA Graph replay via triton.testing.do_bench_cudagraph(rep=150), no profiler; three samples per shape with alternating implementation order, taking the median of each implementation's samples. Each sample is the mean of ten graph replays. Compilation, input generation, router score/top-k computation and communication are excluded.

Tokens	CuTe MoE latency (us)	Native Triton MoE latency (us)	Triton / CuTe	Latency reduction
1	22.690	79.698	3.512x	71.53%
8	91.890	335.377	3.650x	72.60%
32	269.418	771.207	2.862x	65.07%
256	604.272	1402.029	2.320x	56.90%
1024	844.461	1642.910	1.946x	48.60%
2048	1280.327	3015.663	2.355x	57.54%

CuTe is 1.946x–3.650x faster than the existing LMDeploy Triton implementation in this sweep. All outputs are finite; maximum output NRMSE versus Triton is 0.003859 (the benchmark checks < 0.02). This is a numerical smoke check, not a model-quality evaluation or bitwise-equivalence claim.

This is a single-GPU local-MoE benchmark at a TP8 shard shape, not eight-rank TP8 service throughput or EP dispatch/combine performance. Benchmark scripts and raw results remain local, excluded from the PR's code diff. Do not conflate this native-Triton speedup with the separate SGLang Marlin comparison below.

SGLang Marlin comparison

• Hardware: one NVIDIA H200; workload: Kimi TP8-shard shape, 384 experts, hidden size 7168, per-shard intermediate size 256, top-k 8, INT4 group-size 32, BF16 activations/scales.
• Software: PyTorch 2.13.0+cu129, Triton 3.7.1, CuTe DSL 4.4.2, CUDA Python 13.3.1. Reference: local SGLang Marlin checkout based on 9e692c9216c3b5e5c443fecf6b995700eb68d2e4. That checkout has other local changes; the Marlin wrapper, GEMM, repacking and scale-permutation Python files are unmodified.
• Both implementations use identical synthetic weights, activations and precomputed routes, with seed 20260908. Marlin weight repacking is setup-only.
• Timing: warm CUDA Graph replay using triton.testing.do_bench_cudagraph(rep=150), no profiler; median of three runs, each reporting the mean of ten graph replays. The measured region includes route layout preparation, gather, gate/up GEMM, SiLU, down GEMM and weighted reduction. Router score/top-k computation, compilation, repacking and communication are excluded.

Tokens	CuTe MoE latency (us)	SGLang Marlin MoE latency (us)	Marlin / CuTe	CuTe latency change
1	22.750	24.005	1.055x	-5.23%
8	91.843	80.129	0.872x	+14.62%
32	269.476	213.688	0.793x	+26.11%
256	603.588	483.793	0.802x	+24.76%

Marlin / CuTe > 1 means CuTe is faster. CuTe currently wins only at one token in this sweep. Maximum output NRMSE versus Marlin across the completed runs is 0.008765; results are not bitwise identical.

These are single-GPU MoE microbenchmarks at a TP8 shard shape, not eight-rank TP8 service measurements or end-to-end throughput. A separate 1024-token attempt did not complete the Marlin timing and remained in CUDA synchronization; its benchmark process was stopped. No completed 1024/2048-token Marlin comparison or speedup is reported, and no root cause is asserted for that incomplete run. Benchmark scripts and raw results remain local and are not included in the code diff.

Validation

• 37 tests passed and one CuTe test module skipped with the device-capability probe mocked to SM80. This checks A100 provider-selection and collection behavior only; it is not a real A100 numerical run.
• The legacy hidden_dim=1 builder test explicitly checks auto/triton. Explicit cute still rejects unsupported hardware; CuTe numerical tests skip non-Hopper devices before importing optional CuTe dependencies. The default is now auto following review.
• On 2026-09-16, 141 tests passed on H200 in the local working tree after the review changes, covering provider configuration, Triton/CuTe W4A16 kernels and DeepEP regressions. New coverage includes DeepEP-aware compatibility, strict missing-DeepEP errors, NaN-poisoned empty split-K partitions, and full MoE CUDA Graph replay at K=1056. Before changing the kernel, all 12 targeted split/graph cases also passed: the existing epilogue already writes zero accumulators for empty K partitions. The kernel now explicitly clamps their tile count to zero and documents this invariant; no full-buffer zeroing launch was added. Historical benchmark numbers above have not been rerun for this change.
• Ruff, docformatter and git diff --check pass.
• The 2026-09-16 mocked-SM80 provider/collection check passed 45 tests, with the CuTe numerical module skipped. This is not a real A100 GPU run.
• Full upstream CI and the dependency/platform matrix still require verification. No fresh TP8 service run or end-to-end SGLang speedup is claimed.

Remaining CI / compatibility issues

• The previous PR CI run had 12 failures: one builder regression fixed by the test update above, plus 11 response-format/tool-parser failures also present on main. See the previous PR job and the main-branch job.
• The earlier passing base job used XGrammar 0.2.3; both failing jobs used 0.2.6. These observations point to a separate dependency-compatibility issue. This CuTe PR does not change XGrammar pins or parser behavior, and fixing its builder regression alone does not guarantee a green full CI run.
• The default-provider review concern is addressed by selecting auto; explicit cute continues to require a compatible environment.

[Copilot]

🟡 Changes recommended

Address the split-K uninitialized target issue and ensure auto checks DeepEP availability before selecting CuTe.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds a Hopper CuTe DSL W4A16 compressed-tensors MoE backend with configurable provider selection and extensive kernel, routing, and compatibility tests.

Changes:

• Adds CuTe WGMMA/TMA kernels with INT4 dequantization and split-K support.
• Adds auto|triton|cute provider selection, defaulting to CuTe.
• Adds numerical, graph, EP, routing, and provider-selection tests.

File summaries

File	Description
tests/pytorch/kernel/test_compressed_tensors_w4a16.py	Legacy provider and builder coverage
tests/pytorch/kernel/test_compressed_tensors_w4a16_cute.py	CuTe numerical, graph, routing, and EP tests
tests/pytorch/config/test_w4a16_moe_backend.py	Provider compatibility and selection tests
lmdeploy/pytorch/kernels/cuda/compressed_tensors_w4a16_cute.py	CuTe W4A16 kernels and MoE execution paths
lmdeploy/pytorch/envs.py	Backend environment configuration
lmdeploy/pytorch/backends/cuda/moe/compressed_tensors.py	Provider selection and CuTe/DeepEP integration

Review details

Suppressed comments (1)

lmdeploy/pytorch/backends/cuda/moe/compressed_tensors.py:390

• When spec.ep_size > 1, _supports_cute can still return True on a system where use_deepep is false, because it only probes the CuTe kernel import. auto then selects _build_fused_moe_cute, whose DeepEP constructor raises before the Triton branch is selected, contradicting the documented compatibility-based provider selection. Include the DeepEP availability check in _supports_cute (while keeping explicit cute strict).

    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9:
        return False
    try:
        from lmdeploy.pytorch.kernels.cuda.compressed_tensors_w4a16_cute import fused_moe_w4a16_cute  # noqa: F401

• Files reviewed: 6/6 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[grimoire]

Better take a look