[Bugfix] Fix int32 overflow in concat_mla_q flat warp index - #45384
waynehacking8 wants to merge 1 commit into
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
ConcatMLAQKernel computed blockIdx.x * blockDim.x in 32-bit arithmetic, which wraps once the launch exceeds 2^32 threads (more than 2^27 token-heads, e.g. 128 heads x ~1.05M tokens). Wrapped warps alias earlier token/head slots, so they redundantly rewrite early rows of q_out while their own rows are never written, producing incorrect q_out values. Compute the flat warp id in 64-bit and compare against the 64-bit token-head count. Also widen the launcher's total_warps (int x int, latent overflow past 2^31 token-heads) and bounds-check the grid size. Fixes vllm-project#45373 Co-authored-by: Claude Signed-off-by: Wayne Chiu <waynehacking8@gmail.com>
572d15c to
745c0ee
Compare
1 participant
Purpose
Fix #45373:
ConcatMLAQKernelcomputesblockIdx.x * blockDim.xin 32-bit arithmetic. Once the launch exceeds 2^32 threads (> 2^27 token-heads — e.g. the reporter's 128 heads × 140×7491 tokens = 134.2M token-heads → 4.296e9 threads), the product wraps and later warps alias earlier token/head slots: they redundantly rewrite earlyq_outrows while their own rows are never written → incorrectq_out.Changes (both sites):
concat_mla_q.cuh): computeflat_warp_idin 64-bit; compare against the 64-bit token-head count.token_id/head_idstayint(each individually bounded bynum_tokens/num_heads).cache_kernels.cu):total_warps = num_tokens * num_headswasint × int— a second, latent overflow past 2^31 token-heads. Widened toint64_twith aSTD_TORCH_CHECKon the final grid size.Test Plan
The trigger requires a ~154 GB
q_out(134M token-heads × 576 × fp16), which exceeds any single GPU, so validation is a standalone nvcc harness (sm_120, CUDA 13.0) compiling the fixed header from this tree plus an index-math probe replicating the old arithmetic:The old math wraps exactly at block 2^32/256 = 16,777,216, aliasing onto flat id 7 — the mechanism in the issue. The fixed math is correct at all probed blocks.
Leg B — no-regression at feasible scale: the fixed
ConcatMLAQKernel<__half, 512>run on 8192×128 = 1,048,576 token-heads (~1.2 GB, below the wrap threshold), compared element-wise against a CPU reference: 0/603,979,776 mismatches (byte-identical). On sm_120 + CUDA 13 this exercises the 256-bit PTX path (VLLM_256B_PTX_ENABLED=1).clang-format pre-commit hook passes on both files.
Test Result
Leg A: old arithmetic wraps at trigger scale, fixed arithmetic correct. Leg B: byte-identical output at non-overflow scale. (Honest framing: the >150 GB end-to-end trigger cannot be run on a single GPU; correctness at trigger scale is established by the device-level index-math probe + the unchanged data path.)
Duplicate check
No open PR addresses this: searched
45373,concat_mla_q in:title, and the file history (git log origin/main -- csrc/libtorch_stable/concat_mla_q.cuh— last touched by the libtorch_stable migration, no overflow fix). #36743 optimizes the same kernel for ROCm (perf, not a fix) — flagging for conflict awareness.(Optional) Documentation Update
None.
AI assistance disclosure: prepared with AI assistance (Claude); the change was reviewed, compiled, and validated locally by the submitter.