Repository navigation
[Kernel][Build] Bump DeepGEMM to the 26/09/30 release and adapt the sparse indexer - #61149
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Hi @zyongye, I was looking at this as well for my own PR, (#59482), and came to a lot of the same conclusions. One thing that I also found as an issue in DeepGEMM was in cc @Harry-Chen |
|
This pull request has merge conflicts that must be resolved before it can be |
523b752 to
fdbba7a
Compare
mlopart is never used in tree. But thank you for the catch. We should fix it. Thanks |
fdbba7a to
8988e5e
Compare
…parse indexer Pin vllm-project/DeepGEMM dev 6ce1c5e: the upstream 26/09/30 release (vllm-project/DeepGEMM#30), the vendoring import fixes vllm-project/DeepGEMM#31 and vllm-project/DeepGEMM#32, and the paged sparse MQA page-stride relaxation vllm-project/DeepGEMM#33. Vendor the new deep_gemm/locality_domain.py. DeepGEMM#32 makes the MLOPart bridge and the epilogue module use the package's own name, so they also work when vendored as vllm.third_party.deep_gemm. DeepGEMM#33 lowers the page-stride guard from 512 to 256 bytes; the DeepSeek-V4.1 sparse-logits block stride stays 512-aligned, which still satisfies it. The MQA logits API changed: clean_logits/logits_dtype are gone, max_seqlen_k is required, prefill logits are compressed to each row's K span, and SM100 takes only MXFP8/MXFP4 with BF16 weights and returns BF16 logits. On SM100 the FP8 indexer K writers store their existing power-of-two scale as four packed UE8M0 bytes (a lossless MXFP8 encoding), FP8 Q passes unit MX scales with its scale still folded into the weights, and BF16 logits are upcast row-aware for the FP32 top-k kernels. check_wheel_deepgemm.py now imports the vendored package. build_deepgemm_C.py passes --force to setup.py build_ext: DeepGEMM's setup.py only tracks python_api.cpp, so an incremental build after a header-only pin change kept a stale extension. CMake already decides when this step reruns. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
8988e5e to
6f63ba7
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #94360 for commit |
CI selector (shadow): no narrower answer, every step runs: 195 test steps (271 jobs) instead of 195 (271 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Why: fail-open: tools/build_deepgemm_C.py is a docker-image build input (docker/Dockerfile COPY); the CI image is rebuilt from it; running everything Today's rules run everything for this diff. Selector would run (195)
Would skip (today's rules run them) (0)none Would add (today's rules do not run them) (0)none AMD mirrors: would skip (0)none AMD mirrors: would add (0)none 23 changed files · base |
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #94360. |
…parse indexer (vllm-project#61149) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <AndreasKaratzas@users.noreply.github.com>
…parse indexer (vllm-project#61149) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Since vllm-project#61149 the CUDA prefill MQA logits are row-relative (row_starts is None), so gvr_2's read-from-column-0 contract holds for every chunk, not only single-request ones. resolve_prefill_backend now gates on row_starts is None plus GVR2_PREFILL_MIN_ROWS instead of num_reqs; XPU, which still passes per-row starts, stays on top_k_per_row_prefill. FlashInfer 0.7.2 (now pinned) compiles hint-free gvr_2 engines, so pass pre_idx=None with top_k instead of a cached arange(k) hint buffer, and feature-detect that signature. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Jared Wen <w13431838023@gmail.com>
|
Hey @zyongye @JaredforReal, I've been looking at the indexer numerics after this bump and I'm a bit worried about non-DSv4 models. With the old pin, the SM100 kernel summed over heads in FP32 and gave us FP32 logits. The 26/09/30 release dropped That's fine for DSv4, since they trained for it:
However, this might not be the case for other existing models, or new upcoming ones. For instance, DSv3.2 and GLM-5 both use FP32 index scores in their official code: DeepSeek's V3.2 release inference code, and the slime indexer (written by Zilin Zhu @ Z.ai) for GLM-5 through 5.3 and GLM-5.3-Flash (that's their RL stack, but it's what the policy was trained against). So we're now running them with different numerics than their own code uses. This is also already shaping follow-up work in vLLM. #62240 (DeepSelect on BF16 prefill logits), #62172 ( I don't think the evals in the description would catch this either. GSM8K prompts are shorter than Is there an experiment that stressed this that I'm missing? If it's not basically identical, I'd recommend to keep BF16 only for DSv4 and bring back an FP32 path for the rest. FWIW, SGLang already does this: BF16 index logits/weights only for the DSv4.1 index layers on SM100, FP32 for everything else (index_logits.py). Their pinned DeepGEMM still takes |
Overview
Bump the vendored DeepGEMM pin to the tip of
vllm-project/DeepGEMMdev(6ce1c5e): upstream's 26/09/30 public release (vllm-project/DeepGEMM#30, deepseek-ai/DeepGEMM#462) plus vendoring import fixes (vllm-project/DeepGEMM#31, vllm-project/DeepGEMM#32) and a relaxed paged sparse MQA page-stride guard (vllm-project/DeepGEMM#33). Adapt vLLM to the breaking MQA-logits API.Claims
Breaking changes in DeepGEMM
1e1842a→6ce1c5edeep_gemm/__init__.pyimports a newdeep_gemm/locality_domain.pyimport vllm.third_party.deep_gemmfailsdeepgemm.cmake_C.pyregistered the locality-domain shim asdeep_gemm._C.locality_domainvllm.third_party.deep_gemm, and vLLM then silently treats DeepGEMM as unavailabletools/check_wheel_deepgemm.pynow imports the package so CI catches thisdeep_gemm.utils.mlopartand the epilogue module was nameddeep_gemm.epiloguelocalize()(not called by vLLM) would fail when vendored_C.pypasses its__package__to the newinit(library_root_path, python_package_name)fp8_fp4_paged_sparse_mqa_logitspage-stride guard lowered from 512 to 256 B (vllm-project/DeepGEMM#33)fp8_mqa_logits,fp8_paged_mqa_logits,fp8_gemm_nt_skip_head_mid,get_bf16_mega_gate_configremovedclean_logits/logits_dtyperemoved fromfp8_fp4_(paged_)mqa_logits;max_seqlen_know required (> 0)vllm/utils/deep_gemm.py; prefill chunk metadata carries a hostmax_seqlen_kbound (no device sync)fp8_fp4_mqa_logitsreturns compressed logits[M, max_seqlen_k](rowistarts at K positionks[i]) on every arch0, endske - ks; output indices are unchanged (already relative toks)bf16_mega_gatedefaultscoring_func→'sqrtsoftplus'ubatch_utils.py), not combined with Mega MoE todayepilogue=kwarg), locality domains,shutdownSM100 FP8 indexer
Every CUDA FP8 indexer writer already stores an exact power-of-two scale
2^ceil(log2(amax/448))per 128-wide token. On SM100 (mqa_logits_use_mx_scales()) the writers store that scale's exponent as four equal UE8M0 bytes instead of the fp32, which is an exact MXFP8 encoding of the same values (same cache size, same FP8 bytes):indexer_k_quant_and_cache(scale_fmt="mxfp8")Indexer(V3.2 via that class, DeepSeek MTP), HY-V4, DeepSeek-V3.2 under PCP_fused_norm_rope_kernel)GlmMoeDsa), Dots3Note, V3.2 MTPindexer_k_norm_rope_storeReaders need no change:
cp_gather_indexer_k_quant_cache(prefill, all of the above) and the DCP/PCP gathers copy the 4 scale bytes verbatim, and DeepGEMM reads the paged cache directly. The 512-wide MLA caches, MXFP4 caches, and the MiniMax-M3 / Qwen4 indexers are untouched.FP8 Q keeps folding its per-(token, head) power-of-two scale into
weights;mqa_logits_mx_inputspasses unit MX Q scales (exact, sincerelu(s·x) = s·relu(x)fors > 0) and casts weights to BF16. The BF16 logits are upcast to FP32 for the existing FP32-only top-k kernels (as_fp32_logits); routing BF16 logits straight to a BF16-capable top-k (DeepSelect / LiteTopK, #59921/#59922) is a follow-up.SM90 and SM120 keep the FP8 format, FP32 weights and FP32 logits. Consequently the FP8 indexer cache's scale bytes now differ between SM100 and other architectures: moving that KV cache between SM100 and non-SM100 nodes (cross-generation PD disaggregation, shared KV offload) is not supported for these models. Same-architecture setups are unaffected.
Validation
All on 4x GB200 (SM100), CUDA 13, torch 2.13; baseline =
mainwith the old pin1e1842a, same build otherwise.GLM-5.3-Flash (FP8, TP4 + EP, MTP=3,
--max-model-len 131072 --max-num-batched-tokens 131072 --kv-cache-dtype fp8_e4m3 --reasoning-parser glm47,reasoning_effort=low), on the rebased branch vsmainbfa7e5b. MRCR prompts are far longer than the indexer's coverage (2048 pools × 4 = 8K tokens), so they exercise top-k selection; long prefills also engage the TP prefill row sharding from #54951 (confirmed with a temporary log, e.g. 129K-row steps split[47458, 18791, 42577, 19957]across ranks):main+1e1842atests/evals/mrcr), match ratio (2 / 4 / 8 needles)GLM-5.3 serving throughput tracks the MTP acceptance length, which varies per run (the 1337 tok/s run had 1.31 vs ~1.7–1.8 elsewhere); runs with matching acceptance agree (1476 @ 1.70 vs 1471 @ 1.71).
DeepSeek-V4-Flash (before the rebase onto current
main; baseline 26a488a), TP2 + EP,--kv-cache-dtype fp8 --attention_config.indexer_kv_dtype=fp8,--max-model-len 8192; GSM8K viatests/evals/gsm8k/gsm8k_eval.py(1319 questions, 5-shot, 2 rounds). GSM8K prompts fit inside the indexer's coverage (512 × 4 = 2K tokens), so this checks the SM100 path end to end rather than top-k selection:main+1e1842aServing — same server,
vllm bench serve --dataset-name random --random-input-len 4096 --random-output-len 512 --num-prompts 256 --max-concurrency 64 --ignore-eos, one warm-up + 3 runs each, sequential on the same GPUs:main+1e1842aThe ~0.8% gap matches the two extra small kernels per indexer layer on SM100 (weights → BF16, row-aware logits → FP32, ~3 µs together). Feeding BF16 straight from the Q producers and into a BF16 top-k removes both.
Unit tests (GB200):
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.pyfails on this GB200 withunspecified launch failureidentically onmain+1e1842a(pre-existing, unrelated). Not run (no hardware): SM90 and SM120; CI's H100 DeepGEMM job covers the SM90 attention tests. CI has no SM100 DeepGEMM job, so the MX path is validated only by the runs above.Pin delta
1448881→6ce1c5e: the end-to-end runs above used1448881. The two later commits change no kernel math (vllm-project/DeepGEMM#33 only lowers the host and device stride asserts, vllm-project/DeepGEMM#32 only module naming). Re-checked on6ce1c5eon GB200 (the GPUs were shared with another job, so no end-to-end rerun):Details
tools/build_deepgemm_C.pynow passes--forcetosetup.py build_ext. DeepGEMM'ssetup.pylists onlycsrc/python_api.cpp, so after a pin change that touched only headers (config.hppin Support MLOPart imports from vendored package names DeepGEMM#32) an incremental build silently kept the old extension and the import failed withdeep_gemm::init() expected at most 1 argument(s). CMake already reruns this step on DeepGEMM header changes andpython_api.cppis the only translation unit, so forcing adds no work.1e1842aor adapts the indexer to the new MQA API ([Bugfix][Kernel] Bump DeepGEMM pin for SM107 JIT compatibility #59503 bumps to an older SM107 fix already contained in the current pin; [Build] DeepGEMM migration to torch Stable ABI #59482 is the Stable ABI migration).Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code