[PR #459] [MERGED] Sparse MQA logits follow-ups for nv_dev: 16-byte paged KV stride, workspace reuse under graph capture #463

New issue

Closed

opened 2026-10-26 21:18:08 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:08 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/459
Author: @lfr-0531
Created: 9/23/2026
Status: ✅ Merged
Merged: 9/24/2026
Merged by: @RayWang96

Base: nv_dev ← Head: pr/nv-dev-sparse-paged-fixes

📝 Commits (2)

422debb Reuse a warmed sparse MQA logits workspace during graph capture
337b1f8 Accept 16-byte aligned KV pages in the paged sparse MQA logits

📊 Changes

2 files changed (+14 additions, -4 deletions)

View changed files

📝 csrc/apis/attention.hpp (+13 -3)
📝 deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh (+1 -1)

📄 Description

Summary

Two small follow-ups to #447 for the DSv4.1 sparse MQA logits, found while integrating nv_dev into TensorRT-LLM (index cache pages are 64 x (64 + 4) = 4352-byte page-footer pages, and CUDA graphs are captured on a side stream). No kernel algorithm changes.

Accept 16-byte aligned KV pages in the paged sparse MQA logits. fp8_fp4_paged_sparse_mqa_logits required fused_kv_cache.stride(0) % 512 == 0 (host and device asserts). The paged accessor only computes page = base + page_idx * kv_page_stride_bytes and copies the KV / SF blocks with 16-byte cp.async chunks, so 16-byte alignment is the actual requirement (the cache base pointer is now checked for it as well). Caches whose page stride is not padded to 512 bytes can be read in place.
Reuse a warmed sparse MQA logits workspace during graph capture. The per-stream workspace asserted that its first allocation never happens under capture. Frameworks that capture on a dedicated side stream but replay on the stream they warmed up hit that assert on the first captured metadata call. When the current stream is capturing and has no workspace yet, reuse the first workspace warmed on the device (the workspace is per-launch scratch, so this assumes replays are serialized with that stream's eager launches); the assert stays for the genuinely cold case.

Validation

GB300 (SM100), CUDA 13.4, PyTorch 2.14 nv26.08: TensorRT-LLM's CSA2 (DSv4.1 sparse attention) unit suite, 453 cases including paged sparse decode with CUDA-graph capture and replay against native 4352-byte pages: 404 passed / 48 skipped / 1 known harness-only failure, identical to the previous pinned DeepGEMM build. With nv_dev head alone the 7 paged decode cases fail on the 512-byte stride assert.
DeepGEMM's own paged sparse tests pad pages to 512 bytes and are unaffected.

Supersedes #458 (the nv_dev + main 26/09 merge), which #447 has landed.

🤖 Generated with Claude Code

_{🔄 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/459 **Author:** [@lfr-0531](https://github.com/lfr-0531) **Created:** 9/23/2026 **Status:** ✅ Merged **Merged:** 9/24/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `pr/nv-dev-sparse-paged-fixes` --- ### 📝 Commits (2) - [`422debb`](https://github.com/deepseek-ai/DeepGEMM/commit/422debbdf9dcf69421bf4e316944a9c975145ab7) Reuse a warmed sparse MQA logits workspace during graph capture - [`337b1f8`](https://github.com/deepseek-ai/DeepGEMM/commit/337b1f85a9f2add7ff5875266771c0277d5dedc0) Accept 16-byte aligned KV pages in the paged sparse MQA logits ### 📊 Changes **2 files changed** (+14 additions, -4 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/attention.hpp` (+13 -3) 📝 `deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh` (+1 -1) </details> ### 📄 Description ## Summary Two small follow-ups to #447 for the DSv4.1 sparse MQA logits, found while integrating `nv_dev` into TensorRT-LLM (index cache pages are 64 x (64 + 4) = 4352-byte page-footer pages, and CUDA graphs are captured on a side stream). No kernel algorithm changes. 1. **Accept 16-byte aligned KV pages in the paged sparse MQA logits.** `fp8_fp4_paged_sparse_mqa_logits` required `fused_kv_cache.stride(0) % 512 == 0` (host and device asserts). The paged accessor only computes `page = base + page_idx * kv_page_stride_bytes` and copies the KV / SF blocks with 16-byte `cp.async` chunks, so 16-byte alignment is the actual requirement (the cache base pointer is now checked for it as well). Caches whose page stride is not padded to 512 bytes can be read in place. 2. **Reuse a warmed sparse MQA logits workspace during graph capture.** The per-stream workspace asserted that its first allocation never happens under capture. Frameworks that capture on a dedicated side stream but replay on the stream they warmed up hit that assert on the first captured metadata call. When the current stream is capturing and has no workspace yet, reuse the first workspace warmed on the device (the workspace is per-launch scratch, so this assumes replays are serialized with that stream's eager launches); the assert stays for the genuinely cold case. ## Validation * GB300 (SM100), CUDA 13.4, PyTorch 2.14 nv26.08: TensorRT-LLM's CSA2 (DSv4.1 sparse attention) unit suite, 453 cases including paged sparse decode with CUDA-graph capture and replay against native 4352-byte pages: 404 passed / 48 skipped / 1 known harness-only failure, identical to the previous pinned DeepGEMM build. With `nv_dev` head alone the 7 paged decode cases fail on the 512-byte stride assert. * DeepGEMM's own paged sparse tests pad pages to 512 bytes and are unaffected. Supersedes #458 (the nv_dev + main 26/09 merge), which #447 has landed. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:18:08 +00:00