Skip to content

[Kernel] Enable CUTLASS MXFP4 W4A4 MoE on SM12x - #51158

Open
ayrnb wants to merge 2 commits into
vllm-project:mainfrom
ayrnb:feat/sm120-mxfp4-w4a4-moe
Open

[Kernel] Enable CUTLASS MXFP4 W4A4 MoE on SM12x#51158
ayrnb wants to merge 2 commits into
vllm-project:mainfrom
ayrnb:feat/sm120-mxfp4-w4a4-moe

Conversation

@ayrnb

@ayrnb ayrnb commented Aug 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Purpose

Enable CUTLASS MXFP4 W4A4 MoE on SM12x (GeForce / RTX PRO Blackwell). Previously, MXFP4 experts-quant and grouped-GEMM kernels were compiled and gated for SM10x/11x only, so MXFP4 W4A4 MoE silently fell back to Marlin W4A16 on consumer / prosumer Blackwell parts.

Changes

  • mxfp4_blockwise_moe_kernel.cu: Add Mxfp4GroupGemmArchConfig<Arch> trait.
    SM10x/11x keeps its dedicated 1-SM MXFP4 schedule unchanged. SM12x uses
    KernelScheduleAuto / EpilogueScheduleAuto / EpilogueTileAuto —
    the CUTLASS SM120 builder auto-selects the cooperative Ptr-Array
    block-scaled schedule from the pointer-typed StrideA and picks an
    epilogue tile that fits the tighter SMEM budget.

  • mxfp4_experts_quant.cu: Widen capability guard >= 100 && < 120 →
    >= 100 && < 130. SM100/SM120 share the same block-scaled SF layout.

  • CMakeLists.txt: Add both MXFP4 sources to FP4_SM120_SRCS.

  • test_mxfp4_moe.py: Replace hardcoded is_device_capability_family(100)
    skip with runtime mxfp4_experts_quant_supported() query, same predicate
    CutlassExpertsMxfp4 uses.

Test Result

Environment: RTX PRO 5000 (SM120), CUDA 13.0, Qwen3-30B-A3B-Instruct-2507 MXFP4

# serve
vllm serve Qwen3-30B-A3B-Instruct-2507-MXFP4 \
    --tensor-parallel-size 1 \
    --max-num-seqs 256 \
    --attention-backend FLEX_ATTENTION \
    --linear-backend humming \
    --max-num-batched-tokens 65536 \
    --gpu-memory-utilization 0.85 \
    --max-model-len 4096 \
    --no-enable-prefix-caching

# client (prefill)
vllm bench serve --backend vllm --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 1 \
    --num-prompts 256 \
    --base-url http://127.0.0.1:8000/ \
    --max-concurrency $BS --request-rate $BS

# client (decode)
vllm bench serve --backend vllm --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 128 \
    --num-prompts 256 \
    --base-url http://127.0.0.1:8000/ \
    --max-concurrency $BS --request-rate $BS

Throughput

Prefill

Concurrency TTFT_mean(ms) Speedup
W4A16 W4A4
1 79.06 69.53 1.14x
2 82.60 72.45 1.14x
4 91.72 79.64 1.15x
8 133.43 108.31 1.23x
16 582.27 353.79 1.65x
32 1719.19 1598.68 1.08x
64 3339.33 3124.35 1.07x

E2E

Concurrency Req/s Out tok/s Total tok/s TTFT (ms) TPOT (ms)
1 w4a4 0.61 78.54 706.88 69.21 12.28
origin 0.68 87.66 788.97 82.59 10.84
2 w4a4 0.75 96.38 867.43 51.40 20.50
origin 0.83 105.78 952.03 44.47 18.69
4 w4a4 0.91 116.23 1046.05 90.22 33.94
origin 0.96 123.24 1109.20 82.83 32.02
6 w4a4 0.98 125.28 1127.52 123.79 46.98
origin 1.02 130.55 1174.96 116.50 45.10
16 w4a4 1.13 145.22 1307.00 273.55 108.24
origin 1.15 147.02 1323.19 270.33 106.91
32 w4a4 2.17 277.67 2499.03 305.72 111.20
origin 2.18 279.36 2514.21 305.81 110.47
64 w4a4 3.94 504.60 4541.38 344.31 116.62
origin 3.99 510.92 4598.29 340.71 115.03

Accuracy

  • tests/kernels/moe/test_mxfp4_moe.py
image
  • python tests/evals/gsm8k/gsm8k_eval.py
    Running GSM8K evaluation: 1319 questions, 5-shot
    Evaluating: 100%|███████████████████████████████████████████████████████████████████████████████████| 1319/1319 [02:13<00:00, 9.88it/s]
    Results:
    Accuracy: 0.875
    Invalid responses: 0.000
    Total latency: 133.454 s
    Questions per second: 9.884
    Total output tokens: 200389
    Output tokens per second: 1501.558
Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

@ayrnb
[Kernel] Enable CUTLASS MXFP4 W4A4 MoE on SM12x
19932fd 
The MXFP4 experts-quant and grouped-GEMM kernels were compiled and gated
for SM10x/11x only, so MXFP4 W4A4 MoE fell back to Marlin W4A16 on
GeForce / RTX PRO Blackwell. CUTLASS resolves an identical block-scaled
scale-factor layout for arch::Sm100 and arch::Sm120, so a single scale
swizzle serves both families and both kernels can cover SM12x.

- Parameterise the grouped GEMM on ArchTag through a new
  Mxfp4GroupGemmArchConfig trait. SM10x/11x keeps its dedicated 1-SM
  MXFP4 Ptr-Array schedule and fixed epilogue tile unchanged. SM12x has
  no MXFP4-specific Ptr-Array schedule tag, so it lets the builder derive
  the cooperative Ptr-Array block-scaled schedule from the pointer-typed
  StrideA and pick an epilogue tile that fits the tighter SMEM budget of
  a TMEM-less part.
- Widen the experts-quant capability guard to include SM12x.
- Compile both MXFP4 sources for SM12x.
- Gate the MXFP4 MoE kernel tests on the compiled-capability query rather
  than a hardcoded SM family, so they exercise every Blackwell variant
  the wheel was built for instead of silently skipping.

Signed-off-by: ayrnb <641876696@qq.com>
claude[bot]
claude Bot reviewed Aug 5, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

Harry-Chen
Harry-Chen reviewed Aug 5, 2026
View reviewed changes
Comment thread CMakeLists.txt Outdated
@Harry-Chen

Harry-Chen commented Aug 5, 2026

Copy link
Copy Markdown
Member

Thanks! Have you compared the performance before / after switching implementation?

@ayrnb
Lint code
cb5e1f3 
Signed-off-by: ayrnb <641876696@qq.com>
@ayrnb

ayrnb commented Aug 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Thanks! Have you compared the performance before / after switching implementation?

Yes — measured on RTX PRO 5000 (SM120), CUDA 13.0, Qwen3-30B-A3B-Instruct-2507 MXFP4. I've updated the PR description with the full tables.

TL;DR: prefill TTFT improves 7-65%. The decode workload shows a small regression (~1-2%). Still planning to tune it further, especially for decode.

@randomvariable randomvariable mentioned this pull request Aug 19, 2026
Open
randomvariable added a commit to randomvariable/vllm that referenced this pull request Aug 19, 2026
@randomvariable OMP Agent
feat(sm120): GB10 general-purpose serving ports (issue #48, part 1)
29fdcd4 
Port five upstream consumer-Blackwell capabilities so non-DSv4
workloads can serve on GB10 (sm_121, same family as RTX 5090 /
RTX PRO 6000 sm_120):

1. TRTLLM fp8 MoE device gate -> SM_12x (upstream vllm-project#43911).
2. OAI Triton MoE capability window -> (9,0) <= cap < (13,0),
   covering SM120/SM121; kernels are pure Triton JIT with no
   sm90/sm10x-only instructions (upstream vllm-project#41028).
3. CUTLASS grouped GEMM w8a8 for SM120 (upstream vllm-project#43814): new
   grouped_mm_c3x_sm120.cu in FP4_SM120_SRCS, ENABLE_CUTLASS_MOE_SM120
   dispatch + group_gemm_supported range in entry and _custom_ops.py.
   The CMake flag already existed (fork); this wires the kernel it
   promised.
4. NVFP4 grouped MoE pingpong schedule at per-expert M >= 512
   (upstream vllm-project#46481): template<bool UsePingpong> on the sm120 entry;
   the fork's batch-invariance static_assert now pins each
   instantiation (cooperative stays the default/batch-invariant path).
5. CUTLASS MXFP4 W4A4 MoE on SM12x (upstream vllm-project#51158): mxfp4 kernels
   join FP4_SM120_SRCS; kernel/test files byte-match the PR head.

Deferred (tracked on #48): vllm-project#46329 NVFP4 KV-cache enablement — fork
files diverge 300-700 lines (nvfp4_ds_mla ABI, envs, flashinfer
backend); needs a hand-merge plus GB10 GPU validation. vllm-project#50288's
V-scale-swizzle fix is already contained in 46329's kernel.

Verified: py files compile; pre-commit clean (clang-format applied);
test_mxfp4_moe skips in this venv (no GPU/triton_kernels) — GPU
validation rides the next Spark image build.

Co-authored-by: OMP Agent <noreply@omp.local>
Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
randomvariable added a commit to randomvariable/vllm that referenced this pull request Aug 23, 2026
@randomvariable OMP Agent
feat(sm120): GB10 general-purpose serving ports (issue #48, part 1)
e1ba7a0 
Port five upstream consumer-Blackwell capabilities so non-DSv4
workloads can serve on GB10 (sm_121, same family as RTX 5090 /
RTX PRO 6000 sm_120):

1. TRTLLM fp8 MoE device gate -> SM_12x (upstream vllm-project#43911).
2. OAI Triton MoE capability window -> (9,0) <= cap < (13,0),
   covering SM120/SM121; kernels are pure Triton JIT with no
   sm90/sm10x-only instructions (upstream vllm-project#41028).
3. CUTLASS grouped GEMM w8a8 for SM120 (upstream vllm-project#43814): new
   grouped_mm_c3x_sm120.cu in FP4_SM120_SRCS, ENABLE_CUTLASS_MOE_SM120
   dispatch + group_gemm_supported range in entry and _custom_ops.py.
   The CMake flag already existed (fork); this wires the kernel it
   promised.
4. NVFP4 grouped MoE pingpong schedule at per-expert M >= 512
   (upstream vllm-project#46481): template<bool UsePingpong> on the sm120 entry;
   the fork's batch-invariance static_assert now pins each
   instantiation (cooperative stays the default/batch-invariant path).
5. CUTLASS MXFP4 W4A4 MoE on SM12x (upstream vllm-project#51158): mxfp4 kernels
   join FP4_SM120_SRCS; kernel/test files byte-match the PR head.

Deferred (tracked on #48): vllm-project#46329 NVFP4 KV-cache enablement — fork
files diverge 300-700 lines (nvfp4_ds_mla ABI, envs, flashinfer
backend); needs a hand-merge plus GB10 GPU validation. vllm-project#50288's
V-scale-swizzle fix is already contained in 46329's kernel.

Verified: py files compile; pre-commit clean (clang-format applied);
test_mxfp4_moe skips in this venv (no GPU/triton_kernels) — GPU
validation rides the next Spark image build.

Co-authored-by: OMP Agent <noreply@omp.local>
Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
randomvariable added a commit to randomvariable/vllm that referenced this pull request Aug 25, 2026
@randomvariable OMP Agent
feat(sm120): GB10 general-purpose serving ports (issue #48, part 1)
0629ead 
Port five upstream consumer-Blackwell capabilities so non-DSv4
workloads can serve on GB10 (sm_121, same family as RTX 5090 /
RTX PRO 6000 sm_120):

1. TRTLLM fp8 MoE device gate -> SM_12x (upstream vllm-project#43911).
2. OAI Triton MoE capability window -> (9,0) <= cap < (13,0),
   covering SM120/SM121; kernels are pure Triton JIT with no
   sm90/sm10x-only instructions (upstream vllm-project#41028).
3. CUTLASS grouped GEMM w8a8 for SM120 (upstream vllm-project#43814): new
   grouped_mm_c3x_sm120.cu in FP4_SM120_SRCS, ENABLE_CUTLASS_MOE_SM120
   dispatch + group_gemm_supported range in entry and _custom_ops.py.
   The CMake flag already existed (fork); this wires the kernel it
   promised.
4. NVFP4 grouped MoE pingpong schedule at per-expert M >= 512
   (upstream vllm-project#46481): template<bool UsePingpong> on the sm120 entry;
   the fork's batch-invariance static_assert now pins each
   instantiation (cooperative stays the default/batch-invariant path).
5. CUTLASS MXFP4 W4A4 MoE on SM12x (upstream vllm-project#51158): mxfp4 kernels
   join FP4_SM120_SRCS; kernel/test files byte-match the PR head.

Deferred (tracked on #48): vllm-project#46329 NVFP4 KV-cache enablement — fork
files diverge 300-700 lines (nvfp4_ds_mla ABI, envs, flashinfer
backend); needs a hand-merge plus GB10 GPU validation. vllm-project#50288's
V-scale-swizzle fix is already contained in 46329's kernel.

Verified: py files compile; pre-commit clean (clang-format applied);
test_mxfp4_moe skips in this venv (no GPU/triton_kernels) — GPU
validation rides the next Spark image build.

Co-authored-by: OMP Agent <noreply@omp.local>
Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
randomvariable added a commit to randomvariable/vllm that referenced this pull request Aug 29, 2026
@randomvariable OMP Agent
feat(sm120): GB10 general-purpose serving ports (issue #48, part 1)
b16d9fc 
Port five upstream consumer-Blackwell capabilities so non-DSv4
workloads can serve on GB10 (sm_121, same family as RTX 5090 /
RTX PRO 6000 sm_120):

1. TRTLLM fp8 MoE device gate -> SM_12x (upstream vllm-project#43911).
2. OAI Triton MoE capability window -> (9,0) <= cap < (13,0),
   covering SM120/SM121; kernels are pure Triton JIT with no
   sm90/sm10x-only instructions (upstream vllm-project#41028).
3. CUTLASS grouped GEMM w8a8 for SM120 (upstream vllm-project#43814): new
   grouped_mm_c3x_sm120.cu in FP4_SM120_SRCS, ENABLE_CUTLASS_MOE_SM120
   dispatch + group_gemm_supported range in entry and _custom_ops.py.
   The CMake flag already existed (fork); this wires the kernel it
   promised.
4. NVFP4 grouped MoE pingpong schedule at per-expert M >= 512
   (upstream vllm-project#46481): template<bool UsePingpong> on the sm120 entry;
   the fork's batch-invariance static_assert now pins each
   instantiation (cooperative stays the default/batch-invariant path).
5. CUTLASS MXFP4 W4A4 MoE on SM12x (upstream vllm-project#51158): mxfp4 kernels
   join FP4_SM120_SRCS; kernel/test files byte-match the PR head.

Deferred (tracked on #48): vllm-project#46329 NVFP4 KV-cache enablement — fork
files diverge 300-700 lines (nvfp4_ds_mla ABI, envs, flashinfer
backend); needs a hand-merge plus GB10 GPU validation. vllm-project#50288's
V-scale-swizzle fix is already contained in 46329's kernel.

Verified: py files compile; pre-commit clean (clang-format applied);
test_mxfp4_moe skips in this venv (no GPU/triton_kernels) — GPU
validation rides the next Spark image build.

Co-authored-by: OMP Agent <noreply@omp.local>
Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
randomvariable added a commit to randomvariable/vllm that referenced this pull request Aug 29, 2026
@randomvariable OMP Agent
feat(sm120): GB10 general-purpose serving ports (issue #48, part 1)
c4d4dcd 
Port five upstream consumer-Blackwell capabilities so non-DSv4
workloads can serve on GB10 (sm_121, same family as RTX 5090 /
RTX PRO 6000 sm_120):

1. TRTLLM fp8 MoE device gate -> SM_12x (upstream vllm-project#43911).
2. OAI Triton MoE capability window -> (9,0) <= cap < (13,0),
   covering SM120/SM121; kernels are pure Triton JIT with no
   sm90/sm10x-only instructions (upstream vllm-project#41028).
3. CUTLASS grouped GEMM w8a8 for SM120 (upstream vllm-project#43814): new
   grouped_mm_c3x_sm120.cu in FP4_SM120_SRCS, ENABLE_CUTLASS_MOE_SM120
   dispatch + group_gemm_supported range in entry and _custom_ops.py.
   The CMake flag already existed (fork); this wires the kernel it
   promised.
4. NVFP4 grouped MoE pingpong schedule at per-expert M >= 512
   (upstream vllm-project#46481): template<bool UsePingpong> on the sm120 entry;
   the fork's batch-invariance static_assert now pins each
   instantiation (cooperative stays the default/batch-invariant path).
5. CUTLASS MXFP4 W4A4 MoE on SM12x (upstream vllm-project#51158): mxfp4 kernels
   join FP4_SM120_SRCS; kernel/test files byte-match the PR head.

Deferred (tracked on #48): vllm-project#46329 NVFP4 KV-cache enablement — fork
files diverge 300-700 lines (nvfp4_ds_mla ABI, envs, flashinfer
backend); needs a hand-merge plus GB10 GPU validation. vllm-project#50288's
V-scale-swizzle fix is already contained in 46329's kernel.

Verified: py files compile; pre-commit clean (clang-format applied);
test_mxfp4_moe skips in this venv (no GPU/triton_kernels) — GPU
validation rides the next Spark image build.

Co-authored-by: OMP Agent <noreply@omp.local>
Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Harry-Chen Harry-Chen Harry-Chen left review comments
@claude claude[bot] claude[bot] left review comments
@mgoin mgoin Awaiting requested review from mgoin mgoin is a code owner
@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner
@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner
@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner
@zyongye zyongye Awaiting requested review from zyongye zyongye is a code owner
@AndreasKaratzas AndreasKaratzas Awaiting requested review from AndreasKaratzas AndreasKaratzas is a code owner
@LucasWilkinson LucasWilkinson Awaiting requested review from LucasWilkinson LucasWilkinson is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

ci/build nvidia

Projects

NVIDIA
Status: No status

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ayrnb @Harry-Chen