[PR #354] [CLOSED] Fix OOB shared-memory read in paged MQA logits metadata for empty batches #387

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/354
Author: @mingyangHao
Created: 6/4/2026
Status: ❌ Closed

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

📝 Commits (1)

6bdf229 Fix OOB shared-memory read in paged MQA logits metadata for empty batches

📊 Changes

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

View changed files

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

📄 Description

In the non-varlen branch of smxx_paged_mqa_logits_metadata, when the total amount of work is zero (every context_len maps to 0 KV segments, e.g. an all-empty decode/warmup batch), seg_starts is 0 for every SM and the binary search saturates to q_idx == batch_size. The subsequent

num_segs_q = prefix_sum[q_idx] - prefix_sum[q_idx - 1]

then reads prefix_sum[batch_size], which is out of bounds of the prefix_sum shared-memory buffer (it only holds kAlignedBatchSize entries), causing CUDA_ERROR_ILLEGAL_ADDRESS.

The varlen branch already guards the equivalent empty-work case. Clamp num_segs_q to 0 when q_idx >= batch_size so empty SMs fall through to the atom_idx == 0 / kv_split_idx == 0 path and emit the batch_size * num_next_n_atoms end sentinel.

Repro: call get_paged_mqa_logits_metadata with an all-zero context_lens tensor; compute-sanitizer reports "Invalid shared read" in smxx_paged_mqa_logits_metadata, which this patch eliminates. Surfaced by TensorRT-LLM DeepSeek-V4 sparse-MLA indexer warmup (empty KV) on Blackwell.

_{🔄 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/354 **Author:** [@mingyangHao](https://github.com/mingyangHao) **Created:** 6/4/2026 **Status:** ❌ Closed **Base:** `nv_dev` ← **Head:** `fix/paged-mqa-logits-metadata-empty-oob` --- ### 📝 Commits (1) - [`6bdf229`](https://github.com/deepseek-ai/DeepGEMM/commit/6bdf229bfb15e60e755f4f6c6d1d1a2d57aa8bdc) Fix OOB shared-memory read in paged MQA logits metadata for empty batches ### 📊 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 In the non-varlen branch of `smxx_paged_mqa_logits_metadata`, when the total amount of work is zero (every `context_len` maps to 0 KV segments, e.g. an all-empty decode/warmup batch), `seg_starts` is 0 for every SM and the binary search saturates to `q_idx == batch_size`. The subsequent num_segs_q = prefix_sum[q_idx] - prefix_sum[q_idx - 1] then reads `prefix_sum[batch_size]`, which is out of bounds of the `prefix_sum` shared-memory buffer (it only holds `kAlignedBatchSize` entries), causing CUDA_ERROR_ILLEGAL_ADDRESS. The varlen branch already guards the equivalent empty-work case. Clamp `num_segs_q` to 0 when `q_idx >= batch_size` so empty SMs fall through to the `atom_idx == 0` / `kv_split_idx == 0` path and emit the `batch_size * num_next_n_atoms` end sentinel. Repro: call `get_paged_mqa_logits_metadata` with an all-zero `context_lens` tensor; compute-sanitizer reports "Invalid __shared__ read" in `smxx_paged_mqa_logits_metadata`, which this patch eliminates. Surfaced by TensorRT-LLM DeepSeek-V4 sparse-MLA indexer warmup (empty KV) on Blackwell. --- <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:08 +00:00