Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 by lucifer1004 · Pull Request #30272 · sgl-project/sglang

lucifer1004 · 2026-07-06T14:48:24Z

Motivation

DeepSeek-V4-Flash (FP8 checkpoint, MXFP4 experts) could not be served on SM120 (Blackwell desktop/workstation, e.g. RTX PRO 6000). The failures encountered while bringing up TP2/TP4 on 2–4 GPUs were:

The first MoE forward crashed because Marlin dereferenced masked -1 expert blocks under EP.
TP2 could not allocate a viable KV cache because FP32 loader containers inflated expert scales by about 16 GB/rank at EP2.
The first extend crashed because the sparse-prefill kernel supports SM90a/SM100f, not SM120.
Prefill OOMed because the SM120 wrapper used the split-K decode kernel for arbitrary extend batches, allocating roughly 1 MB of scratch per token.
The indexer OOMed because the Torch fallback materialized gathered index KV in FP32 for the full chunk, while the TileLang kernel failed lowering.
The working Marlin MoE path was W4A16 even though DSv4 experts and the available FlashInfer SM120 kernel are natively MXFP4 × MXFP8.

This PR makes DSv4-Flash serve on SM120 at TP2/EP2 and makes the native FlashInfer W4A8 MoE path the SM120 default.

Modifications

FlashInfer MXFP4 MoE (SM120 default)

Add a DSv4 adapter for FlashInfer CUTLASS MoE. On SM120 it keeps routed-expert weights in packed MXFP4, interleaves only the E8M0 scale metadata once after loading, and dynamically quantizes activations to MXFP8. It does not create a W4A16 or dequantized expert-weight copy.
Extend the FlashInfer runner with the MXFP8 × MXFP4 call contract, including activation scales, global weight scale, SwiGLU clamp, and normalized top-k tensors.
Route --moe-runner-backend auto to flashinfer_mxfp4 for DeepSeek V4 on SM120. Explicit backend selections still take precedence.
Retain the SM90 FlashInfer mixed-input W4A16 path and add SM90/SM120 adapter coverage.
Bump flashinfer_python, flashinfer_cubin, and the Docker JIT-cache version from 0.6.12 to 0.6.14, which supplies the SM120 MXFP8 activation-scaling and sparse-MLA APIs used here.

Marlin MXFP4 MoE correctness and memory

marlin_template.h: count and skip expert_id == -1 blocks regardless of is_ep; the standard dispatcher masks non-local experts even when no expert_map is passed.
fused_marlin_moe.py: zero-initialize the aliased intermediate cache because skipped rows remain unwritten but are consumed by the weighted sum.
marlin_utils_fp4.py: release loader-format scale parameters after Marlin preparation.
mxfp4_marlin_moe.py: register loader scales at their checkpoint E8M0 dtype rather than inflating them to FP32.

Attention

deepseek_v4_backend.py: disable the unsupported sparse-prefill path on SM120 and route extend batches through FlashInfer sparse MLA.
flash_mla_sm120.py: call FlashInfer's dispatching entry. Decode-sized batches use split-K; larger extends use the generic paged-attention kernel without per-token split-K scratch.

Indexer

tilelang_kernel.py: make fp8_paged_mqa_logits staging buffers two-dimensional so TileLang's GEMM layout lowering can determine the row size.
server_args.py: enable the TileLang indexer by default on SM120 while retaining the Torch fallback.
indexer.py: chunk the Torch reference query dimension at 1024 tokens to bound fallback memory.

Accuracy Tests

Marlin EP standalone reproduction at DSv4 expert shapes (E=64/rank, H=4096, I=2048, top-k=6): the masked-ID crash is fixed, three runs are bit-stable, and output equals the is_ep=true reference.
E8M0 loader containers produce Marlin weights and scales bit-identical to the previous FP32-container preparation.
SM120 adapter tests verify that post-load preparation preserves the original packed weight parameters and that adapter output matches a direct FlashInfer call. SM90 coverage checks its architecture-specific preparation path.
TileLang indexer versus the Torch reference at decode (n=32, split_kv=8) and prefill (n=2048, split_kv=1): max absolute logit difference 2.3e-5, max relative difference 3.5e-7.
End-to-end TP2/EP2 SM120 full GSM8K: 1,319 samples, exact match 0.955269 (strict-match=0.956027). A 64-sample FlashInfer qualification scored 0.96875, matching the pre-change local Marlin baseline.

Speed Tests and Profiling

random-8k1k serving benchmark: 8,000 input / 1,000 output tokens, four requests per concurrency unit, streaming completions through AIPerf, DeepSeek-V4-Flash TP2/EP2 on 2× RTX PRO 6000 (SM120). The baseline is the successful local Marlin Inferlab run from before this FlashInfer work; the final column comes from the final FlashInfer Inferlab run with the same model, topology, hardware, and client workload. Server launch and JIT startup are outside the measurement window.

concurrency	FlashInfer output tok/s	Marlin output tok/s	change	FlashInfer mean TTFT	Marlin mean TTFT	change
1	53.9	50.9	+5.8%	1.00 s	1.17 s	-14.4%
2	97.8	92.8	+5.5%	1.79 s	2.10 s	-15.1%
4	167.1	156.7	+6.7%	2.91 s	3.42 s	-15.0%
8	255.7	236.8	+8.0%	4.91 s	5.78 s	-15.1%
16	306.5	280.2	+9.4%	10.47 s	12.07 s	-13.3%
32	309.5	283.3	+9.3%	52.13 s	57.47 s	-9.3%

All 252 requests in the final benchmark succeeded. Full GSM8K also completed successfully, and Inferlab verified server cleanup.

Load-end model memory at TP2/EP2 is approximately 75.06 GB/rank, down from 91.99 GB/rank before the E8M0 loader-container fix, leaving roughly 17 GB/rank more KV-pool headroom.

Checklist

Format and repository hooks with prek
SM90 and SM120 FlashInfer MXFP4 adapter tests added
Full GSM8K and TP2/EP2 serving benchmark completed
Documentation: n/a

CI States

Latest PR Test (Base): ⏳ Run #29632757211
Latest PR Test (Extra): 🚫 Run #29632757138

gemini-code-assist · 2026-07-06T14:48:27Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

Carry the SM120 enablement onto the latest upstream layout and make the D-LLM override test exercise a non-FlashInfer starting backend.

lucifer1004 · 2026-07-14T15:30:25Z

@b8zhong Resolved

Fridge003 · 2026-07-18T00:27:35Z

/rerun-failed-ci

Brings in sgl-project#30272 (SM120 DSv4 flashinfer_mxfp4 MoE runner + TP2 memfix, previously carried as a local patch) and the sglang.kernels relocation (RFC sgl-project#29630). Conflict notes: - server_args SM120 block: upstream's out-of-box TileLang-indexer defaults kept, but behind is_set guards so a DeepGEMM build with SM120 attention support can opt back into fp8_paged_mqa_logits; the FLASHMLA_SPARSE_PREFILL force-off is dropped (superseded by the arch gate at the backend dispatch site). - deepseek_v4 hc_combine: fused kernel kept, wrapped in upstream's use_symmetric_memory allocation for the NCCL symmetric all-reduce. - flash_mla_sm120 / mhc imports repointed to sglang.kernels.ops. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…gl-project#30272) into deepseek-base-optimization Brings the SM120 DSv4 PR branch (a979b6b, already merged with sgl-project/sglang main 7a03d30 which includes sgl-project#30272) into the EPv2 integration branch. Conflict notes: - deep_gemm.py: this branch carried early cherry-picks of the PR's SM120 fixes (inline TMA-align, pre-dedup swiglu dispatch); replaced by the PR's refined equivalents (_tma_align_packed_ue8m0 helper, _varlen clamp fallback). The EPv2 pre/post permute paths are kept as-is. - configurer.py: kept this branch's DEEPGEMM_BLACKWELL including SM120 (the EPv2/DeepEP fp8-dispatch gates read it); UE8M0 value is identical either way on SM120. - token_dispatcher/base.py + server_args.py: EPV2 and upstream ASCEND_TP enum/choice entries are unioned; the epv2 validation block is kept ahead of upstream's ascend_tp normalization. - metadata.py: fp4-indexer gating combined with the SM120 JIT-metadata gate and M-chunking. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…sgl-project#30272)

…device actually has Boot 11 failed on both architectures of this rig for two different reasons, and neither was a missing capability -- both were a dispatch that named one architecture where it meant "has this kernel". TP0 (5090, sm120) entered the sparse-prefill branch and hit "Sparse Attention Forward Kernel is only supported on SM90a and SM100f". Upstream already fixed exactly this in the merged PR sgl-project#30272 by adding "and not _is_sm120" to the predicate; our base predates it (git log -S on that file is empty, so this is an old base, not a fork regression). Adopted rather than rebuilt, but phrased as flash_mla_sparse_fwd_supported(device) -- the kernel states its own domain in its error text, so the gate says major in (9, 10) and covers Ampere with the same line. The dense branch it falls to is handed the SAME swa_page_indices / swa_topk_lengths the indexer produced: the kernel changes, the attention support does not. Two independent confirmations that this is lossless rather than an approximation: upstream sgl-project#32464 calls the SM120 dense route intended production behaviour, and our own test kit compares both branches against one _pure_torch_dsv4_combined_reference. TP1/TP2 (3080, sm86) died in dsv4/dequant_k_cache.py on "type fp8e4nv not supported in this architecture". The V4 KV layout is fp8-e4m3 BY ARCHITECTURE (448 fp8 nope + 64 bf16 rope + 7 ue8m0 scale bytes in a uint8 pool); --kv-cache-dtype cannot move it, the pool asserts store_dtype == uint8. Triton's NVIDIA backend admits fp8e4nv only from capability 8.9. The briefing named one file. The bug has four call sites, and fixing only the named one would have MOVED the failure rather than removed it -- the two KV *writers* run before any attention: * dequant_k_cache.py (read) -> manual decode * flash_mla_sm120_triton.py -> manual decode; this is the one that actually becomes load-bearing on sm86 once the sparse branch is left * index_buf_accessor.py (write) -> turned out to be a pure byte copy with no arithmetic between load and store, so both ends are viewed as uint8 unconditionally: byte-identical everywhere, no gate needed * quant_k_cache.py (write) -> a real encode. Rather than hand-roll round-to-nearest-even in Triton, the fallback stages the already-clamped float32 and lets torch do the cast, which sgl-project#262 proved bit-exact on sm86. f32 not bf16 staging: bf16 has 8 mantissa bits and would double-round. Third finding, from sgl-project#33194's report and confirmed here: the compile failure is triggered by the kernel ARGUMENT, not the load -- an fp8-typed tensor fails at the def line even if the pointer is never dereferenced. An in-kernel branch alone cannot fix it, so the host passes the uint8 view and a tl.constexpr selects the decode. nope_cache_view returns view and flag together, following sgl-project#192's pairing invariant: a caller cannot take one without the other. Also removed the module-level SGLANG_SM120_FLASHMLA_BACKEND read in flash_mla_sm120.py -- the exact sgl-project#343 anti-pattern, and its default resolves to a flashinfer kernel gated to cc 120/121, so on a mixed group every rank got whichever card imported first. Resolution is per device now; an explicitly set value still applies everywhere (a statement about the launch, not a probe), and is downgraded loudly, not silently, where the kernel cannot exist. TESTS (desk, CUDA_VISIBLE_DEVICES=99, no GPU touched) test_dsv4_fp8_triton_compat_417.py 12 tests / 262 subtests green. Runs the production @triton.jit decode under TRITON_INTERPRET=1 against a pure Python IEEE-754 reference that never touches a torch fp8 dtype (sgl-project#262's rule: a torch reference would share torch's own Ampere bug and pass vacuously), with an e5m2 cross-probe so a comparison of nothing fails. All 256 codes, signed zero at 0x80, both NaN codes, plus the whole paged kernel vs the torch ref. CAN-FAIL: bits & 7 -> bits & 6 gives 128 failed / 10 passed. test_dsv4_arch_dispatch_417.py 16 tests / 28 subtests green, incl. the two cards of this rig disagreeing inside one process. CAN-FAIL: _FLASH_MLA_CUDA_MAJORS -> (8, 9, 10) gives 7 failed / 15 passed. Regression: unit/layers/attention + unit/utils, 44 failed / 362 passed on this branch vs 44 failed / 334 passed at a56f33a (throwaway detached worktree). Failing set diff EMPTY; all 44 pre-existing "No CUDA GPUs are available". The +28 are these tests. ruff 1 finding on the touched files, the identical 1 at the base (F841, pre-existing). codespell, black, isort clean. mypy clean on both new modules. NOT proved at the desk, by construction: that these kernels COMPILE on a real sm86 card. The Triton interpreter never invokes the NVIDIA backend, which is where fp8e4nv is rejected. Gates for the card window are written down in docs/dev/PLAN_417_dsv4_arch_paths.md along with the full per-architecture coverage matrix, the upstream survey, and the open risks -- chief among them that the indexer (dsv4/indexer.py:667, DeepGEMM metadata) is NOT addressed here and may well be where boot 12 stops, and that sgl-project#33194's chain ends at FP4 routed experts, a wall we expect to miss because we serve GGUF.

…sgl-project#30272)

lucifer1004 requested review from AniZpZ, BBuf, Edwardf0t1, FlamingoPg, HaiShaw, b8zhong and ch-wan as code owners July 6, 2026 14:48

lucifer1004 requested review from 1am9trash, DarkSharpness, Fridge003, HydraQYH, Qiaolin-Yu, YAMY1234, Ying1123, celve, hebiao064, hubertlu-tw, ispobock, kkHuang-amd, merrymercy, rainj-me and yuan-luo as code owners July 6, 2026 14:48

github-actions Bot added deepseek jit-kernel labels Jul 6, 2026

lucifer1004 force-pushed the dsv4 branch from 0dca42e to 195d4fe Compare July 6, 2026 14:58

b8zhong reviewed Jul 7, 2026

View reviewed changes

Comment thread python/sglang/jit_kernel/csrc/gemm/marlin_moe/marlin_template.h Outdated

b8zhong self-assigned this Jul 7, 2026

b8zhong reviewed Jul 7, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/attention/dsa/tilelang_kernel.py Outdated

b8zhong reviewed Jul 7, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/attention/dsv4/indexer.py Outdated

chore(dsv4): merge upstream main

1754606

Carry the SM120 enablement onto the latest upstream layout and make the D-LLM override test exercise a non-FlashInfer starting backend.

fix(ci): add SM120 test entry point

c3b1ea3

robbiemu mentioned this pull request Jul 14, 2026

[Feature] Support serialized static-FP8 MXFP4 MoE on the resident FlashInfer backend (SM120/SM121) #31235

Closed

2 tasks

amd-bot mentioned this pull request Jul 15, 2026

[CI Monitor] Daily Report - 2026-07-15 bingxche/sglang-ci-bot#136

Open

lucifer1004 and others added 4 commits July 14, 2026 21:24

Merge remote-tracking branch 'upstream/main' into dsv4

553644a

fix(quantization): guard FlashInfer adapter imports

31d7d27

Merge remote-tracking branch 'upstream/main' into dsv4

3d0aaeb

Merge branch 'main' into dsv4

eabe8f8

b8zhong mentioned this pull request Jul 15, 2026

SM120 Performance Optimization Plan #19637

Open

68 tasks

Fridge003 reviewed Jul 16, 2026

View reviewed changes

lucifer1004 added 3 commits July 16, 2026 16:19

fix(quantization): simplify mxfp4 architecture check

c3f9bd6

Merge remote-tracking branch 'origin/dsv4' into dsv4

ddb0c45

test(mxfp4): restore SM90 test comments

bdad8c2

Fridge003 added the run-ci-extra label Jul 18, 2026

Merge branch 'main' into dsv4

a966bb8

Fridge003 changed the title ~~SM120 DSv4 TP2 enablement~~ Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 Jul 18, 2026

Merge branch 'main' into dsv4

754b78f

Fridge003 merged commit faf6894 into sgl-project:main Jul 18, 2026
237 of 293 checks passed

TTThanos mentioned this pull request Jul 20, 2026

[Feature] Native SM120 (Blackwell) support for flash_mla_sparse_fwd (DeepSeek-V4 sparse-attention prefill) #31578

Open

Zhylkaaa pushed a commit to Zhylkaaa/sglang that referenced this pull request Jul 29, 2026

Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 (…

2d4a3bd

…sgl-project#30272)

jinzhenfan pushed a commit to jinzhenfan/sglang that referenced this pull request Jul 29, 2026

Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 (…

1e14b46

…sgl-project#30272)

hassellof mentioned this pull request Aug 3, 2026

[Perf] DeepSeek-V4 dsv4 prefill on SM120 (4x RTX PRO 6000, TP=4): ~2-7K tok/s vs ~12.5K reported via vLLM/marlin on identical hardware — expected? tuning guidance? #33422

Open

2 tasks

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 (…

7517854

…sgl-project#30272)

jakki-amd pushed a commit to jakki-amd/sglang that referenced this pull request Sep 9, 2026

Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 (…

edb4361

…sgl-project#30272)

Navigation Menu

Uh oh!

Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2 - #30272

Implement SM120 DeepSeek V4 flashinfer_mxfp4 moe runner backend + TP2#30272
Fridge003 merged 25 commits into
sgl-project:mainfrom
lucifer1004:dsv4

lucifer1004 commented Jul 6, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jul 6, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lucifer1004 commented Jul 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Fridge003 commented Jul 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

lucifer1004 commented Jul 6, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

gemini-code-assist Bot commented Jul 6, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lucifer1004 commented Jul 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Fridge003 commented Jul 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

lucifer1004 commented Jul 6, 2026 •

edited by github-actions Bot

Loading