[Bugfix][GLM-5.3-Flash] Address kpool tail blocks by the padded indexer stride in the NVIDIA prefill seed kernel (#57477)
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>
Jared Wen committed
Sep 20, 2026 at 18:15 UTC
db1bfdd4fb0dd7b8226402ee00abc7a987561b7c
2 files changed
+23
-5
tests/kernels/test_kpool_decode_update_batched.py
+10
-3
@@ -384,9 +384,16 @@ def test_leading_invalid_tail_slot():
384
_assert_eq(r_ref, r_kern)
385
386
387
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
388
-def test_amd_prefill_seed_honors_padded_tail_block_stride():
389
- """The tail shares a padded indexer allocation in production."""
387
+def test_prefill_seed_honors_padded_tail_block_stride():
388
+ """The tail shares a padded indexer allocation in production.
389
+
390
+ ``get_kv_cache_config_from_groups`` aliases each tail tensor onto its
391
+ indexer tensor with the indexer's block stride (38016 B for GLM-5.3-Flash
392
+ vs a dense 2048 B tail block), so a seed kernel that addresses blocks
393
+ densely writes into an unrelated indexer block and leaves the request's
394
+ tail block untouched. Runs on every platform; the NVIDIA kernel had this
395
+ bug while the AMD kernel did not.
396
+ """
397
kpool = 4
398
num_blocks = 6
399
logical_block_elems = 2 * kpool * HEAD_DIM
vllm/models/glm5next/nvidia/ops/kpool_compress.py
+13
-2
@@ -383,6 +383,8 @@ def _kpool_tail_seed_kernel(
383
tslot_ptr,
384
tail_ptr,
385
n_tokens,
386
+ TAIL_BLOCK_ELEMS: tl.constexpr,
387
+ KPOOL_HEAD: tl.constexpr,
388
HEAD_DIM: tl.constexpr,
389
KPOOL: tl.constexpr,
390
BLOCK_D: tl.constexpr,
@@ -393,6 +395,11 @@ def _kpool_tail_seed_kernel(
395
ahead belongs to a different tail block (or is past the batch / padding,
396
slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is
397
``tail[block, {0:K, 1:score}, pos % KPOOL, :]``.
398
+
399
+ The tail cache aliases the indexer cache with the indexer's (padded) block
400
+ stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` /
401
+ ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense
402
+ ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array.
403
"""
404
i = tl.program_id(0)
405
t = tl.load(tslot_ptr + i).to(tl.int64)
@@ -409,11 +416,11 @@ def _kpool_tail_seed_kernel(
416
return
417
offs = tl.arange(0, BLOCK_D)
418
m = offs < HEAD_DIM
412
- base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM
419
+ base = blk * TAIL_BLOCK_ELEMS + (t % KPOOL) * HEAD_DIM
420
k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)
421
s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)
422
tl.store(tail_ptr + base + offs, k, mask=m)
416
- tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m)
423
+ tl.store(tail_ptr + base + KPOOL_HEAD + offs, s, mask=m)
424
425
426
def kpool_seed_tail_cache(
@@ -426,6 +433,8 @@ def kpool_seed_tail_cache(
433
) -> None:
434
"""Seed the paged tail cache from a prefill batch (see the kernel)."""
435
assert tail_kv_cache.dtype == torch.bfloat16
436
+ assert tail_kv_cache.ndim == 4 and tail_kv_cache.shape[1] == 2
437
+ assert tail_kv_cache.stride(3) == 1 and tail_kv_cache.stride(2) == head_dim
438
assert key.dtype == torch.bfloat16
439
n = tslot.shape[0]
440
if n == 0:
@@ -436,6 +445,8 @@ def kpool_seed_tail_cache(
445
tslot,
446
tail_kv_cache,
447
n,
448
+ TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),
449
+ KPOOL_HEAD=tail_kv_cache.stride(1),
450
HEAD_DIM=head_dim,
451
KPOOL=kpool,
452
BLOCK_D=triton.next_power_of_2(head_dim),