From db1bfdd4fb0dd7b8226402ee00abc7a987561b7c Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sun, 20 Sep 2026 18:15:04 +0800 Subject: [PATCH] [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 Co-authored-by: Claude Fable 5.1 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/kernels/test_kpool_decode_update_batched.py | 13 ++++++++++--- vllm/models/glm5next/nvidia/ops/kpool_compress.py | 15 +++++++++++++-- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index afc0c1f836..37ba4b58c8 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -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 diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index 75ccfa8c2a..416d3b5c48 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -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),