[SM120] Support 32-state pages in FP8 paged-MQA logits by luoyuctl · Pull Request #14 · vllm-project/DeepGEMM

luoyuctl · 2026-09-17T09:30:19Z

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.

Permit block_kv in {32, 64} for SM120 at the API and FP8 launcher boundaries.
Reuse the existing split_kv / block_kv grouping: 128 / 32 = 4 groups, compared with 128 / 64 = 2.
Use two KV pipeline stages only for page32 + 64 heads + (varlen or next_n >= 2), to stay within the SM120 shared-memory budget. Other configurations retain three stages.
Add 72 parameterized correctness cases to the existing SM120 MQA test suite.

The device layer is unchanged: #10 vendored DeepGEMM-sm120 v0.1.5, which derives BLOCK_KV = min(PAGE_KV, 64). This PR now contains only host/API changes plus the test matrix.

Relationship to #10

#10 (vendored DeepGEMM-sm120 v0.1.5) landed first and provides the device-side page32 support. This PR is rebased on top of it and contains no device-header edits; it only relaxes the API/launcher page-size checks, adds the SMEM stage heuristic, and adds the test matrix.

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

Note: the results below were recorded on the pre-rebase head. The rebased commit only removes device-header edits, but the SM120 suite still needs to be re-run against vendored v0.1.5 before merge.

Recorded earlier in this development session on NVIDIA RTX PRO 5000 72GB Blackwell (SM120), for this page32 implementation based on a6bbb8000161c0dc3a85a0300a905f76898a7913:

pytest tests/test_sm120_mqa.py -k fp8_paged_mqa_page_geometry
72 passed, 11 deselected in 212.78s

pytest tests/test_sm120_mqa.py -k 'not fp8_paged_mqa_page_geometry and not test_filter'
10 passed, 73 deselected in 9.92s

The imported test_filter helper is excluded from pytest collection in the second command. The added matrix covers page32/64, heads16/32/64, head dimensions32/64/128, uniform next_n=1/2/3 and 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):

batch / context	page32	page64
1 / 4096	3.82	3.82
8 / 4096	5.31	5.08
1 / 32768	4.97	4.62

Submission-time local checks: git diff --check HEAD^ HEAD and 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

Compute-sanitizer, large-load profiling and full-model evaluation of this exact dependency combination remain outstanding.
A separate backport onto an older DeepGEMM 2.6.1 serving image, with vLLM and FlashInfer patches, passed startup and a short chat on 8x SM120. That is NOT an end-to-end validation of this fork commit and is not presented as one.
Human contributor review and reproduction of the relevant tests remain outstanding before merge. Marking this PR ready invites maintainer review; it does not certify completed human review or kernel correctness.

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:

File "/usr/local/lib/python3.12/dist-packages/vllm/v1/executor/multiproc_executor.py", line 103, in _wait_for_response
    response = self.aggregate(self.get_response())
File "/usr/local/lib/python3.12/dist-packages/vllm/v1/executor/multiproc_executor.py", line 437, in get_response
    raise RuntimeError(
RuntimeError: Worker failed with error 'Assertion error (/workspace/.deps/deepgemm-src/csrc/apis/attention.hpp:320): (arch_major == 10 and (block_kv == 32 or block_kv == 64 or block_kv == 128)) or (arch_major == 9 and (block_kv == 32 or block_kv == 64)) or (arch_major == 12 and ((is_fp4 and (block_kv == 32 or block_kv == 64)) or (not is_fp4 and block_kv == 64)))', please check the stack trace above for the root cause

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 _C binary.

The saved build commands used a system-site-packages venv, libdw-dev, and NVIDIA wheel header/library paths:

git submodule update --init --recursive
uv venv --system-site-packages .venv
# Install pytest / required build dependencies in this venv if absent.
# libdw-dev is an OS build prerequisite for the DeepJIT-based fork.
export CPATH="$(find /usr/local/lib/python3.12/dist-packages/nvidia -type d -name include | paste -sd:)"
export LIBRARY_PATH="$(find /usr/local/lib/python3.12/dist-packages/nvidia -type d -name lib | paste -sd:)"
DG_USE_LOCAL_VERSION=0 DG_FORCE_BUILD=1 MAX_JOBS=4 \
  .venv/bin/python setup.py build_ext --inplace
.venv/bin/python -m pytest tests/test_sm120_mqa.py \
  -k fp8_paged_mqa_page_geometry -q
.venv/bin/python -m pytest tests/test_sm120_mqa.py \
  -k 'not fp8_paged_mqa_page_geometry and not test_filter' -q

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.

Item	Fork kernel test	Serving backport smoke test
GPU	1 allocated NVIDIA RTX PRO 5000 72GB Blackwell, SM120	8x same GPU, TP8 + expert parallelism
DeepGEMM	`a6bbb8000161c0dc3a85a0300a905f76898a7913` + the page32 change as pre-rebase commit `471d62c4ef85e259c33cf66b25f48cbbf8e758ef` (rebased onto #10 as `40592bdb7c4b68c0b41af5657049572ea15dc9fb`)	DeepGEMM 2.6.1 lineage, `8b1392b` + a separate page32 backport
Environment provenance	Saved test/build Pod manifests use the same base serving image; saved build log confirms CPython 3.12 and `-std=c++20`. The temporary Pod was deleted; a complete contemporaneous environment dump was not retained.	Versions below were read from the still-running successful Pod, not inferred from the build host
Python	3.12 confirmed by saved build log; exact patch version not independently retained for this test Pod	3.12.3
PyTorch	Shared base image; exact package version was not retained in the four archived test logs	`2.13.0+cu130`
CUDA toolkit / nvcc	Shared base image; not independently retained in the archived test logs	CUDA 13.0, nvcc `V13.0.88`
NVIDIA driver	Historical test-time driver version was not separately archived	`595.58.03`, read from the running node during this documentation update
vLLM	Not a full-model test	`0.1.dev20904+g179dd0fa9` + geometry backports
FlashInfer	Not used for the standalone DeepGEMM kernel tests	`flashinfer-python 0.6.18`, `flashinfer-jit-cache 0.6.18+cu130`, with local decode/prefill patches

Successful 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.

luoyuctl · 2026-09-17T09:50:58Z

@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!

lucifer1004 · 2026-09-17T10:04:31Z

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 sm120_fp8_paged_mqa_logits.cuh. #10 replaces that header with a byte-identical vendored copy from DeepGEMM-sm120 v0.1.3, which renamed BLOCK_KV→PAGE_KV and restructured page addressing: the vendored kernel fixes the compute tile at BLOCK_KV=64 and supports physical PAGE_KV ∈ {64, 128, 256} via within-page TMA offsets. Your two static_asserts reference the old constants and would not survive the rebase textually or semantically.

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 (tile_idx * BLOCK_KV / PAGE_KV, % PAGE_KV) cannot express — so PAGE_KV=32 genuinely requires a BLOCK_KV=32 / 4-group kernel configuration, not just relaxed asserts. That makes it a device-layer feature, and under the vendoring scheme the device layer's single source of truth is DeepGEMM-sm120 (a scheduled watcher flags fork trees that match no release tag, so landing kernel edits directly here would create exactly the divergence we just removed).

Proposed sequencing:

feat(sm120): vendor DeepGEMM-sm120 v0.1.0 device layer #10 lands first (it's review-complete and mergeable).
The page32 kernel support (BLOCK_KV=32 template path with the group/TMA-warp accounting from your change) goes into DeepGEMM-sm120, where CI cross-compiles it against both fork lineages; I can do the port myself with credit to you, or you're welcome to open the PR there directly — whichever you prefer. Released as v0.1.4.
[SM120] Support 32-state pages in FP8 paged-MQA logits #14 then rebases onto dev + vendored v0.1.4 and shrinks to the parts that genuinely belong to this fork: the API/launcher relaxation (block_kv ∈ {32, 64}, the num_kv_stages SMEM heuristic) and your test matrix — no device-header edit needed.

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>

luoyuctl · 2026-09-17T10:14:09Z

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>

lucifer1004 · 2026-09-17T10:27:37Z

Update: the device-side page32 support is now done — DeepGEMM-sm120 v0.1.4 derives BLOCK_KV = min(PAGE_KV, 64) in the FP8 paged MQA logits kernel (mirroring the FP4 sibling), so a 32-row page gets a 32-row compute tile with 4 KV groups at SPLIT_KV=128. #10 has re-vendored at v0.1.4 (ea3b2f8), device-only — the fork's host launcher still asserts page 64, deliberately leaving the host/API enablement and the test matrix to this PR.

Once #10 merges, your rebase shrinks to: API + launcher relaxation (block_kv ∈ {32, 64}), the num_kv_stages SMEM heuristic, and your tests — no device-header edits. Two notes for the rebase: (1) the vendored kernel's constants are renamed (PAGE_KV for the physical page; BLOCK_KV is the derived compute tile), so your two static_asserts are already covered by the vendored compile-time checks; (2) on the nv_dev lineage we needed the same num_kv_stages=2 fallback for page32 + 64 heads + paired queries — the three-stage layout misses the 99 KiB budget by exactly 4 bytes there too, so your heuristic transfers as-is.

* feat(sm120): vendor DeepGEMM-sm120 v0.1.3 device layer Replace the in-tree SM120 device headers with byte-identical vendored copies from lucifer1004/DeepGEMM-sm120 v0.1.3 (merged superset of both lineages). Marker-free vendoring: no manifests, tooling, or header comments enter this fork; provenance and drift control live on the canonical repo's side (per-tag manifest, check_vendor.py --fork, scheduled downstream watcher). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * feat(sm120): adapt host launchers to vendored kernel signatures The vendored DeepGEMM-sm120 v0.1.0 device layer widens several kernel ABIs; adapt the SM120 host glue with neutral mappings: - sm120_bf16_gemm: pass the epilogue operator as a runtime argument (EpilogueArgs, same marshalling as SM100) and split the old int64_t stride_cd_m/stride_cd_batch into u32 stride_d_m/stride_c_m/ stride_d_batch; stride_c_m=0 keeps C sharing D's row stride. - sm120_fp8_fp4_gemm_1d1d: pass shape_cd_m (= shape_m), the runtime epilogue argument, and stride_c_m (= 0) after stride_cd_batch. - sm120_split_k_reduce: the kWithAccumulation template flag is gone; instantiate the 3-arg template and map accumulation onto D in place to the new runtime C operand (gmem_c = D with D's own strides, with_alpha = false). - Masked m-grouped GEMM (bf16 + fp8/fp4): disable the TMA-store epilogue (swizzle_cd_mode = 0) when m % BLOCK_M != 0. The vendored kernels dropped the old masked-boundary scalar fallback, so a boundary tile's full-tile TMA store would spill into the next group's rows; the scalar store path carries the masked row bounds. This mirrors the proven nv_dev-lineage launchers. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * fix(sm120): guard k-grouped GEMM TMA store against unaligned M The vendored v0.1.0 GEMM kernels have no boundary-tile scalar fallback in the TMA-store epilogue, and the k-grouped D descriptor is one flat 2D map over all groups, so TMA can only clamp at the outermost M dim: with m % BLOCK_M != 0, a group's partial tail M tile TMA-store spills into the next group's slab. Mirror the masked-launcher guard (and the nv_dev-lineage reference glue) in both k-grouped launchers: force swizzle_cd_mode = 0 (group-bounded scalar store path) for unaligned m. Repro evidence (num_sms=2, groups=2, m=48, n=128, ks=[8192, 128]; the asymmetric K makes group 0's tile store last, so the spill overwrites group 1's head rows deterministically): - before: bf16/bf16-out 5/5 iterations contaminated (d[1][0:16] = 0.0, expected 768); fp8/fp32-out 5/5 contaminated (d[1][0:16] = 0.0, expected 771) - bf16/fp32-out was already clean (0/5): the bf16 kernel gates its TMA-store epilogue on sizeof(cd_dtype_t) <= 2, so fp32 outputs take the scalar path; the fp8 kernel has no such dtype gate - after: all three configurations 0/5 contaminated Adds test_sm120_kgroup_unaligned_m_tail_tile_isolation to tests/test_sm120_bf16.py (fp32 and bf16 outputs) and tests/test_sm120_fp8_fp4.py (NT and TN layouts), asserting per-group correctness including slab head rows plus flat-storage guard regions. Full SM120 suite: 48 passed (46 baseline nodes + 2 new). Memcheck on the two affected files: 20 passed, 0 sanitizer errors. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * feat(sm120): vendor DeepGEMM-sm120 v0.1.4 device layer 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 (#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> * fix(sm120): vendor DeepGEMM-sm120 v0.1.5 generic-proxy WAR fences Vendor lucifer1004/DeepGEMM-sm120 v0.1.5 (b31a688): 14 fence_view_async_shared() sites across 8 SM120 kernel headers, ordering generic-proxy smem reads (ldmatrix/ld_shared) before the empty-barrier arrive that hands each stage back to the TMA producer. Ports deepseek-ai#453 (bf16 GEMM) and fixes the same hazard found by audit in bmk_bnk_mn, tf32_hc_prenorm, fp8_fp4_gemm_1d1d, and the fp8/fp4 dense+paged MQA logits kernels. Sparse MQA is exempt (legacy cp.async producer). Perf A/B on RTX PRO 6000: no regression beyond run-to-run noise. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --------- Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>

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. Device-side page32 support now comes from the vendored DeepGEMM-sm120 v0.1.5 layer, which derives BLOCK_KV = min(PAGE_KV, 64). This commit keeps only the fork-local host-path changes: - Accept page_kv == 32 at the API and FP8 launcher boundaries. - Use two KV pipeline stages for page32 + 64 heads with varlen or next_n >= 2. Four 32-row KV groups with three stages exceed the 99 KiB SMEM budget by 4 bytes there; the 64-row path is unchanged. - Add the page32/64 correctness matrix to the SM120 MQA suite. 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>

luoyuctl · 2026-09-22T02:30:45Z

Thanks @lucifer1004 — #10 has merged, so this branch is now rebased onto dev (940e98f) as 40592bd, and the PR is down to 3 files / +61/-4:

Device layer untouched. The two DG_STATIC_ASSERTs are dropped and deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh is byte-identical to what feat(sm120): vendor DeepGEMM-sm120 v0.1.0 device layer #10 vendored — the diff against dev is empty for the whole deep_gemm/ tree. Confirmed BLOCK_KV = PAGE_KV < 64 ? PAGE_KV : 64 handles page32 as you described.
Kept: the API relaxation (block_kv in {32, 64}), the launcher relaxation and the num_kv_stages heuristic. Both of your notes hold, and the heuristic is required: at page32 + 64 heads + paired queries, three stages need 101380 B against the 101376 B (99 KiB) budget — over by exactly 4 B — while two stages fit in 80900 B. page64 (93188 B) and heads32 (84996 B) keep three stages unchanged.
The test matrix is unchanged; the numbers currently in the description were recorded on the pre-rebase head and are now marked as such. Will re-run on SM120 against v0.1.5 before merge.

One follow-up worth considering for the canonical repo: the dropped kNumTMAThreads / 32 >= kNumGroups assert is implied today by SPLIT_KV == BLOCK_KV * kNumGroups + kNumTMAThreads == 128 combined with the host's split_kv == 128, but the vendored kernel would have no guard if SPLIT_KV ever changes. Happy to file it in DeepGEMM-sm120 if you think it's worth it.

@zyongye could you take a look when you have time? The scope is now just the host/API relaxation, the SMEM stage heuristic and the test matrix.

#447) * Public Release 26/09 (#432) * Public release 26/09 * Update News * feat(sm120): add sparse MQA support on upstream main APIs Port contiguous and paged sparse kernels to DeepJIT while preserving SM100 dispatch, stream-local metadata, entry-balanced scheduling and register guards. Restore histogram synchronization and capture-safe UE8M0 validation, with independent SM120 numerical and CUDA graph regression fixtures. * feat(sm120): restore native GEMM paths on main APIs Restore the original BF16/FP8/FP4 kernels, scheduler, heuristics, and split-K paths from 139f504 through DeepJIT adapters. Support alpha, independent C/D strides, GPU-only KPSUM, and bounded padding cleanup. Separately repair inherited mixed-tail, producer handoff, grouped row ownership, and TMA output-width defects. Preserve the original TMA/MMA pipelines and dispatch choices except for the documented safety eligibility guard. This is a bounded restoration checkpoint, not full migration or universal performance parity. BF16 odd-N, FP4 K-grouped, and remaining public integrations are unfinished. * feat(sm120): preserve nv_dev features alongside main API SM120 kernels * fix(kernels): address migration review and restore grouped TMA pipeline Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * fix(sm120): harden GEMM execution and complete HeadSplits output mapping - Wait for in-flight TMA descriptor reads before overwriting the per-CTA K-grouped BF16 descriptor, and add full-SM randomized alias/separate-C regressions. - Require packed INT32 SFB on the SM120 skip-head path instead of silently reading FP32 scales as UE8M0. - Apply HeadSplits index mapping with logical-column bounds in scalar stores and split-K reduction, replacing the former capability guards while keeping alpha/C/stride contracts. - Store BF16 batched einsum and HC prenorm directly to supported strided outputs, keeping the existing fallback for unsupported layouts. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * test(sm120): collect SM120 tests into dedicated files Restore tests/test_{attention,bf16,fp8_fp4,hyperconnection}.py to their upstream byte state and consolidate all SM120-specific coverage into tests/test_sm120_*.py plus the shared tests/sm120_exercise.py helper. This eliminates per-release merge conflicts with upstream's wholesale rewrites of the shared test files. Switch gating from collection-time string skipif (which initialized CUDA during pytest collection on non-SM120 machines and broke the sanitizer runner's direct-call path) to upstream's call-time test_filter convention, and give each standalone file its own __main__ runner matching upstream script style. Moved test bodies are AST-identical; test counts and results are unchanged (741 passed both before and after on SM120, memcheck/racecheck subsets clean). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * feat(sm120): vendor device layer from DeepGEMM-sm120 v0.1.3 Adopt the standalone DeepGEMM-sm120 repository as the single source of truth for the SM120 device headers: add the CUDA>=13 compile-time guard (silent block_scale drop on pre-13 ptxas), relocate tensor_map_replace_global_dim_in_smem into common/sm120_utils.cuh so the kernels no longer depend on ptx/tma.cuh providing it, and inline the FP4 smem pack factor in layout/sparse_mqa_logits.cuh. Marker-free vendoring: the files carry no canonical-source comments or other foreign markers; provenance and drift control live on the canonical repo's side (per-tag manifest, check_vendor.py --fork, scheduled downstream watcher). Verified: 741 SM120 tests pass on sm_120a from a fresh JIT cache. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * fix(sm120): gate ptxas -O2 for sparse MQA logits on CUDA 13.3+ CUDA 13.3's ptxas crashes (SIGSEGV; DeepJIT compile subprocess exit 139) when compiling the non-warp-specialized MXFP4 sparse MQA logits variants at -O3 with any --register-usage-level (lucifer1004/DeepGEMM-sm120#2); CUDA 13.2, the validated toolchain, is unaffected. Add a toolchain-gated per-kernel override via DeepJIT CompilerOptions::extra_nvcc_flags: when nvcc reports >= 13.3, the non-warp-specialized instantiations compile with ptxas -O2, which avoids the crash and reproduces the 13.2 register allocation (REG:101 for the reference instantiation). The warp-specialized variant is not affected by the crash and keeps the default flags (its compiled reg==64 contract also verified under -O2). Validated on sm_120a: sparse MQA suite 9/9 under CUDA 13.3.73 with the gate active (previously 3 failing tests), and 9/9 under CUDA 13.2.86 with the gate inactive (unchanged default path). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * feat(sm120): FP8 paged MQA logits PAGE_KV=32 support 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> * perf(sm90): clamp paged MQA scheduler range once at construction The stale/OOB-metadata protection added during the migration lived in the per-task hot path (compound has_next_task/exist_q_atom_idx checks per fetch), regressing paged MQA logits on SM90: ~1% on upstream shapes and 13-16% at next_n=4. Move the protection to the constructor: clamp end_q_atom_idx to the in-batch sentinel (zeroing end_kv_idx at the sentinel) and collapse out-of-range starts, then make the per-task end test a single ordered 64-bit compare on (q_atom, kv) packed keys. The ordered compare also keeps the hang fix: an advance overshooting the end is caught by >=, never by exact equality. context_lens/indices are only dereferenced when current_q_atom_idx is provably in batch (end is clamped at or below the sentinel, and the sentinel range always has end_kv == 0). Patch from local review. Validated on H200: test_paged_mqa_logits 24/24 cases pass; next_n=4 (H=64, D=128, page 64) 140.6us vs 154.6us pre-patch at L=8192 and 1093.3us vs 1233.1us at L=65536 (-9%/-11%), upstream-shape sweep at parity or better. Cross-compiled for sm_90a. Reviewer's full 105-case SM90 suite passes on their side. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * perf(layout): pow2 shift fast path in k-grouped SF pack row accounting pack_fp32_into_ue8m0's per-group loop ran three runtime integer divisions per group per thread on the psum path (align of previous end, align of group K, ceil_div to SF rows). gran_k and k_alignment are powers of two at every call site (host-asserted), so detect that once per block and use shift forms instead; the generic division path is kept for any future non-pow2 alignment. The earlier suggestion to drop the align entirely on SM90/100 was unsafe: psum tail K sizes are not k_alignment-multiples, and without the align the row count comes up short (B300 fp8 psum accumulate regression, diff 0.024 vs 0.008 tolerance). Validated: pack output bitwise identical to the division form on a 128-group psum shape (B300); kernel time 38.7 -> 22.1 us/call (-43%) on the same shape. SM100 k-grouped fp8/bf16 contiguous suites pass on B300 (fresh JIT cache), SM120 k-grouped regressions pass on sm_120a, and the TU cross-compiles for sm_103a/sm_120a. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * docs(layout): note why the pow2 shift forms are exact on SM90/100 psum end offsets are k_alignment-multiples and gran_k divides k_alignment there, so aligned_group_k % gran_k == 0 and the rounding shift is exact division; the ceil semantics matter only on SM120. Comment-only change. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * perf(layout): move gran_k/k_alignment into the JIT template for SF pack Follow-up to 3aca951: the runtime pow2 branch regressed non-pow2 k_alignment (reported by RayWang96: k_alignment=384, a third of their k-grouped cases, 91 -> 124 us vs main on B300). Instead of branching at runtime, pass gran_k/k_alignment as JIT template parameters: the per-group align/ceil_div fold to shifts (pow2) or multiply-high sequences (non-pow2) at compile time, covering both alignment classes with no runtime branch. Runtime kernel args drop accordingly. Validated on B300: pack output bitwise identical to the pre-change form at k_alignment 128 and 384; pack kernel 38.8 -> 14.0 us/call at 128 (-64%) and 65.2 -> 32.3 us/call at 384 (-50%); k-grouped fp8/bf16 contiguous suites pass from a fresh JIT cache. SM120 k-grouped regressions pass on sm_120a; the TU cross-compiles for sm_103a/sm_120a. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> * fix(sm120): vendor v0.1.5 generic-proxy WAR fences from DeepGEMM-sm120 Vendor lucifer1004/DeepGEMM-sm120 v0.1.5 (b31a688): 14 fence_view_async_shared() sites across 8 SM120 kernel headers, ordering generic-proxy smem reads (ldmatrix/ld_shared) before the empty-barrier arrive that hands each stage back to the TMA producer. Ports #453 (bf16 GEMM) and fixes the same hazard found by audit in bmk_bnk_mn, tf32_hc_prenorm, fp8_fp4_gemm_1d1d, and the fp8/fp4 dense+paged MQA logits kernels. Sparse MQA is exempt (legacy cp.async producer). Perf A/B on RTX PRO 6000 (bf16 GEMM, fp8 dense/paged MQA, 3 runs): all deltas within run-to-run noise. Fork suite: 235 passed (GPU4, fresh JIT cache). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --------- Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Zhean Xu <94977922+zheanxu@users.noreply.github.com>

luoyuctl mentioned this pull request Sep 17, 2026

[Bugfix][DSv4.1] Fail fast on SM120 when FlashInfer lacks runtime page sizes vllm-project/vllm#57292

Open

luoyuctl marked this pull request as ready for review September 17, 2026 09:49

This was referenced Sep 17, 2026

feat(sm120): vendor DeepGEMM-sm120 v0.1.0 device layer #10

Merged

Integrate main APIs and SM120 kernels while preserving nv_dev features deepseek-ai/DeepGEMM#447

Merged

luoyuctl force-pushed the fix/sm120-fp8-paged-mqa-page32 branch from 471d62c to 40592bd Compare September 22, 2026 02:26

zyongye approved these changes Sep 22, 2026

View reviewed changes

zyongye merged commit 17ca18f into vllm-project:dev Sep 22, 2026
1 check passed

luoyuctl mentioned this pull request Sep 29, 2026

[Bug]: DeepSeek-V4.1-Flash cannot run on SM120 (RTX PRO 6000): compressed-layer page_block_size=32 has no FlashInfer sparse-MLA decode kernel instantiation (only pbs=64) vllm-project/vllm#59203

Open

xzwgit mentioned this pull request Sep 30, 2026

[Bugfix][Kernel] Bump DeepGEMM pin for SM120 32-state page support (DeepGEMM#14) vllm-project/vllm#59385

Closed

briansrls mentioned this pull request Oct 1, 2026

V4.1 on GB10: 64-token kernel pages on SM12x, FlashInfer 0.7.1rc1, and FLASHINFER_VERSION passed to the build gunb-ai/gunbc#12954

Merged

yyihuang mentioned this pull request Oct 7, 2026

[SM120] Support 128-state pages in FP8 paged-MQA logits #28

Open

aagallo mentioned this pull request Oct 8, 2026

[model] DeepSeek-V4.1-Flash on ml.g7e.48xlarge aws-samples/sagemaker-genai-hosting-examples#265

Merged

Navigation Menu

Repository navigation

[SM120] Support 32-state pages in FP8 paged-MQA logits - #14

[SM120] Support 32-state pages in FP8 paged-MQA logits#14
zyongye merged 1 commit into
vllm-project:devfrom
luoyuctl:fix/sm120-fp8-paged-mqa-page32

luoyuctl commented Sep 17, 2026 •

edited

Loading

Uh oh!

luoyuctl commented Sep 17, 2026

Uh oh!

lucifer1004 commented Sep 17, 2026

Uh oh!

luoyuctl commented Sep 17, 2026

Uh oh!

lucifer1004 commented Sep 17, 2026

Uh oh!

luoyuctl commented Sep 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

luoyuctl commented Sep 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Relationship to #10

Validation

Validation limits and remaining review items

Failure evidence and reproduction

Standalone regression recipe

Environment and evidence boundaries (2026-09-17)

AI assistance

Uh oh!

luoyuctl commented Sep 17, 2026

Uh oh!

lucifer1004 commented Sep 17, 2026

Uh oh!

luoyuctl commented Sep 17, 2026

Uh oh!

lucifer1004 commented Sep 17, 2026

Uh oh!

luoyuctl commented Sep 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

luoyuctl commented Sep 17, 2026 •

edited

Loading