| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton kernels for DeepseekV4 paged K-cache management and sparse-attention index |
| 4 | preparation. |
| 5 | |
| 6 | - quantize_and_insert_k_cache: quantize bf16 K to UE8M0 FP8 and insert into |
| 7 | the paged cache. |
| 8 | - dequantize_and_gather_k_cache: gather and dequantize FP8 K from the paged |
| 9 | cache for sparse/SWA prefill. |
| 10 | - compute_global_topk_indices_and_lens: map local topk indices to global KV |
| 11 | cache slots and count valid entries. |
| 12 | - combine_topk_swa_indices: concatenate topk compressed indices with SWA |
| 13 | window indices for sparse prefill. |
| 14 | """ |
| 15 | |
| 16 | from dataclasses import dataclass |
| 17 | from typing import Any |
| 18 | |
| 19 | import torch |
| 20 | |
| 21 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 22 | get_fp8_min_max, |
| 23 | ) |
| 24 | from vllm.model_executor.warmup.jit_warmup import ( |
| 25 | VllmJitKernel, |
| 26 | _WarmupInputRows, |
| 27 | zip_inputs, |
| 28 | ) |
| 29 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 30 | TritonPointerInputVariant, |
| 31 | TritonWarmupTensor, |
| 32 | ) |
| 33 | from vllm.platforms import current_platform |
| 34 | from vllm.triton_utils import tl, triton |
| 35 | from vllm.utils.import_utils import has_cutedsl |
| 36 | from vllm.utils.math_utils import next_power_of_2 |
| 37 | |
| 38 | # Per-token byte width of the two paged fp8 records. Both put a page's whole |
| 39 | # data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the |
| 40 | # width the KV-cache spec publishes -- identifies the record: |
| 41 | # V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64 |
| 42 | # dims each plus a pad byte. |
| 43 | # V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 |
| 44 | # scales of 32 dims each. FlashMLA's ``ModelType::V41``. |
| 45 | # V4.1 NVFP4 (288 B): 512 e2m1 values packed two per byte (even element in |
| 46 | # the low nibble), then 32 e4m3 scales of 16 dims each. |
| 47 | # FlashMLA's ``ModelType::V41_FP4``, compressed cache only. |
| 48 | V4_BYTES_PER_TOKEN = 584 |
| 49 | V41_BYTES_PER_TOKEN = 528 |
| 50 | V41_QUANT_BLOCK = 32 |
| 51 | V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16 |
| 52 | V41_NVFP4_BYTES_PER_TOKEN = 288 |
| 53 | V41_NVFP4_QUANT_BLOCK = 16 |
| 54 | V41_NVFP4_NUM_SCALES = 512 // V41_NVFP4_QUANT_BLOCK # 32 |
| 55 | |
| 56 | |
| 57 | @triton.jit |
| 58 | def quantize_and_insert_k_kernel( |
| 59 | # Input tensors |
| 60 | k_ptr, # [num_tokens, 512] bf16 |
| 61 | slot_mapping_ptr, # [num_tokens] int64 |
| 62 | # Output tensor |
| 63 | k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) |
| 64 | # Dimensions |
| 65 | num_tokens, |
| 66 | input_dim: tl.constexpr, # 512 |
| 67 | fp8_dim: tl.constexpr, # 448 |
| 68 | bf16_dim: tl.constexpr, # 64 |
| 69 | scale_dim: tl.constexpr, # 8 |
| 70 | quant_block: tl.constexpr, # 64 (quantization block size) |
| 71 | cache_block_size: tl.constexpr, # 64 (paged cache block size) |
| 72 | token_data_size: tl.constexpr, # 576 bytes per token data |
| 73 | block_stride: tl.constexpr, # total bytes per block (padded) |
| 74 | fp8_max: tl.constexpr, |
| 75 | n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding) |
| 76 | use_fnuz: tl.constexpr = False, |
| 77 | ): |
| 78 | """Quantize K tensor and insert into paged K cache. |
| 79 | |
| 80 | K Cache block layout (block_size=64 tokens): |
| 81 | - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16 |
| 82 | - [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales |
| 83 | - [64*576 + 64*8, block_stride): Padding |
| 84 | |
| 85 | One program per token. |
| 86 | |
| 87 | ``use_fnuz=True`` selects FNUZ (``tl.float8e4b8``); default OCP |
| 88 | (``tl.float8e4nv``) matches every production caller. |
| 89 | """ |
| 90 | pid = tl.program_id(0) |
| 91 | |
| 92 | if pid >= num_tokens: |
| 93 | return |
| 94 | |
| 95 | # Get slot mapping |
| 96 | slot_idx = tl.load(slot_mapping_ptr + pid) |
| 97 | if slot_idx == -1: |
| 98 | return |
| 99 | |
| 100 | block_idx = slot_idx // cache_block_size |
| 101 | pos_in_block = slot_idx % cache_block_size |
| 102 | |
| 103 | # Input pointer for this token |
| 104 | input_row_ptr = k_ptr + pid * input_dim |
| 105 | |
| 106 | # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks |
| 107 | # (e.g. >= 57K at block_stride ~37K). Matches gather path below. |
| 108 | cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride |
| 109 | |
| 110 | # Token data pointer: token data is stored contiguously at start of block |
| 111 | # Each token's data is at offset pos_in_block * token_data_size |
| 112 | token_data_ptr = cache_block_ptr + pos_in_block * token_data_size |
| 113 | |
| 114 | # Scale pointer: scales are stored after ALL token data in the block |
| 115 | # Scale for this token is at offset (64 * 576) + pos_in_block * 8 |
| 116 | token_scale_ptr = ( |
| 117 | cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim |
| 118 | ) |
| 119 | |
| 120 | # Token data layout: [0:448] fp8, [448:576] bf16 |
| 121 | token_fp8_ptr = token_data_ptr |
| 122 | token_bf16_ptr = token_data_ptr + fp8_dim |
| 123 | |
| 124 | # ========== Quantize and store FP8 portion (first 448 elements) ========== |
| 125 | # Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent) |
| 126 | for qblock_idx in tl.static_range(n_quant_blocks): |
| 127 | qblock_start = qblock_idx * quant_block |
| 128 | |
| 129 | if qblock_start < fp8_dim: |
| 130 | offsets = qblock_start + tl.arange(0, quant_block) |
| 131 | mask = offsets < fp8_dim |
| 132 | |
| 133 | # Load bf16 input |
| 134 | x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0) |
| 135 | |
| 136 | # Compute absmax scale (same as CUDA kernel) |
| 137 | abs_x = tl.abs(x) |
| 138 | block_max = tl.max(abs_x, axis=0) |
| 139 | block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4) |
| 140 | |
| 141 | # UE8M0: Round scale UP to next power of 2 |
| 142 | # scale = 2^ceil(log2(block_max / fp8_max)) |
| 143 | raw_scale = block_max / fp8_max |
| 144 | log_scale = tl.log2(raw_scale) |
| 145 | exponent = tl.ceil(log_scale) # Round UP to next integer exponent |
| 146 | scale = tl.exp2(exponent) # scale = 2^exponent (power of 2) |
| 147 | |
| 148 | # Quantize to fp8: fp8_value = bf16_value / scale |
| 149 | x_scaled = x / scale |
| 150 | x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max) |
| 151 | |
| 152 | # Convert to fp8 (FNUZ on gfx942, OCP elsewhere), then bitcast to uint8. |
| 153 | if use_fnuz: |
| 154 | x_fp8 = x_clamped.to(tl.float8e4b8) |
| 155 | else: |
| 156 | x_fp8 = x_clamped.to(tl.float8e4nv) |
| 157 | x_uint8 = x_fp8.to(tl.uint8, bitcast=True) |
| 158 | |
| 159 | # Store as uint8 (1 byte each) |
| 160 | tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask) |
| 161 | |
| 162 | # UE8M0 scale encoding: stored_value = exponent + 127 (bias) |
| 163 | # During dequant: scale = 2^(stored_value - 127) |
| 164 | encoded_scale = exponent + 127.0 |
| 165 | encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0) |
| 166 | tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8)) |
| 167 | |
| 168 | # Padding scale at index 7 |
| 169 | tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8)) |
| 170 | |
| 171 | # ========== Store BF16 portion (last 64 elements, no quantization) ========== |
| 172 | bf16_input_offset = fp8_dim |
| 173 | |
| 174 | # Process bf16 in chunks of 16 |
| 175 | bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) |
| 176 | for i in tl.static_range(bf16_dim // 16): |
| 177 | chunk_offsets = i * 16 + tl.arange(0, 16) |
| 178 | bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets) |
| 179 | tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) |
| 180 | |
| 181 | |
| 182 | @triton.jit |
| 183 | def _quantize_and_insert_k_mxfp8_kernel( |
| 184 | k_ptr, # [num_tokens, 512] bf16 |
| 185 | slot_mapping_ptr, # [num_tokens] int64 |
| 186 | k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) |
| 187 | num_tokens, |
| 188 | head_dim: tl.constexpr, # 512, RoPE dims included |
| 189 | scale_dim: tl.constexpr, # 16 |
| 190 | quant_block: tl.constexpr, # 32 dims per UE8M0 scale |
| 191 | cache_block_size: tl.constexpr, # paged cache block size, in tokens |
| 192 | block_stride: tl.constexpr, # total bytes per block (padded) |
| 193 | fp8_max: tl.constexpr, |
| 194 | use_fnuz: tl.constexpr = False, |
| 195 | ): |
| 196 | """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache. |
| 197 | |
| 198 | One program per token. Unlike the V4 record there is no bf16 tail, so every |
| 199 | dim is scaled. |
| 200 | """ |
| 201 | pid = tl.program_id(0) |
| 202 | if pid >= num_tokens: |
| 203 | return |
| 204 | |
| 205 | slot_idx = tl.load(slot_mapping_ptr + pid) |
| 206 | if slot_idx == -1: |
| 207 | return |
| 208 | |
| 209 | block_idx = slot_idx // cache_block_size |
| 210 | pos_in_block = slot_idx % cache_block_size |
| 211 | |
| 212 | # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks. |
| 213 | cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride |
| 214 | token_data_ptr = cache_block_ptr + pos_in_block * head_dim |
| 215 | token_scale_ptr = ( |
| 216 | cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim |
| 217 | ) |
| 218 | |
| 219 | d = tl.arange(0, head_dim) |
| 220 | x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32) |
| 221 | |
| 222 | tiles = tl.reshape(x, (scale_dim, quant_block)) |
| 223 | # Floor the absmax like the CUDA encoder so an all-zero tile still lands on |
| 224 | # a finite scale. |
| 225 | amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4) |
| 226 | exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max))) |
| 227 | scaled = tl.clamp( |
| 228 | tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max |
| 229 | ) |
| 230 | if use_fnuz: # noqa: SIM108 |
| 231 | fp8 = scaled.to(tl.float8e4b8) |
| 232 | else: |
| 233 | fp8 = scaled.to(tl.float8e4nv) |
| 234 | tl.store( |
| 235 | token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,)) |
| 236 | ) |
| 237 | |
| 238 | # UE8M0 encoding: stored_value = exponent + 127 (bias). |
| 239 | encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0) |
| 240 | tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8)) |
| 241 | |
| 242 | |
| 243 | def quantize_and_insert_k_cache( |
| 244 | k: torch.Tensor, # [num_tokens, 512] bf16 |
| 245 | k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 |
| 246 | slot_mapping: torch.Tensor, # [num_tokens] int64 |
| 247 | block_size: int = 64, |
| 248 | is_ue8m0: bool = True, |
| 249 | use_fnuz: bool = False, |
| 250 | bytes_per_token: int = V4_BYTES_PER_TOKEN, |
| 251 | ): |
| 252 | """Quantize K tensor and insert into paged K cache. |
| 253 | |
| 254 | ``bytes_per_token`` picks the record (see the module header): the V4 one, |
| 255 | or V4.1's all-dims MXFP8 one. |
| 256 | |
| 257 | ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on |
| 258 | platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3, |
| 259 | which is used by OCP-encoded caches even on gfx942. |
| 260 | """ |
| 261 | assert k.dim() == 2 and k.shape[1] == 512, ( |
| 262 | f"K must be [num_tokens, 512], got {k.shape}" |
| 263 | ) |
| 264 | assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}" |
| 265 | assert is_ue8m0, "Only support ue8m0 quantization." |
| 266 | |
| 267 | # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to |
| 268 | # padding. Always use slot_mapping.shape[0] as the token count. |
| 269 | num_tokens = slot_mapping.shape[0] |
| 270 | block_stride = k_cache.stride(0) # bytes per block |
| 271 | |
| 272 | if bytes_per_token == V41_BYTES_PER_TOKEN: |
| 273 | if use_fnuz and not current_platform.is_fp8_fnuz(): |
| 274 | raise ValueError("use_fnuz=True requires a platform using FNUZ FP8") |
| 275 | _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)]( |
| 276 | k, |
| 277 | slot_mapping, |
| 278 | k_cache, |
| 279 | num_tokens, |
| 280 | head_dim=512, |
| 281 | scale_dim=V41_NUM_SCALES, |
| 282 | quant_block=V41_QUANT_BLOCK, |
| 283 | cache_block_size=block_size, |
| 284 | block_stride=block_stride, |
| 285 | fp8_max=get_fp8_min_max()[1] |
| 286 | if use_fnuz |
| 287 | else torch.finfo(torch.float8_e4m3fn).max, |
| 288 | use_fnuz=use_fnuz, |
| 289 | ) |
| 290 | return |
| 291 | |
| 292 | TOKEN_FP8_DIM = 448 |
| 293 | TOKEN_BF16_DIM = 64 |
| 294 | TOKEN_SCALE_DIM = 8 |
| 295 | QUANT_BLOCK_SIZE = 64 |
| 296 | if use_fnuz: |
| 297 | if not current_platform.is_fp8_fnuz(): |
| 298 | raise ValueError("use_fnuz=True requires a platform using FNUZ FP8") |
| 299 | _, FP8_MAX = get_fp8_min_max() |
| 300 | else: |
| 301 | FP8_MAX = torch.finfo(torch.float8_e4m3fn).max |
| 302 | TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 |
| 303 | |
| 304 | grid = (num_tokens,) |
| 305 | |
| 306 | quantize_and_insert_k_kernel[grid]( |
| 307 | k, |
| 308 | slot_mapping, |
| 309 | k_cache, |
| 310 | num_tokens, |
| 311 | input_dim=512, |
| 312 | fp8_dim=TOKEN_FP8_DIM, |
| 313 | bf16_dim=TOKEN_BF16_DIM, |
| 314 | scale_dim=TOKEN_SCALE_DIM, |
| 315 | quant_block=QUANT_BLOCK_SIZE, |
| 316 | cache_block_size=block_size, |
| 317 | token_data_size=TOKEN_DATA_SIZE, |
| 318 | block_stride=block_stride, |
| 319 | fp8_max=FP8_MAX, |
| 320 | n_quant_blocks=8, |
| 321 | use_fnuz=use_fnuz, |
| 322 | ) |
| 323 | |
| 324 | |
| 325 | @triton.jit |
| 326 | def _dequantize_and_gather_k_kernel( |
| 327 | out_ptr, |
| 328 | out_stride0, |
| 329 | out_stride1, |
| 330 | k_cache_ptr, |
| 331 | seq_lens_ptr, |
| 332 | block_table_ptr, |
| 333 | offset, |
| 334 | gather_lens_ptr, |
| 335 | # Constants |
| 336 | max_blocks_per_seq: tl.constexpr, |
| 337 | fp8_dim: tl.constexpr, # 448 |
| 338 | bf16_dim: tl.constexpr, # 64 |
| 339 | scale_dim: tl.constexpr, # 8 |
| 340 | quant_block: tl.constexpr, # 64 (quantization block size) |
| 341 | cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size) |
| 342 | token_data_size: tl.constexpr, # 576 bytes per token data |
| 343 | block_stride: tl.constexpr, # total bytes per block (padded) int32 |
| 344 | output_dim: tl.constexpr, # 512 |
| 345 | fp8_max: tl.constexpr, |
| 346 | n_quant_blocks: tl.constexpr, # 7 real blocks |
| 347 | use_fnuz: tl.constexpr = False, |
| 348 | ): |
| 349 | batch_idx = tl.program_id(0) |
| 350 | worker_id = tl.program_id(1) |
| 351 | num_workers = tl.num_programs(1) |
| 352 | |
| 353 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 354 | if gather_lens_ptr is not None: # noqa: SIM108 |
| 355 | gather_len = tl.load(gather_lens_ptr + batch_idx) |
| 356 | else: |
| 357 | # Gather all tokens |
| 358 | gather_len = seq_len |
| 359 | start_pos = seq_len - gather_len |
| 360 | |
| 361 | for i in range(worker_id, gather_len, num_workers): |
| 362 | # Calculate the actual token index in the sequence |
| 363 | pos = start_pos + i |
| 364 | |
| 365 | # Calculate which block and position within block |
| 366 | block_in_seq = pos // cache_block_size |
| 367 | pos_in_block = pos % cache_block_size |
| 368 | |
| 369 | # Get physical block index from block table |
| 370 | block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq |
| 371 | physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 |
| 372 | |
| 373 | # int64: physical_block_idx * block_stride can exceed 2^31 with many |
| 374 | # KV-cache blocks (e.g. >= 57K at block_stride ~37K). |
| 375 | cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride |
| 376 | |
| 377 | # Token data pointer |
| 378 | token_data_ptr = cache_block_ptr + pos_in_block * token_data_size |
| 379 | |
| 380 | # Scale pointer: after all token data |
| 381 | token_scale_ptr = ( |
| 382 | cache_block_ptr |
| 383 | + cache_block_size * token_data_size |
| 384 | + pos_in_block * scale_dim |
| 385 | ) |
| 386 | |
| 387 | # Token data layout: [0:448] fp8, [448:576] bf16 |
| 388 | token_fp8_ptr = token_data_ptr |
| 389 | token_bf16_ptr = token_data_ptr + fp8_dim |
| 390 | |
| 391 | # Output pointer for this token (flattened) |
| 392 | output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 |
| 393 | |
| 394 | # ========== Dequantize FP8 portion using UE8M0 ========== |
| 395 | for qblock_idx in tl.static_range(n_quant_blocks): |
| 396 | qblock_start = qblock_idx * quant_block |
| 397 | |
| 398 | if qblock_start < fp8_dim: |
| 399 | offsets = qblock_start + tl.arange(0, quant_block) |
| 400 | mask = offsets < fp8_dim |
| 401 | |
| 402 | # Load quantized fp8 values (stored as uint8) |
| 403 | x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0) |
| 404 | |
| 405 | # Bitcast uint8 back to fp8 (FNUZ on gfx942, OCP elsewhere). |
| 406 | if use_fnuz: |
| 407 | x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) |
| 408 | else: |
| 409 | x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) |
| 410 | |
| 411 | # Convert fp8 to float32 for computation |
| 412 | x_float = x_fp8.to(tl.float32) |
| 413 | |
| 414 | # Load and decode UE8M0 scale |
| 415 | # UE8M0: scale = 2^(stored_value - 127) |
| 416 | encoded_scale = tl.load(token_scale_ptr + qblock_idx) |
| 417 | exponent = encoded_scale.to(tl.float32) - 127.0 |
| 418 | scale = tl.exp2(exponent) |
| 419 | |
| 420 | # Dequantize: bf16_value = fp8_value * scale |
| 421 | x_dequant = x_float * scale |
| 422 | |
| 423 | # Store as bf16 |
| 424 | tl.store(output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask) |
| 425 | |
| 426 | # ========== Copy BF16 portion directly ========== |
| 427 | bf16_output_offset = fp8_dim # After 448 elements in output |
| 428 | |
| 429 | # Read bf16 from cache |
| 430 | bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) |
| 431 | |
| 432 | # Process in chunks of 16 |
| 433 | for j in tl.static_range(bf16_dim // 16): |
| 434 | chunk_offsets = j * 16 + tl.arange(0, 16) |
| 435 | bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets) |
| 436 | tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) |
| 437 | |
| 438 | |
| 439 | @triton.jit |
| 440 | def _dequantize_and_gather_k_mxfp8_kernel( |
| 441 | out_ptr, |
| 442 | out_stride0, |
| 443 | out_stride1, |
| 444 | k_cache_ptr, |
| 445 | seq_lens_ptr, |
| 446 | block_table_ptr, |
| 447 | offset, |
| 448 | gather_lens_ptr, |
| 449 | max_blocks_per_seq: tl.constexpr, |
| 450 | head_dim: tl.constexpr, # 512 |
| 451 | scale_dim: tl.constexpr, # 16 |
| 452 | quant_block: tl.constexpr, # 32 |
| 453 | cache_block_size: tl.constexpr, |
| 454 | block_stride: tl.constexpr, |
| 455 | use_fnuz: tl.constexpr = False, |
| 456 | ): |
| 457 | """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place.""" |
| 458 | batch_idx = tl.program_id(0) |
| 459 | worker_id = tl.program_id(1) |
| 460 | num_workers = tl.num_programs(1) |
| 461 | |
| 462 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 463 | if gather_lens_ptr is not None: # noqa: SIM108 |
| 464 | gather_len = tl.load(gather_lens_ptr + batch_idx) |
| 465 | else: |
| 466 | gather_len = seq_len |
| 467 | start_pos = seq_len - gather_len |
| 468 | |
| 469 | d = tl.arange(0, head_dim) |
| 470 | s = tl.arange(0, scale_dim) |
| 471 | for i in range(worker_id, gather_len, num_workers): |
| 472 | pos = start_pos + i |
| 473 | block_in_seq = pos // cache_block_size |
| 474 | pos_in_block = pos % cache_block_size |
| 475 | |
| 476 | block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq |
| 477 | physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 |
| 478 | |
| 479 | # int64: physical_block_idx * block_stride can exceed 2^31. |
| 480 | cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride |
| 481 | token_data_ptr = cache_block_ptr + pos_in_block * head_dim |
| 482 | token_scale_ptr = ( |
| 483 | cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim |
| 484 | ) |
| 485 | |
| 486 | x_uint8 = tl.load(token_data_ptr + d) |
| 487 | if use_fnuz: |
| 488 | x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) |
| 489 | else: |
| 490 | x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) |
| 491 | tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block)) |
| 492 | |
| 493 | # UE8M0: scale = 2^(stored_value - 127). |
| 494 | encoded = tl.load(token_scale_ptr + s) |
| 495 | scale = tl.exp2(encoded.to(tl.float32) - 127.0) |
| 496 | dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) |
| 497 | |
| 498 | output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 |
| 499 | tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) |
| 500 | |
| 501 | |
| 502 | @triton.jit |
| 503 | def _dequantize_and_gather_k_nvfp4_kernel( |
| 504 | out_ptr, |
| 505 | out_stride0, |
| 506 | out_stride1, |
| 507 | k_cache_ptr, |
| 508 | seq_lens_ptr, |
| 509 | block_table_ptr, |
| 510 | offset, |
| 511 | gather_lens_ptr, |
| 512 | max_blocks_per_seq: tl.constexpr, |
| 513 | head_dim: tl.constexpr, # 512 |
| 514 | scale_dim: tl.constexpr, # 32 |
| 515 | quant_block: tl.constexpr, # 16 |
| 516 | cache_block_size: tl.constexpr, |
| 517 | block_stride: tl.constexpr, |
| 518 | ): |
| 519 | """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace.""" |
| 520 | batch_idx = tl.program_id(0) |
| 521 | worker_id = tl.program_id(1) |
| 522 | num_workers = tl.num_programs(1) |
| 523 | |
| 524 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 525 | if gather_lens_ptr is not None: # noqa: SIM108 |
| 526 | gather_len = tl.load(gather_lens_ptr + batch_idx) |
| 527 | else: |
| 528 | gather_len = seq_len |
| 529 | start_pos = seq_len - gather_len |
| 530 | |
| 531 | packed_bytes: tl.constexpr = head_dim // 2 |
| 532 | d = tl.arange(0, head_dim) |
| 533 | for i in range(worker_id, gather_len, num_workers): |
| 534 | pos = start_pos + i |
| 535 | block_in_seq = pos // cache_block_size |
| 536 | pos_in_block = pos % cache_block_size |
| 537 | |
| 538 | physical_block_idx = tl.load( |
| 539 | block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq |
| 540 | ) |
| 541 | page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride |
| 542 | |
| 543 | packed = tl.load( |
| 544 | page + pos_in_block * packed_bytes + tl.arange(0, packed_bytes) |
| 545 | ) |
| 546 | # Even element in the low nibble, odd in the high nibble. |
| 547 | codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32)) |
| 548 | # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. |
| 549 | mag_code = codes & 7 |
| 550 | e = (mag_code >> 1).to(tl.float32) |
| 551 | m = (mag_code & 1).to(tl.float32) |
| 552 | mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0)) |
| 553 | vals = tl.where(codes >= 8, -mag, mag) |
| 554 | |
| 555 | sf = tl.load( |
| 556 | page |
| 557 | + cache_block_size * packed_bytes |
| 558 | + pos_in_block * scale_dim |
| 559 | + tl.arange(0, scale_dim) |
| 560 | ) |
| 561 | scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) |
| 562 | tiles = tl.reshape(vals, (scale_dim, quant_block)) |
| 563 | dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) |
| 564 | |
| 565 | output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 |
| 566 | tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) |
| 567 | |
| 568 | |
| 569 | def dequantize_and_gather_k_cache_triton( |
| 570 | # [num_reqs, max_num_tokens, head_size] |
| 571 | out: torch.Tensor, |
| 572 | # [num_blocks, block_size, head_bytes] |
| 573 | k_cache: torch.Tensor, |
| 574 | # [num_reqs] |
| 575 | seq_lens: torch.Tensor, |
| 576 | # [num_reqs] |
| 577 | gather_lens: torch.Tensor | None, |
| 578 | # [num_reqs, max_blocks_per_seq] |
| 579 | block_table: torch.Tensor, |
| 580 | block_size: int, |
| 581 | offset: int, |
| 582 | use_fnuz: bool = False, |
| 583 | ) -> None: |
| 584 | num_reqs = seq_lens.shape[0] |
| 585 | NUM_WORKERS = 128 |
| 586 | if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN: |
| 587 | # Keep short gathers unchanged; cap total workers for multi-request chunks. |
| 588 | num_workers = min( |
| 589 | max(128, 2048 // max(1, num_reqs)), max(128, out.shape[1] - offset) |
| 590 | ) |
| 591 | _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, num_workers)]( |
| 592 | out, |
| 593 | out.stride(0), |
| 594 | out.stride(1), |
| 595 | k_cache, |
| 596 | seq_lens, |
| 597 | block_table, |
| 598 | offset, |
| 599 | gather_lens, |
| 600 | max_blocks_per_seq=block_table.shape[-1], |
| 601 | head_dim=512, |
| 602 | scale_dim=V41_NVFP4_NUM_SCALES, |
| 603 | quant_block=V41_NVFP4_QUANT_BLOCK, |
| 604 | cache_block_size=block_size, |
| 605 | block_stride=k_cache.stride(0), |
| 606 | ) |
| 607 | return |
| 608 | |
| 609 | if k_cache.shape[-1] == V41_BYTES_PER_TOKEN: |
| 610 | _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)]( |
| 611 | out, |
| 612 | out.stride(0), |
| 613 | out.stride(1), |
| 614 | k_cache, |
| 615 | seq_lens, |
| 616 | block_table, |
| 617 | offset, |
| 618 | gather_lens, |
| 619 | max_blocks_per_seq=block_table.shape[-1], |
| 620 | head_dim=512, |
| 621 | scale_dim=V41_NUM_SCALES, |
| 622 | quant_block=V41_QUANT_BLOCK, |
| 623 | cache_block_size=block_size, |
| 624 | block_stride=k_cache.stride(0), |
| 625 | use_fnuz=use_fnuz, |
| 626 | ) |
| 627 | return |
| 628 | |
| 629 | TOKEN_FP8_DIM = 448 |
| 630 | TOKEN_BF16_DIM = 64 |
| 631 | TOKEN_SCALE_DIM = 8 |
| 632 | QUANT_BLOCK_SIZE = 64 |
| 633 | FP8_MAX = 448.0 |
| 634 | TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 |
| 635 | |
| 636 | _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( |
| 637 | out, |
| 638 | out.stride(0), |
| 639 | out.stride(1), |
| 640 | k_cache, |
| 641 | seq_lens, |
| 642 | block_table, |
| 643 | offset, |
| 644 | gather_lens, |
| 645 | max_blocks_per_seq=block_table.shape[-1], |
| 646 | fp8_dim=TOKEN_FP8_DIM, |
| 647 | bf16_dim=TOKEN_BF16_DIM, |
| 648 | scale_dim=TOKEN_SCALE_DIM, |
| 649 | quant_block=QUANT_BLOCK_SIZE, |
| 650 | cache_block_size=block_size, |
| 651 | token_data_size=TOKEN_DATA_SIZE, |
| 652 | block_stride=k_cache.stride(0), |
| 653 | output_dim=512, |
| 654 | fp8_max=FP8_MAX, |
| 655 | n_quant_blocks=7, |
| 656 | use_fnuz=use_fnuz, |
| 657 | ) |
| 658 | |
| 659 | |
| 660 | def dequantize_and_gather_k_cache( |
| 661 | # [num_reqs, max_num_tokens, head_size] |
| 662 | out: torch.Tensor, |
| 663 | # [num_blocks, block_size, head_bytes] |
| 664 | k_cache: torch.Tensor, |
| 665 | # [num_reqs] |
| 666 | seq_lens: torch.Tensor, |
| 667 | # [num_reqs] |
| 668 | gather_lens: torch.Tensor | None, |
| 669 | # [num_reqs, max_blocks_per_seq] |
| 670 | block_table: torch.Tensor, |
| 671 | block_size: int, |
| 672 | offset: int, |
| 673 | use_fnuz: bool = False, |
| 674 | ) -> None: |
| 675 | """Dequantize and gather a paged DSv4 K cache. |
| 676 | |
| 677 | The record is read off ``k_cache.shape[-1]``; see the module header. Only |
| 678 | the fp8 records have a CuteDSL gather, so NVFP4 always takes the Triton |
| 679 | path. |
| 680 | |
| 681 | ``use_fnuz`` MUST match the encoder of the specific cache being read: |
| 682 | ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), |
| 683 | ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder |
| 684 | writes FNUZ on gfx942 and OCP on gfx950). |
| 685 | """ |
| 686 | if has_cutedsl() and k_cache.shape[-1] != V41_NVFP4_BYTES_PER_TOKEN: |
| 687 | # lazily import, otherwise some tests fail due to CUDA driver init failure. |
| 688 | from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( |
| 689 | _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, |
| 690 | ) |
| 691 | |
| 692 | mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN |
| 693 | _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( |
| 694 | out=out, |
| 695 | k_cache=k_cache, |
| 696 | seq_lens=seq_lens, |
| 697 | gather_lens=gather_lens, |
| 698 | block_table=block_table, |
| 699 | block_size=block_size, |
| 700 | offset=offset, |
| 701 | fp8_dim=512 if mxfp8 else 448, |
| 702 | group_size=V41_QUANT_BLOCK if mxfp8 else 64, |
| 703 | ) |
| 704 | return |
| 705 | |
| 706 | dequantize_and_gather_k_cache_triton( |
| 707 | out, |
| 708 | k_cache, |
| 709 | seq_lens, |
| 710 | gather_lens, |
| 711 | block_table, |
| 712 | block_size, |
| 713 | offset, |
| 714 | use_fnuz=use_fnuz, |
| 715 | ) |
| 716 | |
| 717 | |
| 718 | def compute_global_topk_indices_and_lens( |
| 719 | topk_indices: torch.Tensor, |
| 720 | token_to_req_indices: torch.Tensor, |
| 721 | block_table: torch.Tensor, |
| 722 | block_size: int, |
| 723 | is_valid_token: torch.Tensor, |
| 724 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 725 | """Map local topk indices to global KV cache slots and count valid entries. |
| 726 | |
| 727 | Fuses three operations into a single kernel: |
| 728 | 1. Block-table lookup (local index → global slot id) |
| 729 | 2. Valid-entry counting (topk_lens per token) |
| 730 | 3. Masking padding tokens to length 0 |
| 731 | """ |
| 732 | num_tokens = topk_indices.shape[0] |
| 733 | global_topk_indices = torch.empty_like(topk_indices) |
| 734 | topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) |
| 735 | _compute_global_topk_indices_and_lens_kernel[(num_tokens,)]( |
| 736 | global_topk_indices, |
| 737 | global_topk_indices.stride(0), |
| 738 | topk_lens, |
| 739 | topk_indices, |
| 740 | topk_indices.stride(0), |
| 741 | topk_indices.shape[-1], |
| 742 | token_to_req_indices, |
| 743 | block_table, |
| 744 | block_table.stride(0), |
| 745 | block_size, |
| 746 | is_valid_token, |
| 747 | TRITON_BLOCK_SIZE=1024, |
| 748 | ) |
| 749 | return global_topk_indices, topk_lens |
| 750 | |
| 751 | |
| 752 | @triton.jit |
| 753 | def _compute_global_topk_indices_and_lens_kernel( |
| 754 | global_topk_indices_ptr, |
| 755 | global_topk_indices_stride: tl.constexpr, |
| 756 | topk_lens_ptr, |
| 757 | topk_indices_ptr, |
| 758 | topk_indices_stride: tl.constexpr, |
| 759 | topk: tl.constexpr, |
| 760 | token_to_req_indices_ptr, |
| 761 | block_table_ptr, |
| 762 | block_table_stride: tl.constexpr, |
| 763 | block_size: tl.constexpr, |
| 764 | is_valid_token_ptr, |
| 765 | TRITON_BLOCK_SIZE: tl.constexpr, |
| 766 | ): |
| 767 | token_idx = tl.program_id(0) |
| 768 | is_valid_token = tl.load(is_valid_token_ptr + token_idx) |
| 769 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 770 | |
| 771 | count = tl.zeros((), dtype=tl.int32) |
| 772 | for i in range(0, topk, TRITON_BLOCK_SIZE): |
| 773 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 774 | mask = offset < topk |
| 775 | |
| 776 | local_idx = tl.load( |
| 777 | topk_indices_ptr + token_idx * topk_indices_stride + offset, |
| 778 | mask=mask, |
| 779 | other=-1, |
| 780 | ) |
| 781 | is_valid = local_idx >= 0 |
| 782 | |
| 783 | block_indices = local_idx // block_size |
| 784 | block_numbers = tl.load( |
| 785 | block_table_ptr + req_idx * block_table_stride + block_indices, |
| 786 | mask=mask & is_valid, |
| 787 | ) |
| 788 | block_offsets = local_idx % block_size |
| 789 | |
| 790 | slot_ids = block_numbers * block_size + block_offsets |
| 791 | slot_ids = tl.where(is_valid, slot_ids, -1) |
| 792 | tl.store( |
| 793 | global_topk_indices_ptr + token_idx * global_topk_indices_stride + offset, |
| 794 | slot_ids, |
| 795 | mask=mask, |
| 796 | ) |
| 797 | count += tl.sum(is_valid.to(tl.int32), axis=0) |
| 798 | |
| 799 | # Zero out length for padding tokens. |
| 800 | tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) |
| 801 | |
| 802 | |
| 803 | # FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see |
| 804 | # flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is |
| 805 | # 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both. |
| 806 | # The extra slots stay as -1 sentinels and `combined_lens` caps the valid |
| 807 | # range via `topk_length`, so padding is a no-op at kernel level. |
| 808 | _SPARSE_PREFILL_TOPK_ALIGNMENT = 128 |
| 809 | |
| 810 | |
| 811 | def combine_topk_swa_indices( |
| 812 | topk_indices: torch.Tensor, |
| 813 | query_start_loc: torch.Tensor, |
| 814 | seq_lens: torch.Tensor, |
| 815 | gather_lens: torch.Tensor, |
| 816 | window_size: int, |
| 817 | compress_ratio: int, |
| 818 | topk: int, |
| 819 | M: int, |
| 820 | N: int, |
| 821 | out: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 822 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 823 | num_tokens = topk_indices.shape[0] |
| 824 | # The width is fixed per model so the caller-provided workspace stays |
| 825 | # valid on every batch. |
| 826 | combined_topk = ( |
| 827 | (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) |
| 828 | // _SPARSE_PREFILL_TOPK_ALIGNMENT |
| 829 | * _SPARSE_PREFILL_TOPK_ALIGNMENT |
| 830 | ) |
| 831 | if out is None: |
| 832 | combined_indices = torch.full( |
| 833 | (num_tokens, combined_topk), |
| 834 | fill_value=-1, |
| 835 | dtype=torch.int32, |
| 836 | device=topk_indices.device, |
| 837 | ) |
| 838 | combined_lens = torch.empty( |
| 839 | num_tokens, dtype=torch.int32, device=topk_indices.device |
| 840 | ) |
| 841 | else: |
| 842 | combined_indices, combined_lens = out |
| 843 | |
| 844 | _COMBINE_TOPK_SWA_INDICES_KERNEL( |
| 845 | combined_indices, |
| 846 | combined_lens, |
| 847 | topk_indices, |
| 848 | query_start_loc, |
| 849 | seq_lens, |
| 850 | gather_lens, |
| 851 | M, |
| 852 | N, |
| 853 | TOP_K=topk, |
| 854 | COMPRESS_RATIO=compress_ratio, |
| 855 | WINDOW_SIZE=window_size, |
| 856 | ) |
| 857 | return combined_indices, combined_lens |
| 858 | |
| 859 | |
| 860 | _COMBINE_TOPK_SWA_NUM_WORKERS = 256 |
| 861 | |
| 862 | |
| 863 | # Representative pointer alignment variants for Triton pointer specialization. |
| 864 | _COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs( |
| 865 | dict( |
| 866 | topk_indices=True, |
| 867 | query_start_loc=True, |
| 868 | seq_lens=True, |
| 869 | gather_lens=True, |
| 870 | ), |
| 871 | dict( |
| 872 | topk_indices=True, |
| 873 | query_start_loc=False, |
| 874 | seq_lens=False, |
| 875 | gather_lens=True, |
| 876 | ), |
| 877 | dict( |
| 878 | topk_indices=False, |
| 879 | query_start_loc=False, |
| 880 | seq_lens=False, |
| 881 | gather_lens=False, |
| 882 | ), |
| 883 | ) |
| 884 | |
| 885 | |
| 886 | def _hf_config_int(vllm_config: Any, name: str, default: int) -> int: |
| 887 | model_config = getattr(vllm_config, "model_config", None) |
| 888 | hf_config = getattr(model_config, "hf_config", None) |
| 889 | return int(getattr(hf_config, name, default) or default) |
| 890 | |
| 891 | |
| 892 | def _combine_topk_swa_warmup_inputs(vllm_config: Any) -> _WarmupInputRows: |
| 893 | """One warmup row per v4.1 layer type (compress ratio) in the model. |
| 894 | |
| 895 | Every layer slices the shared ``index_topk``-wide indices buffer, so the |
| 896 | padded width is the same for all rows; SWA-only layers (ratio 0) pass |
| 897 | ``TOP_K=0``. |
| 898 | """ |
| 899 | model_config = getattr(vllm_config, "model_config", None) |
| 900 | hf_config = getattr(model_config, "hf_config", None) |
| 901 | compress_ratios = getattr(hf_config, "compress_ratios", None) or [0] |
| 902 | topk = _hf_config_int(vllm_config, "index_topk", 0) |
| 903 | return zip_inputs( |
| 904 | *( |
| 905 | dict(compress_ratio=ratio, topk=topk if ratio > 0 else 0, topk_width=topk) |
| 906 | for ratio in sorted({int(ratio) for ratio in compress_ratios}) |
| 907 | ) |
| 908 | ) |
| 909 | |
| 910 | |
| 911 | def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int: |
| 912 | scheduler_config = getattr(vllm_config, "scheduler_config", None) |
| 913 | return int(getattr(scheduler_config, name, default) or default) |
| 914 | |
| 915 | |
| 916 | class CombineTopkSwaIndicesKernel( |
| 917 | VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"] |
| 918 | ): |
| 919 | @dataclass(frozen=True) |
| 920 | class CompileKey: |
| 921 | TOP_K: int |
| 922 | COMPRESS_RATIO: int |
| 923 | WINDOW_SIZE: int |
| 924 | PADDED_TOP_K: int |
| 925 | input_variant: TritonPointerInputVariant |
| 926 | |
| 927 | @staticmethod |
| 928 | @triton.jit( |
| 929 | do_not_specialize=[ |
| 930 | "combined_indices_stride", |
| 931 | "topk_indices_stride", |
| 932 | "M", |
| 933 | "N", |
| 934 | ] |
| 935 | ) |
| 936 | def kernel( |
| 937 | combined_indices_ptr, |
| 938 | combined_indices_stride, |
| 939 | combined_lens_ptr, |
| 940 | topk_indices_ptr, |
| 941 | topk_indices_stride, |
| 942 | query_start_loc_ptr, |
| 943 | seq_lens_ptr, |
| 944 | gather_lens_ptr, |
| 945 | M, |
| 946 | N, |
| 947 | TOP_K: tl.constexpr, |
| 948 | COMPRESS_RATIO: tl.constexpr, |
| 949 | WINDOW_SIZE: tl.constexpr, |
| 950 | PADDED_TOP_K: tl.constexpr, |
| 951 | ): |
| 952 | batch_idx = tl.program_id(0) |
| 953 | worker_id = tl.program_id(1) |
| 954 | num_workers = tl.num_programs(1) |
| 955 | |
| 956 | # query_start_loc is a global tensor; rebase to chunk-local offsets |
| 957 | # by subtracting the chunk's starting value. |
| 958 | base = tl.load(query_start_loc_ptr) |
| 959 | query_start = tl.load(query_start_loc_ptr + batch_idx) - base |
| 960 | query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base |
| 961 | query_len = query_end - query_start |
| 962 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 963 | gather_len = tl.load(gather_lens_ptr + batch_idx) |
| 964 | start_pos = seq_len - query_len |
| 965 | # The SWA portion of the gathered buffer starts from position |
| 966 | # (seq_len - gather_len), not position 0. We need this offset |
| 967 | # to correctly index into the gathered buffer. |
| 968 | gather_start = seq_len - gather_len |
| 969 | |
| 970 | for token_idx in range(query_start + worker_id, query_end, num_workers): |
| 971 | # topk_len is fully determined by the query token's absolute position: |
| 972 | # the indexer emits min((pos + 1) // compress_ratio, topk_tokens) |
| 973 | # valid entries. Caller passes TOP_K=0 for SWA-only layers to zero |
| 974 | # this out. |
| 975 | token_idx_in_query = token_idx - query_start |
| 976 | pos = start_pos + token_idx_in_query |
| 977 | if COMPRESS_RATIO > 0: |
| 978 | topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) |
| 979 | else: |
| 980 | # SWA-only layers pass TOP_K=0; skip the division entirely |
| 981 | # (integer div by constexpr 0 is UB and yields garbage, which |
| 982 | # turns the store index below negative). |
| 983 | topk_len = 0 |
| 984 | swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0) |
| 985 | # gather_len already excludes context below the request's replay |
| 986 | # start (SWA bounded replay), so the window cannot start before |
| 987 | # the gathered buffer does. |
| 988 | swa_start = tl.maximum(swa_start, gather_start) |
| 989 | swa_len = pos - swa_start + 1 |
| 990 | |
| 991 | offset = tl.arange(0, PADDED_TOP_K) |
| 992 | mask = offset < topk_len |
| 993 | topk_indices = tl.load( |
| 994 | topk_indices_ptr + token_idx * topk_indices_stride + offset, |
| 995 | mask=mask, |
| 996 | ) |
| 997 | tl.store( |
| 998 | combined_indices_ptr + token_idx * combined_indices_stride + offset, |
| 999 | topk_indices + M * batch_idx, |
| 1000 | mask=mask, |
| 1001 | ) |
| 1002 | # Index into gathered buffer: N + (position - gather_start) |
| 1003 | # For positions [swa_start, pos], the buffer indices are: |
| 1004 | # [N + swa_start - gather_start, N + pos - gather_start] |
| 1005 | swa_offset = tl.arange(0, WINDOW_SIZE) |
| 1006 | tl.store( |
| 1007 | combined_indices_ptr |
| 1008 | + token_idx * combined_indices_stride |
| 1009 | + topk_len |
| 1010 | + swa_offset, |
| 1011 | M * batch_idx + N + swa_start + swa_offset - gather_start, |
| 1012 | mask=swa_offset < swa_len, |
| 1013 | ) |
| 1014 | |
| 1015 | combined_len = topk_len + swa_len |
| 1016 | tl.store(combined_lens_ptr + token_idx, combined_len) |
| 1017 | |
| 1018 | def dispatch( # type: ignore[override] |
| 1019 | self, |
| 1020 | *, |
| 1021 | topk_width: int, |
| 1022 | topk_indices: bool, |
| 1023 | query_start_loc: bool, |
| 1024 | seq_lens: bool, |
| 1025 | gather_lens: bool, |
| 1026 | topk: int, |
| 1027 | compress_ratio: int, |
| 1028 | WINDOW_SIZE: int, |
| 1029 | ) -> CompileKey: |
| 1030 | padded_topk = next_power_of_2(topk_width) |
| 1031 | input_variant = TritonPointerInputVariant.from_alignment( |
| 1032 | topk_indices=topk_indices, |
| 1033 | query_start_loc=query_start_loc, |
| 1034 | seq_lens=seq_lens, |
| 1035 | gather_lens=gather_lens, |
| 1036 | ) |
| 1037 | return self.CompileKey( |
| 1038 | TOP_K=topk, |
| 1039 | COMPRESS_RATIO=compress_ratio, |
| 1040 | WINDOW_SIZE=WINDOW_SIZE, |
| 1041 | PADDED_TOP_K=padded_topk, |
| 1042 | input_variant=input_variant, |
| 1043 | ) |
| 1044 | |
| 1045 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 1046 | if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0: |
| 1047 | return [] |
| 1048 | |
| 1049 | window_size = _hf_config_int(vllm_config, "sliding_window", 128) |
| 1050 | return self._trace_dispatch(self.dispatch)( |
| 1051 | _combine_topk_swa_warmup_inputs(vllm_config), |
| 1052 | _COMBINE_TOPK_SWA_POINTER_INPUTS, |
| 1053 | WINDOW_SIZE=window_size, |
| 1054 | ) |
| 1055 | |
| 1056 | def compile(self, compile_key: CompileKey) -> None: |
| 1057 | warmup = getattr(self.kernel, "warmup", None) |
| 1058 | assert warmup is not None |
| 1059 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 1060 | input_variant = compile_key.input_variant |
| 1061 | warmup( |
| 1062 | int32_ptr, |
| 1063 | 1, # do not specialize combined_indices_stride |
| 1064 | int32_ptr, |
| 1065 | input_variant.pointer("topk_indices", torch.int32), |
| 1066 | 1, # do not specialize topk_indices_stride |
| 1067 | input_variant.pointer("query_start_loc", torch.int32), |
| 1068 | input_variant.pointer("seq_lens", torch.int32), |
| 1069 | input_variant.pointer("gather_lens", torch.int32), |
| 1070 | 1, # do not specialize M |
| 1071 | 1, # do not specialize N |
| 1072 | TOP_K=compile_key.TOP_K, |
| 1073 | COMPRESS_RATIO=compile_key.COMPRESS_RATIO, |
| 1074 | WINDOW_SIZE=compile_key.WINDOW_SIZE, |
| 1075 | PADDED_TOP_K=compile_key.PADDED_TOP_K, |
| 1076 | grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS), |
| 1077 | ) |
| 1078 | |
| 1079 | def __call__( |
| 1080 | self, |
| 1081 | combined_indices: torch.Tensor, |
| 1082 | combined_lens: torch.Tensor, |
| 1083 | topk_indices: torch.Tensor, |
| 1084 | query_start_loc: torch.Tensor, |
| 1085 | seq_lens: torch.Tensor, |
| 1086 | gather_lens: torch.Tensor, |
| 1087 | M: int, |
| 1088 | N: int, |
| 1089 | *, |
| 1090 | TOP_K: int, |
| 1091 | COMPRESS_RATIO: int, |
| 1092 | WINDOW_SIZE: int, |
| 1093 | ) -> None: |
| 1094 | num_reqs = seq_lens.shape[0] |
| 1095 | self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)]( |
| 1096 | combined_indices, |
| 1097 | combined_indices.stride(0), |
| 1098 | combined_lens, |
| 1099 | topk_indices, |
| 1100 | topk_indices.stride(0), |
| 1101 | query_start_loc, |
| 1102 | seq_lens, |
| 1103 | gather_lens, |
| 1104 | M, |
| 1105 | N, |
| 1106 | TOP_K=TOP_K, |
| 1107 | COMPRESS_RATIO=COMPRESS_RATIO, |
| 1108 | WINDOW_SIZE=WINDOW_SIZE, |
| 1109 | PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]), |
| 1110 | ) |
| 1111 | |
| 1112 | |
| 1113 | _COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel() |
| 1114 | |
| 1115 | |
| 1116 | def build_flashinfer_mixed_sparse_indices( |
| 1117 | decode_swa_indices: torch.Tensor, |
| 1118 | decode_compressed_indices: torch.Tensor | None, |
| 1119 | decode_compressed_topk_lens: torch.Tensor | None, |
| 1120 | prefill_topk_indices: torch.Tensor, |
| 1121 | query_start_loc: torch.Tensor, |
| 1122 | seq_lens: torch.Tensor, |
| 1123 | token_to_req_indices: torch.Tensor, |
| 1124 | swa_block_table: torch.Tensor, |
| 1125 | swa_block_size: int, |
| 1126 | compressed_block_table: torch.Tensor | None, |
| 1127 | compressed_block_size: int, |
| 1128 | window_size: int, |
| 1129 | compress_ratio: int, |
| 1130 | topk: int, |
| 1131 | decode_compressed_indices_are_local: bool = False, |
| 1132 | decode_is_valid_token: torch.Tensor | None = None, |
| 1133 | swa_block_span: int | None = None, |
| 1134 | compressed_block_span: int | None = None, |
| 1135 | *, |
| 1136 | replay_start: torch.Tensor, |
| 1137 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1138 | """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. |
| 1139 | |
| 1140 | Produces ``sparse_indices`` of shape ``[num_tokens, swa_total_width + |
| 1141 | padded_topk]`` (the first ``swa_total_width`` columns are SWA slot ids, the |
| 1142 | rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length |
| 1143 | per token). Decode tokens read precomputed SWA/compressed indices; prefill |
| 1144 | tokens derive their SWA window from the position and translate local |
| 1145 | compressed indices to global slots via the block tables. |
| 1146 | |
| 1147 | ``replay_start`` ([num_reqs], SWA bounded replay) lower-bounds every |
| 1148 | prefill token's window: positions below it hold no window KV. |
| 1149 | """ |
| 1150 | assert decode_swa_indices.dtype == torch.int32 |
| 1151 | assert decode_swa_indices.dim() == 2 |
| 1152 | swa_total_width = decode_swa_indices.shape[-1] |
| 1153 | assert swa_total_width >= window_size |
| 1154 | if decode_compressed_topk_lens is not None: |
| 1155 | assert decode_compressed_topk_lens.dtype == torch.int32 |
| 1156 | assert prefill_topk_indices.dtype == torch.int32 |
| 1157 | assert prefill_topk_indices.dim() == 2 |
| 1158 | assert query_start_loc.dtype == torch.int32 |
| 1159 | assert seq_lens.dtype == torch.int32 |
| 1160 | assert token_to_req_indices.dtype == torch.int32 |
| 1161 | assert swa_block_table.dtype == torch.int32 |
| 1162 | |
| 1163 | num_decode_tokens = decode_swa_indices.shape[0] |
| 1164 | num_prefill_tokens = prefill_topk_indices.shape[0] |
| 1165 | num_tokens = num_decode_tokens + num_prefill_tokens |
| 1166 | assert token_to_req_indices.shape[0] >= num_tokens |
| 1167 | if decode_compressed_topk_lens is not None: |
| 1168 | assert decode_compressed_topk_lens.shape[0] >= num_decode_tokens |
| 1169 | |
| 1170 | decode_compressed_topk = 0 |
| 1171 | if decode_compressed_indices is None: |
| 1172 | decode_compressed_indices = prefill_topk_indices |
| 1173 | else: |
| 1174 | assert decode_compressed_indices.dtype == torch.int32 |
| 1175 | assert decode_compressed_indices.dim() == 2 |
| 1176 | assert decode_compressed_indices.shape[0] == num_decode_tokens |
| 1177 | decode_compressed_topk = decode_compressed_indices.shape[-1] |
| 1178 | if decode_compressed_topk > 0 and decode_compressed_indices_are_local: |
| 1179 | assert decode_is_valid_token is not None |
| 1180 | assert decode_is_valid_token.dtype == torch.bool |
| 1181 | assert decode_is_valid_token.shape[0] >= num_decode_tokens |
| 1182 | else: |
| 1183 | decode_is_valid_token = token_to_req_indices |
| 1184 | |
| 1185 | if compressed_block_table is None: |
| 1186 | compressed_block_table = swa_block_table |
| 1187 | assert compressed_block_table.dtype == torch.int32 |
| 1188 | has_decode_compressed_lens = decode_compressed_topk_lens is not None |
| 1189 | if decode_compressed_topk_lens is None: |
| 1190 | decode_compressed_topk_lens = token_to_req_indices |
| 1191 | |
| 1192 | # The FlashInfer TRTLLM-gen sparse-MLA kernels require every per-token topk |
| 1193 | # index row to start on a 16-byte boundary: the kernel loads the compressed |
| 1194 | # indices with 128-bit (16-byte) vectorized loads, so a misaligned row would |
| 1195 | # fault or read across rows. 16 bytes = 4 int32 indices, so round the topk |
| 1196 | # width (and hence the row stride, since the SWA columns are fixed-width) up |
| 1197 | # to a multiple of 4. The extra columns are filled with -1 (invalid) and bounded |
| 1198 | # by ``sparse_topk_lens``, so padding never changes the attention result. |
| 1199 | padded_topk = max(topk, decode_compressed_topk) |
| 1200 | padded_topk = (padded_topk + 3) // 4 * 4 |
| 1201 | sparse_indices = torch.empty( |
| 1202 | (num_tokens, swa_total_width + padded_topk), |
| 1203 | dtype=torch.int32, |
| 1204 | device=decode_swa_indices.device, |
| 1205 | ) |
| 1206 | sparse_topk_lens = torch.empty( |
| 1207 | num_tokens, dtype=torch.int32, device=decode_swa_indices.device |
| 1208 | ) |
| 1209 | if num_tokens == 0: |
| 1210 | return sparse_indices, sparse_topk_lens |
| 1211 | |
| 1212 | window_block_size = triton.next_power_of_2(max(swa_total_width, 1)) |
| 1213 | topk_block_size = triton.next_power_of_2(max(padded_topk, 1)) |
| 1214 | max_block_size = max(window_block_size, topk_block_size) |
| 1215 | num_warps = 4 if max_block_size >= 256 else 1 |
| 1216 | |
| 1217 | # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV. |
| 1218 | swa_span = swa_block_size if swa_block_span is None else swa_block_span |
| 1219 | compressed_span = ( |
| 1220 | compressed_block_size |
| 1221 | if compressed_block_span is None |
| 1222 | else compressed_block_span |
| 1223 | ) |
| 1224 | _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)]( |
| 1225 | sparse_indices, |
| 1226 | sparse_indices.stride(0), |
| 1227 | sparse_topk_lens, |
| 1228 | decode_swa_indices, |
| 1229 | decode_swa_indices.stride(0), |
| 1230 | decode_compressed_indices, |
| 1231 | decode_compressed_indices.stride(0), |
| 1232 | decode_compressed_topk_lens, |
| 1233 | decode_is_valid_token, |
| 1234 | prefill_topk_indices, |
| 1235 | prefill_topk_indices.stride(0), |
| 1236 | query_start_loc, |
| 1237 | seq_lens, |
| 1238 | token_to_req_indices, |
| 1239 | replay_start, |
| 1240 | swa_block_table, |
| 1241 | swa_block_table.stride(0), |
| 1242 | swa_block_size, |
| 1243 | swa_span, |
| 1244 | compressed_block_table, |
| 1245 | compressed_block_table.stride(0), |
| 1246 | compressed_block_size, |
| 1247 | compressed_span, |
| 1248 | NUM_DECODE_TOKENS=num_decode_tokens, |
| 1249 | WINDOW_SIZE=window_size, |
| 1250 | SWA_INDEX_WIDTH=swa_total_width, |
| 1251 | COMPRESS_RATIO=compress_ratio, |
| 1252 | TOP_K=topk, |
| 1253 | PADDED_TOP_K=padded_topk, |
| 1254 | PREFILL_TOPK_STRIDE=prefill_topk_indices.shape[-1], |
| 1255 | DECODE_COMPRESSED_TOPK=decode_compressed_topk, |
| 1256 | DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local, |
| 1257 | HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens, |
| 1258 | WINDOW_BLOCK_SIZE=window_block_size, |
| 1259 | TOPK_BLOCK_SIZE=topk_block_size, |
| 1260 | num_warps=num_warps, |
| 1261 | ) |
| 1262 | return sparse_indices, sparse_topk_lens |
| 1263 | |
| 1264 | |
| 1265 | @triton.jit |
| 1266 | def _remap_flashinfer_index(values, block_size, block_span): |
| 1267 | # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so |
| 1268 | # packed pages (#44577) need block*block_size+off -> block*block_span+off. |
| 1269 | # The launcher declares these indices with sparse_indices_are_storage_offsets=True. |
| 1270 | is_valid = values >= 0 |
| 1271 | safe_values = tl.where(is_valid, values, 0) |
| 1272 | values = (safe_values // block_size) * block_span |
| 1273 | values += safe_values % block_size |
| 1274 | return tl.where(is_valid, values, -1) |
| 1275 | |
| 1276 | |
| 1277 | @triton.jit( |
| 1278 | do_not_specialize=[ |
| 1279 | "sparse_indices_stride", |
| 1280 | "decode_swa_stride", |
| 1281 | "decode_compressed_stride", |
| 1282 | "prefill_topk_stride", |
| 1283 | "swa_block_table_stride", |
| 1284 | "swa_block_size", |
| 1285 | "swa_block_span", |
| 1286 | "compressed_block_table_stride", |
| 1287 | "compressed_block_size", |
| 1288 | "compressed_block_span", |
| 1289 | "NUM_DECODE_TOKENS", |
| 1290 | "PREFILL_TOPK_STRIDE", |
| 1291 | ] |
| 1292 | ) |
| 1293 | def _build_flashinfer_mixed_sparse_indices_kernel( |
| 1294 | sparse_indices_ptr, |
| 1295 | sparse_indices_stride, |
| 1296 | sparse_topk_lens_ptr, |
| 1297 | decode_swa_indices_ptr, |
| 1298 | decode_swa_stride, |
| 1299 | decode_compressed_indices_ptr, |
| 1300 | decode_compressed_stride, |
| 1301 | decode_compressed_topk_lens_ptr, |
| 1302 | decode_is_valid_token_ptr, |
| 1303 | prefill_topk_indices_ptr, |
| 1304 | prefill_topk_stride, |
| 1305 | query_start_loc_ptr, |
| 1306 | seq_lens_ptr, |
| 1307 | token_to_req_indices_ptr, |
| 1308 | replay_start_ptr, |
| 1309 | swa_block_table_ptr, |
| 1310 | swa_block_table_stride, |
| 1311 | swa_block_size, |
| 1312 | swa_block_span, |
| 1313 | compressed_block_table_ptr, |
| 1314 | compressed_block_table_stride, |
| 1315 | compressed_block_size, |
| 1316 | compressed_block_span, |
| 1317 | NUM_DECODE_TOKENS, |
| 1318 | WINDOW_SIZE: tl.constexpr, |
| 1319 | SWA_INDEX_WIDTH: tl.constexpr, |
| 1320 | COMPRESS_RATIO: tl.constexpr, |
| 1321 | TOP_K: tl.constexpr, |
| 1322 | PADDED_TOP_K: tl.constexpr, |
| 1323 | PREFILL_TOPK_STRIDE, |
| 1324 | DECODE_COMPRESSED_TOPK: tl.constexpr, |
| 1325 | DECODE_COMPRESSED_INDICES_ARE_LOCAL: tl.constexpr, |
| 1326 | HAS_DECODE_COMPRESSED_LENS: tl.constexpr, |
| 1327 | WINDOW_BLOCK_SIZE: tl.constexpr, |
| 1328 | TOPK_BLOCK_SIZE: tl.constexpr, |
| 1329 | ): |
| 1330 | token_idx = tl.program_id(0) |
| 1331 | |
| 1332 | if token_idx < NUM_DECODE_TOKENS: |
| 1333 | for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): |
| 1334 | offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) |
| 1335 | mask = offset < SWA_INDEX_WIDTH |
| 1336 | values = tl.load( |
| 1337 | decode_swa_indices_ptr + token_idx * decode_swa_stride + offset, |
| 1338 | mask=mask, |
| 1339 | other=-1, |
| 1340 | ) |
| 1341 | values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) |
| 1342 | tl.store( |
| 1343 | sparse_indices_ptr + token_idx * sparse_indices_stride + offset, |
| 1344 | values, |
| 1345 | mask=mask, |
| 1346 | ) |
| 1347 | |
| 1348 | compressed_len = tl.zeros((), dtype=tl.int32) |
| 1349 | for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE): |
| 1350 | offset = i + tl.arange(0, TOPK_BLOCK_SIZE) |
| 1351 | mask = offset < PADDED_TOP_K |
| 1352 | values = tl.load( |
| 1353 | decode_compressed_indices_ptr |
| 1354 | + token_idx * decode_compressed_stride |
| 1355 | + offset, |
| 1356 | mask=offset < DECODE_COMPRESSED_TOPK, |
| 1357 | other=-1, |
| 1358 | ) |
| 1359 | if DECODE_COMPRESSED_INDICES_ARE_LOCAL: |
| 1360 | token_valid = tl.load(decode_is_valid_token_ptr + token_idx) |
| 1361 | is_valid = values >= 0 |
| 1362 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 1363 | block_indices = values // compressed_block_size |
| 1364 | block_numbers = tl.load( |
| 1365 | compressed_block_table_ptr |
| 1366 | + req_idx * compressed_block_table_stride |
| 1367 | + block_indices, |
| 1368 | mask=mask & is_valid, |
| 1369 | other=-1, |
| 1370 | ) |
| 1371 | block_offsets = values % compressed_block_size |
| 1372 | values = block_numbers * compressed_block_size + block_offsets |
| 1373 | values = tl.where(is_valid, values, -1) |
| 1374 | compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0) |
| 1375 | values = _remap_flashinfer_index( |
| 1376 | values, compressed_block_size, compressed_block_span |
| 1377 | ) |
| 1378 | tl.store( |
| 1379 | sparse_indices_ptr |
| 1380 | + token_idx * sparse_indices_stride |
| 1381 | + SWA_INDEX_WIDTH |
| 1382 | + offset, |
| 1383 | values, |
| 1384 | mask=mask, |
| 1385 | ) |
| 1386 | |
| 1387 | if DECODE_COMPRESSED_TOPK == 0: |
| 1388 | compressed_len = tl.zeros((), dtype=tl.int32) |
| 1389 | elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL: |
| 1390 | if HAS_DECODE_COMPRESSED_LENS: |
| 1391 | compressed_len = tl.load(decode_compressed_topk_lens_ptr + token_idx) |
| 1392 | else: |
| 1393 | compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) |
| 1394 | |
| 1395 | tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len) |
| 1396 | return |
| 1397 | |
| 1398 | prefill_idx = token_idx - NUM_DECODE_TOKENS |
| 1399 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 1400 | query_start = tl.load(query_start_loc_ptr + req_idx) |
| 1401 | query_end = tl.load(query_start_loc_ptr + req_idx + 1) |
| 1402 | query_len = query_end - query_start |
| 1403 | seq_len = tl.load(seq_lens_ptr + req_idx) |
| 1404 | start_pos = seq_len - query_len |
| 1405 | token_idx_in_query = token_idx - query_start |
| 1406 | pos = start_pos + token_idx_in_query |
| 1407 | swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0) |
| 1408 | # SWA bounded replay: no window KV exists below the request's replay start. |
| 1409 | swa_start_pos = tl.maximum(swa_start_pos, tl.load(replay_start_ptr + req_idx)) |
| 1410 | swa_len = pos - swa_start_pos + 1 |
| 1411 | if COMPRESS_RATIO > 0: |
| 1412 | topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) |
| 1413 | else: |
| 1414 | # SWA-only layers (ratio 0): no indexer topk; div by constexpr 0 is UB. |
| 1415 | topk_len = 0 |
| 1416 | |
| 1417 | for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): |
| 1418 | offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) |
| 1419 | mask = offset < SWA_INDEX_WIDTH |
| 1420 | pos_offset = swa_start_pos + offset |
| 1421 | block_indices = pos_offset // swa_block_size |
| 1422 | block_numbers = tl.load( |
| 1423 | swa_block_table_ptr + req_idx * swa_block_table_stride + block_indices, |
| 1424 | mask=mask & (offset < swa_len), |
| 1425 | other=-1, |
| 1426 | ) |
| 1427 | block_offsets = pos_offset % swa_block_size |
| 1428 | slot_ids = block_numbers * swa_block_size + block_offsets |
| 1429 | slot_ids = tl.where(offset < swa_len, slot_ids, -1) |
| 1430 | slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span) |
| 1431 | tl.store( |
| 1432 | sparse_indices_ptr + token_idx * sparse_indices_stride + offset, |
| 1433 | slot_ids, |
| 1434 | mask=mask, |
| 1435 | ) |
| 1436 | |
| 1437 | for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE): |
| 1438 | offset = i + tl.arange(0, TOPK_BLOCK_SIZE) |
| 1439 | mask = offset < PADDED_TOP_K |
| 1440 | local_idx = tl.load( |
| 1441 | prefill_topk_indices_ptr + prefill_idx * prefill_topk_stride + offset, |
| 1442 | mask=(offset < PREFILL_TOPK_STRIDE) & (offset < topk_len), |
| 1443 | other=-1, |
| 1444 | ) |
| 1445 | is_valid = local_idx >= 0 |
| 1446 | block_indices = local_idx // compressed_block_size |
| 1447 | block_numbers = tl.load( |
| 1448 | compressed_block_table_ptr |
| 1449 | + req_idx * compressed_block_table_stride |
| 1450 | + block_indices, |
| 1451 | mask=mask & is_valid, |
| 1452 | other=-1, |
| 1453 | ) |
| 1454 | block_offsets = local_idx % compressed_block_size |
| 1455 | slot_ids = block_numbers * compressed_block_size + block_offsets |
| 1456 | slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1) |
| 1457 | slot_ids = _remap_flashinfer_index( |
| 1458 | slot_ids, compressed_block_size, compressed_block_span |
| 1459 | ) |
| 1460 | tl.store( |
| 1461 | sparse_indices_ptr |
| 1462 | + token_idx * sparse_indices_stride |
| 1463 | + SWA_INDEX_WIDTH |
| 1464 | + offset, |
| 1465 | slot_ids, |
| 1466 | mask=mask, |
| 1467 | ) |
| 1468 | |
| 1469 | tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len) |