[PR #346] [CLOSED] fix: use dynamic shared memory in smxx_paged_mqa_logits_metadata to support large batch sizes #380

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/346
Author: @Oxygen56
Created: 5/31/2026
Status: ❌ Closed

Base: main ← Head: fix/issue-322

📝 Commits (2)

36abc9c Fix Mega MoE kernel shared memory calculation error
b4ab2b6 fix: use dynamic shared memory in smxx_paged_mqa_logits_metadata to support large batch sizes

📊 Changes

3 files changed (+11 additions, -9 deletions)

View changed files

📝 csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp (+2 -2)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+1 -1)
📝 deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh (+8 -6)

📄 Description

Summary

Fixes #322. The smxx_paged_mqa_logits_metadata kernel used static __shared__ arrays sized by kAlignedBatchSize, which caused shared memory overflow at batch sizes >= 4096. The three arrays (varlen_atom_token_start, varlen_atom_context_len, prefix_sum) consumed 3 * kAlignedBatchSize * 4 bytes, exceeding the 48KB limit.

Changes

deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh: Replaced three static __shared__ arrays with an extern __shared__ pool, using pointer arithmetic to partition it. Updated varlen_num_atoms_shared accesses to use pointer dereferencing.
csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp: Fixed the smem_size calculation from aligned_batch_size * sizeof(int) to (3 * aligned_batch_size + 1) * sizeof(int) to match the actual allocation (the FP4 version already had the correct calculation).

Background

The dynamic shared memory approach works because the construct_launch_config function in the JIT runtime already calls cudaFuncSetAttribute(..., cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size), which raises the per-block shared memory limit beyond the default 48KB to support the requested allocation (up to the GPU's capacity of ~227KB).

Tested with the reproducer from the issue at batch_size=8192.

Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com

_{🔄 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/346 **Author:** [@Oxygen56](https://github.com/Oxygen56) **Created:** 5/31/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/issue-322` --- ### 📝 Commits (2) - [`36abc9c`](https://github.com/deepseek-ai/DeepGEMM/commit/36abc9c3dd95c9b7412613c72c271b6bb6ad577b) Fix Mega MoE kernel shared memory calculation error - [`b4ab2b6`](https://github.com/deepseek-ai/DeepGEMM/commit/b4ab2b6c0f16831b2edb2cf3d2f698b205e2b4d2) fix: use dynamic shared memory in smxx_paged_mqa_logits_metadata to support large batch sizes ### 📊 Changes **3 files changed** (+11 additions, -9 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp` (+2 -2) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+1 -1) 📝 `deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh` (+8 -6) </details> ### 📄 Description ## Summary Fixes #322. The `smxx_paged_mqa_logits_metadata` kernel used static `__shared__` arrays sized by `kAlignedBatchSize`, which caused shared memory overflow at batch sizes >= 4096. The three arrays (`varlen_atom_token_start`, `varlen_atom_context_len`, `prefix_sum`) consumed `3 * kAlignedBatchSize * 4` bytes, exceeding the 48KB limit. ## Changes 1. **`deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh`**: Replaced three static `__shared__` arrays with an `extern __shared__` pool, using pointer arithmetic to partition it. Updated `varlen_num_atoms_shared` accesses to use pointer dereferencing. 2. **`csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp`**: Fixed the `smem_size` calculation from `aligned_batch_size * sizeof(int)` to `(3 * aligned_batch_size + 1) * sizeof(int)` to match the actual allocation (the FP4 version already had the correct calculation). ## Background The dynamic shared memory approach works because the `construct_launch_config` function in the JIT runtime already calls `cudaFuncSetAttribute(..., cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)`, which raises the per-block shared memory limit beyond the default 48KB to support the requested allocation (up to the GPU's capacity of ~227KB). Tested with the reproducer from the issue at batch_size=8192. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --- <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:04 +00:00