Conversation
The SM12x FP8 paged-MQA path asserted block_kv == 64, which leaves the DeepSeek-V4.1 sparse indexer without a usable page on SM120: vLLM disables the MXFP4 indexer cache on SM12x (dsa_indexer_uses_fp4()), so the FP8 cache is the only option, and V4.1 mixes compress_ratios 1 and 2, so a 64-token storage block yields num_states = 32 for the ratio-2 layers. Accept block_kv == 32 by splitting the KV range into kNumGroups = 4 groups (64 / 32), matching the existing SM90 block_kv=32 port, and drop num_kv_stages to 2 for the (block_kv=32, num_heads=64) shape with varlen or next_n >= 2 so the kernel stays within the 99 KiB SMEM budget. Verified on RTX PRO 5000 (sm_120): 72/72 new parameterized cases pass (page 32/64 x heads 16/32/64 x next_n 1/2/3 + varlen x head_dim 32/64/128, including random physical pages, padded strides, empty history, partial pages, CUDA graph replay) and the 10 pre-existing SM120 cases still pass. Graph-replay latency for the page-32 path is within 4.5%-7.6% of page-64 at heads=32. Signed-off-by: zack <51604064+luoyuctl@users.noreply.github.com>
|
@lucifer1004 @zyongye Could you take a look at this change when you have time? DeepSeek-V4.1's ratio-2 layers produce 32-state indexer pages with a 64-token kernel block, but the SM120 FP8 paged-MQA path currently rejects them. This patch extends the existing kernel rather than adding a fallback. It keeps the page64 path unchanged and reduces KV pipeline stages only for page32 configurations that would exceed the shared-memory budget. We recorded 72 new correctness cases and 10 existing tests passing on RTX PRO 5000 (SM120), including CUDA graph replay. The PR description includes the test environment, reproduction steps, and remaining validation gaps. @lucifer1004, I'd especially appreciate your review of the TMA/barrier handling and advice on coordinating this with #10. @zyongye, your feedback on the scope and merge order for this fork would also be helpful. Thanks! |
|
Hi @luoyuctl — I'm the author of #10 (the SM120 device-layer vendoring). Happy to coordinate; here's how the pieces fit. Where the conflict is. #14's device-header edit lands on the fork's pre-vendoring Why page32 still needs kernel work. With 32-row pages, a 64-row compute tile would straddle two non-contiguous physical pages, which the current TMA addressing ( Proposed sequencing:
Same feature is relevant to the nv_dev-lineage fork (deepseek-ai#447, DSv4.1 indexer pages), which picks it up by vendoring the same tag — another reason to route the kernel change through the canonical repo rather than per-fork. Happy to review either way. |
Mirror the FP4 sibling's derivation in the FP8 kernel: BLOCK_KV = min(PAGE_KV, 64), so a 32-row physical page gets a 32-row compute tile (a 64-row tile would straddle two non-contiguous pages, which the within-page TMA offset addressing cannot express). SPLIT_KV stays 128 = BLOCK_KV * groups, so PAGE_KV=32 yields 4 KV groups (4 TMA warps each own one group). Motivated by DSv4.1 indexer caches, which mix 64- and 32-state pages (vllm-project/DeepGEMM#14). Gate coverage: instantiate FP8 paged MQA at PAGE_KV 32/64/256. Compile gate passes against both deployment lineages (deepseek main, vllm dev). Runtime enablement stays with the forks' host glue. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
Thanks @lucifer1004, that sequencing makes sense. Please go ahead with the page32 kernel port in DeepGEMM-sm120, and thanks for offering to credit our contribution. Once #10 and the updated device layer are integrated into this fork, we’ll rebase #14, drop the device-header changes, and keep the API/launcher updates and test matrix. We’ll rerun the SM120 tests against that version and ask for another review. Thanks for helping coordinate this! |
FP8 paged MQA logits gains PAGE_KV=32 (BLOCK_KV derived as min(PAGE_KV, 64), mirroring the FP4 sibling). Device-only update: this fork's host launcher still restricts paged FP8 to page 64, so page32 stays inert until host glue opts in (vllm-project#14). Validated on sm_120a: test_sm120_mqa.py + test_sm120_fp8_fp4.py 23/23 passed from a fresh JIT cache (only the pre-existing test_filter collection quirk remains, deepseek-ai#446). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Vendor DeepGEMM-sm120 v0.1.4, which derives BLOCK_KV = min(PAGE_KV, 64) in the FP8 paged kernel (mirroring the FP4 sibling): a 32-row page gets a 32-row compute tile and 4 KV groups at SPLIT_KV=128, instead of a 64-row tile straddling two non-contiguous physical pages. DSv4.1 indexer caches mix 64- and 32-state pages (vllm-project#14). Host glue: relax the FP8 launcher and fused-cache API gates to admit block_kv=32 for arch 12, derive tile_kv = min(block_kv, 64) like the FP4 launcher, and drop to two KV stages for page32 + 64 heads + paired queries, where three stages exceed the 99 KiB SMEM budget by 4 bytes. Tests: the paged-MQA contract matrix gains (fp8, page32), and a focused case covers the two-stage fallback (page32, 64 heads, paired/varlen). Validated on sm_120a from a fresh JIT cache: paged MQA suite 20/20 passed, incl. graph replay and legacy-API cross-checks. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
|
Update: the device-side page32 support is now done — DeepGEMM-sm120 v0.1.4 derives Once #10 merges, your rebase shrinks to: API + launcher relaxation ( |
2 participants
Summary
Enable FP8 paged-MQA logits with 32-state KV pages on SM120, without switching to an FP4 indexer cache or introducing a fallback kernel.
DeepSeek-V4.1 mixes compression ratios 1 and 2. With a 64-token kernel block, these yield indexer pages of 64 and 32 states respectively. The SM120 FP8 path currently rejects the latter in both the API and host launcher.
block_kv in {32, 64}for SM120 at the API and FP8 launcher boundaries.split_kv / block_kvgrouping:128 / 32 = 4groups, compared with128 / 64 = 2.next_n >= 2), to stay within the SM120 shared-memory budget. Other configurations retain three stages.Duplicate-work check / overlap
Checked open SM120 / paged-MQA PRs before submission. #10 migrates the SM120 device layer and changes page addressing, but its FP8 paged kernel still restricts PAGE_KV to 64/128/256; it does not enable FP8 page32. This narrow API/launcher/kernel change addresses that missing case. The device-header edits will need reconciliation if #10 lands first; I am happy to coordinate rather than maintain divergent kernels.
Related integration discussion: vllm-project/vllm#56461, vllm-project/vllm#56702, vllm-project/vllm#57292. This PR alone does not provide complete model support: vLLM cache geometry and FlashInfer dual-cache prefill support are separate dependencies.
Validation
Recorded earlier in this development session on NVIDIA RTX PRO 5000 72GB Blackwell (SM120), for this page32 implementation based on
a6bbb8000161c0dc3a85a0300a905f76898a7913:The imported
test_filterhelper is excluded from pytest collection in the second command. The added matrix covers page32/64, heads16/32/64, head dimensions32/64/128, uniformnext_n=1/2/3and varlen, shuffled physical pages, padded page strides, zero/partial contexts, a float64 reference, repeat execution and CUDA graph replay.Recorded metadata+logits graph-replay medians at 32 heads, in microseconds (not end-to-end model throughput):
Submission-time local checks:
git diff --check HEAD^ HEADand Python AST parsing of the test file passed. GPU tests were not rerun while preparing this draft. No repository pre-commit configuration is present.Validation limits and remaining review items
Failure evidence and reproduction
The following is an excerpt from the archived old serving-image error (2026-09-17); it is not a traceback captured against this fork's current base. Paths/line numbers belong to that image. The final assertion was propagated through the worker RPC:
Process/timestamp prefixes were removed for readability. The recovered excerpt is the RPC propagation tail; it is not presented as the full original worker call stack. The independent fork code has the same SM120 FP8 page32 rejection; its line numbers differ.
Standalone regression recipe
On an SM120 CUDA machine with a matching PyTorch/CUDA development environment, check out this PR, initialize its submodules, and build the extension and its matching JIT headers together. Merely replacing the CUDA header leaves the API/launcher restrictions inside the old
_Cbinary.The saved build commands used a system-site-packages venv,
libdw-dev, and NVIDIA wheel header/library paths:Adjust the NVIDIA wheel search root for a different Python installation. The recorded results are 72 and 10 passes, respectively, as listed above. For a pre-fix comparison, retain the added test file but rebuild the parent revision's extension and matching headers in a separate checkout; page32 is expected to fail the original assertion. That exact parent-revision pytest rerun was not recorded, so it is a reproduction recipe, not an additional claimed test result.
Environment and evidence boundaries (2026-09-17)
These are two different validation tracks; they must not be combined into a claim that the current vLLM PR head was tested end-to-end.
a6bbb8000161c0dc3a85a0300a905f76898a7913+ the page32 change now committed as471d62c4ef85e259c33cf66b25f48cbbf8e758ef8b1392b+ a separate page32 backport-std=c++20. The temporary Pod was deleted; a complete contemporaneous environment dump was not retained.2.13.0+cu130V13.0.88595.58.03, read from the running node during this documentation update0.1.dev20904+g179dd0fa9+ geometry backportsflashinfer-python 0.6.18,flashinfer-jit-cache 0.6.18+cu130, with local decode/prefill patchesSuccessful serving image manifest digest, confirmed from the Pod's
imageID:sha256:406cb11244831d00bbe05e90834a94177e3befe505d61570f22e2e1eb6cc75f5.This is the manifest digest; the previously quoted
17bff2ab...is an image-config digest and must not be used as the pull digest.The successful deployment used
--tensor-parallel-size 8,--enable-expert-parallel,--kv-cache-dtype fp8,--block-size 64,--max-model-len 262144,--max-num-seqs 64,--max-num-batched-tokens 8192,--gpu-memory-utilization 0.95, EPLB, and DSpark K=5 with adaptive verification disabled. These are experiment settings, not a claim that every setting is necessary or optimal.AI assistance
AI assistance was used for implementation, review, test execution and preparation of this PR. The implementation is submitted for review; human review and validation remain required before merge. No completed human sign-off on kernel correctness is asserted.
Commit-message correction
The existing commit message mistakenly writes
64 / 32when describing four KV groups. The implementation reusessplit_kv / block_kv, which is128 / 32 = 4. The source and summary above are correct; this note corrects the prose without rewriting the tested commit.