iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][GLM-5.3-Flash] Address kpool tail blocks by the padded indexer stride in the NVIDIA prefill seed kernel (#57477)

Browse Source 
Signed-off-by: Jared Wen <w13431838023@gmail.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Jared Wen
2026-09-20 18:15:04 +08:00
committed by GitHub
co-authored by Claude Fable 5.1 mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 0748d3bd57
commit db1bfdd4fb
2 changed files with 23 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/test_kpool_decode_update_batched.py
+10 -3
View File
@@ -384,9 +384,16 @@ def test_leading_invalid_tail_slot():
_assert_eq(r_ref, r_kern)
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
def test_amd_prefill_seed_honors_padded_tail_block_stride():
"""The tail shares a padded indexer allocation in production."""
def test_prefill_seed_honors_padded_tail_block_stride():
"""The tail shares a padded indexer allocation in production.
``get_kv_cache_config_from_groups`` aliases each tail tensor onto its
indexer tensor with the indexer's block stride (38016 B for GLM-5.3-Flash
vs a dense 2048 B tail block), so a seed kernel that addresses blocks
densely writes into an unrelated indexer block and leaves the request's
tail block untouched. Runs on every platform; the NVIDIA kernel had this
bug while the AMD kernel did not.
"""
kpool = 4
num_blocks = 6
logical_block_elems = 2 * kpool * HEAD_DIM
vllm/models/glm5next/nvidia/ops/kpool_compress.py
+13 -2
View File
@@ -383,6 +383,8 @@ def _kpool_tail_seed_kernel(
tslot_ptr,
tail_ptr,
n_tokens,
TAIL_BLOCK_ELEMS: tl.constexpr,
KPOOL_HEAD: tl.constexpr,
HEAD_DIM: tl.constexpr,
KPOOL: tl.constexpr,
BLOCK_D: tl.constexpr,
@@ -393,6 +395,11 @@ def _kpool_tail_seed_kernel(
ahead belongs to a different tail block (or is past the batch / padding,
slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is
``tail[block, {0:K, 1:score}, pos % KPOOL, :]``.
The tail cache aliases the indexer cache with the indexer's (padded) block
stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` /
``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense
``[num_blocks, 2, KPOOL, HEAD_DIM]`` array.
"""
i = tl.program_id(0)
t = tl.load(tslot_ptr + i).to(tl.int64)
@@ -409,11 +416,11 @@ def _kpool_tail_seed_kernel(
return
offs = tl.arange(0, BLOCK_D)
m = offs < HEAD_DIM
base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM
base = blk * TAIL_BLOCK_ELEMS + (t % KPOOL) * HEAD_DIM
k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)
s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)
tl.store(tail_ptr + base + offs, k, mask=m)
tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m)
tl.store(tail_ptr + base + KPOOL_HEAD + offs, s, mask=m)
def kpool_seed_tail_cache(
@@ -426,6 +433,8 @@ def kpool_seed_tail_cache(
) -> None:
"""Seed the paged tail cache from a prefill batch (see the kernel)."""
assert tail_kv_cache.dtype == torch.bfloat16
assert tail_kv_cache.ndim == 4 and tail_kv_cache.shape[1] == 2
assert tail_kv_cache.stride(3) == 1 and tail_kv_cache.stride(2) == head_dim
assert key.dtype == torch.bfloat16
n = tslot.shape[0]
if n == 0:
@@ -436,6 +445,8 @@ def kpool_seed_tail_cache(
tslot,
tail_kv_cache,
n,
TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),
KPOOL_HEAD=tail_kv_cache.stride(1),
HEAD_DIM=head_dim,
KPOOL=kpool,
BLOCK_D=triton.next_power_of_2(head_dim),