Repository navigation
Conversation
…2 overflow _apply_write_kernel computes row_idx * row_stride + start_idx where row_idx is loaded from an int32 staging buffer and row_stride is a host scalar (gpu.stride(0)) that Triton specializes to int32. The multiply therefore runs in int32 and wraps once row_idx * stride >= 2**31, yielding a negative (or re-wrapped) offset: an out-of-bounds store, i.e. a GPU memory access fault or a silent write to the wrong row. The same class of bug affects the sampler penalties path: PenaltiesState builds its idx_mapping with an explicit dtype=torch.int32 (everywhere else the mapping is np.intp/int64 and Triton promotes), so _bincount_kernel and _penalties_kernel compute req_state_idx * stride in int32. Practical trigger: RequestState.all_token_ids is (max_num_reqs, max_model_len) int32, so at max_model_len=1M any row index >= 2148 overflows; MRV2's free_indices.pop() assigns the first request the highest row, so any engine with max_num_reqs >= 2149 faults on request #1 (or during warmup, when sampler warmup exercises the bincount path eagerly). This affects every platform, not just ROCm. Promote the row index to int64 at load in _apply_write_kernel, _penalties_kernel and _bincount_kernel, matching the existing house pattern in _zero_kv_blocks_kernel. Width-promoted index math costs one widened multiply per program. Regression tests allocate a >2^31-element int32 tensor and are skipped unless the device has enough free memory. Both tests abort the process with a GPU memory access fault when the casts are reverted. Closes vllm-project#57030 Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm
requested review from
WoosukKwon,
njhill and
yewentao256
as code owners
September 29, 2026 07:39
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR does
Closes #57030.
Casts the row index to
tl.int64at load in the three Model Runner V2 kernels whose stride products can leave int32:vllm/v1/worker/gpu/buffer_utils.py_apply_write_kernel(row_idx * row_stride + start_idx)row_idx * max_model_lenvllm/v1/worker/gpu/sample/penalties.py_bincount_kernel(req_state_idx * all_token_ids_stride)req_state_idx * max_model_lenvllm/v1/worker/gpu/sample/penalties.py_penalties_kernel(same load, same mapping, cast for symmetry)req_state_idx * stridesrow_idxcomes from an int32 staging buffer androw_strideis a host scalar (gpu.stride(0)) that Triton specializes to int32, so the multiply runs in int32. Same for the penalties path, wherePenaltiesState.apply_staged_writesbuildsidx_mappingwith an explicitdtype=torch.int32(everywhere else the mapping isnp.intp/int64 and Triton's usuali32 * i64 -> i64promotion hides the problem).Once
row_idx * stride >= 2**31the offset wraps → out-of-bounds store/load (GPU memory access fault) or, when the re-wrapped offset lands back inside the allocation, silent corruption on the wrong row. WithRequestState.all_token_ids = (max_num_reqs, max_model_len)int32 and MRV2'sfree_indices.pop()assigning the first request the highest row, any engine at 1M context withmax_num_reqs >= 2149dies on request #1 — and the penalties variant additionally trips during MRV2's kernel warmup (sampler warmup stages penalty requests bottom-up frommax_num_reqs - 1). This is not ROCm-specific: any platform running MRV2 with largemax_model_len * max_num_reqsis exposed; the issue is only labelled ROCm because that's where it was reported.The pattern matches the existing house fix in
_zero_kv_blocks_kernel(vllm/v1/worker/utils.py), which already does.to(tl.int64). Cost: one widened multiply per program.Verification (8x MI355X, gfx950, ROCm 7.2, torch 2.12.0+git6bbd260, main @
6cbbea8721)Repo is pristine on this branch except for this PR (
git statusclean before/after).StagedWriteTensor((65536, 40000), int32), write row 65535Memory access fault by GPU node-2 ... Reason: Unknown(SIGABRT + GPU coredump) — byte-identical to the reporter's crash[1,2,3,4]at row 65535all_token_ids-shaped(4096, 1_000_000), row 30002.0e9 < 2**31)(5000, 1M), row 4295bincountat row 2559 of(2560, 1M)(realpenalties.bincount)Threshold, for reviewers wondering why the defaults don't trip it: overflow iff
row_idx * max_model_len >= 2**31→ row ≥ 2148 at 1M ctx → needsmax_num_reqs ≥ 2149(or ~716 for mropeprefill_positions,num_dims=3). At ≤512k context with defaultmax_num_seqsthe product stays below2**31, which is why the reporter saw "≤512k fine, 1M dies every time".Tests
New
tests/v1/worker/test_gpu_int64_row_offsets.py:test_staged_write_high_row— write through_apply_write_kernelat row 2559 of a(2560, 1M)int32 tensor (10 GiB), assert values land at the target row and not clobber row 0 (the silent-corruption signature).test_bincount_high_row— same-shapedall_token_idsthrough the realpenalties.bincount, assert prompt bits / output counts land on the right row.Both are gated on
current_platform.is_cuda_alike()+ ≥11 GiB free device memory (torch.accelerator.get_memory_info), since the tensors must exceed 2³¹ elements by construction — no smaller shape demonstrates the bug.Switchback evidence on this branch: with the casts reverted, each test independently aborts the process with
Fatal Python error: Aborted(GPU memory access fault, exit 134); with the fix,2 passed. Full-repo sanity:tests/kernels/core/test_uva.py24 passed, 4 skipped.Note on #57524
#57524 targets the same bug. This PR is an independent fix preferred by maintainers: it is based on current
main(#57524 no longer applies — itstest_uva.pyhunk anchors at an EOF that has since moved), gates the bincount test onis_cuda_alike()instead ofis_cuda()(which isFalseunderRocmPlatform, so that test could never run on the platform named in the issue), and covers_penalties_kernelas well.