[PR #353] [MERGED] Fix OOB in paged MQA logits metadata scheduler #386

New issue

Closed

opened 2026-10-26 21:17:06 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:06 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/353
Author: @Barry-Delaney
Created: 6/4/2026
Status: ✅ Merged
Merged: 6/4/2026
Merged by: @RayWang96

Base: nv_dev ← Head: fix/paged-mqa-metadata-oob

📝 Commits (1)

07a044c Fix OOB in paged MQA logits metadata scheduler

📊 Changes

1 file changed (+1 additions, -1 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh (+1 -1)

📄 Description

When all context_lens are zero (e.g. warmup decode with kv_lens=1 and indexer_compress_ratio=4 giving gen_seq_lens=0), the binary search pushes lo to batch_size. The subsequent prefix_sum[q_idx] access is OOB when q_idx == batch_size == kAlignedBatchSize, since prefix_sum only has indices [0, batch_size).

Fix: clamp q_idx to batch_size - 1. When the binary search lands past all queries (total work = 0), num_segs_q correctly evaluates to 0, and the existing guards (num_segs_q > 0 ? ...) produce the sentinel values (atom_idx=0, kv_split_idx=0).

_{🔄 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/353 **Author:** [@Barry-Delaney](https://github.com/Barry-Delaney) **Created:** 6/4/2026 **Status:** ✅ Merged **Merged:** 6/4/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `fix/paged-mqa-metadata-oob` --- ### 📝 Commits (1) - [`07a044c`](https://github.com/deepseek-ai/DeepGEMM/commit/07a044c4cab056ce8a18ea32c4418f229c33c8f9) Fix OOB in paged MQA logits metadata scheduler ### 📊 Changes **1 file changed** (+1 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh` (+1 -1) </details> ### 📄 Description When all context_lens are zero (e.g. warmup decode with kv_lens=1 and indexer_compress_ratio=4 giving gen_seq_lens=0), the binary search pushes lo to batch_size. The subsequent prefix_sum[q_idx] access is OOB when q_idx == batch_size == kAlignedBatchSize, since prefix_sum only has indices [0, batch_size). Fix: clamp q_idx to batch_size - 1. When the binary search lands past all queries (total work = 0), num_segs_q correctly evaluates to 0, and the existing guards (num_segs_q > 0 ? ...) produce the sentinel values (atom_idx=0, kv_split_idx=0). --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:17:06 +00:00