deepseek-ai/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00
Code Issues 162 Projects Packages Wiki Activity

[PR #340] Add SM90 FP8 paged MQA logits support for next_n=3 #373

New issue
Open
opened 2026-10-26 21:16:59 +00:00 by bsakaguchi · 0 comments
bsakaguchi commented 2026-10-26 21:16:59 +00:00
Owner
Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/340
Author: @yangsiqt
Created: 5/28/2026
Status: 🔄 Open

Base: main ← Head: fix-sm90-nextn3-paged-mqa-logits

📝 Commits (2)

  • be546e7 Support SM90 paged MQA logits NextN=3
  • 8ceef6b Test SM90 paged MQA logits NextN=3

📊 Changes

4 files changed (+197 additions, -104 deletions)

View changed files

📝 csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp (+13 -6)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh (+180 -94)
📝 deep_gemm/include/deep_gemm/scheduler/sm90_paged_mqa_logits.cuh (+3 -3)
📝 tests/test_attention.py (+1 -1)

📄 Description

Summary

This PR adds SM90/Hopper support for FP8 paged MQA logits with next_n=3,
which is required by vLLM MTP speculative decoding when
num_speculative_tokens=2. The unmodified SM90 path rejects this case with
the existing next_n == 1 or next_n == 2 assertion.

The change is limited to the SM90 non-varlen path. Existing next_n=1/2,
SM100 behavior, and the public Python API remain unchanged.

Related issue: https://github.com/vllm-project/vllm/issues/43457

Implementation

The SM90 kernel uses a two-token WGMMA atom. next_n=3 is implemented as a
pair plus a single-token tail instead of introducing a three-token WGMMA shape:

tokens 0 and 1 -> existing two-token WGMMA path
token 2        -> single-token TailWGMMA path
  • Metadata treats the three tokens as one logical scheduler atom, so the pair
    and tail computations reuse each loaded KV tile.
  • A four-stage Q pipeline prefetches the pair and tail TMA stages while keeping
    the existing producer/consumer barrier lifecycle.
  • The tail uses a two-token TMA box with the unused row zero-filled, but only
    the valid third-token row is reduced and stored.
  • Scheduler helpers retain the original batch row for block-table lookup and
    map the atom to the correct output token offset.

SM90 varlen next_n=3 and next_n>3 remain out of scope.

Validation

GPU: NVIDIA H20 96 GB (SM90)
CUDA toolkit / NVCC: 12.9
PyTorch: 2.8.0+cu129
Baseline: 559d79fb6994a58b8a15b4b93bf13ccc16edf247
Candidate: 8ceef6bad7b4359234cecf87a48c63b6f0c8c0db
  • Unmodified main reproduces the expected assertion for next_n=3.
  • The patched kernel passes B=16,17,256,257, H=64, D=128, L=512,
    FP32 logits, and 20 repeated calls per shape. Repeated outputs are bitwise
    identical; simulated-reference differences are 7.73e-9 to 8.21e-9.
  • tests/test_attention.py::test_paged_mqa_logits passes all 36/36 SM90 cases:
    next_n=1/2/3, FP32/BF16 logits, H=32/64, D=128, B=256/4096, and
    average KV lengths 8192/65536 subject to the test's 32M-token cap.
  • compute-sanitizer --tool memcheck reports 0 errors for
    B=16,17,32,33, next_n=3, L=512, with three repeated calls.

Standalone Performance

CUDA-event benchmark: B=256, H=64, D=128, BLOCK_KV=64, FP32 logits,
20 warmups, 100 timed iterations, five repeats. Values are median milliseconds
per call; JIT compilation is excluded.

L next_n Upstream main Patched Change vs main
4096 1 0.06982 ms 0.06994 ms +0.16%
4096 2 0.13241 ms 0.13172 ms -0.53%
4096 3 unsupported 0.19984 ms N/A
8192 1 0.13204 ms 0.13155 ms -0.37%
8192 2 0.25780 ms 0.25631 ms -0.58%
8192 3 unsupported 0.38742 ms N/A

T(next_n=3) / T(next_n=2) is 1.517x at L=4096 and 1.512x at
L=8192; existing next_n=1/2 performance remains within 1% of upstream.

NCU Profile

A separate same-machine comparison used one H20, CUDA 12.8,
PyTorch 2.11.0+cu128, and Nsight Compute 2025.1. The table reports FP32 NCU
Duration and patched next_n=3 hardware metrics.

L main n2 patched n2 n2 change patched n3 n3/n2 SM DRAM Occupancy
8192 0.279840 ms 0.278016 ms -0.65% 0.421952 ms 1.518x 94.80% 18.91% 31.01%
32768 1.079552 ms 1.068640 ms -1.01% 1.604800 ms 1.502x 97.50% 19.10% 31.16%

For next_n=3, NCU reports 96 registers/thread and 185.4 KB dynamic shared
memory per block. The kernel is compute/pipeline-bound rather than DRAM-bound;
its occupancy, eligible-warps, and barrier-stall profile matches the prior
implementation. Same-run main-versus-patched ratios are used here because
absolute NCU durations from different profiler or toolchain configurations are
not directly comparable. DeepGEMM recommends NVCC 12.9 or newer for best
performance; the primary CUDA-event results above use CUDA 12.9.

Integration Context

An earlier revision was exercised end to end with vLLM on 2x H20 using TP=2
and num_speculative_tokens=2: the server initialized, /v1/completions
returned 200 OK, no DeepGEMM next_n assertion occurred, and speculative
decode metrics were emitted. This is included as consumer-path validation;
the refreshed commit's correctness and performance conclusions use the
standalone results above.

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/340 **Author:** [@yangsiqt](https://github.com/yangsiqt) **Created:** 5/28/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix-sm90-nextn3-paged-mqa-logits` --- ### 📝 Commits (2) - [`be546e7`](https://github.com/deepseek-ai/DeepGEMM/commit/be546e7d09127ac4fedb18208acdb20be8795fcf) Support SM90 paged MQA logits NextN=3 - [`8ceef6b`](https://github.com/deepseek-ai/DeepGEMM/commit/8ceef6bad7b4359234cecf87a48c63b6f0c8c0db) Test SM90 paged MQA logits NextN=3 ### 📊 Changes **4 files changed** (+197 additions, -104 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp` (+13 -6) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh` (+180 -94) 📝 `deep_gemm/include/deep_gemm/scheduler/sm90_paged_mqa_logits.cuh` (+3 -3) 📝 `tests/test_attention.py` (+1 -1) </details> ### 📄 Description ## Summary This PR adds SM90/Hopper support for FP8 paged MQA logits with `next_n=3`, which is required by vLLM MTP speculative decoding when `num_speculative_tokens=2`. The unmodified SM90 path rejects this case with the existing `next_n == 1 or next_n == 2` assertion. The change is limited to the SM90 non-varlen path. Existing `next_n=1/2`, SM100 behavior, and the public Python API remain unchanged. Related issue: https://github.com/vllm-project/vllm/issues/43457 ## Implementation The SM90 kernel uses a two-token WGMMA atom. `next_n=3` is implemented as a pair plus a single-token tail instead of introducing a three-token WGMMA shape: ```text tokens 0 and 1 -> existing two-token WGMMA path token 2 -> single-token TailWGMMA path ``` - Metadata treats the three tokens as one logical scheduler atom, so the pair and tail computations reuse each loaded KV tile. - A four-stage Q pipeline prefetches the pair and tail TMA stages while keeping the existing producer/consumer barrier lifecycle. - The tail uses a two-token TMA box with the unused row zero-filled, but only the valid third-token row is reduced and stored. - Scheduler helpers retain the original batch row for block-table lookup and map the atom to the correct output token offset. SM90 varlen `next_n=3` and `next_n>3` remain out of scope. ## Validation ```text GPU: NVIDIA H20 96 GB (SM90) CUDA toolkit / NVCC: 12.9 PyTorch: 2.8.0+cu129 Baseline: 559d79fb6994a58b8a15b4b93bf13ccc16edf247 Candidate: 8ceef6bad7b4359234cecf87a48c63b6f0c8c0db ``` - Unmodified main reproduces the expected assertion for `next_n=3`. - The patched kernel passes `B=16,17,256,257`, `H=64`, `D=128`, `L=512`, FP32 logits, and 20 repeated calls per shape. Repeated outputs are bitwise identical; simulated-reference differences are `7.73e-9` to `8.21e-9`. - `tests/test_attention.py::test_paged_mqa_logits` passes all 36/36 SM90 cases: `next_n=1/2/3`, FP32/BF16 logits, `H=32/64`, `D=128`, `B=256/4096`, and average KV lengths `8192/65536` subject to the test's 32M-token cap. - `compute-sanitizer --tool memcheck` reports `0 errors` for `B=16,17,32,33`, `next_n=3`, `L=512`, with three repeated calls. ## Standalone Performance CUDA-event benchmark: `B=256`, `H=64`, `D=128`, `BLOCK_KV=64`, FP32 logits, 20 warmups, 100 timed iterations, five repeats. Values are median milliseconds per call; JIT compilation is excluded. | L | next_n | Upstream main | Patched | Change vs main | |---:|---:|---:|---:|---:| | 4096 | 1 | 0.06982 ms | 0.06994 ms | +0.16% | | 4096 | 2 | 0.13241 ms | 0.13172 ms | -0.53% | | 4096 | 3 | unsupported | 0.19984 ms | N/A | | 8192 | 1 | 0.13204 ms | 0.13155 ms | -0.37% | | 8192 | 2 | 0.25780 ms | 0.25631 ms | -0.58% | | 8192 | 3 | unsupported | 0.38742 ms | N/A | `T(next_n=3) / T(next_n=2)` is `1.517x` at `L=4096` and `1.512x` at `L=8192`; existing `next_n=1/2` performance remains within 1% of upstream. ## NCU Profile A separate same-machine comparison used one H20, CUDA 12.8, PyTorch 2.11.0+cu128, and Nsight Compute 2025.1. The table reports FP32 NCU Duration and patched `next_n=3` hardware metrics. | L | main n2 | patched n2 | n2 change | patched n3 | n3/n2 | SM | DRAM | Occupancy | |---:|---:|---:|---:|---:|---:|---:|---:|---:| | 8192 | 0.279840 ms | 0.278016 ms | -0.65% | 0.421952 ms | 1.518x | 94.80% | 18.91% | 31.01% | | 32768 | 1.079552 ms | 1.068640 ms | -1.01% | 1.604800 ms | 1.502x | 97.50% | 19.10% | 31.16% | For `next_n=3`, NCU reports 96 registers/thread and 185.4 KB dynamic shared memory per block. The kernel is compute/pipeline-bound rather than DRAM-bound; its occupancy, eligible-warps, and barrier-stall profile matches the prior implementation. Same-run main-versus-patched ratios are used here because absolute NCU durations from different profiler or toolchain configurations are not directly comparable. DeepGEMM recommends NVCC 12.9 or newer for best performance; the primary CUDA-event results above use CUDA 12.9. ## Integration Context An earlier revision was exercised end to end with vLLM on 2x H20 using TP=2 and `num_speculative_tokens=2`: the server initialized, `/v1/completions` returned `200 OK`, no DeepGEMM `next_n` assertion occurred, and speculative decode metrics were emitted. This is included as consumer-path validation; the refreshed commit's correctness and performance conclusions use the standalone results above. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
public-sync-20261014
0930-release
nv_dev
fix-mega-moe-fence
public-release-2609
public-release-2607
public-release-2606
mega-update
public-release-260416
ko3n1g/ci/build-wheels
nv_dev_f8e8fb5
nv_dev_950d31a
nv_dev_67fc648
nv_dev_c491439
nv_dev_4ff3f54
v2.1.1.post3
v2.1.1.post2
v2.1.1.post1
v2.1.1
v2.1.0
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepGEMM#373
No description provided.