[Bug]: _compute_slot_mapping_kernel reads block_table out of bounds for cache groups whose block table is narrower than the sequence

[zzw09773]

Summary

_compute_slot_mapping_kernel in vllm/v1/worker/block_table.py masks its block-table load by
token validity and CP ownership, but not by whether the computed column index is inside the
row. Any KV cache group whose block table does not span the sequence in raw token positions
therefore performs an unbounded out-of-bounds read, scaling with sequence length.

On 8× MI325X (gfx942, ROCm 7.2.3) this eventually crosses into an unmapped page and takes down
every rank with Memory access fault by GPU node-N ... Reason: Unknown.

The missing guard

vllm/v1/worker/block_table.py @ main 94d96e24:

block_indices = (                                    # :463-466
    virtual_block_indices * BLOCKS_PER_KV_BLOCK
    + local_block_offsets // block_size
)
block_numbers = tl.load(                             # :467-471
    block_table_ptr + row_offset + block_indices,
    mask=mask & is_local,                            # :469  ← no range term
    other=0,
).to(tl.int64)
slot_offsets = local_block_offsets % block_size
slot_ids = block_numbers * block_size + slot_offsets
slot_ids = tl.where(is_local, slot_ids, PAD_ID)      # :474  ← no range term

The model-side QSA slot helpers in #53896 do bound this (valid &= logical_blocks < block_table.shape[1]
and a clamp before indexing); the generic kernel is the one path we found that does not.

Two independent instances, both from in-flight model PRs

Neither spec class is on main yet; both are added by open PRs, and each declares a
one-block-per-request side cache that the kernel then indexes with raw token positions:

PR	spec	max_num_blocks_per_req()	effective row width	index computed
#53896 (Qwen3.8-Flash-Next)	CircularBufferSpec	1	128 tokens (after the 128-token alignment in get_block_table_width)	pos // 8
#53906 (GLM-5.3-Flash)	KpoolTailSpec	1	128 tokens	pos // index_kpool

Both reach the kernel because vllm/v1/worker/gpu_model_runner.py:7387-7390 assigns
SlotMappingMode.NONE only to KVCacheSpecKind.MAMBA and TOKEN_TO_KV_SLOT to everything else.
(That classification rule deserves its own discussion; we have raised it on both PRs rather than here.)

Observed

Instrumenting the kernel to print when max(block_indices) >= block_table_stride:

[slotmap!!OOR] kvbs=8 bs=8 bpk=1 stride=16 rows=16

A 200-token request triggers it 456 times. At ~95k tokens the read is ~11,875 int32 elements
(≈47 KB) past a ~1 KB tensor, and that is where it started faulting for us. The threshold is not
a clean constant — it depends on what the caching allocator has mapped after the tensor, which is
why it is latency-to-detect and layout-dependent. Enabling speculative decoding (which adds KV
cache groups and shifts the allocation layout) moved the threshold: with MTP
num_speculative_tokens=3 the fault appeared at ~95k tokens, with =1 at ~100k, and without
MTP the same model ran 250k without faulting — the read was still landing in mapped memory.

Fix

+        in_range = block_indices < block_table_stride
         block_numbers = tl.load(
             block_table_ptr + row_offset + block_indices,
-            mask=mask & is_local,
+            mask=mask & is_local & in_range,
             other=0,
         ).to(tl.int64)
         slot_offsets = local_block_offsets % block_size
         slot_ids = block_numbers * block_size + slot_offsets
-        slot_ids = tl.where(is_local, slot_ids, PAD_ID)
+        slot_ids = tl.where(is_local & in_range, slot_ids, PAD_ID)

For a group whose block table does span the sequence, in_range is always true and behaviour is
unchanged. For a narrow group the entry becomes PAD_SLOT_ID instead of a wild read — and in both
instances above the generic output is discarded anyway, because each builder recomputes its own
circular mapping (see the KpoolTailMetadataBuilder docstring in #53906, which states the generic
per-group kernel output "cannot express a 1-block-per-request circular buffer").

Result

With the two-line guard, Qwen3.8-Flash-Next with MTP num_speculative_tokens=3 at
--max-model-len 262144 — the configuration that faulted at ~95k above — runs needle-in-haystack
100k / 150k / 250k all 3/3 with zero GPU faults. GLM-5.3-Flash without MTP (which had not faulted
before the guard; its read stayed inside mapped memory) is unchanged at 100k / 250k 3/3.

Independent of ROCm?

The missing mask is platform-independent. Whether it manifests on CUDA depends on the allocator
layout and how far past the row the read goes; we only have gfx942 data.

Reproduction

1. Serve either model from the PRs above on 8 GPUs with a long context.
2. Add a debug print in _compute_slot_mapping_kernel when
   max(block_indices) >= block_table_stride (one-shot; the condition is true on essentially
   every step, so an unconditional device print will flood).
3. Any request longer than the side cache's aligned row width (128 tokens for both specs) trips it.

What we verified / did not verify

• Verified: the missing mask on main @ 94d96e24 (no in_range token anywhere in the file);
  the OOR instrumentation counts; that the two-line guard eliminates the faults on our hardware.
• Not verified: that the out-of-bounds read is the only cause of every fault we saw. It was
  for one configuration (GLM without MTP, and Qwen with MTP); a second GLM configuration
  (with MTP) still faults after this fix, for an unrelated reason in aiter's MLA decode dispatch,
  filed separately.
• Not verified: any CUDA-platform manifestation.

[github-actions]

CC @hongxiayang @tjtanaa @vllmellm @giuseppegrossi for ROCm-related issue

[frank-suwen]

Hi, thanks for the detailed investigation and hardware validation. Are you planning to submit the kernel-side guard yourself, or would it be helpful if I prepared a separate PR? I reviewed current main and did not find an overlapping implementation PR.

I also noticed that the Model Runner V2 _compute_slot_mappings_kernel() in vllm/v1/worker/gpu/block_table.py appears to contain an analogous unbounded block-table load. If useful, I can investigate covering both that path and the reported ComputeSlotMappingKernel in vllm/v1/worker/block_table.py with the proposed bounds mask and PAD_SLOT_ID behavior, together with focused synthetic GPU regression tests.

I can independently validate the kernel behavior on NVIDIA CUDA while treating your MI325X results as the ROCm end-to-end evidence. Please let me know how you would prefer to coordinate and whether both runner paths should be included in scope. Thanks!

[zzw09773]

Thanks — please go ahead and prepare the PR; that is the better outcome. We only have the ROCm side, and our fix lives in a locally patched image rather than a branch.

Scope: yes, please include both paths. We did not look at the V2 runner's _compute_slot_mappings_kernel in vllm/v1/worker/gpu/block_table.py; if it carries the same unbounded block-table load it should get the same guard.

What we did locally (for reference, not a proposal for the exact form): in ComputeSlotMappingKernel we added in_range = block_indices < block_table_stride, folded it into the existing mask (mask & is_local & in_range), and used tl.where(is_local & in_range, ..., PAD_SLOT_ID) for the slot ids, so out-of-range lookups become PAD_SLOT_ID rather than reads past the row. That is what has been serving on MI325X since; it is the reason the 262k-context runs with MTP (num_speculative_tokens=3) on Qwen3.8-Flash-Next stopped faulting.

Validation: once the PR is up, ping us and we will run it on the 8×MI325X setup — the long-context MTP case that originally faulted (≥~95k tokens with CircularBufferSpec's one-block table) plus the short-context regression, and report back on the PR. Your CUDA validation plus synthetic tests covering both kernels sounds like the right split.

[zzw09773]

One more note for the V2 path, from reading main @ fd57c4b (vllm/v1/worker/gpu/block_table.py, _compute_slot_mappings_kernel, lines 333-344):

block_indices = local_positions // kernel_block_size
block_offsets = local_positions % kernel_block_size
block_numbers = tl.load(
    block_table_ptr + req_state_idx * block_table_stride + block_indices,
    mask=is_local,
    other=0,
)
slot_ids = block_numbers * kernel_block_size + block_offsets
if CP_SIZE != 1:
    slot_ids = tl.where(is_local, slot_ids, PAD_ID)

Two things worth building into the guard:

1. block_indices is never compared against block_table_stride (the row width), so the same class of out-of-row read exists here as in the V1 kernel. Adding block_indices < block_table_stride to the load mask closes it.
2. The other=0 on that load is a trap once a range mask is added: a masked-out lane yields block_numbers = 0, and slot_ids = 0 * kernel_block_size + block_offsets is a legitimate slot inside block 0, not a padding slot. In the common CP_SIZE == 1 path there is no tl.where(..., PAD_ID) afterwards (is_local is the constant True), so the value would be stored as-is. The guard therefore needs to write PAD_ID explicitly for out-of-range lanes, the way the V1 kernel does at its tl.where(is_local, slot_ids, PAD_ID).

Neither of these is a new finding beyond what you described — just the two spots we would check first when reviewing the PR. We have not run the V2 runner on this model (it selects the V1 runner), so this is from reading only.