On unified-memory devices (e.g. NVIDIA GB10 / Grace Blackwell, 121 GB shared between host and device) the CUDA caching allocator's reserved pool competes directly with system RAM. In that setting the sparse lightning indexer's dense-logits buffer becomes a practical problem during long prefills.
fp8_mqa_logits_triton() (vllm/models/deepseek_v4/nvidia/ops/sm12x_mqa.py, line ~140; two sibling variants in the same file do the same) allocates its output with a fresh torch.empty((num_q, seq_len_kv), dtype=torch.float32) on every call. During a chunked prefill this happens once per chunk per compressed layer, and seq_len_kv (the compressed context) grows monotonically as the prefill advances. Consequences:
- the caching allocator cannot reuse previously freed blocks (each request is larger than every cached block), so it keeps mapping new segments;
memory_allocated stays flat while memory_reserved ratchets upward and is not returned;- on unified memory that reserve is system RAM: on a GB10 the ratchet is ~2.6–3 GiB for a single 32K-token needle prompt and scales with context length, pushing the machine toward its memory watchdog.
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True does not help (measured: identical growth).
Observed with torch.cuda.memory._record_memory_history(): in one prefill window, 105 allocations up to 256 MiB attributed to that stack, with fresh segment_map events of ~254 MiB each at runtime.
A possible low-risk fix: keep one flat, grow-only fp32 buffer per device and hand the kernel buffer[:M*N].view(M, N). The view is contiguous with strides (N, 1) — byte-identical layout to a fresh torch.empty(M, N) — and the kernel already takes strides as arguments, so the generated Triton code and reduction order are unchanged. Since split_indexer_prefill_chunks (vllm/v1/attention/backends/mla/indexer.py) already bounds M * N * 4 to VLLM_SPARSE_INDEXER_MAX_LOGITS_MB, such a pool converges to that bound and stops growing.
Happy to provide the full memory-history traces if useful.
Disclosure: this report was prepared with AI assistance (Claude); the measurements and code reading were verified by a human before filing.
On unified-memory devices (e.g. NVIDIA GB10 / Grace Blackwell, 121 GB shared between host and device) the CUDA caching allocator's reserved pool competes directly with system RAM. In that setting the sparse lightning indexer's dense-logits buffer becomes a practical problem during long prefills.
fp8_mqa_logits_triton()(vllm/models/deepseek_v4/nvidia/ops/sm12x_mqa.py, line ~140; two sibling variants in the same file do the same) allocates its output with a freshtorch.empty((num_q, seq_len_kv), dtype=torch.float32)on every call. During a chunked prefill this happens once per chunk per compressed layer, andseq_len_kv(the compressed context) grows monotonically as the prefill advances. Consequences:memory_allocatedstays flat whilememory_reservedratchets upward and is not returned;PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Truedoes not help (measured: identical growth).Observed with
torch.cuda.memory._record_memory_history(): in one prefill window, 105 allocations up to 256 MiB attributed to that stack, with freshsegment_mapevents of ~254 MiB each at runtime.A possible low-risk fix: keep one flat, grow-only fp32 buffer per device and hand the kernel
buffer[:M*N].view(M, N). The view is contiguous with strides(N, 1)— byte-identical layout to a freshtorch.empty(M, N)— and the kernel already takes strides as arguments, so the generated Triton code and reduction order are unchanged. Sincesplit_indexer_prefill_chunks(vllm/v1/attention/backends/mla/indexer.py) already boundsM * N * 4toVLLM_SPARSE_INDEXER_MAX_LOGITS_MB, such a pool converges to that bound and stops growing.Happy to provide the full memory-history traces if useful.
Disclosure: this report was prepared with AI assistance (Claude); the measurements and code reading were verified by a human before filing.