[PR #344] [CLOSED] Fix an OOB access in the the paged MQA logits kernel #376

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/344
Author: @tirthasheshpatel
Created: 5/29/2026
Status: ❌ Closed

Base: main ← Head: fix-paged-mqa-logits-oob-context-lens

📝 Commits (1)

a8cf3fa fix(paged_mqa_logits): guard get_num_kv against out-of-batch q indices

📊 Changes

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

View changed files

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

📄 Description

The fp8_paged_mqa_logits kernel splits work across SMs. When there's less work than the number of SMs, some SMs get no work and are handed a past-the-end sentinel index (batch_size). The scheduler constructor ends up reading context_lens[batch_size] in get_num_kv which is an OOB read. Usually harmless (lands in padded memory), but if context_lens sits at a page edge and the read hits unmapped memory, we run into an IMA.

Observed the IMA here: https://github.com/sgl-project/sglang/issues/26541

Bisected to 7f2a703 (#304, "Mega MoE + FP4 Indexer release"), which introduced the new scheduler and out-of-batch atom emission.

A standalone repro:

"""IMA repro for DeepGEMM commit 7f2a703 (PagedMQALogitsScheduler OOB).

Covers decode (next_n=1) and speculative decode (next_n=2,3). Each case places
context_lens against an unmapped guard page; the scheduler's out-of-batch read
(context_lens[>= batch_size*next_n]) then page-faults on the unpatched kernel.
"""

import torch, deep_gemm
from cuda.bindings.driver import (
        CUmemAccess_flags, CUmemAccessDesc, CUmemAllocationGranularity_flags,
        CUmemAllocationProp, CUmemAllocationType, CUmemLocationType,
        cuCtxGetDevice, cuMemAddressReserve, cuMemCreate,
        cuMemGetAllocationGranularity, cuMemMap, cuMemSetAccess, cuMemcpyDtoD,
    )


def guarded_context_lens(batch_size, next_n, ctx):
    """
    A normal torch.empty((1,1), int32) won't fault on an OOB in CUDA.
    PyTorch's caching allocator rounds every allocation up to a large block
    (>=512 B, carved from 2 MiB+ segments), and puts the tensor at the start
    of the block so `context_lens[1]` lands inside mapped pool memory and
    just returns garbage. To force the fault, this function allocates a
    tensor at the tail of the page with nothing mapped immediately after
    it.
    """

    n = batch_size * next_n
    torch.cuda.synchronize()
    prop = CUmemAllocationProp()
    prop.type = CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
    prop.location.type = CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
    prop.location.id = int(cuCtxGetDevice()[1])
    page = int(cuMemGetAllocationGranularity(
        prop, CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM)[1])

    va = cuMemAddressReserve(2 * page, page, 0, 0)[1]
    cuMemMap(va, page, 0, cuMemCreate(page, prop, 0)[1], 0)
    access = CUmemAccessDesc()
    access.location = prop.location
    access.flags = CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
    cuMemSetAccess(va, page, [access], 1)

    data_ptr = int(va) + page - n * 4
    seed = torch.full((n,), ctx, dtype=torch.int32, device="cuda")
    cuMemcpyDtoD(data_ptr, seed.data_ptr(), n * 4)
    torch.cuda.synchronize()

    class _CAI:
        __cuda_array_interface__ = {
            "shape": (batch_size, next_n), "typestr": "<i4",
            "data": (data_ptr, False), "version": 3,
        }
    return torch.as_tensor(_CAI(), device="cuda")


def run(next_n):
    batch_size, num_heads, head_dim, page_sz, max_blocks = 1, 32, 128, 64, 32

    context_lens = guarded_context_lens(batch_size, next_n, 1024)
    q = torch.randn((batch_size, next_n, num_heads, head_dim),
                    dtype=torch.bfloat16, device="cuda").to(torch.float8_e4m3fn)
    kv_cache = torch.zeros((4096, page_sz, 1, head_dim + 4), dtype=torch.uint8, device="cuda")
    weights = torch.randn((batch_size * next_n, num_heads), dtype=torch.float32, device="cuda")
    block_table = torch.zeros((batch_size, max_blocks), dtype=torch.int32, device="cuda")

    schedule_meta = deep_gemm.get_paged_mqa_logits_metadata(
        context_lens, page_sz, deep_gemm.get_num_sms())
    deep_gemm.fp8_paged_mqa_logits(
        q, kv_cache, weights, context_lens, block_table, schedule_meta,
        max_blocks * page_sz, clean_logits=False)
    torch.cuda.synchronize()


if __name__ == "__main__":
    # next_n = 1 => no spec-decode
    # next_n = 2, 3 => spec-decode
    next_n = 1  # also fails for 2, and 3
    run(next_n)

cc @mattteochen @nvjullin @nvpohanh

_{🔄 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/344 **Author:** [@tirthasheshpatel](https://github.com/tirthasheshpatel) **Created:** 5/29/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix-paged-mqa-logits-oob-context-lens` --- ### 📝 Commits (1) - [`a8cf3fa`](https://github.com/deepseek-ai/DeepGEMM/commit/a8cf3fa16db9235c8f517602fa82c127f16e22c2) fix(paged_mqa_logits): guard get_num_kv against out-of-batch q indices ### 📊 Changes **1 file changed** (+2 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh` (+2 -0) </details> ### 📄 Description The `fp8_paged_mqa_logits` kernel splits work across SMs. When there's less work than the number of SMs, some SMs get no work and are handed a past-the-end sentinel index (`batch_size`). The scheduler constructor ends up reading `context_lens[batch_size]` in `get_num_kv` which is an OOB read. Usually harmless (lands in padded memory), but if `context_lens` sits at a page edge and the read hits unmapped memory, we run into an IMA. Observed the IMA here: https://github.com/sgl-project/sglang/issues/26541 Bisected to 7f2a703 (#304, "Mega MoE + FP4 Indexer release"), which introduced the new scheduler and out-of-batch atom emission. A standalone repro: ```python """IMA repro for DeepGEMM commit 7f2a703 (PagedMQALogitsScheduler OOB). Covers decode (next_n=1) and speculative decode (next_n=2,3). Each case places context_lens against an unmapped guard page; the scheduler's out-of-batch read (context_lens[>= batch_size*next_n]) then page-faults on the unpatched kernel. """ import torch, deep_gemm from cuda.bindings.driver import ( CUmemAccess_flags, CUmemAccessDesc, CUmemAllocationGranularity_flags, CUmemAllocationProp, CUmemAllocationType, CUmemLocationType, cuCtxGetDevice, cuMemAddressReserve, cuMemCreate, cuMemGetAllocationGranularity, cuMemMap, cuMemSetAccess, cuMemcpyDtoD, ) def guarded_context_lens(batch_size, next_n, ctx): """ A normal torch.empty((1,1), int32) won't fault on an OOB in CUDA. PyTorch's caching allocator rounds every allocation up to a large block (>=512 B, carved from 2 MiB+ segments), and puts the tensor at the start of the block so `context_lens[1]` lands inside mapped pool memory and just returns garbage. To force the fault, this function allocates a tensor at the tail of the page with nothing mapped immediately after it. """ n = batch_size * next_n torch.cuda.synchronize() prop = CUmemAllocationProp() prop.type = CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED prop.location.type = CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE prop.location.id = int(cuCtxGetDevice()[1]) page = int(cuMemGetAllocationGranularity( prop, CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM)[1]) va = cuMemAddressReserve(2 * page, page, 0, 0)[1] cuMemMap(va, page, 0, cuMemCreate(page, prop, 0)[1], 0) access = CUmemAccessDesc() access.location = prop.location access.flags = CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE cuMemSetAccess(va, page, [access], 1) data_ptr = int(va) + page - n * 4 seed = torch.full((n,), ctx, dtype=torch.int32, device="cuda") cuMemcpyDtoD(data_ptr, seed.data_ptr(), n * 4) torch.cuda.synchronize() class _CAI: __cuda_array_interface__ = { "shape": (batch_size, next_n), "typestr": "<i4", "data": (data_ptr, False), "version": 3, } return torch.as_tensor(_CAI(), device="cuda") def run(next_n): batch_size, num_heads, head_dim, page_sz, max_blocks = 1, 32, 128, 64, 32 context_lens = guarded_context_lens(batch_size, next_n, 1024) q = torch.randn((batch_size, next_n, num_heads, head_dim), dtype=torch.bfloat16, device="cuda").to(torch.float8_e4m3fn) kv_cache = torch.zeros((4096, page_sz, 1, head_dim + 4), dtype=torch.uint8, device="cuda") weights = torch.randn((batch_size * next_n, num_heads), dtype=torch.float32, device="cuda") block_table = torch.zeros((batch_size, max_blocks), dtype=torch.int32, device="cuda") schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( context_lens, page_sz, deep_gemm.get_num_sms()) deep_gemm.fp8_paged_mqa_logits( q, kv_cache, weights, context_lens, block_table, schedule_meta, max_blocks * page_sz, clean_logits=False) torch.cuda.synchronize() if __name__ == "__main__": # next_n = 1 => no spec-decode # next_n = 2, 3 => spec-decode next_n = 1 # also fails for 2, and 3 run(next_n) ``` cc @mattteochen @nvjullin @nvpohanh --- <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:00 +00:00