| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """kpool (key-pooling) Triton kernels for the sparse-attention indexer. |
| 4 | |
| 5 | The cache stores POOLS (1 entry per ``pool_size`` consecutive tokens) rather |
| 6 | than individual tokens. ``compress_ratio == pool_size`` on the kv_cache_spec |
| 7 | makes the metadata builder emit pool-granular slot_mapping / seq_lens / |
| 8 | cu_seq_lens / page_table for free; this file supplies the compress-write |
| 9 | kernel (replacing ``indexer_k_quant_and_cache``) and the pool-level topk |
| 10 | helpers (select pools -> expand to tokens -> append tail). |
| 11 | """ |
| 12 | |
| 13 | from __future__ import annotations |
| 14 | |
| 15 | import torch |
| 16 | |
| 17 | from vllm.platforms import current_platform |
| 18 | from vllm.triton_utils import tl, triton |
| 19 | |
| 20 | # The GLM-5.3-Flash indexer head dimension is fixed at 128. |
| 21 | INDEX_HEAD_DIM = 128 |
| 22 | FP8_DTYPE = current_platform.fp8_dtype() |
| 23 | FP8_MAX = torch.finfo(FP8_DTYPE).max |
| 24 | |
| 25 | |
| 26 | @triton.jit |
| 27 | def _cache_k_offset( |
| 28 | token_offset, |
| 29 | dim_offset, |
| 30 | head_dim: tl.constexpr, |
| 31 | preshuffle: tl.constexpr, |
| 32 | ): |
| 33 | if preshuffle: |
| 34 | return ( |
| 35 | (token_offset // 16) * 16 * head_dim |
| 36 | + (dim_offset // 16) * 16 * 16 |
| 37 | + (token_offset % 16) * 16 |
| 38 | + dim_offset % 16 |
| 39 | ) |
| 40 | return token_offset * head_dim + dim_offset |
| 41 | |
| 42 | |
| 43 | # Hadamard-128 rotation |
| 44 | |
| 45 | |
| 46 | @triton.jit |
| 47 | def _hadamard128_stage(x, GROUPS: tl.constexpr, STRIDE: tl.constexpr): |
| 48 | x3 = tl.reshape(x, (GROUPS, 2, STRIDE)) |
| 49 | x3 = tl.trans(x3, 0, 2, 1) |
| 50 | a, b = tl.split(x3) |
| 51 | x3 = tl.join(a + b, a - b) |
| 52 | x3 = tl.trans(x3, 0, 2, 1) |
| 53 | return tl.reshape(x3, (128,)) |
| 54 | |
| 55 | |
| 56 | @triton.jit |
| 57 | def _hadamard128(x): |
| 58 | x = _hadamard128_stage(x, 64, 1) |
| 59 | x = _hadamard128_stage(x, 32, 2) |
| 60 | x = _hadamard128_stage(x, 16, 4) |
| 61 | x = _hadamard128_stage(x, 8, 8) |
| 62 | x = _hadamard128_stage(x, 4, 16) |
| 63 | x = _hadamard128_stage(x, 2, 32) |
| 64 | x = _hadamard128_stage(x, 1, 64) |
| 65 | return x * 0.08838834764831845 # 1/sqrt(128) |
| 66 | |
| 67 | |
| 68 | # Fused pool compression and cache write. |
| 69 | |
| 70 | |
| 71 | @triton.jit |
| 72 | def _kpool_softmax_rotate_write_cache_kernel( |
| 73 | buf_fp8_ptr, |
| 74 | buf_fp32_ptr, |
| 75 | slot_k_ptr, |
| 76 | slot_score_ptr, |
| 77 | ape_ptr, |
| 78 | loc_ptr, |
| 79 | write_mask_ptr, |
| 80 | compressed_k_ptr, |
| 81 | compressed_scale_ptr, |
| 82 | slot_k_stride_0, |
| 83 | slot_k_stride_1, |
| 84 | slot_score_stride_0, |
| 85 | slot_score_stride_1, |
| 86 | ape_stride_0, |
| 87 | PAGE_SIZE: tl.constexpr, |
| 88 | BUF_NUMEL_PER_PAGE: tl.constexpr, |
| 89 | POOL_SIZE: tl.constexpr, |
| 90 | HEAD_DIM: tl.constexpr, |
| 91 | S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, |
| 92 | FP8_MAX: tl.constexpr, |
| 93 | PRESHUFFLE: tl.constexpr, |
| 94 | ROUND_SCALE: tl.constexpr, |
| 95 | HAS_WRITE_MASK: tl.constexpr, |
| 96 | RETURN_COMPRESSED: tl.constexpr, |
| 97 | WRITE_CACHE: tl.constexpr, |
| 98 | BLOCK_D: tl.constexpr, |
| 99 | ): |
| 100 | """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k -> |
| 101 | Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``.""" |
| 102 | row = tl.program_id(0) |
| 103 | do_write = True |
| 104 | if HAS_WRITE_MASK: |
| 105 | do_write = tl.load(write_mask_ptr + row) |
| 106 | |
| 107 | offs = tl.arange(0, BLOCK_D) |
| 108 | mask = (offs < HEAD_DIM) & do_write |
| 109 | |
| 110 | # --- Pass 1: per-dim max over the pool (softmax numerical stability) --- |
| 111 | max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32) |
| 112 | for slot in tl.static_range(0, POOL_SIZE): |
| 113 | score = tl.load( |
| 114 | slot_score_ptr |
| 115 | + row * slot_score_stride_0 |
| 116 | + slot * slot_score_stride_1 |
| 117 | + offs, |
| 118 | mask=mask, |
| 119 | other=0.0, |
| 120 | ).to(tl.float32) |
| 121 | score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to( |
| 122 | tl.float32 |
| 123 | ) |
| 124 | max_score = tl.maximum(max_score, score) |
| 125 | |
| 126 | # --- Pass 2: softmax-weighted sum of K --- |
| 127 | acc = tl.full((BLOCK_D,), 0.0, tl.float32) |
| 128 | denom = tl.full((BLOCK_D,), 0.0, tl.float32) |
| 129 | for slot in tl.static_range(0, POOL_SIZE): |
| 130 | score = tl.load( |
| 131 | slot_score_ptr |
| 132 | + row * slot_score_stride_0 |
| 133 | + slot * slot_score_stride_1 |
| 134 | + offs, |
| 135 | mask=mask, |
| 136 | other=0.0, |
| 137 | ).to(tl.float32) |
| 138 | score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to( |
| 139 | tl.float32 |
| 140 | ) |
| 141 | prob = tl.exp(score - max_score) |
| 142 | denom += prob |
| 143 | k = tl.load( |
| 144 | slot_k_ptr + row * slot_k_stride_0 + slot * slot_k_stride_1 + offs, |
| 145 | mask=mask, |
| 146 | other=0.0, |
| 147 | ).to(tl.float32) |
| 148 | acc += k * prob |
| 149 | |
| 150 | x = acc / denom |
| 151 | x = tl.where(do_write, x, 0.0).to(tl.bfloat16).to(tl.float32) |
| 152 | |
| 153 | # Match the unfused pooled-K path's bf16 precision before quantization. |
| 154 | x = _hadamard128(x).to(tl.bfloat16).to(tl.float32) |
| 155 | |
| 156 | # --- per-vector absmax fp8 quant --- |
| 157 | fp8_max_inv = 1.0 / FP8_MAX |
| 158 | absmax = tl.max(tl.abs(x), axis=0) |
| 159 | absmax = tl.maximum(absmax, 1e-4) |
| 160 | if ROUND_SCALE: |
| 161 | scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv))) |
| 162 | else: |
| 163 | scale = absmax * fp8_max_inv |
| 164 | quantized = x / scale |
| 165 | quantized = tl.minimum(tl.maximum(quantized, -FP8_MAX), FP8_MAX) |
| 166 | |
| 167 | if WRITE_CACHE: |
| 168 | loc = tl.load(loc_ptr + row, mask=do_write, other=0) |
| 169 | loc_page_index = loc // PAGE_SIZE |
| 170 | loc_token_offset_in_page = loc % PAGE_SIZE |
| 171 | out_k_offsets = loc_page_index * BUF_NUMEL_PER_PAGE + _cache_k_offset( |
| 172 | loc_token_offset_in_page, |
| 173 | offs, |
| 174 | HEAD_DIM, |
| 175 | PRESHUFFLE, |
| 176 | ) |
| 177 | out_s_offset = ( |
| 178 | loc_page_index * BUF_NUMEL_PER_PAGE // 4 |
| 179 | + S_OFFSET_NBYTES_IN_PAGE // 4 |
| 180 | + loc_token_offset_in_page |
| 181 | ) |
| 182 | tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask) |
| 183 | tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) |
| 184 | |
| 185 | if RETURN_COMPRESSED: |
| 186 | tl.store( |
| 187 | compressed_k_ptr + row * HEAD_DIM + offs, |
| 188 | quantized, |
| 189 | mask=offs < HEAD_DIM, |
| 190 | ) |
| 191 | tl.store(compressed_scale_ptr + row, scale) |
| 192 | |
| 193 | |
| 194 | def kpool_compress_and_write_cache( |
| 195 | kv_cache: torch.Tensor, |
| 196 | slot_k: torch.Tensor, |
| 197 | slot_score: torch.Tensor, |
| 198 | ape: torch.Tensor, |
| 199 | loc: torch.Tensor, |
| 200 | pool_size: int, |
| 201 | head_dim: int = INDEX_HEAD_DIM, |
| 202 | write_mask: torch.Tensor | None = None, |
| 203 | round_scale: bool = True, |
| 204 | return_compressed: bool = False, |
| 205 | write_cache: bool = True, |
| 206 | ): |
| 207 | """Compress ``pool_size`` tokens into one fp8 K and write at ``loc``. |
| 208 | |
| 209 | Args: |
| 210 | kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8. |
| 211 | slot_k: ``[n_pools, pool_size, head_dim]`` bf16 — raw per-token K. |
| 212 | slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score. |
| 213 | ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias. |
| 214 | loc: ``[n_pools]`` int64 — flat physical slot per pool. |
| 215 | pool_size: Number of tokens compressed into one cache entry. |
| 216 | head_dim: Indexer head dimension. |
| 217 | write_mask: ``[n_pools]`` bool — pools to write, or None for all. |
| 218 | round_scale: Round each fp8 scale down to a power of two. |
| 219 | return_compressed: Also return the compressed K and scales. |
| 220 | write_cache: Write the compressed result into ``kv_cache``. |
| 221 | |
| 222 | """ |
| 223 | assert slot_k.ndim == 3 |
| 224 | assert slot_score.shape == slot_k.shape |
| 225 | assert ape.shape == slot_k.shape[1:] |
| 226 | assert slot_k.shape[2] == head_dim |
| 227 | assert slot_k.dtype == torch.bfloat16 |
| 228 | assert ape.dtype == torch.float32 |
| 229 | assert kv_cache.dtype == torch.uint8 |
| 230 | assert loc.dtype == torch.int64 |
| 231 | assert write_cache or return_compressed |
| 232 | |
| 233 | page_size = kv_cache.shape[1] |
| 234 | buf = kv_cache |
| 235 | slot_k = slot_k.contiguous() |
| 236 | slot_score = slot_score.contiguous() |
| 237 | ape = ape.contiguous() |
| 238 | loc = loc.contiguous() |
| 239 | if write_mask is None: |
| 240 | write_mask = torch.empty((1,), dtype=torch.bool, device=slot_k.device) |
| 241 | has_write_mask = False |
| 242 | else: |
| 243 | assert write_mask.shape == (slot_k.shape[0],) |
| 244 | write_mask = write_mask.contiguous() |
| 245 | has_write_mask = True |
| 246 | assert not return_compressed |
| 247 | |
| 248 | if slot_k.shape[0] == 0: |
| 249 | if return_compressed: |
| 250 | return ( |
| 251 | torch.empty( |
| 252 | (0, head_dim), |
| 253 | dtype=FP8_DTYPE, |
| 254 | device=slot_k.device, |
| 255 | ), |
| 256 | torch.empty((0,), dtype=torch.float32, device=slot_k.device), |
| 257 | ) |
| 258 | return None |
| 259 | |
| 260 | buf_fp8 = buf.view(FP8_DTYPE) |
| 261 | buf_fp32 = buf.view(torch.float32) |
| 262 | # bytes per page (last dim of kv_cache) viewed as uint8 |
| 263 | buf_numel_per_page = buf.stride(0) |
| 264 | s_offset_nbytes_in_page = page_size * head_dim |
| 265 | |
| 266 | if return_compressed: |
| 267 | compressed_k = torch.empty( |
| 268 | (slot_k.shape[0], head_dim), |
| 269 | dtype=FP8_DTYPE, |
| 270 | device=slot_k.device, |
| 271 | ) |
| 272 | compressed_scale = torch.empty( |
| 273 | (slot_k.shape[0],), dtype=torch.float32, device=slot_k.device |
| 274 | ) |
| 275 | else: |
| 276 | compressed_k = buf_fp8 |
| 277 | compressed_scale = buf_fp32 |
| 278 | |
| 279 | if page_size > 1: |
| 280 | assert page_size % 16 == 0, "ROCm preshuffle requires 16-token tiles" |
| 281 | |
| 282 | _kpool_softmax_rotate_write_cache_kernel[(slot_k.shape[0],)]( |
| 283 | buf_fp8, |
| 284 | buf_fp32, |
| 285 | slot_k, |
| 286 | slot_score, |
| 287 | ape, |
| 288 | loc, |
| 289 | write_mask, |
| 290 | compressed_k, |
| 291 | compressed_scale, |
| 292 | slot_k.stride(0), |
| 293 | slot_k.stride(1), |
| 294 | slot_score.stride(0), |
| 295 | slot_score.stride(1), |
| 296 | ape.stride(0), |
| 297 | PAGE_SIZE=page_size, |
| 298 | BUF_NUMEL_PER_PAGE=buf_numel_per_page, |
| 299 | POOL_SIZE=slot_k.shape[1], |
| 300 | HEAD_DIM=head_dim, |
| 301 | S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page, |
| 302 | FP8_MAX=FP8_MAX, |
| 303 | PRESHUFFLE=page_size > 1, |
| 304 | ROUND_SCALE=round_scale, |
| 305 | HAS_WRITE_MASK=has_write_mask, |
| 306 | RETURN_COMPRESSED=return_compressed, |
| 307 | WRITE_CACHE=write_cache, |
| 308 | BLOCK_D=triton.next_power_of_2(head_dim), |
| 309 | ) |
| 310 | |
| 311 | if return_compressed: |
| 312 | return compressed_k, compressed_scale |
| 313 | return None |
| 314 | |
| 315 | |
| 316 | # Seed each request's incomplete pool into its paged tail during prefill. |
| 317 | |
| 318 | |
| 319 | @triton.jit |
| 320 | def _kpool_tail_seed_kernel( |
| 321 | key_ptr, |
| 322 | score_ptr, |
| 323 | tslot_ptr, |
| 324 | tail_ptr, |
| 325 | n_tokens, |
| 326 | TAIL_BLOCK_ELEMS: tl.constexpr, |
| 327 | KPOOL_HEAD: tl.constexpr, |
| 328 | HEAD_DIM: tl.constexpr, |
| 329 | KPOOL: tl.constexpr, |
| 330 | RING: tl.constexpr, |
| 331 | BLOCK_D: tl.constexpr, |
| 332 | ): |
| 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, |
| 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 |
| 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). |
| 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 |
| 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( |
| 361 | tail_ptr + block_base + KPOOL_HEAD + (t % RING) * HEAD_DIM + offs, s, mask=m |
| 362 | ) |
| 363 | |
| 364 | |
| 365 | def kpool_seed_tail_cache( |
| 366 | tail_kv_cache: torch.Tensor, |
| 367 | key: torch.Tensor, |
| 368 | gate_score: torch.Tensor, |
| 369 | tslot: torch.Tensor, |
| 370 | kpool: int, |
| 371 | head_dim: int = INDEX_HEAD_DIM, |
| 372 | ) -> None: |
| 373 | """Seed the paged tail cache from a prefill batch (see the kernel).""" |
| 374 | assert tail_kv_cache.dtype == torch.bfloat16 |
| 375 | assert key.dtype == torch.bfloat16 |
| 376 | n = tslot.shape[0] |
| 377 | if n == 0: |
| 378 | return |
| 379 | _kpool_tail_seed_kernel[(n,)]( |
| 380 | key, |
| 381 | gate_score, |
| 382 | tslot, |
| 383 | tail_kv_cache, |
| 384 | n, |
| 385 | TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0), |
| 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 | |
| 393 | |
| 394 | # Update each request's tail during decode and write completed pools. |
| 395 | |
| 396 | |
| 397 | @triton.jit |
| 398 | def _kpool_decode_update_batched_kernel( |
| 399 | buf_fp8_ptr, |
| 400 | buf_fp32_ptr, |
| 401 | tail_kv_ptr, |
| 402 | tail_slot_mapping_ptr, # [B, NEXT_N] int32 |
| 403 | key_ptr, # [B, NEXT_N, HEAD_DIM] bf16 |
| 404 | key_stride_b, |
| 405 | key_stride_t, |
| 406 | slot_score_ptr, # [B, NEXT_N, HEAD_DIM] bf16 |
| 407 | ss_stride_b, |
| 408 | ss_stride_t, |
| 409 | ape_ptr, |
| 410 | ape_stride_0, |
| 411 | slot_mapping_ptr, # [B, NEXT_N] int32 |
| 412 | positions_ptr, # [B, NEXT_N] int32 |
| 413 | NEXT_N, # runtime token count per request (no .item() needed) |
| 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, |
| 421 | S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, |
| 422 | FP8_MAX: tl.constexpr, |
| 423 | PRESHUFFLE: tl.constexpr, |
| 424 | ROUND_SCALE: tl.constexpr, |
| 425 | BLOCK_D: tl.constexpr, |
| 426 | ): |
| 427 | """One program per request; iterates its NEXT_N verify tokens in order. |
| 428 | |
| 429 | Replaces the caller's per-token sequential launch loop. The intra-request |
| 430 | iteration MUST stay in position order: a pool-completion at token t* reads |
| 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 |
| 434 | programs are independent (distinct tail blocks). RING >= POOL_SIZE. |
| 435 | """ |
| 436 | req = tl.program_id(0) |
| 437 | offs = tl.arange(0, BLOCK_D) |
| 438 | dim_mask = offs < HEAD_DIM |
| 439 | |
| 440 | for t in tl.range(0, NEXT_N): |
| 441 | idx = req * NEXT_N + t |
| 442 | cache_loc = tl.load(slot_mapping_ptr + idx) |
| 443 | pos = tl.load(positions_ptr + idx) |
| 444 | safe_pos = tl.maximum(pos, 0) |
| 445 | pos_valid = (cache_loc >= 0) & (pos >= 0) |
| 446 | |
| 447 | 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 |
| 452 | # negative sentinel -- reading it from token 0 would poison every |
| 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) // 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 |
| 460 | # the token-granular tail slot -- not on `pos_valid`, which keys off the |
| 461 | # POOL-granular `slot_mapping` and is therefore only true on the pool's |
| 462 | # last token. Gating the stash on pos_valid dropped every intra-pool |
| 463 | # token, so a decode-built pool compressed 3 stale ring entries (the |
| 464 | # prefill-seeded prompt tail, frozen forever) plus the current token. |
| 465 | stash_valid = (pos >= 0) & (tail_slot >= 0) |
| 466 | |
| 467 | key = tl.load( |
| 468 | key_ptr + req * key_stride_b + t * key_stride_t + offs, |
| 469 | mask=dim_mask, |
| 470 | other=0.0, |
| 471 | ).to(tl.float32) |
| 472 | score_current = tl.load( |
| 473 | slot_score_ptr + req * ss_stride_b + t * ss_stride_t + offs, |
| 474 | mask=dim_mask, |
| 475 | other=0.0, |
| 476 | ).to(tl.float32) |
| 477 | |
| 478 | if pos_valid & (slot == POOL_SIZE - 1): |
| 479 | pool_logical_start = safe_pos - slot |
| 480 | |
| 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) % RING |
| 485 | score_buf = tl.load( |
| 486 | tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs, |
| 487 | mask=dim_mask, |
| 488 | other=0.0, |
| 489 | ).to(tl.float32) |
| 490 | score = tl.where(is_current, score_current, score_buf) |
| 491 | score += tl.load( |
| 492 | ape_ptr + pool_slot * ape_stride_0 + offs, |
| 493 | mask=dim_mask, |
| 494 | other=0.0, |
| 495 | ).to(tl.float32) |
| 496 | max_score = tl.maximum(max_score, score) |
| 497 | |
| 498 | acc = tl.full((BLOCK_D,), 0.0, tl.float32) |
| 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) % RING |
| 503 | score_buf = tl.load( |
| 504 | tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs, |
| 505 | mask=dim_mask, |
| 506 | other=0.0, |
| 507 | ).to(tl.float32) |
| 508 | score = tl.where(is_current, score_current, score_buf) |
| 509 | score += tl.load( |
| 510 | ape_ptr + pool_slot * ape_stride_0 + offs, |
| 511 | mask=dim_mask, |
| 512 | other=0.0, |
| 513 | ).to(tl.float32) |
| 514 | prob = tl.exp(score - max_score) |
| 515 | denom += prob |
| 516 | k_buf = tl.load( |
| 517 | tail_kv_ptr + block_base + phys * HEAD_DIM + offs, |
| 518 | mask=dim_mask, |
| 519 | other=0.0, |
| 520 | ).to(tl.float32) |
| 521 | k = tl.where(is_current, key, k_buf) |
| 522 | acc += k * prob |
| 523 | |
| 524 | x = (acc / denom).to(tl.bfloat16).to(tl.float32) |
| 525 | x = _hadamard128(x).to(tl.bfloat16).to(tl.float32) |
| 526 | |
| 527 | fp8_max_inv = 1.0 / FP8_MAX |
| 528 | absmax = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-4) |
| 529 | if ROUND_SCALE: |
| 530 | scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv))) |
| 531 | else: |
| 532 | scale = absmax * fp8_max_inv |
| 533 | quantized = tl.minimum(tl.maximum(x / scale, -FP8_MAX), FP8_MAX) |
| 534 | |
| 535 | loc = cache_loc.to(tl.int64) |
| 536 | loc_page_index = loc // PAGE_SIZE |
| 537 | loc_token_offset_in_page = loc % PAGE_SIZE |
| 538 | out_k_offsets = loc_page_index * BUF_NUMEL_PER_PAGE + _cache_k_offset( |
| 539 | loc_token_offset_in_page, |
| 540 | offs, |
| 541 | HEAD_DIM, |
| 542 | PRESHUFFLE, |
| 543 | ) |
| 544 | out_s_offset = ( |
| 545 | loc_page_index * BUF_NUMEL_PER_PAGE // 4 |
| 546 | + S_OFFSET_NBYTES_IN_PAGE // 4 |
| 547 | + loc_token_offset_in_page |
| 548 | ) |
| 549 | tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask) |
| 550 | tl.store(buf_fp32_ptr + out_s_offset, scale) |
| 551 | |
| 552 | # Stash the current token AFTER any completion read so the completion |
| 553 | # uses prior stashes (and the current token's own key/score via |
| 554 | # is_current), then leaves this token for future pools. Order matches |
| 555 | # the per-token kernel: completion read first, stash second. |
| 556 | update_mask = dim_mask & stash_valid |
| 557 | tl.store( |
| 558 | tail_kv_ptr + block_base + phys_slot * HEAD_DIM + offs, |
| 559 | key, |
| 560 | mask=update_mask, |
| 561 | ) |
| 562 | tl.store( |
| 563 | tail_kv_ptr + block_base + KPOOL_HEAD + phys_slot * HEAD_DIM + offs, |
| 564 | score_current, |
| 565 | mask=update_mask, |
| 566 | ) |
| 567 | |
| 568 | |
| 569 | def kpool_decode_update_and_maybe_write_cache_batched( |
| 570 | kv_cache: torch.Tensor, |
| 571 | tail_kv_cache: torch.Tensor, |
| 572 | tail_slot_mapping: torch.Tensor, |
| 573 | key: torch.Tensor, |
| 574 | slot_score: torch.Tensor, |
| 575 | ape: torch.Tensor, |
| 576 | slot_mapping: torch.Tensor, |
| 577 | positions: torch.Tensor, |
| 578 | pool_size: int, |
| 579 | head_dim: int = INDEX_HEAD_DIM, |
| 580 | round_scale: bool = True, |
| 581 | ) -> None: |
| 582 | """Batched decode-step kpool update for spec verify (``next_n > 1``). |
| 583 | |
| 584 | One launch replaces the caller's per-token loop. Inputs are grouped per |
| 585 | request: ``[num_requests, next_n, ...]``. Each program handles one |
| 586 | request's ``next_n`` tokens in position order (see the kernel docstring for |
| 587 | why ordering is required for pool-completion correctness). |
| 588 | |
| 589 | Plain decode (``next_n == 1``) is handled here too — the kernel collapses |
| 590 | to a single-iteration loop. |
| 591 | |
| 592 | Args: |
| 593 | kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8. |
| 594 | tail_kv_cache: paged tail cache ``[num_blocks, 2, pool_size, head_dim]`` |
| 595 | bf16 (K at half 0, gate score at half 1). |
| 596 | tail_slot_mapping: ``[num_requests, next_n]`` int32. |
| 597 | key: ``[num_requests, next_n, head_dim]`` bf16. |
| 598 | slot_score: ``[num_requests, next_n, head_dim]`` bf16. |
| 599 | ape: ``[pool_size, head_dim]`` fp32. |
| 600 | slot_mapping: ``[num_requests, next_n]`` int32. |
| 601 | positions: ``[num_requests, next_n]`` int32. |
| 602 | pool_size: Number of tokens compressed into one cache entry. |
| 603 | head_dim: Indexer head dimension. |
| 604 | round_scale: Round each fp8 scale down to a power of two. |
| 605 | |
| 606 | """ |
| 607 | num_requests, next_n = key.shape[0], key.shape[1] |
| 608 | if num_requests == 0 or next_n == 0: |
| 609 | return |
| 610 | assert tail_kv_cache.ndim == 4 |
| 611 | assert tail_kv_cache.shape[1] == 2 |
| 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 |
| 617 | assert slot_score.shape == key.shape |
| 618 | assert ape.shape == (pool_size, head_dim) |
| 619 | assert tail_slot_mapping.shape == (num_requests, next_n) |
| 620 | assert slot_mapping.shape == (num_requests, next_n) |
| 621 | assert positions.shape == (num_requests, next_n) |
| 622 | assert key.dtype == torch.bfloat16 |
| 623 | assert slot_score.dtype == torch.bfloat16 |
| 624 | assert ape.dtype == torch.float32 |
| 625 | assert kv_cache.dtype == torch.uint8 |
| 626 | |
| 627 | page_size = kv_cache.shape[1] |
| 628 | buf = kv_cache |
| 629 | buf_fp8 = buf.view(FP8_DTYPE) |
| 630 | buf_fp32 = buf.view(torch.float32) |
| 631 | |
| 632 | # The kernel indexes the int tensors as ``req * next_n + t`` (row-major), |
| 633 | # so they must be contiguous. Callers pass either a view of a contiguous |
| 634 | # slice or a freshly scattered tensor, making these no-ops; the calls guard |
| 635 | # against a future caller handing over a strided view. |
| 636 | tail_slot_mapping = tail_slot_mapping.contiguous() |
| 637 | slot_mapping = slot_mapping.contiguous() |
| 638 | positions = positions.contiguous() |
| 639 | |
| 640 | if page_size > 1: |
| 641 | assert page_size % 16 == 0, "ROCm preshuffle requires 16-token tiles" |
| 642 | |
| 643 | _kpool_decode_update_batched_kernel[(num_requests,)]( |
| 644 | buf_fp8, |
| 645 | buf_fp32, |
| 646 | tail_kv_cache, |
| 647 | tail_slot_mapping, |
| 648 | key, |
| 649 | key.stride(0), |
| 650 | key.stride(1), |
| 651 | slot_score, |
| 652 | slot_score.stride(0), |
| 653 | slot_score.stride(1), |
| 654 | ape, |
| 655 | ape.stride(0), |
| 656 | slot_mapping, |
| 657 | positions, |
| 658 | next_n, |
| 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, |
| 666 | S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim, |
| 667 | FP8_MAX=FP8_MAX, |
| 668 | PRESHUFFLE=page_size > 1, |
| 669 | ROUND_SCALE=round_scale, |
| 670 | BLOCK_D=triton.next_power_of_2(head_dim), |
| 671 | ) |
| 672 | |
| 673 | |
| 674 | # Pool-level top-k helpers. |
| 675 | |
| 676 | |
| 677 | def history_group_budget_for_topk(topk: int, pool_size: int) -> int: |
| 678 | """Number of pools to select so that expanding yields ``topk`` tokens.""" |
| 679 | assert topk % pool_size == 0 |
| 680 | return topk // pool_size |
| 681 | |
| 682 | |
| 683 | def expand_pools_to_tokens( |
| 684 | group_ids: torch.Tensor, |
| 685 | group_valid: torch.Tensor, |
| 686 | topk: int, |
| 687 | pool_size: int, |
| 688 | page_table: torch.Tensor | None = None, |
| 689 | topk_offsets: torch.Tensor | None = None, |
| 690 | ) -> torch.Tensor: |
| 691 | """Expand selected full-pool ids to a strict-width token topk tensor.""" |
| 692 | assert group_ids.ndim == 2 |
| 693 | assert group_valid.shape == group_ids.shape |
| 694 | assert topk % pool_size == 0 |
| 695 | assert group_ids.shape[1] == history_group_budget_for_topk(topk, pool_size) |
| 696 | assert page_table is None or topk_offsets is None |
| 697 | |
| 698 | device = group_ids.device |
| 699 | offsets = torch.arange(pool_size, device=device, dtype=torch.int64) |
| 700 | token_ids = group_ids.to(torch.int64).unsqueeze(-1) * pool_size + offsets |
| 701 | token_ids = token_ids.reshape(group_ids.shape[0], topk) |
| 702 | valid = ( |
| 703 | group_valid.unsqueeze(-1) |
| 704 | .expand(-1, -1, pool_size) |
| 705 | .reshape(group_ids.shape[0], topk) |
| 706 | ) |
| 707 | |
| 708 | if page_table is not None: |
| 709 | assert page_table.ndim == 2 |
| 710 | safe_ids = token_ids.clamp(min=0, max=page_table.shape[1] - 1) |
| 711 | output = torch.gather(page_table, dim=1, index=safe_ids).to(torch.int32) |
| 712 | elif topk_offsets is not None: |
| 713 | if topk_offsets.ndim == 2: |
| 714 | assert topk_offsets.shape[1] == 1 |
| 715 | topk_offsets = topk_offsets.squeeze(1) |
| 716 | output = (token_ids + topk_offsets.to(torch.int64).unsqueeze(1)).to(torch.int32) |
| 717 | else: |
| 718 | output = token_ids.to(torch.int32) |
| 719 | |
| 720 | return torch.where(valid, output, torch.full_like(output, -1)) |
| 721 | |
| 722 | |
| 723 | def append_tail_to_topk( |
| 724 | topk_result: torch.Tensor, |
| 725 | seq_lens: torch.Tensor, |
| 726 | pool_lens: torch.Tensor, |
| 727 | pool_size: int, |
| 728 | page_table: torch.Tensor | None = None, |
| 729 | topk_offsets: torch.Tensor | None = None, |
| 730 | ) -> torch.Tensor: |
| 731 | """Append non-pooled tail tokens after expanded history tokens. |
| 732 | |
| 733 | ``index_kpool_always_select_tail`` keeps the (incomplete) trailing pool so |
| 734 | the most recent tokens are always attended to. |
| 735 | """ |
| 736 | assert topk_result.dtype == torch.int32 |
| 737 | assert seq_lens.ndim == 1 |
| 738 | assert pool_lens.ndim == 1 |
| 739 | |
| 740 | tail_pool = pool_size - 1 |
| 741 | if tail_pool == 0: |
| 742 | return topk_result |
| 743 | |
| 744 | rows, n_cols = topk_result.shape |
| 745 | out_cols = n_cols + tail_pool |
| 746 | out = torch.empty( |
| 747 | (rows, out_cols), dtype=topk_result.dtype, device=topk_result.device |
| 748 | ) |
| 749 | |
| 750 | # tail tokens: [pool_len*pool_size, seq_len) for each row. |
| 751 | pool_len = pool_lens.to(torch.int32) |
| 752 | tail_start = pool_len * pool_size |
| 753 | seq_len = seq_lens.to(torch.int32) |
| 754 | tail_count = seq_len - tail_start # in [0, pool_size) |
| 755 | |
| 756 | cols = torch.arange(out_cols, device=topk_result.device)[None, :] |
| 757 | history_len = n_cols |
| 758 | is_history = cols < history_len |
| 759 | tail_off = cols - history_len |
| 760 | is_tail = (tail_off >= 0) & (tail_off < tail_count[:, None]) |
| 761 | |
| 762 | # safe_hist must be per-row [rows, out_cols] so the gather reads each row's |
| 763 | # OWN history. cols is [1, out_cols]; if used directly, gather (which does |
| 764 | # NOT broadcast the index) would read only row 0 of topk_result, making every |
| 765 | # query inherit row 0's history (empty for the first token) and lose all its |
| 766 | # selected tokens — only the per-row tail would survive. This only manifests |
| 767 | # for multi-row sparse PREFILL (decode has 1 row, so it reads its own row 0). |
| 768 | safe_hist = torch.minimum(cols, torch.full_like(cols, n_cols - 1)).expand( |
| 769 | rows, out_cols |
| 770 | ) |
| 771 | history_val = torch.gather(topk_result, 1, safe_hist) |
| 772 | |
| 773 | tail_raw = tail_start[:, None] + tail_off |
| 774 | tail_val = tail_raw.to(torch.int32) |
| 775 | if page_table is not None: |
| 776 | safe_tail = tail_raw.clamp(min=0, max=page_table.shape[1] - 1) |
| 777 | tail_val = torch.gather(page_table, 1, safe_tail).to(torch.int32) |
| 778 | elif topk_offsets is not None: |
| 779 | tail_val = (tail_raw + topk_offsets.to(torch.int64).unsqueeze(1)).to( |
| 780 | torch.int32 |
| 781 | ) |
| 782 | |
| 783 | out = torch.where(is_history, history_val, -1) |
| 784 | out = torch.where(is_tail, tail_val, out) |
| 785 | return out |
| 786 | |
| 787 | |
| 788 | @triton.jit |
| 789 | def _expand_pools_and_append_tail_kernel( |
| 790 | pool_ids_ptr, # [rows, n_groups], int (any int dtype) |
| 791 | seq_lens_ptr, # [rows], int32 (token-granular seq_len) |
| 792 | out_ptr, # [rows, out_cols], int32 |
| 793 | topk, # n_groups * pool_size |
| 794 | out_cols, # topk + pool_size - 1 |
| 795 | POOL_SIZE: tl.constexpr, |
| 796 | BLOCK_COLS: tl.constexpr, |
| 797 | pid_s0, |
| 798 | out_s0, |
| 799 | ): |
| 800 | # Fuses expand_pools_to_tokens + append_tail_to_topk (identity path) into a |
| 801 | # single kernel. Each program writes one (row, column-tile) of the output. |
| 802 | row = tl.program_id(0) |
| 803 | tile = tl.program_id(1) |
| 804 | cols = tile * BLOCK_COLS + tl.arange(0, BLOCK_COLS) |
| 805 | mask = cols < out_cols |
| 806 | |
| 807 | seq_len = tl.load(seq_lens_ptr + row) |
| 808 | pool_len = seq_len // POOL_SIZE |
| 809 | tail_start = pool_len * POOL_SIZE |
| 810 | tail_count = seq_len - tail_start # in [0, POOL_SIZE) |
| 811 | |
| 812 | # History region [0, topk): expand selected pool g = cols // POOL_SIZE. |
| 813 | is_history = cols < topk |
| 814 | g = cols // POOL_SIZE |
| 815 | o = cols % POOL_SIZE |
| 816 | pid = tl.load(pool_ids_ptr + row * pid_s0 + g, mask=mask & is_history, other=-1) |
| 817 | hist_val = (pid * POOL_SIZE + o).to(tl.int32) |
| 818 | hist_out = tl.where(pid >= 0, hist_val, -1) |
| 819 | |
| 820 | # Tail region [topk, out_cols): the request's trailing incomplete pool. |
| 821 | tail_off = cols - topk |
| 822 | is_tail = (tail_off >= 0) & (tail_off < tail_count) |
| 823 | tail_val = (tail_start + tail_off).to(tl.int32) |
| 824 | tail_out = tl.where(is_tail, tail_val, -1) |
| 825 | |
| 826 | result = tl.where(is_history, hist_out, tail_out) |
| 827 | tl.store(out_ptr + row * out_s0 + cols, result, mask=mask) |
| 828 | |
| 829 | |
| 830 | def expand_pools_and_append_tail( |
| 831 | pool_ids: torch.Tensor, |
| 832 | seq_lens: torch.Tensor, |
| 833 | pool_size: int, |
| 834 | ) -> torch.Tensor: |
| 835 | """Fuse ``expand_pools_to_tokens`` + ``append_tail_to_topk`` (identity path). |
| 836 | |
| 837 | Produces the same ``[rows, topk + pool_size - 1]`` int32 output as calling |
| 838 | the two functions in sequence when neither ``page_table`` nor |
| 839 | ``topk_offsets`` is passed — the only path used by the GLM-5.3-Flash indexer. |
| 840 | The kernel derives ``pool_len = seq_len // pool_size`` internally, so the |
| 841 | caller no longer needs to precompute it. Replaces ~25 elementwise kernels |
| 842 | with one Triton launch. |
| 843 | """ |
| 844 | rows, n_groups = pool_ids.shape |
| 845 | topk = n_groups * pool_size |
| 846 | out_cols = topk + pool_size - 1 |
| 847 | out = torch.empty((rows, out_cols), dtype=torch.int32, device=pool_ids.device) |
| 848 | BLOCK_COLS = 128 |
| 849 | n_tiles = triton.cdiv(out_cols, BLOCK_COLS) |
| 850 | _expand_pools_and_append_tail_kernel[(rows, n_tiles)]( |
| 851 | pool_ids, |
| 852 | seq_lens, |
| 853 | out, |
| 854 | topk, |
| 855 | out_cols, |
| 856 | POOL_SIZE=pool_size, |
| 857 | BLOCK_COLS=BLOCK_COLS, |
| 858 | pid_s0=pool_ids.stride(0), |
| 859 | out_s0=out.stride(0), |
| 860 | ) |
| 861 | return out |