| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit test for the batched kpool decode-update kernel. |
| 4 | |
| 5 | Validates ``kpool_decode_update_and_maybe_write_cache_batched`` against an |
| 6 | independent pure-torch reference that replicates the per-request, in-position |
| 7 | order semantics: stash each token into a paged tail ring; on pool completion |
| 8 | (``pos % pool_size == pool_size-1``) softmax(gate+ape)-weighted sum + Hadamard-128 |
| 9 | + per-vector fp8 absmax quant + write to the indexer K cache. Covers |
| 10 | no-completion, completion-at-end, completion-mid-batch, non-uniform padding, |
| 11 | plain decode, plus a randomized fuzz pass. |
| 12 | |
| 13 | The kernel iterates each request's ``next_n`` tokens in position order inside |
| 14 | one program (grid = num_requests) to preserve the pool-completion |
| 15 | read-after-stash dependency; the reference mirrors that ordering. |
| 16 | |
| 17 | ``test_decode_writer_matches_prefill_writer`` is deliberately NOT |
| 18 | reference-based: it checks the decode writer against the *prefill* writer |
| 19 | (``kpool_compress_and_write_cache``), the invariant that actually matters in |
| 20 | production. A hand-written reference can drift to match a buggy kernel -- that |
| 21 | is exactly how the stash-gating bug (intra-pool tokens never entering the tail |
| 22 | ring, because the stash was gated on the pool-granular ``slot_mapping``) stayed |
| 23 | green here. |
| 24 | """ |
| 25 | |
| 26 | import math |
| 27 | |
| 28 | import pytest |
| 29 | import torch |
| 30 | |
| 31 | from vllm.platforms import current_platform |
| 32 | |
| 33 | if current_platform.is_rocm(): |
| 34 | from vllm.models.glm5next.amd.ops.kpool_compress import ( |
| 35 | kpool_compress_and_write_cache, |
| 36 | kpool_decode_update_and_maybe_write_cache_batched, |
| 37 | kpool_seed_tail_cache, |
| 38 | ) |
| 39 | else: |
| 40 | from vllm.models.glm5next.nvidia.ops.kpool_compress import ( |
| 41 | kpool_compress_and_write_cache, |
| 42 | kpool_decode_update_and_maybe_write_cache_batched, |
| 43 | kpool_seed_tail_cache, |
| 44 | ) |
| 45 | |
| 46 | HEAD_DIM = 128 |
| 47 | POOL_SIZE = 16 |
| 48 | PAGE_SIZE = 64 |
| 49 | NUM_BLOCKS = 32 |
| 50 | ROUND_SCALE = True |
| 51 | FP8_DTYPE = current_platform.fp8_dtype() |
| 52 | FP8_MAX = torch.finfo(FP8_DTYPE).max |
| 53 | |
| 54 | |
| 55 | def _make_caches(): |
| 56 | kv = torch.zeros( |
| 57 | NUM_BLOCKS, PAGE_SIZE, HEAD_DIM + 4, dtype=torch.uint8, device="cuda" |
| 58 | ) |
| 59 | tail = torch.zeros( |
| 60 | NUM_BLOCKS, 2, POOL_SIZE, HEAD_DIM, dtype=torch.bfloat16, device="cuda" |
| 61 | ) |
| 62 | return kv, tail |
| 63 | |
| 64 | |
| 65 | def _tail_slot_for(blocks, pos): |
| 66 | """tail_slot = block*POOL + pos%POOL; each request owns a distinct tail block.""" |
| 67 | blk = torch.tensor(blocks, device=pos.device, dtype=torch.int32).unsqueeze(1) |
| 68 | return (blk * POOL_SIZE + pos % POOL_SIZE).to(torch.int32) |
| 69 | |
| 70 | |
| 71 | def _seed_prior(tail, blocks, n_prior, seed=42): |
| 72 | if n_prior <= 0: |
| 73 | return |
| 74 | g = torch.Generator(device=tail.device).manual_seed(seed) |
| 75 | prior_k = torch.randn( |
| 76 | len(blocks), |
| 77 | n_prior, |
| 78 | HEAD_DIM, |
| 79 | dtype=torch.bfloat16, |
| 80 | device=tail.device, |
| 81 | generator=g, |
| 82 | ) |
| 83 | prior_s = torch.randn( |
| 84 | len(blocks), |
| 85 | n_prior, |
| 86 | HEAD_DIM, |
| 87 | dtype=torch.bfloat16, |
| 88 | device=tail.device, |
| 89 | generator=g, |
| 90 | ) |
| 91 | for i, blk in enumerate(blocks): |
| 92 | tail[blk, 0, :n_prior, :] = prior_k[i] |
| 93 | tail[blk, 1, :n_prior, :] = prior_s[i] |
| 94 | |
| 95 | |
| 96 | def _hadamard128_torch(x: torch.Tensor) -> torch.Tensor: |
| 97 | """Reference Hadamard-128 on the last dim (must be 128).""" |
| 98 | n = x.shape[-1] |
| 99 | assert n == 128 |
| 100 | h = torch.tensor([[1.0, 1.0], [1.0, -1.0]], dtype=torch.float32, device=x.device) |
| 101 | while h.shape[0] < n: |
| 102 | h = torch.cat([torch.cat([h, h], dim=1), torch.cat([h, -h], dim=1)], dim=0) |
| 103 | h = h / math.sqrt(n) |
| 104 | return x @ h |
| 105 | |
| 106 | |
| 107 | def _torch_reference( |
| 108 | kv: torch.Tensor, |
| 109 | tail: torch.Tensor, |
| 110 | tail_slot: torch.Tensor, |
| 111 | key: torch.Tensor, |
| 112 | score: torch.Tensor, |
| 113 | ape: torch.Tensor, |
| 114 | slot_map: torch.Tensor, |
| 115 | pos: torch.Tensor, |
| 116 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 117 | """Independent reference for the batched decode-update kernel. |
| 118 | |
| 119 | For each request, iterate its next_n tokens in order. On each token: |
| 120 | - if pos%POOL == POOL-1 and pos_valid: compress the pool (slots |
| 121 | [pool_start..pool_start+POOL-1], current token via is_current) and write |
| 122 | fp8 K + fp32 scale to kv_cache at cache_loc. |
| 123 | - always stash the current token's K/score into tail[block, pos%POOL]. |
| 124 | """ |
| 125 | kv = kv.clone() |
| 126 | tail = tail.clone() |
| 127 | B, next_n = pos.shape |
| 128 | # The indexer K cache is [num_blocks, PAGE_SIZE, HEAD_DIM+4] uint8 but the |
| 129 | # kernels interpret each page as [HEAD_DIM*PAGE_SIZE bytes of K (token-major) |
| 130 | # | 4*PAGE_SIZE bytes of fp32 scale (token-major)]. Operate on a flat byte |
| 131 | # view so the reference writes K and scale at the exact offsets the kernel |
| 132 | # uses (page_base + tok*HEAD_DIM for K; page_base + HEAD_DIM*PAGE_SIZE + |
| 133 | # tok*4 for scale). |
| 134 | page_bytes = PAGE_SIZE * (HEAD_DIM + 4) |
| 135 | k_region = HEAD_DIM * PAGE_SIZE |
| 136 | tail_slot_cpu = tail_slot.cpu().tolist() |
| 137 | slot_map_cpu = slot_map.cpu().tolist() |
| 138 | pos_cpu = pos.cpu().tolist() |
| 139 | key_cpu = key.float().cpu() |
| 140 | score_cpu = score.float().cpu() |
| 141 | ape_cpu = ape.cpu() |
| 142 | tail_cpu = tail.float().cpu() |
| 143 | kv_flat = kv.view(torch.uint8).reshape(-1).cpu() |
| 144 | |
| 145 | for b in range(B): |
| 146 | for t in range(next_n): |
| 147 | cache_loc = slot_map_cpu[b][t] |
| 148 | p = pos_cpu[b][t] |
| 149 | pos_valid = cache_loc >= 0 and p >= 0 |
| 150 | safe_pos = max(p, 0) |
| 151 | slot = safe_pos % POOL_SIZE |
| 152 | phys_slot = safe_pos % POOL_SIZE |
| 153 | # Per-token block derivation (a leading invalid sentinel must not |
| 154 | # poison the base for the rest of the request); clamped like the |
| 155 | # kernel so an invalid entry can't form a negative base. |
| 156 | block = max(tail_slot_cpu[b][t], 0) // POOL_SIZE |
| 157 | |
| 158 | cur_key = key_cpu[b, t] |
| 159 | cur_score = score_cpu[b, t] |
| 160 | |
| 161 | if pos_valid and slot == POOL_SIZE - 1: |
| 162 | pool_logical_start = safe_pos - slot |
| 163 | pool_scores = [] |
| 164 | pool_ks = [] |
| 165 | for ps in range(POOL_SIZE): |
| 166 | is_current = ps == slot |
| 167 | phys = (pool_logical_start + ps) % POOL_SIZE |
| 168 | if is_current: |
| 169 | s = cur_score |
| 170 | k = cur_key |
| 171 | else: |
| 172 | s = tail_cpu[block, 1, phys] |
| 173 | k = tail_cpu[block, 0, phys] |
| 174 | s = s + ape_cpu[ps] |
| 175 | pool_scores.append(s) |
| 176 | pool_ks.append(k) |
| 177 | pool_scores = torch.stack(pool_scores) # [POOL, D] |
| 178 | pool_ks = torch.stack(pool_ks) # [POOL, D] |
| 179 | max_score = pool_scores.max(dim=0).values |
| 180 | prob = torch.exp(pool_scores - max_score) |
| 181 | denom = prob.sum(dim=0) |
| 182 | acc = (pool_ks * prob).sum(dim=0) |
| 183 | x = (acc / denom).to(torch.bfloat16).to(torch.float32) |
| 184 | x = _hadamard128_torch(x).to(torch.bfloat16).to(torch.float32) |
| 185 | absmax = torch.clamp(x.abs().max(), min=1e-4) |
| 186 | if ROUND_SCALE: |
| 187 | scale = torch.exp2(torch.ceil(torch.log2(absmax / FP8_MAX))) |
| 188 | else: |
| 189 | scale = absmax / FP8_MAX |
| 190 | quantized = torch.clamp(x / scale, -FP8_MAX, FP8_MAX).to(FP8_DTYPE) |
| 191 | # write K and scale at the separated-layout offsets |
| 192 | loc = cache_loc |
| 193 | loc_page_index = loc // PAGE_SIZE |
| 194 | loc_tok = loc % PAGE_SIZE |
| 195 | page_base = loc_page_index * page_bytes |
| 196 | if current_platform.is_rocm(): |
| 197 | dims = torch.arange(HEAD_DIM) |
| 198 | k_off = ( |
| 199 | page_base |
| 200 | + (loc_tok // 16) * 16 * HEAD_DIM |
| 201 | + (dims // 16) * 16 * 16 |
| 202 | + (loc_tok % 16) * 16 |
| 203 | + dims % 16 |
| 204 | ) |
| 205 | else: |
| 206 | k_off = page_base + loc_tok * HEAD_DIM + torch.arange(HEAD_DIM) |
| 207 | s_off = page_base + k_region + loc_tok * 4 |
| 208 | kv_flat[k_off] = quantized.view(torch.uint8) |
| 209 | kv_flat[s_off : s_off + 4] = scale.detach().reshape(1).view(torch.uint8) |
| 210 | |
| 211 | # stash -- gated on the TOKEN-granular tail slot, not on pos_valid. |
| 212 | # pos_valid keys off the pool-granular cache_loc, which is -1 for |
| 213 | # every token that is not the pool's last, so gating the stash on it |
| 214 | # would drop all intra-pool tokens. |
| 215 | if p >= 0 and tail_slot_cpu[b][t] >= 0: |
| 216 | tail_cpu[block, 0, phys_slot] = cur_key |
| 217 | tail_cpu[block, 1, phys_slot] = cur_score |
| 218 | |
| 219 | kv_out = kv_flat.view(NUM_BLOCKS, PAGE_SIZE, HEAD_DIM + 4).to(device="cuda") |
| 220 | return kv_out, tail_cpu.to(torch.bfloat16).to(device="cuda") |
| 221 | |
| 222 | |
| 223 | def _assert_eq(r_ref, r_kern): |
| 224 | kv_ref, tail_ref = r_ref |
| 225 | kv_kern, tail_kern = r_kern |
| 226 | assert torch.equal(kv_ref, kv_kern), ( |
| 227 | "kv_cache differs: max diff " |
| 228 | f"{(kv_ref.int() - kv_kern.int()).abs().max().item()}" |
| 229 | ) |
| 230 | assert torch.equal(tail_ref, tail_kern), ( |
| 231 | "tail_kv_cache differs: max diff " |
| 232 | f"{(tail_ref.float() - tail_kern.float()).abs().max().item()}" |
| 233 | ) |
| 234 | |
| 235 | |
| 236 | def _cache_pool_bytes(kv_cache: torch.Tensor, pool_slot: int) -> torch.Tensor: |
| 237 | """Read a logical pool's K and scale bytes from the platform cache layout.""" |
| 238 | page_size = kv_cache.shape[1] |
| 239 | head_dim = kv_cache.shape[2] - 4 |
| 240 | page_idx, token_offset = divmod(pool_slot, page_size) |
| 241 | flat = kv_cache[page_idx].reshape(-1) |
| 242 | dims = torch.arange(head_dim, device=kv_cache.device) |
| 243 | if current_platform.is_rocm() and page_size > 1: |
| 244 | k_offsets = ( |
| 245 | (token_offset // 16) * 16 * head_dim |
| 246 | + (dims // 16) * 16 * 16 |
| 247 | + (token_offset % 16) * 16 |
| 248 | + dims % 16 |
| 249 | ) |
| 250 | else: |
| 251 | k_offsets = token_offset * head_dim + dims |
| 252 | scale_offset = page_size * head_dim + 4 * token_offset |
| 253 | return torch.cat((flat[k_offsets], flat[scale_offset : scale_offset + 4])) |
| 254 | |
| 255 | |
| 256 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required") |
| 257 | def test_amd_prefill_writer_uses_preshuffled_cache_layout(): |
| 258 | from vllm.models.glm5next.amd.ops.kpool_compress import ( |
| 259 | kpool_compress_and_write_cache as amd_kpool_compress, |
| 260 | ) |
| 261 | |
| 262 | torch.manual_seed(0) |
| 263 | token_offset = 17 |
| 264 | kv = torch.zeros(1, PAGE_SIZE, HEAD_DIM + 4, dtype=torch.uint8, device="cuda") |
| 265 | key = torch.randn(1, POOL_SIZE, HEAD_DIM, dtype=torch.bfloat16, device="cuda") |
| 266 | score = torch.randn_like(key) |
| 267 | ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda") |
| 268 | compressed_k, compressed_scale = amd_kpool_compress( |
| 269 | kv, |
| 270 | key, |
| 271 | score, |
| 272 | ape, |
| 273 | torch.tensor([token_offset], dtype=torch.int64, device="cuda"), |
| 274 | pool_size=POOL_SIZE, |
| 275 | head_dim=HEAD_DIM, |
| 276 | round_scale=ROUND_SCALE, |
| 277 | return_compressed=True, |
| 278 | ) |
| 279 | |
| 280 | dim = torch.arange(HEAD_DIM, device="cuda") |
| 281 | offsets = ( |
| 282 | (token_offset // 16) * 16 * HEAD_DIM |
| 283 | + (dim // 16) * 16 * 16 |
| 284 | + (token_offset % 16) * 16 |
| 285 | + dim % 16 |
| 286 | ) |
| 287 | flat = kv.view(torch.uint8).reshape(-1) |
| 288 | stored_k = flat[offsets].view(compressed_k.dtype) |
| 289 | scale_offset = PAGE_SIZE * HEAD_DIM + token_offset * 4 |
| 290 | stored_scale = flat[scale_offset : scale_offset + 4].view(torch.float32) |
| 291 | |
| 292 | assert torch.equal(stored_k, compressed_k[0]) |
| 293 | assert torch.equal(stored_scale, compressed_scale) |
| 294 | |
| 295 | |
| 296 | def _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos): |
| 297 | kv = kv.clone() |
| 298 | tail = tail.clone() |
| 299 | kpool_decode_update_and_maybe_write_cache_batched( |
| 300 | kv, |
| 301 | tail, |
| 302 | tail_slot, |
| 303 | key, |
| 304 | score, |
| 305 | ape, |
| 306 | slot_map, |
| 307 | pos, |
| 308 | POOL_SIZE, |
| 309 | HEAD_DIM, |
| 310 | round_scale=ROUND_SCALE, |
| 311 | ) |
| 312 | return kv, tail |
| 313 | |
| 314 | |
| 315 | @pytest.mark.parametrize("pool_size", [4, 16]) |
| 316 | @pytest.mark.parametrize("ring_pools", [1, 2]) |
| 317 | def test_decode_writer_matches_prefill_writer(pool_size, ring_pools): |
| 318 | ring = ring_pools * pool_size |
| 319 | n_pools, page, nblk = 8, 64, 4 |
| 320 | n_tok = n_pools * pool_size |
| 321 | dev = "cuda" |
| 322 | torch.manual_seed(0) |
| 323 | k = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 324 | score = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 325 | ape = torch.randn(pool_size, HEAD_DIM, dtype=torch.float32, device=dev) |
| 326 | |
| 327 | kv_prefill = torch.zeros(nblk, page, HEAD_DIM + 4, dtype=torch.uint8, device=dev) |
| 328 | kpool_compress_and_write_cache( |
| 329 | kv_prefill, |
| 330 | k.view(n_pools, pool_size, HEAD_DIM), |
| 331 | score.view(n_pools, pool_size, HEAD_DIM), |
| 332 | ape, |
| 333 | torch.arange(n_pools, dtype=torch.int64, device=dev), |
| 334 | pool_size=pool_size, |
| 335 | head_dim=HEAD_DIM, |
| 336 | round_scale=ROUND_SCALE, |
| 337 | ) |
| 338 | |
| 339 | # One request owning tail block 0, fed one token per decode step. |
| 340 | kv_decode = torch.zeros_like(kv_prefill) |
| 341 | tail = torch.zeros(nblk, 2, ring, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 342 | for t in range(n_tok): |
| 343 | completes = t % pool_size == pool_size - 1 |
| 344 | kpool_decode_update_and_maybe_write_cache_batched( |
| 345 | kv_decode, |
| 346 | tail, |
| 347 | # token-granular: every token has a valid tail slot |
| 348 | torch.tensor([[t % ring]], dtype=torch.int32, device=dev), |
| 349 | k[t].view(1, 1, HEAD_DIM), |
| 350 | score[t].view(1, 1, HEAD_DIM), |
| 351 | ape, |
| 352 | # pool-granular: only the pool's last token carries a cache slot |
| 353 | torch.tensor( |
| 354 | [[t // pool_size if completes else -1]], dtype=torch.int32, device=dev |
| 355 | ), |
| 356 | torch.tensor([[t]], dtype=torch.int32, device=dev), |
| 357 | pool_size, |
| 358 | HEAD_DIM, |
| 359 | round_scale=ROUND_SCALE, |
| 360 | ) |
| 361 | |
| 362 | differing = [ |
| 363 | p |
| 364 | for p in range(n_pools) |
| 365 | if not torch.equal( |
| 366 | _cache_pool_bytes(kv_prefill, p), _cache_pool_bytes(kv_decode, p) |
| 367 | ) |
| 368 | ] |
| 369 | assert not differing, ( |
| 370 | f"decode-written pools differ from prefill-written pools: " |
| 371 | f"{len(differing)}/{n_pools} (pool_size={pool_size}, first={differing[:5]})" |
| 372 | ) |
| 373 | |
| 374 | |
| 375 | @pytest.mark.parametrize("ring_pools", [1, 2]) |
| 376 | def test_rejected_draft_redo_needs_ring_slots(ring_pools): |
| 377 | """With a one-pool ring, the drafts behind a rejected pool-completing draft |
| 378 | overwrote the pool's earlier keys, so its redo compressed wrong keys.""" |
| 379 | pool, spec, page, nblk = 4, 3, 64, 2 |
| 380 | ring = ring_pools * pool |
| 381 | dev = "cuda" |
| 382 | torch.manual_seed(1) |
| 383 | n_tok = 3 * pool |
| 384 | k = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 385 | score = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 386 | ape = torch.randn(pool, HEAD_DIM, dtype=torch.float32, device=dev) |
| 387 | kv_ref = torch.zeros(nblk, page, HEAD_DIM + 4, dtype=torch.uint8, device=dev) |
| 388 | kpool_compress_and_write_cache( |
| 389 | kv_ref, |
| 390 | k.view(3, pool, HEAD_DIM), |
| 391 | score.view(3, pool, HEAD_DIM), |
| 392 | ape, |
| 393 | torch.arange(3, dtype=torch.int64, device=dev), |
| 394 | pool_size=pool, |
| 395 | head_dim=HEAD_DIM, |
| 396 | round_scale=ROUND_SCALE, |
| 397 | ) |
| 398 | |
| 399 | kv = torch.zeros_like(kv_ref) |
| 400 | tail = torch.zeros(nblk, 2, ring, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 401 | |
| 402 | def step(positions, keys, scores): |
| 403 | pos = torch.tensor([positions], dtype=torch.int32, device=dev) |
| 404 | slots = [(p // pool) if p % pool == pool - 1 else -1 for p in positions] |
| 405 | kpool_decode_update_and_maybe_write_cache_batched( |
| 406 | kv, |
| 407 | tail, |
| 408 | pos % ring, |
| 409 | keys.view(1, -1, HEAD_DIM), |
| 410 | scores.view(1, -1, HEAD_DIM), |
| 411 | ape, |
| 412 | torch.tensor([slots], dtype=torch.int32, device=dev), |
| 413 | pos, |
| 414 | pool, |
| 415 | HEAD_DIM, |
| 416 | round_scale=ROUND_SCALE, |
| 417 | ) |
| 418 | |
| 419 | for t in range(7): |
| 420 | step([t], k[t], score[t]) |
| 421 | # Control: verified token 7 completes pool 1 before drafts 8..10 are stashed. |
| 422 | drafts = torch.randn(spec, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 423 | draft_scores = torch.randn(spec, HEAD_DIM, dtype=torch.bfloat16, device=dev) |
| 424 | step( |
| 425 | [7, 8, 9, 10], |
| 426 | torch.cat([k[7:8], drafts]), |
| 427 | torch.cat([score[7:8], draft_scores]), |
| 428 | ) |
| 429 | step([8, 9, 10, 11], k[8:12], score[8:12]) # all drafts rejected |
| 430 | for p in (1, 2): |
| 431 | assert torch.equal(_cache_pool_bytes(kv, p), _cache_pool_bytes(kv_ref, p)), p |
| 432 | |
| 433 | # Draft 7 completes pool 1 and is rejected. With a one-pool ring, drafts |
| 434 | # 8 and 9 overwrite the slots of positions 4 and 5, which are read by the |
| 435 | # redo of 7. |
| 436 | kv.zero_() |
| 437 | tail.zero_() |
| 438 | for t in range(6): |
| 439 | step([t], k[t], score[t]) |
| 440 | step( |
| 441 | [6, 7, 8, 9], |
| 442 | torch.cat([k[6:7], drafts]), |
| 443 | torch.cat([score[6:7], draft_scores]), |
| 444 | ) |
| 445 | step([7, 8, 9, 10], k[7:11], score[7:11]) |
| 446 | pool1_ok = torch.equal(_cache_pool_bytes(kv, 1), _cache_pool_bytes(kv_ref, 1)) |
| 447 | if ring_pools == 1: |
| 448 | assert not pool1_ok, "expected a one-pool ring to corrupt pool 1" |
| 449 | else: |
| 450 | assert pool1_ok |
| 451 | |
| 452 | |
| 453 | def test_leading_invalid_tail_slot(): |
| 454 | """A request whose FIRST token carries an invalid (-1) tail slot while a |
| 455 | later token is a real pool completion. |
| 456 | |
| 457 | The tail block must be derived per token, not from token 0: a leading |
| 458 | invalid sentinel would otherwise poison the base address for the whole |
| 459 | request (out-of-bounds tail reads on the completion). |
| 460 | """ |
| 461 | torch.manual_seed(0) |
| 462 | B, next_n, blocks = 2, 4, [3, 5] |
| 463 | # req 0: token 0 invalid (pos -1), tokens 1..3 valid, completion at pos 15 |
| 464 | # req 1: all valid, no completion |
| 465 | pos = torch.tensor( |
| 466 | [[-1, 13, 14, 15], [4, 5, 6, 7]], dtype=torch.int32, device="cuda" |
| 467 | ) |
| 468 | safe_pos = torch.where(pos >= 0, pos, 0) |
| 469 | tail_slot = _tail_slot_for(blocks, safe_pos) |
| 470 | # leading invalid entry carries the -1 sentinel, as the scatter path emits |
| 471 | tail_slot[0, 0] = -1 |
| 472 | slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") |
| 473 | slot_map[0, 3] = 15 # req 0 completes its pool on the last verify token |
| 474 | |
| 475 | key = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") |
| 476 | score = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") |
| 477 | ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda") |
| 478 | |
| 479 | kv, tail = _make_caches() |
| 480 | _seed_prior(tail, blocks, 13) |
| 481 | r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) |
| 482 | r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) |
| 483 | _assert_eq(r_ref, r_kern) |
| 484 | |
| 485 | |
| 486 | def test_prefill_seed_honors_padded_tail_block_stride(): |
| 487 | """The tail shares a padded indexer allocation in production. |
| 488 | |
| 489 | ``get_kv_cache_config_from_groups`` aliases each tail tensor onto its |
| 490 | indexer tensor with the indexer's block stride (38016 B for GLM-5.3-Flash |
| 491 | vs a dense 2048 B tail block), so a seed kernel that addresses blocks |
| 492 | densely writes into an unrelated indexer block and leaves the request's |
| 493 | tail block untouched. Runs on every platform; the NVIDIA kernel had this |
| 494 | bug while the AMD kernel did not. |
| 495 | """ |
| 496 | kpool = 4 |
| 497 | num_blocks = 6 |
| 498 | logical_block_elems = 2 * kpool * HEAD_DIM |
| 499 | padded_block_elems = logical_block_elems + 256 |
| 500 | sentinel = -123.0 |
| 501 | backing = torch.full( |
| 502 | (num_blocks * padded_block_elems,), |
| 503 | sentinel, |
| 504 | dtype=torch.bfloat16, |
| 505 | device="cuda", |
| 506 | ) |
| 507 | tail = torch.as_strided( |
| 508 | backing, |
| 509 | size=(num_blocks, 2, kpool, HEAD_DIM), |
| 510 | stride=(padded_block_elems, kpool * HEAD_DIM, HEAD_DIM, 1), |
| 511 | ) |
| 512 | |
| 513 | block = 3 |
| 514 | ring_slot = 2 |
| 515 | key = torch.arange(HEAD_DIM, dtype=torch.bfloat16, device="cuda").unsqueeze(0) |
| 516 | score = (key + 256).to(torch.bfloat16) |
| 517 | tail_slot = torch.tensor( |
| 518 | [block * kpool + ring_slot], dtype=torch.int32, device="cuda" |
| 519 | ) |
| 520 | |
| 521 | kpool_seed_tail_cache(tail, key, score, tail_slot, kpool, HEAD_DIM) |
| 522 | torch.accelerator.synchronize() |
| 523 | |
| 524 | assert torch.equal(tail[block, 0, ring_slot], key[0]) |
| 525 | assert torch.equal(tail[block, 1, ring_slot], score[0]) |
| 526 | |
| 527 | compact_offset = (block * 2 * kpool + ring_slot) * HEAD_DIM |
| 528 | assert torch.all(backing[compact_offset : compact_offset + HEAD_DIM] == sentinel) |
| 529 | |
| 530 | |
| 531 | @pytest.mark.parametrize( |
| 532 | "case_id", |
| 533 | [ |
| 534 | "no_completion", |
| 535 | "completion_at_end", |
| 536 | "completion_mid_batch", |
| 537 | "non_uniform_padding", |
| 538 | "plain_decode", |
| 539 | ], |
| 540 | ) |
| 541 | def test_batched_matches_reference(case_id): |
| 542 | torch.manual_seed(0) |
| 543 | if case_id == "no_completion": |
| 544 | B, next_n, blocks = 3, 4, [0, 1, 2] |
| 545 | pos = ( |
| 546 | torch.arange(next_n, device="cuda", dtype=torch.int32) |
| 547 | .unsqueeze(0) |
| 548 | .expand(B, -1) |
| 549 | .contiguous() |
| 550 | ) |
| 551 | slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") |
| 552 | n_prior = 0 |
| 553 | elif case_id == "completion_at_end": |
| 554 | B, next_n, blocks = 2, 4, [0, 1] |
| 555 | pos = torch.tensor( |
| 556 | [[12, 13, 14, 15], [12, 13, 14, 15]], dtype=torch.int32, device="cuda" |
| 557 | ) |
| 558 | slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") |
| 559 | slot_map[:, 3] = torch.tensor( |
| 560 | [15, PAGE_SIZE + 15], dtype=torch.int32, device="cuda" |
| 561 | ) |
| 562 | n_prior = POOL_SIZE - next_n |
| 563 | elif case_id == "completion_mid_batch": |
| 564 | B, next_n, blocks = 3, 4, [0, 1, 2] |
| 565 | pos = torch.tensor([[13, 14, 15, 16]] * B, dtype=torch.int32, device="cuda") |
| 566 | slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") |
| 567 | slot_map[:, 2] = torch.tensor( |
| 568 | [15, PAGE_SIZE + 15, 2 * PAGE_SIZE + 15], dtype=torch.int32, device="cuda" |
| 569 | ) |
| 570 | n_prior = 13 |
| 571 | elif case_id == "non_uniform_padding": |
| 572 | B, next_n, blocks = 2, 4, [0, 1] |
| 573 | pos = torch.tensor( |
| 574 | [[12, 13, 14, 15], [12, 13, -1, -1]], dtype=torch.int32, device="cuda" |
| 575 | ) |
| 576 | slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") |
| 577 | slot_map[0, 3] = 15 |
| 578 | n_prior = POOL_SIZE - 4 |
| 579 | else: # plain_decode |
| 580 | B, next_n, blocks = 4, 1, [0, 1, 2, 3] |
| 581 | pos = torch.tensor([[5], [6], [7], [8]], dtype=torch.int32, device="cuda") |
| 582 | slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda") |
| 583 | n_prior = 0 |
| 584 | |
| 585 | if case_id == "non_uniform_padding": |
| 586 | safe_pos = torch.where(pos >= 0, pos, 0) |
| 587 | tail_slot = torch.where(pos >= 0, _tail_slot_for(blocks, safe_pos), 0) |
| 588 | else: |
| 589 | tail_slot = _tail_slot_for(blocks, pos) |
| 590 | |
| 591 | key = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") |
| 592 | score = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda") |
| 593 | ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda") |
| 594 | |
| 595 | kv, tail = _make_caches() |
| 596 | _seed_prior(tail, blocks, n_prior) |
| 597 | r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) |
| 598 | r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) |
| 599 | _assert_eq(r_ref, r_kern) |
| 600 | |
| 601 | |
| 602 | @pytest.mark.parametrize("seed", list(range(20))) |
| 603 | def test_batched_matches_reference_fuzz(seed): |
| 604 | """Random B / next_n / start positions; covers 0, 1, and multi completion.""" |
| 605 | g = torch.Generator(device="cuda").manual_seed(seed) |
| 606 | B = int(torch.randint(1, 6, (1,), generator=g, device="cuda").item()) |
| 607 | next_n = int(torch.randint(1, 8, (1,), generator=g, device="cuda").item()) |
| 608 | blocks = list(range(B)) |
| 609 | |
| 610 | starts = torch.randint(0, 33, (B,), generator=g, device="cuda", dtype=torch.int32) |
| 611 | pos = starts.unsqueeze(1) + torch.arange( |
| 612 | next_n, device="cuda", dtype=torch.int32 |
| 613 | ).unsqueeze(0) |
| 614 | tail_slot = _tail_slot_for(blocks, pos) |
| 615 | |
| 616 | is_completion = pos % POOL_SIZE == POOL_SIZE - 1 |
| 617 | blk = torch.tensor(blocks, device="cuda", dtype=torch.int32).unsqueeze(1) |
| 618 | pool_slot = blk * PAGE_SIZE + (POOL_SIZE - 1) |
| 619 | slot_map = torch.where(is_completion, pool_slot, torch.full_like(pos, -1)) |
| 620 | |
| 621 | key = torch.randn( |
| 622 | B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda", generator=g |
| 623 | ) |
| 624 | score = torch.randn( |
| 625 | B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda", generator=g |
| 626 | ) |
| 627 | ape = torch.randn( |
| 628 | POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda", generator=g |
| 629 | ) |
| 630 | |
| 631 | kv, tail = _make_caches() |
| 632 | prior_g = torch.Generator(device="cuda").manual_seed(seed + 1000) |
| 633 | for b in range(B): |
| 634 | n_prior = int(starts[b].item()) % POOL_SIZE |
| 635 | if n_prior > 0: |
| 636 | pk = torch.randn( |
| 637 | n_prior, |
| 638 | HEAD_DIM, |
| 639 | dtype=torch.bfloat16, |
| 640 | device="cuda", |
| 641 | generator=prior_g, |
| 642 | ) |
| 643 | ps = torch.randn( |
| 644 | n_prior, |
| 645 | HEAD_DIM, |
| 646 | dtype=torch.bfloat16, |
| 647 | device="cuda", |
| 648 | generator=prior_g, |
| 649 | ) |
| 650 | tail[blocks[b], 0, :n_prior, :] = pk |
| 651 | tail[blocks[b], 1, :n_prior, :] = ps |
| 652 | |
| 653 | r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos) |
| 654 | r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos) |
| 655 | _assert_eq(r_ref, r_kern) |