@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2617fe9383

[Bugfix][GLM-5.3-Flash] kpool corruption with speculative decoding (#58454)

Signed-off-by: Matt Mastracci <matthew@mastracci.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Matt Mastracci committed Sep 25, 2026 at 10:51 UTC 2617fe938355594c48d4512a2ef6b470962aac1a

5 files changed +208 -49

tests/kernels/test_kpool_decode_update_batched.py

+90 -4

@@ -293,8 +293,9 @@ def _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos):

293

294

295 @pytest.mark.parametrize("pool_size", [4, 16])

296 -def test_decode_writer_matches_prefill_writer(pool_size):

297 - """Compare production decode and prefill writers for pool sizes 4 and 16."""

296 +@pytest.mark.parametrize("ring_pools", [1, 2])

297 +def test_decode_writer_matches_prefill_writer(pool_size, ring_pools):

298 + ring = ring_pools * pool_size

299 n_pools, page, nblk = 8, 64, 4

300 n_tok = n_pools * pool_size

301 dev = "cuda"

@@ -317,14 +318,14 @@ def test_decode_writer_matches_prefill_writer(pool_size):

318

319 # One request owning tail block 0, fed one token per decode step.

320 kv_decode = torch.zeros_like(kv_prefill)

320 - tail = torch.zeros(nblk, 2, pool_size, HEAD_DIM, dtype=torch.bfloat16, device=dev)

321 + tail = torch.zeros(nblk, 2, ring, HEAD_DIM, dtype=torch.bfloat16, device=dev)

322 for t in range(n_tok):

323 completes = t % pool_size == pool_size - 1

324 kpool_decode_update_and_maybe_write_cache_batched(

325 kv_decode,

326 tail,

327 # token-granular: every token has a valid tail slot

327 - torch.tensor([[t % pool_size]], dtype=torch.int32, device=dev),

328 + torch.tensor([[t % ring]], dtype=torch.int32, device=dev),

329 k[t].view(1, 1, HEAD_DIM),

330 score[t].view(1, 1, HEAD_DIM),

331 ape,

@@ -351,6 +352,91 @@ def test_decode_writer_matches_prefill_writer(pool_size):

352 )

353

354

355 +@pytest.mark.parametrize("ring_pools", [1, 2])

356 +def test_rejected_draft_redo_needs_ring_slots(ring_pools):

357 + """With a one-pool ring, the drafts behind a rejected pool-completing draft

358 + overwrote the pool's earlier keys, so its redo compressed wrong keys."""

359 + pool, spec, page, nblk = 4, 3, 64, 2

360 + ring = ring_pools * pool

361 + dev = "cuda"

362 + torch.manual_seed(1)

363 + n_tok = 3 * pool

364 + k = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev)

365 + score = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev)

366 + ape = torch.randn(pool, HEAD_DIM, dtype=torch.float32, device=dev)

367 + kv_ref = torch.zeros(nblk, page, HEAD_DIM + 4, dtype=torch.uint8, device=dev)

368 + kpool_compress_and_write_cache(

369 + kv_ref,

370 + k.view(3, pool, HEAD_DIM),

371 + score.view(3, pool, HEAD_DIM),

372 + ape,

373 + torch.arange(3, dtype=torch.int64, device=dev),

374 + pool_size=pool,

375 + head_dim=HEAD_DIM,

376 + round_scale=ROUND_SCALE,

377 + )

378 +

379 + kv = torch.zeros_like(kv_ref)

380 + tail = torch.zeros(nblk, 2, ring, HEAD_DIM, dtype=torch.bfloat16, device=dev)

381 +

382 + def pool_bytes(cache, p):

383 + # Page layout: [page * HEAD_DIM bytes of K rows | page * 4 bytes of scales].

384 + flat = cache[0].reshape(-1)

385 + k_bytes = flat[p * HEAD_DIM : (p + 1) * HEAD_DIM]

386 + s_bytes = flat[page * HEAD_DIM + 4 * p : page * HEAD_DIM + 4 * (p + 1)]

387 + return torch.cat([k_bytes, s_bytes])

388 +

389 + def step(positions, keys, scores):

390 + pos = torch.tensor([positions], dtype=torch.int32, device=dev)

391 + slots = [(p // pool) if p % pool == pool - 1 else -1 for p in positions]

392 + kpool_decode_update_and_maybe_write_cache_batched(

393 + kv,

394 + tail,

395 + pos % ring,

396 + keys.view(1, -1, HEAD_DIM),

397 + scores.view(1, -1, HEAD_DIM),

398 + ape,

399 + torch.tensor([slots], dtype=torch.int32, device=dev),

400 + pos,

401 + pool,

402 + HEAD_DIM,

403 + round_scale=ROUND_SCALE,

404 + )

405 +

406 + for t in range(7):

407 + step([t], k[t], score[t])

408 + # Control: verified token 7 completes pool 1 before drafts 8..10 are stashed.

409 + drafts = torch.randn(spec, HEAD_DIM, dtype=torch.bfloat16, device=dev)

410 + draft_scores = torch.randn(spec, HEAD_DIM, dtype=torch.bfloat16, device=dev)

411 + step(

412 + [7, 8, 9, 10],

413 + torch.cat([k[7:8], drafts]),

414 + torch.cat([score[7:8], draft_scores]),

415 + )

416 + step([8, 9, 10, 11], k[8:12], score[8:12]) # all drafts rejected

417 + for p in (1, 2):

418 + assert torch.equal(pool_bytes(kv, p), pool_bytes(kv_ref, p)), p

419 +

420 + # Draft 7 completes pool 1 and is rejected. With a one-pool ring, drafts

421 + # 8 and 9 overwrite the slots of positions 4 and 5, which are read by the

422 + # redo of 7.

423 + kv.zero_()

424 + tail.zero_()

425 + for t in range(6):

426 + step([t], k[t], score[t])

427 + step(

428 + [6, 7, 8, 9],

429 + torch.cat([k[6:7], drafts]),

430 + torch.cat([score[6:7], draft_scores]),

431 + )

432 + step([7, 8, 9, 10], k[7:11], score[7:11])

433 + pool1_ok = torch.equal(pool_bytes(kv, 1), pool_bytes(kv_ref, 1))

434 + if ring_pools == 1:

435 + assert not pool1_ok, "expected a one-pool ring to corrupt pool 1"

436 + else:

437 + assert pool1_ok

438 +

439 +

440 def test_leading_invalid_tail_slot():

441 """A request whose FIRST token carries an invalid (-1) tail slot while a

442 later token is a real pool completion.

tests/v1/attention/test_kpool_tail_slot_mapping.py

+68 -11

@@ -24,6 +24,8 @@ from types import SimpleNamespace

24 import pytest

25 import torch

27 +from vllm.config import VllmConfig, set_current_vllm_config

28 +from vllm.models.glm5next.common.attention import Glm5NextTailCache

29 from vllm.v1.attention.backend import CommonAttentionMetadata

30 from vllm.v1.attention.backends.mla.indexer import (

31 KpoolTailBackend,

@@ -56,6 +58,30 @@ def test_tail_backend_layout_matches_kernel_pointer_arithmetic():

58 assert content_stride == 1

61 +@pytest.mark.parametrize(

62 + "num_speculative_tokens, ring", [(0, 4), (1, 8), (4, 8), (7, 16), (13, 32)]

63 +)

64 +def test_tail_ring_divides_the_attention_block(num_speculative_tokens, ring):

65 + """With 7 draft tokens the ring was 12 slots. 12 does not divide the

66 + 640-token KDA attention block, so the scheduler block grew to their lcm

67 + and prefix-cache hits were cut down to multiples of 1920."""

68 + with set_current_vllm_config(VllmConfig()):

69 + cache = Glm5NextTailCache(

70 + head_dim=128,

71 + dtype=torch.bfloat16,

72 + prefix="tail",

73 + cache_config=SimpleNamespace(block_size=640),

74 + index_kpool=KPOOL,

75 + )

76 + spec = cache.get_kv_cache_spec(

77 + SimpleNamespace(num_speculative_tokens=num_speculative_tokens)

78 + )

79 +

80 + assert spec.block_size == spec.sliding_window == ring

81 + assert ring >= KPOOL + num_speculative_tokens

82 + assert 640 % ring == 0

83 +

84 +

85 def test_tail_spec_opts_out_of_generic_slot_mapping():

86 """The tail row is one block wide (padded to the block-table alignment), so

87 the generic kernel's ``pos // kpool`` column index runs off the end of the

@@ -296,31 +322,30 @@ def test_builder_reuses_slot_mapping_storage():

322

323

324 class TailRingMirror:

299 - """Mirror of _kpool_tail_seed_kernel / _kpool_decode_update_batched_kernel

300 - addressing: block = tail_slot // kpool, ring offset = pos % kpool; a pool

301 - completing at pos reads ring slots (pool_start + s) % kpool and uses the

302 - current token's own K/score for the last member."""

325 + """Mirror of the tail-ring addressing in _kpool_tail_seed_kernel and

326 + _kpool_decode_update_batched_kernel."""

327

304 - def __init__(self, num_blocks, kpool=KPOOL):

328 + def __init__(self, num_blocks, kpool=KPOOL, ring=None):

329 self.kpool = kpool

306 - self.k = torch.full((num_blocks, kpool, 3), float("nan"))

307 - self.s = torch.full((num_blocks, kpool, 3), float("nan"))

330 + self.ring = ring or kpool

331 + self.k = torch.full((num_blocks, self.ring, 3), float("nan"))

332 + self.s = torch.full((num_blocks, self.ring, 3), float("nan"))

333

334 def stash(self, tail_slot, pos, k, s):

310 - blk, off = tail_slot // self.kpool, pos % self.kpool

335 + blk, off = tail_slot // self.ring, pos % self.ring

336 self.k[blk, off] = k

337 self.s[blk, off] = s

338

339 seed = stash # the seed kernel writes with the same addressing

340

341 def complete(self, tail_slot, pos, k, s):

317 - blk = tail_slot // self.kpool

342 + blk = tail_slot // self.ring

343 start = pos - (self.kpool - 1)

344 kk = torch.stack(

320 - [self.k[blk, (start + i) % self.kpool] for i in range(self.kpool)]

345 + [self.k[blk, (start + i) % self.ring] for i in range(self.kpool)]

346 )

347 ss = torch.stack(

323 - [self.s[blk, (start + i) % self.kpool] for i in range(self.kpool)]

348 + [self.s[blk, (start + i) % self.ring] for i in range(self.kpool)]

349 )

350 kk[-1], ss[-1] = k, s # is_current for the completing token

351 w = torch.softmax(ss, dim=0)

@@ -454,3 +479,35 @@ def test_triton_mapping_reads_strided_block_table():

479 for req, blk in enumerate(own_blocks):

480 start, end = int(qsl[req]), int(qsl[req + 1])

481 assert (got[start:end] // KPOOL == blk).all()

482 +

483 +

484 +@pytest.mark.parametrize("ring_pools", [1, 2])

485 +def test_rejected_completing_draft_needs_ring_slots(ring_pools):

486 + """With a one-pool ring, the drafts behind a rejected pool-completing draft

487 + overwrote the pool's earlier keys, so its redo compressed wrong keys."""

488 + ring_size = ring_pools * KPOOL

489 + truth = TailRingMirror(num_blocks=2, ring=ring_size)

490 + ring = TailRingMirror(num_blocks=2, ring=ring_size)

491 + block = 1

492 +

493 + def slot(pos):

494 + return block * ring_size + pos % ring_size

495 +

496 + for pos in range(4, 7):

497 + truth.stash(slot(pos), pos, *token_kv(0, pos))

498 + ring.stash(slot(pos), pos, *token_kv(0, pos))

499 + expected = truth.complete(slot(7), 7, *token_kv(0, 7))

500 +

501 + # Draft 7 completes the pool and is rejected. With a one-pool ring, drafts

502 + # 8..10 overwrite the slots of positions 4..6, which are read by the redo

503 + # of 7.

504 + for pos in range(7, 11):

505 + k, s = token_kv(9, pos) # draft values

506 + if pos % KPOOL == KPOOL - 1:

507 + ring.complete(slot(pos), pos, k, s)

508 + ring.stash(slot(pos), pos, k, s)

509 + redo = ring.complete(slot(7), 7, *token_kv(0, 7))

510 + if ring_pools == 1:

511 + assert not torch.allclose(redo, expected)

512 + else:

513 + torch.testing.assert_close(redo, expected)

vllm/models/glm5next/amd/ops/kpool_compress.py

+18 -16

@@ -327,37 +327,38 @@ def _kpool_tail_seed_kernel(

327 KPOOL_HEAD: tl.constexpr,

328 HEAD_DIM: tl.constexpr,

329 KPOOL: tl.constexpr,

330 + RING: tl.constexpr,

331 BLOCK_D: tl.constexpr,

332 ):

332 - """Copy token ``i``'s raw K + gate into its request's tail block.

333 + """Copy token ``i``'s raw K + gate into its request's tail ring.

334

335 Token ``i`` is among its request's last KPOOL tokens iff the token KPOOL

336 ahead belongs to a different tail block (or is past the batch / padding,

336 - slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is

337 - ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``.

337 + slot < 0). ``tslot = block * RING + pos % RING``; the destination is

338 + ``tail[block, {0:K, 1:score}, pos % RING, :]``.

339 """

340 i = tl.program_id(0)

341 t = tl.load(tslot_ptr + i).to(tl.int64)

342 if t < 0:

343 return

343 - blk = t // KPOOL # t >= 0 here, so trunc == floor

344 + blk = t // RING # t >= 0 here, so trunc == floor

345 ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to(

346 tl.int64

347 )

348 # Match the torch semantics exactly: a negative ahead slot floors to a

349 # block id that differs from every real block -> token is in the tail.

350 # Only divide non-negative slots (Triton int div truncates, torch floors).

350 - if ahead >= 0 and ahead // KPOOL == blk:

351 + if ahead >= 0 and ahead // RING == blk:

352 return

353 offs = tl.arange(0, BLOCK_D)

354 m = offs < HEAD_DIM

355 block_base = blk * TAIL_BLOCK_ELEMS

355 - base = block_base + (t % KPOOL) * HEAD_DIM

356 + base = block_base + (t % RING) * HEAD_DIM

357 k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)

358 s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)

359 tl.store(tail_ptr + base + offs, k, mask=m)

360 tl.store(

360 - tail_ptr + block_base + KPOOL_HEAD + (t % KPOOL) * HEAD_DIM + offs, s, mask=m

361 + tail_ptr + block_base + KPOOL_HEAD + (t % RING) * HEAD_DIM + offs, s, mask=m

362 )

363

364

@@ -385,6 +386,7 @@ def kpool_seed_tail_cache(

386 KPOOL_HEAD=tail_kv_cache.stride(1),

387 HEAD_DIM=head_dim,

388 KPOOL=kpool,

389 + RING=tail_kv_cache.shape[2],

390 BLOCK_D=triton.next_power_of_2(head_dim),

391 )

392

@@ -412,6 +414,7 @@ def _kpool_decode_update_batched_kernel(

414 PAGE_SIZE: tl.constexpr,

415 BUF_NUMEL_PER_PAGE: tl.constexpr,

416 POOL_SIZE: tl.constexpr,

417 + RING: tl.constexpr,

418 TAIL_BLOCK_ELEMS: tl.constexpr,

419 KPOOL_HEAD: tl.constexpr,

420 HEAD_DIM: tl.constexpr,

@@ -428,10 +431,7 @@ def _kpool_decode_update_batched_kernel(

431 the tail-ring slots that tokens t < t* (same request) just stashed in this

432 same invocation. ``tl.range`` iterates sequentially within the program, so

433 those stashes are visible to the later completion read. Cross-request

431 - programs are independent (distinct tail blocks). With NEXT_N < POOL_SIZE

432 - (the spec-verify case: NEXT_N ~= num_spec+1, POOL_SIZE=16) at most one

433 - completion can occur per request per call, but the ordered loop is correct

434 - for any NEXT_N.

434 + programs are independent (distinct tail blocks). RING >= POOL_SIZE.

435 """

436 req = tl.program_id(0)

437 offs = tl.arange(0, BLOCK_D)

@@ -445,7 +445,7 @@ def _kpool_decode_update_batched_kernel(

445 pos_valid = (cache_loc >= 0) & (pos >= 0)

446

447 slot = safe_pos % POOL_SIZE

448 - phys_slot = safe_pos % POOL_SIZE

448 + phys_slot = safe_pos % RING

449

450 # Derive the tail block from THIS token's tail_slot (the request's block

451 # is constant across a pool, but a padded / invalid entry carries a

@@ -453,7 +453,7 @@ def _kpool_decode_update_batched_kernel(

453 # token's base address). Clamp so an invalid entry can never form an

454 # out-of-bounds base; the accesses below are gated on pos_valid anyway.

455 tail_slot = tl.load(tail_slot_mapping_ptr + idx)

456 - block = tl.maximum(tail_slot, 0).to(tl.int64) // POOL_SIZE

456 + block = tl.maximum(tail_slot, 0).to(tl.int64) // RING

457 block_base = block * TAIL_BLOCK_ELEMS

458

459 # The tail-ring stash must run for EVERY real token, so it is gated on

@@ -481,7 +481,7 @@ def _kpool_decode_update_batched_kernel(

481 max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32)

482 for pool_slot in tl.static_range(0, POOL_SIZE):

483 is_current = pool_slot == slot

484 - phys = (pool_logical_start + pool_slot) % POOL_SIZE

484 + phys = (pool_logical_start + pool_slot) % RING

485 score_buf = tl.load(

486 tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,

487 mask=dim_mask,

@@ -499,7 +499,7 @@ def _kpool_decode_update_batched_kernel(

499 denom = tl.full((BLOCK_D,), 0.0, tl.float32)

500 for pool_slot in tl.static_range(0, POOL_SIZE):

501 is_current = pool_slot == slot

502 - phys = (pool_logical_start + pool_slot) % POOL_SIZE

502 + phys = (pool_logical_start + pool_slot) % RING

503 score_buf = tl.load(

504 tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,

505 mask=dim_mask,

@@ -609,7 +609,8 @@ def kpool_decode_update_and_maybe_write_cache_batched(

609 return

610 assert tail_kv_cache.ndim == 4

611 assert tail_kv_cache.shape[1] == 2

612 - assert tail_kv_cache.shape[2] == pool_size

612 + ring = tail_kv_cache.shape[2]

613 + assert ring >= pool_size and ring % pool_size == 0, (ring, pool_size)

614 assert tail_kv_cache.shape[3] == head_dim

615 assert tail_kv_cache.dtype == torch.bfloat16

616 assert key.ndim == 3 and key.shape[2] == head_dim

@@ -658,6 +659,7 @@ def kpool_decode_update_and_maybe_write_cache_batched(

659 PAGE_SIZE=page_size,

660 BUF_NUMEL_PER_PAGE=buf.stride(0),

661 POOL_SIZE=pool_size,

662 + RING=ring,

663 TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),

664 KPOOL_HEAD=tail_kv_cache.stride(1),

665 HEAD_DIM=head_dim,

vllm/models/glm5next/common/attention.py

+15 -3

@@ -34,6 +34,7 @@ from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool

34 from vllm.platforms import current_platform

35 from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig

36 from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES

37 +from vllm.utils.math_utils import cdiv, next_power_of_2

38 from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec

40 logger = init_logger(__name__)

@@ -160,7 +161,7 @@ class Glm5NextTailCache(DeepseekV32IndexerCache):

161 """Paged circular buffer for the kpool indexer's in-progress (tail) pool.

162

163 Holds the trailing incomplete pool's raw K + gate score: one block of

163 - ``index_kpool`` slots per request, overwritten in place by ``pos % kpool``

164 + ``ring`` slots per request, overwritten in place by ``pos % ring``

165 as decode/spec-decode advances. Prefill seeds it (instead of discarding the

166 tail raw K+gate); the connector transfers it across PD; decode reads it to

167 compress the boundary pool correctly. ``KpoolTailSpec`` /

@@ -190,13 +191,24 @@ class Glm5NextTailCache(DeepseekV32IndexerCache):

191 def get_kv_cache_spec(self, vllm_config: VllmConfig):

192 # The two head slots form [K, gate score] in the generic

193 # [block, head, state, content] cache view.

194 + # Drafts are stashed before acceptance. With a one-pool ring, the

195 + # drafts behind a rejected pool-completing draft overwrite the keys

196 + # read by its redo.

197 + span = self._index_kpool + vllm_config.num_speculative_tokens

198 + ring = self._index_kpool * next_power_of_2(cdiv(span, self._index_kpool))

199 + # ring must divide the attention block size (a multiple of 128).

200 + assert self.cache_config.block_size % ring == 0, (

201 + f"Glm5NextTailCache: cache_config.block_size "

202 + f"({self.cache_config.block_size}) must be a multiple of the "

203 + f"tail ring ({ring})"

204 + )

205 return KpoolTailSpec(

194 - block_size=self._index_kpool,

206 + block_size=ring,

207 num_kv_heads=2,

208 head_size=self.head_dim,

209 head_size_v=0,

210 dtype=torch.bfloat16,

199 - sliding_window=self._index_kpool,

211 + sliding_window=ring,

212 )

213

214 def get_attn_backend(self):

vllm/models/glm5next/nvidia/ops/kpool_compress.py

+17 -15

@@ -387,14 +387,15 @@ def _kpool_tail_seed_kernel(

387 KPOOL_HEAD: tl.constexpr,

388 HEAD_DIM: tl.constexpr,

389 KPOOL: tl.constexpr,

390 + RING: tl.constexpr,

391 BLOCK_D: tl.constexpr,

392 ):

392 - """Copy token ``i``'s raw K + gate into its request's tail block.

393 + """Copy token ``i``'s raw K + gate into its request's tail ring.

394

395 Token ``i`` is among its request's last KPOOL tokens iff the token KPOOL

396 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, :]``.

397 + slot < 0). ``tslot = block * RING + pos % RING``; the destination is

398 + ``tail[block, {0:K, 1:score}, pos % RING, :]``.

399

400 The tail cache aliases the indexer cache with the indexer's (padded) block

401 stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` /

@@ -405,18 +406,18 @@ def _kpool_tail_seed_kernel(

406 t = tl.load(tslot_ptr + i).to(tl.int64)

407 if t < 0:

408 return

408 - blk = t // KPOOL # t >= 0 here, so trunc == floor

409 + blk = t // RING # t >= 0 here, so trunc == floor

410 ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to(

411 tl.int64

412 )

413 # Match the torch semantics exactly: a negative ahead slot floors to a

414 # block id that differs from every real block -> token is in the tail.

415 # Only divide non-negative slots (Triton int div truncates, torch floors).

415 - if ahead >= 0 and ahead // KPOOL == blk:

416 + if ahead >= 0 and ahead // RING == blk:

417 return

418 offs = tl.arange(0, BLOCK_D)

419 m = offs < HEAD_DIM

419 - base = blk * TAIL_BLOCK_ELEMS + (t % KPOOL) * HEAD_DIM

420 + base = blk * TAIL_BLOCK_ELEMS + (t % RING) * HEAD_DIM

421 k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)

422 s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)

423 tl.store(tail_ptr + base + offs, k, mask=m)

@@ -449,6 +450,7 @@ def kpool_seed_tail_cache(

450 KPOOL_HEAD=tail_kv_cache.stride(1),

451 HEAD_DIM=head_dim,

452 KPOOL=kpool,

453 + RING=tail_kv_cache.shape[2],

454 BLOCK_D=triton.next_power_of_2(head_dim),

455 )

456

@@ -476,6 +478,7 @@ def _kpool_decode_update_batched_kernel(

478 PAGE_SIZE: tl.constexpr,

479 BUF_NUMEL_PER_PAGE: tl.constexpr,

480 POOL_SIZE: tl.constexpr,

481 + RING: tl.constexpr,

482 TAIL_BLOCK_ELEMS: tl.constexpr,

483 KPOOL_HEAD: tl.constexpr,

484 HEAD_DIM: tl.constexpr,

@@ -490,10 +493,7 @@ def _kpool_decode_update_batched_kernel(

493 the tail-ring slots that tokens t < t* (same request) just stashed in this

494 same invocation. ``tl.range`` iterates sequentially within the program, so

495 those stashes are visible to the later completion read. Cross-request

493 - programs are independent (distinct tail blocks). With NEXT_N < POOL_SIZE

494 - (the spec-verify case: NEXT_N ~= num_spec+1, POOL_SIZE=16) at most one

495 - completion can occur per request per call, but the ordered loop is correct

496 - for any NEXT_N.

496 + programs are independent (distinct tail blocks). RING >= POOL_SIZE.

497 """

498 req = tl.program_id(0)

499 offs = tl.arange(0, BLOCK_D)

@@ -507,7 +507,7 @@ def _kpool_decode_update_batched_kernel(

507 pos_valid = (cache_loc >= 0) & (pos >= 0)

508

509 slot = safe_pos % POOL_SIZE

510 - phys_slot = safe_pos % POOL_SIZE

510 + phys_slot = safe_pos % RING

511

512 # Derive the tail block from THIS token's tail_slot (the request's block

513 # is constant across a pool, but a padded / invalid entry carries a

@@ -515,7 +515,7 @@ def _kpool_decode_update_batched_kernel(

515 # token's base address). Clamp so an invalid entry can never form an

516 # out-of-bounds base; the accesses below are gated on pos_valid anyway.

517 tail_slot = tl.load(tail_slot_mapping_ptr + idx)

518 - block = tl.maximum(tail_slot, 0).to(tl.int64) // POOL_SIZE

518 + block = tl.maximum(tail_slot, 0).to(tl.int64) // RING

519 block_base = block * TAIL_BLOCK_ELEMS

520

521 # The tail-ring stash must run for EVERY real token, so it is gated on

@@ -543,7 +543,7 @@ def _kpool_decode_update_batched_kernel(

543 max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32)

544 for pool_slot in tl.static_range(0, POOL_SIZE):

545 is_current = pool_slot == slot

546 - phys = (pool_logical_start + pool_slot) % POOL_SIZE

546 + phys = (pool_logical_start + pool_slot) % RING

547 score_buf = tl.load(

548 tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,

549 mask=dim_mask,

@@ -561,7 +561,7 @@ def _kpool_decode_update_batched_kernel(

561 denom = tl.full((BLOCK_D,), 0.0, tl.float32)

562 for pool_slot in tl.static_range(0, POOL_SIZE):

563 is_current = pool_slot == slot

564 - phys = (pool_logical_start + pool_slot) % POOL_SIZE

564 + phys = (pool_logical_start + pool_slot) % RING

565 score_buf = tl.load(

566 tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,

567 mask=dim_mask,

@@ -671,7 +671,8 @@ def kpool_decode_update_and_maybe_write_cache_batched(

671 return

672 assert tail_kv_cache.ndim == 4

673 assert tail_kv_cache.shape[1] == 2

674 - assert tail_kv_cache.shape[2] == pool_size

674 + ring = tail_kv_cache.shape[2]

675 + assert ring >= pool_size and ring % pool_size == 0, (ring, pool_size)

676 assert tail_kv_cache.shape[3] == head_dim

677 assert tail_kv_cache.dtype == torch.bfloat16

678 assert key.ndim == 3 and key.shape[2] == head_dim

@@ -717,6 +718,7 @@ def kpool_decode_update_and_maybe_write_cache_batched(

718 PAGE_SIZE=page_size,

719 BUF_NUMEL_PER_PAGE=buf.stride(0),

720 POOL_SIZE=pool_size,

721 + RING=ring,

722 TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),

723 KPOOL_HEAD=tail_kv_cache.stride(1),

724 HEAD_DIM=head_dim,