deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #311] [CLOSED] fix: correct fused KV cache stride in paged MQA logits #358

New issue

Closed

opened 2026-10-26 21:16:46 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:16:46 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/311
Author: @JasonOA888
Created: 4/19/2026
Status: ❌ Closed

Base: main ← Head: fix/paged-mqa-kv-cache-stride

📝 Commits (1)

510c226 fix: correct fused KV cache stride in paged MQA logits

📊 Changes

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

View changed files

📝 csrc/apis/attention.hpp (+6 -6)
📝 tests/test_attention.py (+12 -8)

📄 Description

Bug

fp8_fp4_paged_mqa_logits decomposes the fused KV cache into separate kv_cache and kv_cache_sf tensors via torch::from_blob. The strides passed to from_blob assume a flat layout (all values then all SF), but the API assertion enforces an interleaved layout ([values | SF] per row):

DG_HOST_ASSERT(fused_kv_cache.stride(1) == fp4_with_sf_bytes);  // FP4: 68
DG_HOST_ASSERT(fused_kv_cache.stride(1) == head_dim_with_sf);  // FP8: 132

What was wrong

kv_cache stride(1):

FP4: was head_dim/2 (64), should be fp4_with_sf_bytes (68)
FP8: was head_dim (128), should be head_dim_with_sf (132)

kv_cache_sf offset:

Was block_kv * head_dim, pointing past all data rows into undefined territory
Should be head_dim (or head_dim/2 for FP4), pointing to the SF of the first row

kv_cache_sf stride(1):

Was 1 (contiguous), but SF values are interleaved every fp4_with_sf_bytes/head_dim_with_sf bytes

Impact

TMA descriptors constructed with wrong strides cause misaligned reads for any row > 0 in the paged KV cache. This silently corrupts attention logits for multi-row blocks.

The test helper had the same flat-layout assumption, masking the bug for small block sizes where corruption fell in the last 1-2 rows.

Fix

Use the correct row stride (fp4_with_sf_bytes / head_dim_with_sf) for both kv_cache and kv_cache_sf from_blob
Fix kv_cache_sf offset to head_dim / head_dim / 2 (start of first row's SF)
Fix test helpers to construct proper interleaved layout

_{🔄 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/311 **Author:** [@JasonOA888](https://github.com/JasonOA888) **Created:** 4/19/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/paged-mqa-kv-cache-stride` --- ### 📝 Commits (1) - [`510c226`](https://github.com/deepseek-ai/DeepGEMM/commit/510c2263c8c11cf4355d76eaf2b3d2821d8d8fdb) fix: correct fused KV cache stride in paged MQA logits ### 📊 Changes **2 files changed** (+18 additions, -14 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/attention.hpp` (+6 -6) 📝 `tests/test_attention.py` (+12 -8) </details> ### 📄 Description ## Bug `fp8_fp4_paged_mqa_logits` decomposes the fused KV cache into separate `kv_cache` and `kv_cache_sf` tensors via `torch::from_blob`. The strides passed to `from_blob` assume a **flat layout** (all values then all SF), but the API assertion enforces an **interleaved layout** (`[values | SF]` per row): ```cpp DG_HOST_ASSERT(fused_kv_cache.stride(1) == fp4_with_sf_bytes); // FP4: 68 DG_HOST_ASSERT(fused_kv_cache.stride(1) == head_dim_with_sf); // FP8: 132 ``` ### What was wrong **kv_cache stride(1):** - FP4: was `head_dim/2` (64), should be `fp4_with_sf_bytes` (68) - FP8: was `head_dim` (128), should be `head_dim_with_sf` (132) **kv_cache_sf offset:** - Was `block_kv * head_dim`, pointing past all data rows into undefined territory - Should be `head_dim` (or `head_dim/2` for FP4), pointing to the SF of the first row **kv_cache_sf stride(1):** - Was `1` (contiguous), but SF values are interleaved every `fp4_with_sf_bytes`/`head_dim_with_sf` bytes ### Impact TMA descriptors constructed with wrong strides cause misaligned reads for any row > 0 in the paged KV cache. This silently corrupts attention logits for multi-row blocks. The test helper had the same flat-layout assumption, masking the bug for small block sizes where corruption fell in the last 1-2 rows. ## Fix - Use the correct row stride (`fp4_with_sf_bytes` / `head_dim_with_sf`) for both `kv_cache` and `kv_cache_sf` from_blob - Fix `kv_cache_sf` offset to `head_dim` / `head_dim / 2` (start of first row's SF) - Fix test helpers to construct proper interleaved layout --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:16:46 +00:00

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

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#358

No description provided.

Rows
Columns