Skip to content

[Bugfix][ROCm] Bound kpool prefill tail-seed stores by NUM_TAIL_BLOCKS #56059

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zzw09773 wants to merge 1 commit into
base: main
Choose a base branch
from
+11 −0
Open

[Bugfix][ROCm] Bound kpool prefill tail-seed stores by NUM_TAIL_BLOCKS #56059

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions vllm/models/glm5next/amd/ops/kpool_compress.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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
Comment on lines +385 to +386

@chatgpt-codex-connector chatgpt-codex-connector Bot Sep 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Bound the decode tail accesses as well

In the reported prefix-cache scenario where this guard sees blk >= tail_kv_cache.shape[0], it only suppresses the prefill seed fault. The same request retains that tail-block mapping during decode, but both AMD and NVIDIA _kpool_decode_update_batched_kernel only require tail_slot >= 0 before using block = tail_slot // POOL_SIZE for tail reads and stores (AMD lines 499-509 and 600-609; NVIDIA lines 502-512 and 603-612). Consequently, the next decode step can access the same out-of-range block and GPU-fault; the decode path needs the corresponding upper-bound handling rather than relying on the new seed-only check.

Useful? React with 👍 / 👎.

ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to(
tl.int64
)
Expand Down Expand Up @@ -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]),
)


Expand Down
4 changes: 4 additions & 0 deletions vllm/models/glm5next/nvidia/ops/kpool_compress.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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
)
Expand Down Expand Up @@ -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]),
)


Expand Down
Loading