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