From 0b964b0a9db402150e64d5eab4cb58a959386ab3 Mon Sep 17 00:00:00 2001 From: "Hsiu-Ying,Kung" Date: Wed, 9 Sep 2026 18:38:44 +0800 Subject: [PATCH] [Bugfix][ROCm] Bound kpool prefill tail-seed stores by NUM_TAIL_BLOCKS Prefill _kpool_tail_seed_kernel addressed tail[blk] with a physical id that can be main-KV magnitude after prefix-cache churn. Decode already refused block >= n_tail; seed did not, and the store GPU-faulted. Measured on gfx942 GLM-5.3-Flash (vllm#56037 Crash D): HIP_LAUNCH_BLOCKING captured ShaderName _kpool_tail_seed_kernel. --- vllm/models/glm5next/amd/ops/kpool_compress.py | 7 +++++++ vllm/models/glm5next/nvidia/ops/kpool_compress.py | 4 ++++ 2 files changed, 11 insertions(+) diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 64e6fd46d73f..1420664ffab0 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -365,6 +365,7 @@ def _kpool_tail_seed_kernel( HEAD_DIM: tl.constexpr, KPOOL: tl.constexpr, BLOCK_D: tl.constexpr, + NUM_TAIL_BLOCKS: tl.constexpr, ): """Copy token ``i``'s raw K + gate into its request's tail block. @@ -378,6 +379,11 @@ def _kpool_tail_seed_kernel( if t < 0: return blk = t // KPOOL # t >= 0 here, so trunc == floor + # Prefill seed used to address tail[blk] with a physical id that can be + # main-KV magnitude after prefix-cache churn. Decode already bounds this; + # skip OOB stores (vllm#56037 Crash D). + if blk >= NUM_TAIL_BLOCKS: + return ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to( tl.int64 ) @@ -423,6 +429,7 @@ def kpool_seed_tail_cache( HEAD_DIM=head_dim, KPOOL=kpool, BLOCK_D=triton.next_power_of_2(head_dim), + NUM_TAIL_BLOCKS=int(tail_kv_cache.shape[0]), ) diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index 70b41021aab0..4c5f70d5eebe 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -378,6 +378,7 @@ def _kpool_tail_seed_kernel( HEAD_DIM: tl.constexpr, KPOOL: tl.constexpr, BLOCK_D: tl.constexpr, + NUM_TAIL_BLOCKS: tl.constexpr, ): """Copy token ``i``'s raw K + gate into its request's tail block. @@ -391,6 +392,8 @@ def _kpool_tail_seed_kernel( if t < 0: return blk = t // KPOOL # t >= 0 here, so trunc == floor + if blk >= NUM_TAIL_BLOCKS: + return ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to( tl.int64 ) @@ -431,6 +434,7 @@ def kpool_seed_tail_cache( HEAD_DIM=head_dim, KPOOL=kpool, BLOCK_D=triton.next_power_of_2(head_dim), + NUM_TAIL_BLOCKS=int(tail_kv_cache.shape[0]), )