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

[GH-ISSUE #322] get_paged_mqa_logits_metadata() does not work with large batch size due to excessive static shared memory impl #160

New issue
Closed
opened 2026-10-26 21:13:56 +00:00 by bsakaguchi · 1 comment
bsakaguchi commented 2026-10-26 21:13:56 +00:00
Owner
Copy link

Originally created by @gau-nernst on GitHub (Apr 29, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/322

import torch

from vllm.third_party.deep_gemm import get_paged_mqa_logits_metadata

batch_size = 8192
block_size = 64
num_sms = torch.cuda.get_device_properties(0).multi_processor_count

context_lens = torch.full((batch_size, 1), 4096, dtype=torch.int32, device="cuda")
indices = torch.arange(batch_size, dtype=torch.int32, device="cuda") // 2

get_paged_mqa_logits_metadata(
    context_lens,
    block_size,
    num_sms,
    indices=indices,
)
torch.cuda.synchronize()

NVCC compilation failed: ptxas error   : Entry function '_ZN9deep_gemm5sched30smxx_paged_mqa_logits_metadataILj8192ELj256ELj152ELb1EEEvjjbPKjS3_Pj' uses too much shared data (0x18004 bytes, 0xc000 max)

Traceback (most recent call last):
  File "/home/thien/debug.py", line 12, in <module>
    get_paged_mqa_logits_metadata(
RuntimeError: Assertion error (/workspace/.deps/deepgemm-src/csrc/apis/../jit_kernels/impls/../../jit/compiler.hpp:228): false and "NVCC compilation failed"

This is caused by static shared memory being linear with batch size

github.com/deepseek-ai/DeepGEMM@891d57b4db/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh (L19-L20)

Would it possible to make it support large batch size instead? Either with dynamic shared memory or reimplement in a way that doesn't need that much shared memory.

The context is that I'm trying to use fp8_fp4_paged_mqa_logits() for prefill case, which requires calling this metadata preparation kernel. Thank you.

Originally created by @gau-nernst on GitHub (Apr 29, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/322 ```python import torch from vllm.third_party.deep_gemm import get_paged_mqa_logits_metadata batch_size = 8192 block_size = 64 num_sms = torch.cuda.get_device_properties(0).multi_processor_count context_lens = torch.full((batch_size, 1), 4096, dtype=torch.int32, device="cuda") indices = torch.arange(batch_size, dtype=torch.int32, device="cuda") // 2 get_paged_mqa_logits_metadata( context_lens, block_size, num_sms, indices=indices, ) torch.cuda.synchronize() ``` ``` NVCC compilation failed: ptxas error : Entry function '_ZN9deep_gemm5sched30smxx_paged_mqa_logits_metadataILj8192ELj256ELj152ELb1EEEvjjbPKjS3_Pj' uses too much shared data (0x18004 bytes, 0xc000 max) Traceback (most recent call last): File "/home/thien/debug.py", line 12, in <module> get_paged_mqa_logits_metadata( RuntimeError: Assertion error (/workspace/.deps/deepgemm-src/csrc/apis/../jit_kernels/impls/../../jit/compiler.hpp:228): false and "NVCC compilation failed" ``` This is caused by static shared memory being linear with batch size https://github.com/deepseek-ai/DeepGEMM/blob/891d57b4db1071624b5c8fa0d1e51cb317fa709f/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh#L19-L20 Would it possible to make it support large batch size instead? Either with dynamic shared memory or reimplement in a way that doesn't need that much shared memory. The context is that I'm trying to use `fp8_fp4_paged_mqa_logits()` for prefill case, which requires calling this metadata preparation kernel. Thank you.
bsakaguchi commented 2026-10-26 21:13:59 +00:00
Author
Owner
Copy link

@toffee-desuwa commented on GitHub (May 30, 2026):

@LyricZhao — I looked into #322 and I think I found the cause. Wanted to ask your preferred direction before anyone writes a fix.

The cause: In paged_mqa_logits.cuh, the kernel smxx_paged_mqa_logits_metadata declares three static shared arrays, each sized by kAlignedBatchSize:

varlen_atom_token_start[kAlignedBatchSize] (line 19)
varlen_atom_context_len[kAlignedBatchSize] (line 20)
prefix_sum[kAlignedBatchSize] (line 57)
That is about 3 * 4 * kAlignedBatchSize bytes of static shared memory. When the batch size is large (roughly 4096+), this goes over the 48 KB static smem limit, so the kernel cannot launch. The host side in attention.hpp does not catch this: there is no batch-size check before launch, and the only smem accounting is for dynamic shared memory, which this kernel does not use. So it cannot detect a static smem overflow. This matches @gau-nernst's report: fine at small batch, breaks at large batch.

Three possible fixes (not sure which you prefer):

Option 1 — Move the three arrays to dynamic shared memory (extern shared), and pass the size through the launch config and the host check, so the existing accounting covers them.

Option 2 — Process the batch in fixed-size chunks, so static smem stays bounded no matter how large the batch is.

Option 3 — Keep it static but add a host-side DG_HOST_ASSERT on the batch-size limit. This is the smallest change: it turns a silent launch failure into a clear error, but it caps the supported batch instead of fixing it.

I'm happy to open a PR for whichever direction fits the codebase best. I just didn't want to decide the design myself on a core kernel. Thanks!

<!-- gh-comment-id:4582475995 --> @toffee-desuwa commented on GitHub (May 30, 2026): @LyricZhao — I looked into [#322](https://github.com/deepseek-ai/DeepGEMM/issues/322) and I think I found the cause. Wanted to ask your preferred direction before anyone writes a fix. The cause: In paged_mqa_logits.cuh, the kernel smxx_paged_mqa_logits_metadata declares three static __shared__ arrays, each sized by kAlignedBatchSize: varlen_atom_token_start[kAlignedBatchSize] (line 19) varlen_atom_context_len[kAlignedBatchSize] (line 20) prefix_sum[kAlignedBatchSize] (line 57) That is about 3 * 4 * kAlignedBatchSize bytes of static shared memory. When the batch size is large (roughly 4096+), this goes over the 48 KB static smem limit, so the kernel cannot launch. The host side in attention.hpp does not catch this: there is no batch-size check before launch, and the only smem accounting is for dynamic shared memory, which this kernel does not use. So it cannot detect a static smem overflow. This matches @gau-nernst's report: fine at small batch, breaks at large batch. Three possible fixes (not sure which you prefer): Option 1 — Move the three arrays to dynamic shared memory (extern __shared__), and pass the size through the launch config and the host check, so the existing accounting covers them. Option 2 — Process the batch in fixed-size chunks, so static smem stays bounded no matter how large the batch is. Option 3 — Keep it static but add a host-side DG_HOST_ASSERT on the batch-size limit. This is the smallest change: it turns a silent launch failure into a clear error, but it caps the supported batch instead of fixing it. I'm happy to open a PR for whichever direction fits the codebase best. I just didn't want to decide the design myself on a core kernel. Thanks!
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#160
No description provided.