| 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 | WarmupIntRange, |
| 26 | kernel_launcher, |
| 27 | zip_inputs, |
| 28 | ) |
| 29 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 30 | LaunchSpec, |
| 31 | TritonPointerInputVariant, |
| 32 | TritonWarmupTensor, |
| 33 | VllmTritonJitKernel, |
| 34 | triton_scalar_specialization_rep, |
| 35 | ) |
| 36 | from vllm.platforms import current_platform |
| 37 | from vllm.triton_utils import tl, triton |
| 38 | from vllm.utils.import_utils import has_cutedsl |
| 39 | from vllm.utils.math_utils import next_power_of_2 |
| 40 | |
| 41 | |
| 42 | @triton.jit |
| 43 | def quantize_and_insert_k_kernel( |
| 44 | # Input tensors |
| 45 | k_ptr, # [num_tokens, 512] bf16 |
| 46 | slot_mapping_ptr, # [num_tokens] int64 |
| 47 | # Output tensor |
| 48 | k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) |
| 49 | # Dimensions |
| 50 | num_tokens, |
| 51 | input_dim: tl.constexpr, # 512 |
| 52 | fp8_dim: tl.constexpr, # 448 |
| 53 | bf16_dim: tl.constexpr, # 64 |
| 54 | scale_dim: tl.constexpr, # 8 |
| 55 | quant_block: tl.constexpr, # 64 (quantization block size) |
| 56 | cache_block_size: tl.constexpr, # 64 (paged cache block size) |
| 57 | token_data_size: tl.constexpr, # 576 bytes per token data |
| 58 | block_stride: tl.constexpr, # total bytes per block (padded) |
| 59 | fp8_max: tl.constexpr, |
| 60 | n_quant_blocks: tl.constexpr, # 8 (7 real + 1 padding) |
| 61 | use_fnuz: tl.constexpr = False, |
| 62 | ): |
| 63 | """Quantize K tensor and insert into paged K cache. |
| 64 | |
| 65 | K Cache block layout (block_size=64 tokens): |
| 66 | - [0, 64*576): Token data, each token has 448 fp8 + 128 bf16 |
| 67 | - [64*576, 64*576 + 64*8): Scales, each token has 8 uint8 scales |
| 68 | - [64*576 + 64*8, block_stride): Padding |
| 69 | |
| 70 | One program per token. |
| 71 | |
| 72 | ``use_fnuz=True`` selects FNUZ (``tl.float8e4b8``); default OCP |
| 73 | (``tl.float8e4nv``) matches every production caller. |
| 74 | """ |
| 75 | pid = tl.program_id(0) |
| 76 | |
| 77 | if pid >= num_tokens: |
| 78 | return |
| 79 | |
| 80 | # Get slot mapping |
| 81 | slot_idx = tl.load(slot_mapping_ptr + pid) |
| 82 | if slot_idx == -1: |
| 83 | return |
| 84 | |
| 85 | block_idx = slot_idx // cache_block_size |
| 86 | pos_in_block = slot_idx % cache_block_size |
| 87 | |
| 88 | # Input pointer for this token |
| 89 | input_row_ptr = k_ptr + pid * input_dim |
| 90 | |
| 91 | # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks |
| 92 | # (e.g. >= 57K at block_stride ~37K). Matches gather path below. |
| 93 | cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride |
| 94 | |
| 95 | # Token data pointer: token data is stored contiguously at start of block |
| 96 | # Each token's data is at offset pos_in_block * token_data_size |
| 97 | token_data_ptr = cache_block_ptr + pos_in_block * token_data_size |
| 98 | |
| 99 | # Scale pointer: scales are stored after ALL token data in the block |
| 100 | # Scale for this token is at offset (64 * 576) + pos_in_block * 8 |
| 101 | token_scale_ptr = ( |
| 102 | cache_block_ptr + cache_block_size * token_data_size + pos_in_block * scale_dim |
| 103 | ) |
| 104 | |
| 105 | # Token data layout: [0:448] fp8, [448:576] bf16 |
| 106 | token_fp8_ptr = token_data_ptr |
| 107 | token_bf16_ptr = token_data_ptr + fp8_dim |
| 108 | |
| 109 | # ========== Quantize and store FP8 portion (first 448 elements) ========== |
| 110 | # Using UE8M0 quantization strategy (scale is power of 2, stored as uint8 exponent) |
| 111 | for qblock_idx in tl.static_range(n_quant_blocks): |
| 112 | qblock_start = qblock_idx * quant_block |
| 113 | |
| 114 | if qblock_start < fp8_dim: |
| 115 | offsets = qblock_start + tl.arange(0, quant_block) |
| 116 | mask = offsets < fp8_dim |
| 117 | |
| 118 | # Load bf16 input |
| 119 | x = tl.load(input_row_ptr + offsets, mask=mask, other=0.0) |
| 120 | |
| 121 | # Compute absmax scale (same as CUDA kernel) |
| 122 | abs_x = tl.abs(x) |
| 123 | block_max = tl.max(abs_x, axis=0) |
| 124 | block_max = tl.maximum(block_max, 1e-4) # Match CUDA: fmaxf(amax, 1e-4) |
| 125 | |
| 126 | # UE8M0: Round scale UP to next power of 2 |
| 127 | # scale = 2^ceil(log2(block_max / fp8_max)) |
| 128 | raw_scale = block_max / fp8_max |
| 129 | log_scale = tl.log2(raw_scale) |
| 130 | exponent = tl.ceil(log_scale) # Round UP to next integer exponent |
| 131 | scale = tl.exp2(exponent) # scale = 2^exponent (power of 2) |
| 132 | |
| 133 | # Quantize to fp8: fp8_value = bf16_value / scale |
| 134 | x_scaled = x / scale |
| 135 | x_clamped = tl.clamp(x_scaled, -fp8_max, fp8_max) |
| 136 | |
| 137 | # Convert to fp8 (FNUZ on gfx942, OCP elsewhere), then bitcast to uint8. |
| 138 | if use_fnuz: |
| 139 | x_fp8 = x_clamped.to(tl.float8e4b8) |
| 140 | else: |
| 141 | x_fp8 = x_clamped.to(tl.float8e4nv) |
| 142 | x_uint8 = x_fp8.to(tl.uint8, bitcast=True) |
| 143 | |
| 144 | # Store as uint8 (1 byte each) |
| 145 | tl.store(token_fp8_ptr + offsets, x_uint8, mask=mask) |
| 146 | |
| 147 | # UE8M0 scale encoding: stored_value = exponent + 127 (bias) |
| 148 | # During dequant: scale = 2^(stored_value - 127) |
| 149 | encoded_scale = exponent + 127.0 |
| 150 | encoded_scale = tl.maximum(tl.minimum(encoded_scale, 255.0), 0.0) |
| 151 | tl.store(token_scale_ptr + qblock_idx, encoded_scale.to(tl.uint8)) |
| 152 | |
| 153 | # Padding scale at index 7 |
| 154 | tl.store(token_scale_ptr + 7, tl.zeros((), dtype=tl.uint8)) |
| 155 | |
| 156 | # ========== Store BF16 portion (last 64 elements, no quantization) ========== |
| 157 | bf16_input_offset = fp8_dim |
| 158 | |
| 159 | # Process bf16 in chunks of 16 |
| 160 | bf16_out_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) |
| 161 | for i in tl.static_range(bf16_dim // 16): |
| 162 | chunk_offsets = i * 16 + tl.arange(0, 16) |
| 163 | bf16_vals = tl.load(input_row_ptr + bf16_input_offset + chunk_offsets) |
| 164 | tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) |
| 165 | |
| 166 | |
| 167 | def quantize_and_insert_k_cache( |
| 168 | k: torch.Tensor, # [num_tokens, 512] bf16 |
| 169 | k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 |
| 170 | slot_mapping: torch.Tensor, # [num_tokens] int64 |
| 171 | block_size: int = 64, |
| 172 | is_ue8m0: bool = True, |
| 173 | use_fnuz: bool = False, |
| 174 | ): |
| 175 | """Quantize K tensor and insert into paged K cache. |
| 176 | |
| 177 | K Cache block layout (block_size=64 tokens): |
| 178 | - First 64 * 576 = 36864 bytes: Token data |
| 179 | - Each token: 448 bytes (fp8) + 128 bytes (bf16) |
| 180 | - Next 64 * 8 = 512 bytes: Scales |
| 181 | - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) |
| 182 | - Padded to multiple of 576 |
| 183 | |
| 184 | ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on |
| 185 | platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3, |
| 186 | which is used by OCP-encoded caches even on gfx942. |
| 187 | """ |
| 188 | assert k.dim() == 2 and k.shape[1] == 512, ( |
| 189 | f"K must be [num_tokens, 512], got {k.shape}" |
| 190 | ) |
| 191 | assert k.dtype == torch.bfloat16, f"K must be bf16, got {k.dtype}" |
| 192 | assert is_ue8m0, "Only support ue8m0 quantization." |
| 193 | |
| 194 | # NOTE: When using DP, slot_mapping.shape[0] can be less than k.shape[0] due to |
| 195 | # padding. Always use slot_mapping.shape[0] as the token count. |
| 196 | num_tokens = slot_mapping.shape[0] |
| 197 | block_stride = k_cache.stride(0) # bytes per block |
| 198 | |
| 199 | TOKEN_FP8_DIM = 448 |
| 200 | TOKEN_BF16_DIM = 64 |
| 201 | TOKEN_SCALE_DIM = 8 |
| 202 | QUANT_BLOCK_SIZE = 64 |
| 203 | if use_fnuz: |
| 204 | if not current_platform.is_fp8_fnuz(): |
| 205 | raise ValueError("use_fnuz=True requires a platform using FNUZ FP8") |
| 206 | _, FP8_MAX = get_fp8_min_max() |
| 207 | else: |
| 208 | FP8_MAX = torch.finfo(torch.float8_e4m3fn).max |
| 209 | TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 |
| 210 | |
| 211 | grid = (num_tokens,) |
| 212 | |
| 213 | quantize_and_insert_k_kernel[grid]( |
| 214 | k, |
| 215 | slot_mapping, |
| 216 | k_cache, |
| 217 | num_tokens, |
| 218 | input_dim=512, |
| 219 | fp8_dim=TOKEN_FP8_DIM, |
| 220 | bf16_dim=TOKEN_BF16_DIM, |
| 221 | scale_dim=TOKEN_SCALE_DIM, |
| 222 | quant_block=QUANT_BLOCK_SIZE, |
| 223 | cache_block_size=block_size, |
| 224 | token_data_size=TOKEN_DATA_SIZE, |
| 225 | block_stride=block_stride, |
| 226 | fp8_max=FP8_MAX, |
| 227 | n_quant_blocks=8, |
| 228 | use_fnuz=use_fnuz, |
| 229 | ) |
| 230 | |
| 231 | |
| 232 | class DequantizeAndGatherKCacheKernel( |
| 233 | VllmTritonJitKernel["DequantizeAndGatherKCacheKernel.CompileKey"] |
| 234 | ): |
| 235 | NUM_WORKERS = 128 |
| 236 | |
| 237 | @dataclass(frozen=True) |
| 238 | class CompileKey: |
| 239 | max_blocks_per_seq: int |
| 240 | cache_block_size: int |
| 241 | block_stride: int |
| 242 | use_fnuz: bool |
| 243 | has_gather_lens: bool |
| 244 | offset: int |
| 245 | |
| 246 | @staticmethod |
| 247 | @triton.jit |
| 248 | def kernel( |
| 249 | out_ptr, |
| 250 | out_stride0, |
| 251 | out_stride1, |
| 252 | k_cache_ptr, |
| 253 | seq_lens_ptr, |
| 254 | block_table_ptr, |
| 255 | offset, |
| 256 | gather_lens_ptr, |
| 257 | # Constants |
| 258 | max_blocks_per_seq: tl.constexpr, |
| 259 | fp8_dim: tl.constexpr, # 448 |
| 260 | bf16_dim: tl.constexpr, # 64 |
| 261 | scale_dim: tl.constexpr, # 8 |
| 262 | quant_block: tl.constexpr, # 64 (quantization block size) |
| 263 | cache_block_size: tl.constexpr, # 64 or 128 (paged cache block size) |
| 264 | token_data_size: tl.constexpr, # 576 bytes per token data |
| 265 | block_stride: tl.constexpr, # total bytes per block (padded) int32 |
| 266 | output_dim: tl.constexpr, # 512 |
| 267 | fp8_max: tl.constexpr, |
| 268 | n_quant_blocks: tl.constexpr, # 7 real blocks |
| 269 | use_fnuz: tl.constexpr = False, |
| 270 | ): |
| 271 | batch_idx = tl.program_id(0) |
| 272 | worker_id = tl.program_id(1) |
| 273 | num_workers = tl.num_programs(1) |
| 274 | |
| 275 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 276 | if gather_lens_ptr is not None: # noqa: SIM108 |
| 277 | gather_len = tl.load(gather_lens_ptr + batch_idx) |
| 278 | else: |
| 279 | # Gather all tokens |
| 280 | gather_len = seq_len |
| 281 | start_pos = seq_len - gather_len |
| 282 | |
| 283 | for i in range(worker_id, gather_len, num_workers): |
| 284 | # Calculate the actual token index in the sequence |
| 285 | pos = start_pos + i |
| 286 | |
| 287 | # Calculate which block and position within block |
| 288 | block_in_seq = pos // cache_block_size |
| 289 | pos_in_block = pos % cache_block_size |
| 290 | |
| 291 | # Get physical block index from block table |
| 292 | block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq |
| 293 | physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 |
| 294 | |
| 295 | # int64: physical_block_idx * block_stride can exceed 2^31 with many |
| 296 | # KV-cache blocks (e.g. >= 57K at block_stride ~37K). |
| 297 | cache_block_ptr = ( |
| 298 | k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride |
| 299 | ) |
| 300 | |
| 301 | # Token data pointer |
| 302 | token_data_ptr = cache_block_ptr + pos_in_block * token_data_size |
| 303 | |
| 304 | # Scale pointer: after all token data |
| 305 | token_scale_ptr = ( |
| 306 | cache_block_ptr |
| 307 | + cache_block_size * token_data_size |
| 308 | + pos_in_block * scale_dim |
| 309 | ) |
| 310 | |
| 311 | # Token data layout: [0:448] fp8, [448:576] bf16 |
| 312 | token_fp8_ptr = token_data_ptr |
| 313 | token_bf16_ptr = token_data_ptr + fp8_dim |
| 314 | |
| 315 | # Output pointer for this token (flattened) |
| 316 | output_row_ptr = ( |
| 317 | out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 |
| 318 | ) |
| 319 | |
| 320 | # ========== Dequantize FP8 portion using UE8M0 ========== |
| 321 | for qblock_idx in tl.static_range(n_quant_blocks): |
| 322 | qblock_start = qblock_idx * quant_block |
| 323 | |
| 324 | if qblock_start < fp8_dim: |
| 325 | offsets = qblock_start + tl.arange(0, quant_block) |
| 326 | mask = offsets < fp8_dim |
| 327 | |
| 328 | # Load quantized fp8 values (stored as uint8) |
| 329 | x_uint8 = tl.load(token_fp8_ptr + offsets, mask=mask, other=0) |
| 330 | |
| 331 | # Bitcast uint8 back to fp8 (FNUZ on gfx942, OCP elsewhere). |
| 332 | if use_fnuz: |
| 333 | x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) |
| 334 | else: |
| 335 | x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) |
| 336 | |
| 337 | # Convert fp8 to float32 for computation |
| 338 | x_float = x_fp8.to(tl.float32) |
| 339 | |
| 340 | # Load and decode UE8M0 scale |
| 341 | # UE8M0: scale = 2^(stored_value - 127) |
| 342 | encoded_scale = tl.load(token_scale_ptr + qblock_idx) |
| 343 | exponent = encoded_scale.to(tl.float32) - 127.0 |
| 344 | scale = tl.exp2(exponent) |
| 345 | |
| 346 | # Dequantize: bf16_value = fp8_value * scale |
| 347 | x_dequant = x_float * scale |
| 348 | |
| 349 | # Store as bf16 |
| 350 | tl.store( |
| 351 | output_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask |
| 352 | ) |
| 353 | |
| 354 | # ========== Copy BF16 portion directly ========== |
| 355 | bf16_output_offset = fp8_dim # After 448 elements in output |
| 356 | |
| 357 | # Read bf16 from cache |
| 358 | bf16_cache_ptr = token_bf16_ptr.to(tl.pointer_type(tl.bfloat16)) |
| 359 | |
| 360 | # Process in chunks of 16 |
| 361 | for j in tl.static_range(bf16_dim // 16): |
| 362 | chunk_offsets = j * 16 + tl.arange(0, 16) |
| 363 | bf16_vals = tl.load(bf16_cache_ptr + chunk_offsets) |
| 364 | tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) |
| 365 | |
| 366 | def dispatch( # type: ignore[override] |
| 367 | self, |
| 368 | *, |
| 369 | max_model_len: int, |
| 370 | block_table_block_size: int, |
| 371 | cache_block_size: int, |
| 372 | offset: int, |
| 373 | **compile_key_fields: bool, |
| 374 | ) -> CompileKey: |
| 375 | token_stride = 576 |
| 376 | scale_stride = 8 |
| 377 | unpadded = cache_block_size * (token_stride + scale_stride) |
| 378 | block_stride = ((unpadded + token_stride - 1) // token_stride) * token_stride |
| 379 | return self.CompileKey( |
| 380 | **compile_key_fields, |
| 381 | max_blocks_per_seq=(max_model_len + block_table_block_size - 1) |
| 382 | // block_table_block_size, |
| 383 | cache_block_size=cache_block_size, |
| 384 | block_stride=block_stride, |
| 385 | offset=triton_scalar_specialization_rep(offset), |
| 386 | ) |
| 387 | |
| 388 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 389 | max_model_len = vllm_config.model_config.max_model_len |
| 390 | block_size = vllm_config.cache_config.block_size |
| 391 | if max_model_len <= 0 or block_size <= 0: |
| 392 | return [] |
| 393 | |
| 394 | compress_ratios = frozenset( |
| 395 | max(1, int(compress_ratio)) |
| 396 | for compress_ratio in vllm_config.model_config.hf_config.compress_ratios |
| 397 | ) |
| 398 | return self._trace_dispatch(self.dispatch)( |
| 399 | zip_inputs( |
| 400 | dict( |
| 401 | cache_block_size=block_size, |
| 402 | has_gather_lens=True, |
| 403 | use_fnuz=current_platform.is_fp8_fnuz(), |
| 404 | offset=1, |
| 405 | enabled=True, |
| 406 | ), |
| 407 | dict( |
| 408 | cache_block_size=block_size, |
| 409 | has_gather_lens=True, |
| 410 | use_fnuz=current_platform.is_fp8_fnuz(), |
| 411 | offset=2, |
| 412 | enabled=True, |
| 413 | ), |
| 414 | dict( |
| 415 | cache_block_size=block_size, |
| 416 | has_gather_lens=True, |
| 417 | use_fnuz=current_platform.is_fp8_fnuz(), |
| 418 | offset=16, |
| 419 | enabled=True, |
| 420 | ), |
| 421 | dict( |
| 422 | cache_block_size=max(block_size // 4, 1), |
| 423 | has_gather_lens=False, |
| 424 | use_fnuz=False, |
| 425 | offset=16, |
| 426 | enabled=4 in compress_ratios, |
| 427 | ), |
| 428 | dict( |
| 429 | cache_block_size=max(block_size // 128, 1), |
| 430 | has_gather_lens=False, |
| 431 | use_fnuz=False, |
| 432 | offset=16, |
| 433 | enabled=128 in compress_ratios, |
| 434 | ), |
| 435 | ), |
| 436 | max_model_len=max_model_len, |
| 437 | block_table_block_size=block_size, |
| 438 | _when=lambda *, enabled: enabled, |
| 439 | ) |
| 440 | |
| 441 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 442 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 443 | return dict( |
| 444 | out=TritonWarmupTensor( |
| 445 | torch.bfloat16, |
| 446 | shape=(1, 1, 512), |
| 447 | strides=(512, 512, 1), |
| 448 | ), |
| 449 | k_cache=TritonWarmupTensor( |
| 450 | torch.uint8, |
| 451 | shape=(1, 1), |
| 452 | strides=(compile_key.block_stride, 1), |
| 453 | ), |
| 454 | seq_lens=int32_ptr, |
| 455 | gather_lens=(int32_ptr if compile_key.has_gather_lens else None), |
| 456 | block_table=TritonWarmupTensor( |
| 457 | torch.int32, |
| 458 | shape=(1, compile_key.max_blocks_per_seq), |
| 459 | ), |
| 460 | block_size=compile_key.cache_block_size, |
| 461 | offset=compile_key.offset, |
| 462 | use_fnuz=compile_key.use_fnuz, |
| 463 | ) |
| 464 | |
| 465 | @kernel_launcher |
| 466 | def __call__( |
| 467 | self, |
| 468 | out: torch.Tensor, |
| 469 | k_cache: torch.Tensor, |
| 470 | seq_lens: torch.Tensor, |
| 471 | gather_lens: torch.Tensor | None, |
| 472 | block_table: torch.Tensor, |
| 473 | block_size: int, |
| 474 | offset: int, |
| 475 | *, |
| 476 | use_fnuz: bool = False, |
| 477 | ) -> LaunchSpec: |
| 478 | num_reqs = seq_lens.shape[0] |
| 479 | return (num_reqs, self.NUM_WORKERS), dict( |
| 480 | out_stride0=out.stride(0), |
| 481 | out_stride1=out.stride(1), |
| 482 | max_blocks_per_seq=block_table.shape[-1], |
| 483 | fp8_dim=448, |
| 484 | bf16_dim=64, |
| 485 | scale_dim=8, |
| 486 | quant_block=64, |
| 487 | cache_block_size=block_size, |
| 488 | token_data_size=576, |
| 489 | block_stride=k_cache.stride(0), |
| 490 | output_dim=512, |
| 491 | fp8_max=448.0, |
| 492 | n_quant_blocks=7, |
| 493 | ) |
| 494 | |
| 495 | |
| 496 | def dequantize_and_gather_k_cache( |
| 497 | # [num_reqs, max_num_tokens, head_size] |
| 498 | out: torch.Tensor, |
| 499 | # [num_blocks, block_size, head_bytes] |
| 500 | k_cache: torch.Tensor, |
| 501 | # [num_reqs] |
| 502 | seq_lens: torch.Tensor, |
| 503 | # [num_reqs] |
| 504 | gather_lens: torch.Tensor | None, |
| 505 | # [num_reqs, max_blocks_per_seq] |
| 506 | block_table: torch.Tensor, |
| 507 | block_size: int, |
| 508 | offset: int, |
| 509 | use_fnuz: bool = False, |
| 510 | ) -> None: |
| 511 | """Dequantize and gather a paged DSv4 K cache. |
| 512 | |
| 513 | ``use_fnuz`` MUST match the encoder of the specific cache being read: |
| 514 | ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), |
| 515 | ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder |
| 516 | writes FNUZ on gfx942 and OCP on gfx950). |
| 517 | """ |
| 518 | if has_cutedsl(): |
| 519 | # lazily import, otherwise some tests fail due to CUDA driver init failure. |
| 520 | from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( |
| 521 | _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, |
| 522 | ) |
| 523 | |
| 524 | _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( |
| 525 | out=out, |
| 526 | k_cache=k_cache, |
| 527 | seq_lens=seq_lens, |
| 528 | gather_lens=gather_lens, |
| 529 | block_table=block_table, |
| 530 | block_size=block_size, |
| 531 | offset=offset, |
| 532 | ) |
| 533 | return |
| 534 | |
| 535 | _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL( |
| 536 | out, |
| 537 | k_cache, |
| 538 | seq_lens, |
| 539 | gather_lens, |
| 540 | block_table, |
| 541 | block_size, |
| 542 | offset, |
| 543 | use_fnuz=use_fnuz, |
| 544 | ) |
| 545 | |
| 546 | |
| 547 | def compute_global_topk_indices_and_lens( |
| 548 | topk_indices: torch.Tensor, |
| 549 | token_to_req_indices: torch.Tensor, |
| 550 | block_table: torch.Tensor, |
| 551 | block_size: int, |
| 552 | is_valid_token: torch.Tensor, |
| 553 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 554 | """Map local topk indices to global KV cache slots and count valid entries. |
| 555 | |
| 556 | Fuses three operations into a single kernel: |
| 557 | 1. Block-table lookup (local index → global slot id) |
| 558 | 2. Valid-entry counting (topk_lens per token) |
| 559 | 3. Masking padding tokens to length 0 |
| 560 | """ |
| 561 | num_tokens = topk_indices.shape[0] |
| 562 | global_topk_indices = torch.empty_like(topk_indices) |
| 563 | topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) |
| 564 | _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL( |
| 565 | global_topk_indices, |
| 566 | topk_lens, |
| 567 | topk_indices, |
| 568 | token_to_req_indices, |
| 569 | block_table, |
| 570 | block_size, |
| 571 | is_valid_token, |
| 572 | ) |
| 573 | return global_topk_indices, topk_lens |
| 574 | |
| 575 | |
| 576 | class ComputeGlobalTopkIndicesAndLensKernel( |
| 577 | VllmTritonJitKernel["ComputeGlobalTopkIndicesAndLensKernel.CompileKey"] |
| 578 | ): |
| 579 | TRITON_BLOCK_SIZE = 1024 |
| 580 | |
| 581 | @dataclass(frozen=True) |
| 582 | class CompileKey: |
| 583 | global_topk_indices_stride: int |
| 584 | topk_indices_stride: int |
| 585 | topk: int |
| 586 | block_table_stride: int |
| 587 | block_size: int |
| 588 | |
| 589 | @staticmethod |
| 590 | @triton.jit |
| 591 | def kernel( |
| 592 | global_topk_indices_ptr, |
| 593 | global_topk_indices_stride: tl.constexpr, |
| 594 | topk_lens_ptr, |
| 595 | topk_indices_ptr, |
| 596 | topk_indices_stride: tl.constexpr, |
| 597 | topk: tl.constexpr, |
| 598 | token_to_req_indices_ptr, |
| 599 | block_table_ptr, |
| 600 | block_table_stride: tl.constexpr, |
| 601 | block_size: tl.constexpr, |
| 602 | is_valid_token_ptr, |
| 603 | TRITON_BLOCK_SIZE: tl.constexpr, |
| 604 | ): |
| 605 | token_idx = tl.program_id(0) |
| 606 | is_valid_token = tl.load(is_valid_token_ptr + token_idx) |
| 607 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 608 | |
| 609 | count = tl.zeros((), dtype=tl.int32) |
| 610 | for i in range(0, topk, TRITON_BLOCK_SIZE): |
| 611 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 612 | mask = offset < topk |
| 613 | |
| 614 | local_idx = tl.load( |
| 615 | topk_indices_ptr + token_idx * topk_indices_stride + offset, |
| 616 | mask=mask, |
| 617 | other=-1, |
| 618 | ) |
| 619 | is_valid = local_idx >= 0 |
| 620 | |
| 621 | block_indices = local_idx // block_size |
| 622 | block_numbers = tl.load( |
| 623 | block_table_ptr + req_idx * block_table_stride + block_indices, |
| 624 | mask=mask & is_valid, |
| 625 | ) |
| 626 | block_offsets = local_idx % block_size |
| 627 | |
| 628 | slot_ids = block_numbers * block_size + block_offsets |
| 629 | slot_ids = tl.where(is_valid, slot_ids, -1) |
| 630 | tl.store( |
| 631 | global_topk_indices_ptr |
| 632 | + token_idx * global_topk_indices_stride |
| 633 | + offset, |
| 634 | slot_ids, |
| 635 | mask=mask, |
| 636 | ) |
| 637 | count += tl.sum(is_valid.to(tl.int32), axis=0) |
| 638 | |
| 639 | # Zero out length for padding tokens. |
| 640 | tl.store(topk_lens_ptr + token_idx, tl.where(is_valid_token, count, 0)) |
| 641 | |
| 642 | def dispatch( # type: ignore[override] |
| 643 | self, |
| 644 | *, |
| 645 | topk_width: int, |
| 646 | block_size: int, |
| 647 | block_table_block_size: int, |
| 648 | max_model_len: int, |
| 649 | ) -> CompileKey: |
| 650 | return self.CompileKey( |
| 651 | global_topk_indices_stride=topk_width, |
| 652 | topk_indices_stride=topk_width, |
| 653 | topk=topk_width, |
| 654 | block_table_stride=(max_model_len + block_table_block_size - 1) |
| 655 | // block_table_block_size, |
| 656 | block_size=block_size, |
| 657 | ) |
| 658 | |
| 659 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 660 | index_topk = vllm_config.model_config.hf_config.index_topk |
| 661 | compress_ratios = vllm_config.model_config.hf_config.compress_ratios |
| 662 | cache_block_size = vllm_config.cache_config.block_size |
| 663 | if index_topk <= 0 or cache_block_size <= 0: |
| 664 | return [] |
| 665 | |
| 666 | max_c128a_topk = ( |
| 667 | ((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128 |
| 668 | ) |
| 669 | active_c128a_topk_widths = WarmupIntRange( |
| 670 | 128, |
| 671 | max_c128a_topk + 1, |
| 672 | advance=lambda width: ( |
| 673 | min(width * 2, max_c128a_topk) |
| 674 | if width < max_c128a_topk |
| 675 | else max_c128a_topk + 1 |
| 676 | ), |
| 677 | ) |
| 678 | trace_dispatch = self._trace_dispatch(self.dispatch) |
| 679 | c4_keys = trace_dispatch( |
| 680 | topk_width=index_topk if 4 in compress_ratios else (), |
| 681 | block_size=max(1, cache_block_size // 4), |
| 682 | block_table_block_size=cache_block_size, |
| 683 | max_model_len=vllm_config.model_config.max_model_len, |
| 684 | ) |
| 685 | c128a_keys = trace_dispatch( |
| 686 | topk_width=active_c128a_topk_widths if 128 in compress_ratios else (), |
| 687 | block_size=max(1, cache_block_size // 128), |
| 688 | block_table_block_size=cache_block_size, |
| 689 | max_model_len=vllm_config.model_config.max_model_len, |
| 690 | ) |
| 691 | return list(dict.fromkeys((*c4_keys, *c128a_keys))) |
| 692 | |
| 693 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 694 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 695 | return dict( |
| 696 | global_topk_indices=TritonWarmupTensor( |
| 697 | torch.int32, |
| 698 | shape=(1, compile_key.topk), |
| 699 | strides=(compile_key.global_topk_indices_stride, 1), |
| 700 | ), |
| 701 | topk_lens=int32_ptr, |
| 702 | topk_indices=TritonWarmupTensor( |
| 703 | torch.int32, |
| 704 | shape=(1, compile_key.topk), |
| 705 | strides=(compile_key.topk_indices_stride, 1), |
| 706 | ), |
| 707 | token_to_req_indices=int32_ptr, |
| 708 | block_table=TritonWarmupTensor( |
| 709 | torch.int32, |
| 710 | shape=(1, compile_key.block_table_stride), |
| 711 | ), |
| 712 | block_size=compile_key.block_size, |
| 713 | is_valid_token=TritonWarmupTensor(torch.bool), |
| 714 | ) |
| 715 | |
| 716 | @kernel_launcher |
| 717 | def __call__( |
| 718 | self, |
| 719 | global_topk_indices: torch.Tensor, |
| 720 | topk_lens: torch.Tensor, |
| 721 | topk_indices: torch.Tensor, |
| 722 | token_to_req_indices: torch.Tensor, |
| 723 | block_table: torch.Tensor, |
| 724 | block_size: int, |
| 725 | is_valid_token: torch.Tensor, |
| 726 | ) -> LaunchSpec: |
| 727 | num_tokens = topk_indices.shape[0] |
| 728 | return (num_tokens,), dict( |
| 729 | global_topk_indices_stride=global_topk_indices.stride(0), |
| 730 | topk_indices_stride=topk_indices.stride(0), |
| 731 | topk=topk_indices.shape[-1], |
| 732 | block_table_stride=block_table.stride(0), |
| 733 | TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE, |
| 734 | ) |
| 735 | |
| 736 | |
| 737 | _DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL = DequantizeAndGatherKCacheKernel() |
| 738 | _COMPUTE_GLOBAL_TOPK_INDICES_AND_LENS_KERNEL = ComputeGlobalTopkIndicesAndLensKernel() |
| 739 | |
| 740 | |
| 741 | # FlashMLA sparse prefill asserts `params.topk % B_TOPK == 0` (see |
| 742 | # flashmla/csrc/sm100/prefill/sparse/fwd/head{64,128}/phase1.cuh). B_TOPK is |
| 743 | # 64 for the h_q=64 kernel and 128 for h_q=128; pad to 128 to satisfy both. |
| 744 | # The extra slots stay as -1 sentinels and `combined_lens` caps the valid |
| 745 | # range via `topk_length`, so padding is a no-op at kernel level. |
| 746 | _SPARSE_PREFILL_TOPK_ALIGNMENT = 128 |
| 747 | |
| 748 | |
| 749 | def combine_topk_swa_indices( |
| 750 | topk_indices: torch.Tensor, |
| 751 | query_start_loc: torch.Tensor, |
| 752 | seq_lens: torch.Tensor, |
| 753 | gather_lens: torch.Tensor, |
| 754 | window_size: int, |
| 755 | compress_ratio: int, |
| 756 | topk: int, |
| 757 | M: int, |
| 758 | N: int, |
| 759 | out: tuple[torch.Tensor, torch.Tensor] | None = None, |
| 760 | left_visible: torch.Tensor | None = None, |
| 761 | right_visible: torch.Tensor | None = None, |
| 762 | max_image_tokens: int = 0, |
| 763 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 764 | num_tokens = topk_indices.shape[0] |
| 765 | # max_image_tokens widens the SWA column region for in-image |
| 766 | # bidirectional visibility; the width is fixed per model so the |
| 767 | # caller-provided workspace stays valid on image-free batches. |
| 768 | combined_topk = ( |
| 769 | (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) |
| 770 | // _SPARSE_PREFILL_TOPK_ALIGNMENT |
| 771 | * _SPARSE_PREFILL_TOPK_ALIGNMENT |
| 772 | ) |
| 773 | if out is None: |
| 774 | combined_indices = torch.full( |
| 775 | (num_tokens, combined_topk), |
| 776 | fill_value=-1, |
| 777 | dtype=torch.int32, |
| 778 | device=topk_indices.device, |
| 779 | ) |
| 780 | combined_lens = torch.empty( |
| 781 | num_tokens, dtype=torch.int32, device=topk_indices.device |
| 782 | ) |
| 783 | else: |
| 784 | combined_indices, combined_lens = out |
| 785 | combined_indices.fill_(-1) |
| 786 | |
| 787 | _COMBINE_TOPK_SWA_INDICES_KERNEL( |
| 788 | combined_indices, |
| 789 | combined_lens, |
| 790 | topk_indices, |
| 791 | query_start_loc, |
| 792 | seq_lens, |
| 793 | gather_lens, |
| 794 | M, |
| 795 | N, |
| 796 | TOP_K=topk, |
| 797 | COMPRESS_RATIO=compress_ratio, |
| 798 | WINDOW_SIZE=window_size, |
| 799 | left_visible=left_visible, |
| 800 | right_visible=right_visible, |
| 801 | max_image_tokens=max_image_tokens, |
| 802 | ) |
| 803 | return combined_indices, combined_lens |
| 804 | |
| 805 | |
| 806 | def _hf_config_int(vllm_config: Any, name: str, default: int) -> int: |
| 807 | model_config = getattr(vllm_config, "model_config", None) |
| 808 | hf_config = getattr(model_config, "hf_config", None) |
| 809 | return int(getattr(hf_config, name, default) or default) |
| 810 | |
| 811 | |
| 812 | def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int: |
| 813 | scheduler_config = getattr(vllm_config, "scheduler_config", None) |
| 814 | return int(getattr(scheduler_config, name, default) or default) |
| 815 | |
| 816 | |
| 817 | class CombineTopkSwaIndicesKernel( |
| 818 | VllmTritonJitKernel["CombineTopkSwaIndicesKernel.CompileKey"] |
| 819 | ): |
| 820 | NUM_WORKERS = 256 |
| 821 | |
| 822 | @dataclass(frozen=True) |
| 823 | class CompileKey: |
| 824 | TOP_K: int |
| 825 | COMPRESS_RATIO: int |
| 826 | WINDOW_SIZE: int |
| 827 | IMAGE_WIDTH: int |
| 828 | PADDED_TOP_K: int |
| 829 | input_variant: TritonPointerInputVariant |
| 830 | |
| 831 | @staticmethod |
| 832 | @triton.jit( |
| 833 | do_not_specialize=[ |
| 834 | "combined_indices_stride", |
| 835 | "topk_indices_stride", |
| 836 | "M", |
| 837 | "N", |
| 838 | ] |
| 839 | ) |
| 840 | def kernel( |
| 841 | combined_indices_ptr, |
| 842 | combined_indices_stride, |
| 843 | combined_lens_ptr, |
| 844 | topk_indices_ptr, |
| 845 | topk_indices_stride, |
| 846 | query_start_loc_ptr, |
| 847 | seq_lens_ptr, |
| 848 | gather_lens_ptr, |
| 849 | left_visible_ptr, |
| 850 | right_visible_ptr, |
| 851 | M, |
| 852 | N, |
| 853 | TOP_K: tl.constexpr, |
| 854 | COMPRESS_RATIO: tl.constexpr, |
| 855 | WINDOW_SIZE: tl.constexpr, |
| 856 | IMAGE_WIDTH: tl.constexpr, |
| 857 | PADDED_TOP_K: tl.constexpr, |
| 858 | ): |
| 859 | batch_idx = tl.program_id(0) |
| 860 | worker_id = tl.program_id(1) |
| 861 | num_workers = tl.num_programs(1) |
| 862 | |
| 863 | # query_start_loc is a global tensor; rebase to chunk-local offsets |
| 864 | # by subtracting the chunk's starting value. |
| 865 | base = tl.load(query_start_loc_ptr) |
| 866 | query_start = tl.load(query_start_loc_ptr + batch_idx) - base |
| 867 | query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base |
| 868 | query_len = query_end - query_start |
| 869 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 870 | gather_len = tl.load(gather_lens_ptr + batch_idx) |
| 871 | start_pos = seq_len - query_len |
| 872 | # The SWA portion of the gathered buffer starts from position |
| 873 | # (seq_len - gather_len), not position 0. We need this offset |
| 874 | # to correctly index into the gathered buffer. |
| 875 | gather_start = seq_len - gather_len |
| 876 | |
| 877 | for token_idx in range(query_start + worker_id, query_end, num_workers): |
| 878 | # topk_len is fully determined by the query token's absolute position: |
| 879 | # both the C4A indexer and the C128A metadata builder emit |
| 880 | # min((pos + 1) // compress_ratio, topk_tokens) valid entries. |
| 881 | # Caller passes TOP_K=0 for SWA-only layers to zero this out. |
| 882 | token_idx_in_query = token_idx - query_start |
| 883 | pos = start_pos + token_idx_in_query |
| 884 | topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) |
| 885 | if IMAGE_WIDTH > 0: |
| 886 | # In-image bidirectional visibility: the window starts up to |
| 887 | # max(left - (window - 1), 0) positions earlier and extends |
| 888 | # `right` positions past the query token (both 0 outside |
| 889 | # image spans, reducing to the plain causal window). |
| 890 | left = tl.load(left_visible_ptr + token_idx) |
| 891 | right = tl.load(right_visible_ptr + token_idx) |
| 892 | else: |
| 893 | left = 0 |
| 894 | right = 0 |
| 895 | left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) |
| 896 | swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) |
| 897 | swa_len = pos + right - swa_start + 1 |
| 898 | |
| 899 | offset = tl.arange(0, PADDED_TOP_K) |
| 900 | mask = offset < topk_len |
| 901 | topk_indices = tl.load( |
| 902 | topk_indices_ptr + token_idx * topk_indices_stride + offset, |
| 903 | mask=mask, |
| 904 | ) |
| 905 | tl.store( |
| 906 | combined_indices_ptr + token_idx * combined_indices_stride + offset, |
| 907 | topk_indices + M * batch_idx, |
| 908 | mask=mask, |
| 909 | ) |
| 910 | # Index into gathered buffer: N + (position - gather_start) |
| 911 | # For positions [swa_start, pos + right], the buffer indices are: |
| 912 | # [N + swa_start - gather_start, N + pos + right - gather_start] |
| 913 | for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE): |
| 914 | swa_offset = i + tl.arange(0, WINDOW_SIZE) |
| 915 | tl.store( |
| 916 | combined_indices_ptr |
| 917 | + token_idx * combined_indices_stride |
| 918 | + topk_len |
| 919 | + swa_offset, |
| 920 | M * batch_idx + N + swa_start + swa_offset - gather_start, |
| 921 | mask=swa_offset < swa_len, |
| 922 | ) |
| 923 | |
| 924 | combined_len = topk_len + swa_len |
| 925 | tl.store(combined_lens_ptr + token_idx, combined_len) |
| 926 | |
| 927 | def dispatch( # type: ignore[override] |
| 928 | self, |
| 929 | *, |
| 930 | topk_width: int, |
| 931 | topk_indices: bool, |
| 932 | query_start_loc: bool, |
| 933 | seq_lens: bool, |
| 934 | gather_lens: bool, |
| 935 | compress_ratio: int, |
| 936 | index_topk: int, |
| 937 | active_topk_width: int, |
| 938 | WINDOW_SIZE: int, |
| 939 | image_width: int = 0, |
| 940 | ) -> CompileKey: |
| 941 | topk_width = active_topk_width if compress_ratio == 128 else index_topk |
| 942 | topk = 0 if compress_ratio == 1 else topk_width |
| 943 | padded_topk = next_power_of_2(topk_width) |
| 944 | input_variant = TritonPointerInputVariant.from_alignment( |
| 945 | topk_indices=topk_indices, |
| 946 | query_start_loc=query_start_loc, |
| 947 | seq_lens=seq_lens, |
| 948 | gather_lens=gather_lens, |
| 949 | ) |
| 950 | return self.CompileKey( |
| 951 | TOP_K=topk, |
| 952 | COMPRESS_RATIO=compress_ratio, |
| 953 | WINDOW_SIZE=WINDOW_SIZE, |
| 954 | IMAGE_WIDTH=image_width, |
| 955 | PADDED_TOP_K=padded_topk, |
| 956 | input_variant=input_variant, |
| 957 | ) |
| 958 | |
| 959 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 960 | if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0: |
| 961 | return [] |
| 962 | |
| 963 | window_size = _hf_config_int(vllm_config, "sliding_window", 128) |
| 964 | hf_config = vllm_config.model_config.hf_config |
| 965 | index_topk = int(getattr(hf_config, "index_topk", 0) or 0) |
| 966 | compress_ratios = tuple( |
| 967 | sorted({max(1, int(ratio)) for ratio in hf_config.compress_ratios}) |
| 968 | ) |
| 969 | max_c128a_topk = ( |
| 970 | ((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128 |
| 971 | ) |
| 972 | active_topk_widths = WarmupIntRange( |
| 973 | 128, |
| 974 | max_c128a_topk + 1, |
| 975 | advance=lambda width: ( |
| 976 | min(width * 2, max_c128a_topk) |
| 977 | if width < max_c128a_topk |
| 978 | else max_c128a_topk + 1 |
| 979 | ), |
| 980 | ) |
| 981 | image_width = ( |
| 982 | _hf_config_int(vllm_config, "vision_max_n_token", 0) |
| 983 | if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0 |
| 984 | else 0 |
| 985 | ) |
| 986 | # Warm both the plain-window variant (batches without image spans) |
| 987 | # and the in-image bidirectional variant. |
| 988 | image_widths = [0, image_width] if image_width > 0 else [0] |
| 989 | return self._trace_dispatch(self.dispatch)( |
| 990 | zip_inputs( |
| 991 | dict( |
| 992 | topk_indices=True, |
| 993 | query_start_loc=True, |
| 994 | seq_lens=True, |
| 995 | gather_lens=True, |
| 996 | ), |
| 997 | dict( |
| 998 | topk_indices=True, |
| 999 | query_start_loc=False, |
| 1000 | seq_lens=False, |
| 1001 | gather_lens=True, |
| 1002 | ), |
| 1003 | dict( |
| 1004 | topk_indices=False, |
| 1005 | query_start_loc=False, |
| 1006 | seq_lens=False, |
| 1007 | gather_lens=False, |
| 1008 | ), |
| 1009 | ), |
| 1010 | compress_ratio=compress_ratios, |
| 1011 | index_topk=index_topk, |
| 1012 | active_topk_width=active_topk_widths, |
| 1013 | WINDOW_SIZE=window_size, |
| 1014 | image_width=image_widths, |
| 1015 | ) |
| 1016 | |
| 1017 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1018 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 1019 | variant = compile_key.input_variant |
| 1020 | has_image = compile_key.IMAGE_WIDTH > 0 |
| 1021 | return dict( |
| 1022 | combined_indices=int32_ptr, |
| 1023 | combined_lens=int32_ptr, |
| 1024 | topk_indices=variant.pointer( |
| 1025 | "topk_indices", torch.int32, shape=(1, compile_key.PADDED_TOP_K) |
| 1026 | ), |
| 1027 | query_start_loc=variant.pointer("query_start_loc", torch.int32), |
| 1028 | seq_lens=variant.pointer("seq_lens", torch.int32), |
| 1029 | gather_lens=variant.pointer("gather_lens", torch.int32), |
| 1030 | M=1, |
| 1031 | N=1, |
| 1032 | TOP_K=compile_key.TOP_K, |
| 1033 | COMPRESS_RATIO=compile_key.COMPRESS_RATIO, |
| 1034 | WINDOW_SIZE=compile_key.WINDOW_SIZE, |
| 1035 | # Reproduce the runtime substitution: no-image batches launch with |
| 1036 | # ``topk_indices`` in the visibility pointer slots (see __call__). |
| 1037 | left_visible=(int32_ptr if has_image else None), |
| 1038 | right_visible=(int32_ptr if has_image else None), |
| 1039 | max_image_tokens=compile_key.IMAGE_WIDTH, |
| 1040 | ) |
| 1041 | |
| 1042 | @kernel_launcher |
| 1043 | def __call__( |
| 1044 | self, |
| 1045 | combined_indices: torch.Tensor, |
| 1046 | combined_lens: torch.Tensor, |
| 1047 | topk_indices: torch.Tensor, |
| 1048 | query_start_loc: torch.Tensor, |
| 1049 | seq_lens: torch.Tensor, |
| 1050 | gather_lens: torch.Tensor, |
| 1051 | M: int, |
| 1052 | N: int, |
| 1053 | *, |
| 1054 | TOP_K: int, |
| 1055 | COMPRESS_RATIO: int, |
| 1056 | WINDOW_SIZE: int, |
| 1057 | left_visible: torch.Tensor | None = None, |
| 1058 | right_visible: torch.Tensor | None = None, |
| 1059 | max_image_tokens: int = 0, |
| 1060 | ) -> LaunchSpec: |
| 1061 | num_reqs = seq_lens.shape[0] |
| 1062 | has_image = left_visible is not None |
| 1063 | image_width = max_image_tokens if has_image else 0 |
| 1064 | return (num_reqs, self.NUM_WORKERS), dict( |
| 1065 | combined_indices_stride=combined_indices.stride(0), |
| 1066 | topk_indices_stride=topk_indices.stride(0), |
| 1067 | left_visible_ptr=left_visible if has_image else topk_indices, |
| 1068 | right_visible_ptr=right_visible if has_image else topk_indices, |
| 1069 | TOP_K=TOP_K, |
| 1070 | COMPRESS_RATIO=COMPRESS_RATIO, |
| 1071 | WINDOW_SIZE=WINDOW_SIZE, |
| 1072 | IMAGE_WIDTH=image_width, |
| 1073 | PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]), |
| 1074 | ) |
| 1075 | |
| 1076 | |
| 1077 | _COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel() |
| 1078 | |
| 1079 | |
| 1080 | def build_flashinfer_mixed_sparse_indices( |
| 1081 | decode_swa_indices: torch.Tensor, |
| 1082 | decode_compressed_indices: torch.Tensor | None, |
| 1083 | decode_compressed_topk_lens: torch.Tensor | None, |
| 1084 | prefill_topk_indices: torch.Tensor, |
| 1085 | query_start_loc: torch.Tensor, |
| 1086 | seq_lens: torch.Tensor, |
| 1087 | token_to_req_indices: torch.Tensor, |
| 1088 | swa_block_table: torch.Tensor, |
| 1089 | swa_block_size: int, |
| 1090 | compressed_block_table: torch.Tensor | None, |
| 1091 | compressed_block_size: int, |
| 1092 | window_size: int, |
| 1093 | compress_ratio: int, |
| 1094 | topk: int, |
| 1095 | decode_compressed_indices_are_local: bool = False, |
| 1096 | decode_is_valid_token: torch.Tensor | None = None, |
| 1097 | swa_block_span: int | None = None, |
| 1098 | compressed_block_span: int | None = None, |
| 1099 | prefill_left_visible: torch.Tensor | None = None, |
| 1100 | prefill_right_visible: torch.Tensor | None = None, |
| 1101 | max_image_tokens: int = 0, |
| 1102 | num_rows: int | None = None, |
| 1103 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1104 | """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. |
| 1105 | |
| 1106 | Produces ``sparse_indices`` of shape ``[num_tokens, swa_total_width + |
| 1107 | padded_topk]`` (the first ``swa_total_width`` columns are SWA slot ids, the |
| 1108 | rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length |
| 1109 | per token). Decode tokens read precomputed SWA/compressed indices; prefill |
| 1110 | tokens derive their SWA window from the position and translate local |
| 1111 | compressed indices to global slots via the block tables. |
| 1112 | |
| 1113 | ``num_rows`` (>= ``num_tokens``) sizes both outputs for a kernel call that |
| 1114 | spans a CUDA-graph-padded batch; the rows past ``num_tokens`` stay unset. |
| 1115 | |
| 1116 | When ``prefill_left_visible``/``prefill_right_visible`` are given (vision |
| 1117 | variant), the SWA column region widens by ``max_image_tokens`` and prefill |
| 1118 | tokens inside an image span get a bidirectionally widened window; decode |
| 1119 | rows are padded with -1 across the extra columns. |
| 1120 | """ |
| 1121 | assert decode_swa_indices.dtype == torch.int32 |
| 1122 | assert decode_swa_indices.dim() == 2 |
| 1123 | swa_index_width = decode_swa_indices.shape[-1] |
| 1124 | assert swa_index_width >= window_size |
| 1125 | has_image = prefill_left_visible is not None |
| 1126 | image_width = max_image_tokens if has_image else 0 |
| 1127 | swa_total_width = swa_index_width + image_width |
| 1128 | if has_image: |
| 1129 | assert prefill_right_visible is not None |
| 1130 | if decode_compressed_topk_lens is not None: |
| 1131 | assert decode_compressed_topk_lens.dtype == torch.int32 |
| 1132 | assert prefill_topk_indices.dtype == torch.int32 |
| 1133 | assert prefill_topk_indices.dim() == 2 |
| 1134 | assert query_start_loc.dtype == torch.int32 |
| 1135 | assert seq_lens.dtype == torch.int32 |
| 1136 | assert token_to_req_indices.dtype == torch.int32 |
| 1137 | assert swa_block_table.dtype == torch.int32 |
| 1138 | |
| 1139 | num_decode_tokens = decode_swa_indices.shape[0] |
| 1140 | num_prefill_tokens = prefill_topk_indices.shape[0] |
| 1141 | num_tokens = num_decode_tokens + num_prefill_tokens |
| 1142 | assert token_to_req_indices.shape[0] >= num_tokens |
| 1143 | if decode_compressed_topk_lens is not None: |
| 1144 | assert decode_compressed_topk_lens.shape[0] >= num_decode_tokens |
| 1145 | |
| 1146 | decode_compressed_topk = 0 |
| 1147 | if decode_compressed_indices is None: |
| 1148 | decode_compressed_indices = prefill_topk_indices |
| 1149 | else: |
| 1150 | assert decode_compressed_indices.dtype == torch.int32 |
| 1151 | assert decode_compressed_indices.dim() == 2 |
| 1152 | assert decode_compressed_indices.shape[0] == num_decode_tokens |
| 1153 | decode_compressed_topk = decode_compressed_indices.shape[-1] |
| 1154 | if decode_compressed_topk > 0 and decode_compressed_indices_are_local: |
| 1155 | assert decode_is_valid_token is not None |
| 1156 | assert decode_is_valid_token.dtype == torch.bool |
| 1157 | assert decode_is_valid_token.shape[0] >= num_decode_tokens |
| 1158 | else: |
| 1159 | decode_is_valid_token = token_to_req_indices |
| 1160 | |
| 1161 | if compressed_block_table is None: |
| 1162 | compressed_block_table = swa_block_table |
| 1163 | assert compressed_block_table.dtype == torch.int32 |
| 1164 | has_decode_compressed_lens = decode_compressed_topk_lens is not None |
| 1165 | if decode_compressed_topk_lens is None: |
| 1166 | decode_compressed_topk_lens = token_to_req_indices |
| 1167 | |
| 1168 | # The FlashInfer TRTLLM-gen sparse-MLA kernels require every per-token topk |
| 1169 | # index row to start on a 16-byte boundary: the kernel loads the compressed |
| 1170 | # indices with 128-bit (16-byte) vectorized loads, so a misaligned row would |
| 1171 | # fault or read across rows. 16 bytes = 4 int32 indices, so round the topk |
| 1172 | # width (and hence the row stride, since the SWA columns are fixed-width) up |
| 1173 | # to a multiple of 4. The extra columns are filled with -1 (invalid) and bounded |
| 1174 | # by ``sparse_topk_lens``, so padding never changes the attention result. |
| 1175 | padded_topk = max(topk, decode_compressed_topk) |
| 1176 | padded_topk = (padded_topk + 3) // 4 * 4 |
| 1177 | num_rows = num_tokens if num_rows is None else num_rows |
| 1178 | assert num_rows >= num_tokens |
| 1179 | sparse_indices = torch.empty( |
| 1180 | (num_rows, swa_total_width + padded_topk), |
| 1181 | dtype=torch.int32, |
| 1182 | device=decode_swa_indices.device, |
| 1183 | ) |
| 1184 | sparse_topk_lens = torch.empty( |
| 1185 | num_rows, dtype=torch.int32, device=decode_swa_indices.device |
| 1186 | ) |
| 1187 | if num_tokens == 0: |
| 1188 | return sparse_indices, sparse_topk_lens |
| 1189 | |
| 1190 | window_block_size = triton.next_power_of_2(max(swa_index_width, 1)) |
| 1191 | topk_block_size = triton.next_power_of_2(max(padded_topk, 1)) |
| 1192 | max_block_size = max(window_block_size, topk_block_size) |
| 1193 | num_warps = 4 if max_block_size >= 256 else 1 |
| 1194 | |
| 1195 | # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV. |
| 1196 | swa_span = swa_block_size if swa_block_span is None else swa_block_span |
| 1197 | compressed_span = ( |
| 1198 | compressed_block_size |
| 1199 | if compressed_block_span is None |
| 1200 | else compressed_block_span |
| 1201 | ) |
| 1202 | # The launch grid follows the output rows, so hand over only the real ones. |
| 1203 | _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL( |
| 1204 | sparse_indices[:num_tokens], |
| 1205 | sparse_topk_lens[:num_tokens], |
| 1206 | decode_swa_indices, |
| 1207 | decode_compressed_indices, |
| 1208 | decode_compressed_topk_lens, |
| 1209 | decode_is_valid_token, |
| 1210 | prefill_topk_indices, |
| 1211 | prefill_left_visible if has_image else token_to_req_indices, |
| 1212 | prefill_right_visible if has_image else token_to_req_indices, |
| 1213 | query_start_loc, |
| 1214 | seq_lens, |
| 1215 | token_to_req_indices, |
| 1216 | swa_block_table, |
| 1217 | swa_block_size, |
| 1218 | swa_span, |
| 1219 | compressed_block_table, |
| 1220 | compressed_block_size, |
| 1221 | compressed_span, |
| 1222 | num_decode_tokens=num_decode_tokens, |
| 1223 | window_size=window_size, |
| 1224 | swa_index_width=swa_index_width, |
| 1225 | image_width=image_width, |
| 1226 | compress_ratio=compress_ratio, |
| 1227 | topk=topk, |
| 1228 | padded_topk=padded_topk, |
| 1229 | prefill_topk_stride=prefill_topk_indices.shape[-1], |
| 1230 | decode_compressed_topk=decode_compressed_topk, |
| 1231 | decode_compressed_indices_are_local=decode_compressed_indices_are_local, |
| 1232 | has_decode_compressed_lens=has_decode_compressed_lens, |
| 1233 | window_block_size=window_block_size, |
| 1234 | topk_block_size=topk_block_size, |
| 1235 | num_warps=num_warps, |
| 1236 | ) |
| 1237 | return sparse_indices, sparse_topk_lens |
| 1238 | |
| 1239 | |
| 1240 | @triton.jit |
| 1241 | def _remap_flashinfer_index(values, block_size, block_span): |
| 1242 | # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so |
| 1243 | # packed pages (#44577) need block*block_size+off -> block*block_span+off. |
| 1244 | # The launcher declares these indices with sparse_indices_are_storage_offsets=True. |
| 1245 | is_valid = values >= 0 |
| 1246 | safe_values = tl.where(is_valid, values, 0) |
| 1247 | values = (safe_values // block_size) * block_span |
| 1248 | values += safe_values % block_size |
| 1249 | return tl.where(is_valid, values, -1) |
| 1250 | |
| 1251 | |
| 1252 | class BuildFlashinferMixedSparseIndicesKernel( |
| 1253 | VllmTritonJitKernel["BuildFlashinferMixedSparseIndicesKernel.CompileKey"] |
| 1254 | ): |
| 1255 | @dataclass(frozen=True) |
| 1256 | class CompileKey: |
| 1257 | window_size: int |
| 1258 | swa_index_width: int |
| 1259 | image_width: int |
| 1260 | compress_ratio: int |
| 1261 | top_k: int |
| 1262 | padded_top_k: int |
| 1263 | decode_compressed_topk: int |
| 1264 | decode_compressed_indices_are_local: bool |
| 1265 | has_decode_compressed_lens: bool |
| 1266 | window_block_size: int |
| 1267 | topk_block_size: int |
| 1268 | |
| 1269 | @staticmethod |
| 1270 | @triton.jit( |
| 1271 | do_not_specialize=[ |
| 1272 | "sparse_indices_stride", |
| 1273 | "decode_swa_stride", |
| 1274 | "decode_compressed_stride", |
| 1275 | "prefill_topk_stride", |
| 1276 | "swa_block_table_stride", |
| 1277 | "swa_block_size", |
| 1278 | "swa_block_span", |
| 1279 | "compressed_block_table_stride", |
| 1280 | "compressed_block_size", |
| 1281 | "compressed_block_span", |
| 1282 | "NUM_DECODE_TOKENS", |
| 1283 | "PREFILL_TOPK_STRIDE", |
| 1284 | ] |
| 1285 | ) |
| 1286 | def kernel( |
| 1287 | sparse_indices_ptr, |
| 1288 | sparse_indices_stride, |
| 1289 | sparse_topk_lens_ptr, |
| 1290 | decode_swa_indices_ptr, |
| 1291 | decode_swa_stride, |
| 1292 | decode_compressed_indices_ptr, |
| 1293 | decode_compressed_stride, |
| 1294 | decode_compressed_topk_lens_ptr, |
| 1295 | decode_is_valid_token_ptr, |
| 1296 | prefill_topk_indices_ptr, |
| 1297 | prefill_topk_stride, |
| 1298 | left_visible_ptr, |
| 1299 | right_visible_ptr, |
| 1300 | query_start_loc_ptr, |
| 1301 | seq_lens_ptr, |
| 1302 | token_to_req_indices_ptr, |
| 1303 | swa_block_table_ptr, |
| 1304 | swa_block_table_stride, |
| 1305 | swa_block_size, |
| 1306 | swa_block_span, |
| 1307 | compressed_block_table_ptr, |
| 1308 | compressed_block_table_stride, |
| 1309 | compressed_block_size, |
| 1310 | compressed_block_span, |
| 1311 | NUM_DECODE_TOKENS, |
| 1312 | WINDOW_SIZE: tl.constexpr, |
| 1313 | SWA_INDEX_WIDTH: tl.constexpr, |
| 1314 | IMAGE_WIDTH: tl.constexpr, |
| 1315 | COMPRESS_RATIO: tl.constexpr, |
| 1316 | TOP_K: tl.constexpr, |
| 1317 | PADDED_TOP_K: tl.constexpr, |
| 1318 | PREFILL_TOPK_STRIDE, |
| 1319 | DECODE_COMPRESSED_TOPK: tl.constexpr, |
| 1320 | DECODE_COMPRESSED_INDICES_ARE_LOCAL: tl.constexpr, |
| 1321 | HAS_DECODE_COMPRESSED_LENS: tl.constexpr, |
| 1322 | WINDOW_BLOCK_SIZE: tl.constexpr, |
| 1323 | TOPK_BLOCK_SIZE: tl.constexpr, |
| 1324 | ): |
| 1325 | token_idx = tl.program_id(0) |
| 1326 | # Total SWA column count; > SWA_INDEX_WIDTH only for in-image |
| 1327 | # bidirectional visibility (vision variant), where the extra columns are |
| 1328 | # -1-padded for decode rows. |
| 1329 | SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH |
| 1330 | |
| 1331 | if token_idx < NUM_DECODE_TOKENS: |
| 1332 | # A decode row wider than the window is DSpark's non-causal block: |
| 1333 | # the block-anchored window of context followed by the whole block, |
| 1334 | # contiguous from column 0. The kernel treats every entry of its |
| 1335 | # active ranges as a key -- the first min(WINDOW_SIZE, pos + 1) |
| 1336 | # columns and [WINDOW_SIZE, sparse_topk_lens) -- so -1 must not |
| 1337 | # appear in either. Fill the causal window with the first visible |
| 1338 | # entries and put the rest from column WINDOW_SIZE on, followed |
| 1339 | # directly by the compressed entries. |
| 1340 | NONCAUSAL: tl.constexpr = SWA_INDEX_WIDTH > WINDOW_SIZE |
| 1341 | compressed_start = SWA_TOTAL_WIDTH |
| 1342 | if NONCAUSAL: |
| 1343 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 1344 | query_start = tl.load(query_start_loc_ptr + req_idx) |
| 1345 | query_len = tl.load(query_start_loc_ptr + req_idx + 1) - query_start |
| 1346 | context_len = tl.load(seq_lens_ptr + req_idx) - query_len |
| 1347 | token_in_query = token_idx - query_start |
| 1348 | window_len = tl.minimum(context_len + token_in_query + 1, WINDOW_SIZE) |
| 1349 | spill_len = ( |
| 1350 | tl.minimum(context_len, WINDOW_SIZE) + query_len - window_len |
| 1351 | ) |
| 1352 | # A CUDA-graph pad token sits outside its empty request, so |
| 1353 | # the kernel never reads its row; keep it the shape of a |
| 1354 | # causal pad row (all -1, WINDOW_SIZE long) rather than let |
| 1355 | # the formulas above go negative. |
| 1356 | is_block_token = (token_in_query >= 0) & (token_in_query < query_len) |
| 1357 | window_len = tl.where(is_block_token, window_len, 0) |
| 1358 | spill_len = tl.where(is_block_token, spill_len, 0) |
| 1359 | compressed_start = WINDOW_SIZE + spill_len |
| 1360 | for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): |
| 1361 | offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) |
| 1362 | mask = offset < SWA_TOTAL_WIDTH |
| 1363 | src = offset |
| 1364 | src_mask = offset < SWA_INDEX_WIDTH |
| 1365 | if NONCAUSAL: |
| 1366 | # Columns from compressed_start on belong to the loops below. |
| 1367 | mask = mask & (offset < compressed_start) |
| 1368 | spill = offset >= WINDOW_SIZE |
| 1369 | src = tl.where(spill, offset - WINDOW_SIZE + window_len, offset) |
| 1370 | src_mask = tl.where( |
| 1371 | spill, |
| 1372 | offset < WINDOW_SIZE + spill_len, |
| 1373 | offset < window_len, |
| 1374 | ) |
| 1375 | values = tl.load( |
| 1376 | decode_swa_indices_ptr + token_idx * decode_swa_stride + src, |
| 1377 | mask=src_mask, |
| 1378 | other=-1, |
| 1379 | ) |
| 1380 | values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) |
| 1381 | tl.store( |
| 1382 | sparse_indices_ptr + token_idx * sparse_indices_stride + offset, |
| 1383 | values, |
| 1384 | mask=mask, |
| 1385 | ) |
| 1386 | |
| 1387 | compressed_len = tl.zeros((), dtype=tl.int32) |
| 1388 | for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE): |
| 1389 | offset = i + tl.arange(0, TOPK_BLOCK_SIZE) |
| 1390 | mask = offset < PADDED_TOP_K |
| 1391 | values = tl.load( |
| 1392 | decode_compressed_indices_ptr |
| 1393 | + token_idx * decode_compressed_stride |
| 1394 | + offset, |
| 1395 | mask=offset < DECODE_COMPRESSED_TOPK, |
| 1396 | other=-1, |
| 1397 | ) |
| 1398 | if DECODE_COMPRESSED_INDICES_ARE_LOCAL: |
| 1399 | token_valid = tl.load(decode_is_valid_token_ptr + token_idx) |
| 1400 | is_valid = values >= 0 |
| 1401 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 1402 | block_indices = values // compressed_block_size |
| 1403 | block_numbers = tl.load( |
| 1404 | compressed_block_table_ptr |
| 1405 | + req_idx * compressed_block_table_stride |
| 1406 | + block_indices, |
| 1407 | mask=mask & is_valid, |
| 1408 | other=-1, |
| 1409 | ) |
| 1410 | block_offsets = values % compressed_block_size |
| 1411 | values = block_numbers * compressed_block_size + block_offsets |
| 1412 | values = tl.where(is_valid, values, -1) |
| 1413 | compressed_len += tl.sum( |
| 1414 | (is_valid & token_valid).to(tl.int32), axis=0 |
| 1415 | ) |
| 1416 | values = _remap_flashinfer_index( |
| 1417 | values, compressed_block_size, compressed_block_span |
| 1418 | ) |
| 1419 | tl.store( |
| 1420 | sparse_indices_ptr |
| 1421 | + token_idx * sparse_indices_stride |
| 1422 | + compressed_start |
| 1423 | + offset, |
| 1424 | values, |
| 1425 | mask=mask, |
| 1426 | ) |
| 1427 | |
| 1428 | if NONCAUSAL: |
| 1429 | # The row's remaining columns, left over by the shifted |
| 1430 | # compressed block. |
| 1431 | for i in range(0, SWA_TOTAL_WIDTH - WINDOW_SIZE, WINDOW_BLOCK_SIZE): |
| 1432 | offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) |
| 1433 | tl.store( |
| 1434 | sparse_indices_ptr |
| 1435 | + token_idx * sparse_indices_stride |
| 1436 | + compressed_start |
| 1437 | + PADDED_TOP_K |
| 1438 | + offset, |
| 1439 | tl.full((WINDOW_BLOCK_SIZE,), -1, dtype=tl.int32), |
| 1440 | mask=offset < SWA_TOTAL_WIDTH - compressed_start, |
| 1441 | ) |
| 1442 | |
| 1443 | if DECODE_COMPRESSED_TOPK == 0: |
| 1444 | compressed_len = tl.zeros((), dtype=tl.int32) |
| 1445 | elif not DECODE_COMPRESSED_INDICES_ARE_LOCAL: |
| 1446 | if HAS_DECODE_COMPRESSED_LENS: |
| 1447 | compressed_len = tl.load( |
| 1448 | decode_compressed_topk_lens_ptr + token_idx |
| 1449 | ) |
| 1450 | else: |
| 1451 | compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) |
| 1452 | |
| 1453 | tl.store( |
| 1454 | sparse_topk_lens_ptr + token_idx, compressed_start + compressed_len |
| 1455 | ) |
| 1456 | return |
| 1457 | |
| 1458 | prefill_idx = token_idx - NUM_DECODE_TOKENS |
| 1459 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 1460 | query_start = tl.load(query_start_loc_ptr + req_idx) |
| 1461 | query_end = tl.load(query_start_loc_ptr + req_idx + 1) |
| 1462 | query_len = query_end - query_start |
| 1463 | seq_len = tl.load(seq_lens_ptr + req_idx) |
| 1464 | start_pos = seq_len - query_len |
| 1465 | token_idx_in_query = token_idx - query_start |
| 1466 | pos = start_pos + token_idx_in_query |
| 1467 | if IMAGE_WIDTH > 0: |
| 1468 | # In-image bidirectional visibility: window starts up to |
| 1469 | # max(left - (window - 1), 0) earlier and extends `right` past pos. |
| 1470 | left = tl.load(left_visible_ptr + token_idx) |
| 1471 | right = tl.load(right_visible_ptr + token_idx) |
| 1472 | else: |
| 1473 | left = 0 |
| 1474 | right = 0 |
| 1475 | left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) |
| 1476 | swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) |
| 1477 | swa_len = pos + right - swa_start_pos + 1 |
| 1478 | topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) |
| 1479 | |
| 1480 | for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): |
| 1481 | offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) |
| 1482 | mask = offset < SWA_TOTAL_WIDTH |
| 1483 | pos_offset = swa_start_pos + offset |
| 1484 | block_indices = pos_offset // swa_block_size |
| 1485 | block_numbers = tl.load( |
| 1486 | swa_block_table_ptr + req_idx * swa_block_table_stride + block_indices, |
| 1487 | mask=mask & (offset < swa_len), |
| 1488 | other=-1, |
| 1489 | ) |
| 1490 | block_offsets = pos_offset % swa_block_size |
| 1491 | slot_ids = block_numbers * swa_block_size + block_offsets |
| 1492 | slot_ids = tl.where(offset < swa_len, slot_ids, -1) |
| 1493 | slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span) |
| 1494 | tl.store( |
| 1495 | sparse_indices_ptr + token_idx * sparse_indices_stride + offset, |
| 1496 | slot_ids, |
| 1497 | mask=mask, |
| 1498 | ) |
| 1499 | |
| 1500 | for i in range(0, PADDED_TOP_K, TOPK_BLOCK_SIZE): |
| 1501 | offset = i + tl.arange(0, TOPK_BLOCK_SIZE) |
| 1502 | mask = offset < PADDED_TOP_K |
| 1503 | local_idx = tl.load( |
| 1504 | prefill_topk_indices_ptr + prefill_idx * prefill_topk_stride + offset, |
| 1505 | mask=(offset < PREFILL_TOPK_STRIDE) & (offset < topk_len), |
| 1506 | other=-1, |
| 1507 | ) |
| 1508 | is_valid = local_idx >= 0 |
| 1509 | block_indices = local_idx // compressed_block_size |
| 1510 | block_numbers = tl.load( |
| 1511 | compressed_block_table_ptr |
| 1512 | + req_idx * compressed_block_table_stride |
| 1513 | + block_indices, |
| 1514 | mask=mask & is_valid, |
| 1515 | other=-1, |
| 1516 | ) |
| 1517 | block_offsets = local_idx % compressed_block_size |
| 1518 | slot_ids = block_numbers * compressed_block_size + block_offsets |
| 1519 | slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1) |
| 1520 | slot_ids = _remap_flashinfer_index( |
| 1521 | slot_ids, compressed_block_size, compressed_block_span |
| 1522 | ) |
| 1523 | tl.store( |
| 1524 | sparse_indices_ptr |
| 1525 | + token_idx * sparse_indices_stride |
| 1526 | + SWA_TOTAL_WIDTH |
| 1527 | + offset, |
| 1528 | slot_ids, |
| 1529 | mask=mask, |
| 1530 | ) |
| 1531 | |
| 1532 | tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len) |
| 1533 | |
| 1534 | def dispatch( # type: ignore[override] |
| 1535 | self, |
| 1536 | *, |
| 1537 | swa_index_width: int, |
| 1538 | image_width: int, |
| 1539 | compress_ratio: int, |
| 1540 | index_topk: int, |
| 1541 | active_topk_width: int, |
| 1542 | has_prefill: bool, |
| 1543 | **compile_key_fields: Any, |
| 1544 | ) -> CompileKey: |
| 1545 | decode_compressed_topk = ( |
| 1546 | active_topk_width |
| 1547 | if compress_ratio == 128 |
| 1548 | else index_topk |
| 1549 | if compress_ratio == 4 |
| 1550 | else 0 |
| 1551 | ) |
| 1552 | topk_width = decode_compressed_topk if has_prefill else 0 |
| 1553 | padded_top_k = max(topk_width, decode_compressed_topk) |
| 1554 | padded_top_k = (padded_top_k + 3) // 4 * 4 |
| 1555 | return self.CompileKey( |
| 1556 | **compile_key_fields, |
| 1557 | compress_ratio=compress_ratio, |
| 1558 | top_k=topk_width, |
| 1559 | decode_compressed_indices_are_local=compress_ratio == 4, |
| 1560 | has_decode_compressed_lens=compress_ratio == 128, |
| 1561 | swa_index_width=swa_index_width, |
| 1562 | image_width=image_width, |
| 1563 | padded_top_k=padded_top_k, |
| 1564 | decode_compressed_topk=decode_compressed_topk, |
| 1565 | window_block_size=next_power_of_2(max(swa_index_width, 1)), |
| 1566 | topk_block_size=next_power_of_2(max(padded_top_k, 1)), |
| 1567 | ) |
| 1568 | |
| 1569 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 1570 | hf_config = vllm_config.model_config.hf_config |
| 1571 | window_size = int(getattr(hf_config, "sliding_window", 128) or 128) |
| 1572 | speculative_config = vllm_config.speculative_config |
| 1573 | if speculative_config is not None and speculative_config.use_dspark(): |
| 1574 | from vllm.v1.attention.backends.mla.compressor_utils import ( |
| 1575 | get_dspark_swa_index_width, |
| 1576 | ) |
| 1577 | |
| 1578 | swa_index_width: int | tuple[int, int] = ( |
| 1579 | window_size, |
| 1580 | get_dspark_swa_index_width( |
| 1581 | window_size, |
| 1582 | speculative_config.num_speculative_tokens, |
| 1583 | ), |
| 1584 | ) |
| 1585 | else: |
| 1586 | swa_index_width = window_size |
| 1587 | max_image_tokens = ( |
| 1588 | int(getattr(hf_config, "vision_max_n_token", 0) or 0) |
| 1589 | if int(getattr(hf_config, "vision_n_layers", 0) or 0) > 0 |
| 1590 | else 0 |
| 1591 | ) |
| 1592 | image_widths = (0, max_image_tokens) if max_image_tokens > 0 else 0 |
| 1593 | index_topk = int(getattr(hf_config, "index_topk", 0) or 0) |
| 1594 | compress_ratios = tuple( |
| 1595 | sorted({max(1, int(ratio)) for ratio in hf_config.compress_ratios}) |
| 1596 | ) |
| 1597 | max_c128a_topk = ( |
| 1598 | ((vllm_config.model_config.max_model_len + 127) // 128 + 127) // 128 * 128 |
| 1599 | ) |
| 1600 | active_topk_widths = WarmupIntRange( |
| 1601 | 128, |
| 1602 | max_c128a_topk + 1, |
| 1603 | advance=lambda width: ( |
| 1604 | min(width * 2, max_c128a_topk) |
| 1605 | if width < max_c128a_topk |
| 1606 | else max_c128a_topk + 1 |
| 1607 | ), |
| 1608 | ) |
| 1609 | return self._trace_dispatch(self.dispatch)( |
| 1610 | compress_ratio=compress_ratios, |
| 1611 | index_topk=index_topk, |
| 1612 | active_topk_width=active_topk_widths, |
| 1613 | has_prefill=(False, True), |
| 1614 | window_size=window_size, |
| 1615 | swa_index_width=swa_index_width, |
| 1616 | image_width=image_widths, |
| 1617 | ) |
| 1618 | |
| 1619 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1620 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 1621 | swa_block_size = 64 |
| 1622 | compressed_block_size = ( |
| 1623 | swa_block_size |
| 1624 | if compile_key.compress_ratio == 1 |
| 1625 | else 256 // compile_key.compress_ratio |
| 1626 | ) |
| 1627 | num_warps = ( |
| 1628 | 4 |
| 1629 | if max(compile_key.window_block_size, compile_key.topk_block_size) >= 256 |
| 1630 | else 1 |
| 1631 | ) |
| 1632 | return dict( |
| 1633 | sparse_indices=TritonWarmupTensor( |
| 1634 | torch.int32, |
| 1635 | shape=( |
| 1636 | 1, |
| 1637 | compile_key.swa_index_width |
| 1638 | + compile_key.image_width |
| 1639 | + compile_key.padded_top_k, |
| 1640 | ), |
| 1641 | ), |
| 1642 | sparse_topk_lens=int32_ptr, |
| 1643 | decode_swa_indices=TritonWarmupTensor( |
| 1644 | torch.int32, |
| 1645 | shape=(1, compile_key.swa_index_width), |
| 1646 | ), |
| 1647 | decode_compressed_indices=int32_ptr, |
| 1648 | decode_compressed_topk_lens=int32_ptr, |
| 1649 | decode_is_valid_token=TritonWarmupTensor( |
| 1650 | torch.bool |
| 1651 | if compile_key.decode_compressed_indices_are_local |
| 1652 | else torch.int32 |
| 1653 | ), |
| 1654 | prefill_topk_indices=int32_ptr, |
| 1655 | prefill_left_visible=int32_ptr, |
| 1656 | prefill_right_visible=int32_ptr, |
| 1657 | query_start_loc=int32_ptr, |
| 1658 | seq_lens=int32_ptr, |
| 1659 | token_to_req_indices=int32_ptr, |
| 1660 | swa_block_table=TritonWarmupTensor(torch.int32, shape=(1, 1)), |
| 1661 | swa_block_size=swa_block_size, |
| 1662 | swa_block_span=1, |
| 1663 | compressed_block_table=TritonWarmupTensor(torch.int32, shape=(1, 1)), |
| 1664 | compressed_block_size=compressed_block_size, |
| 1665 | compressed_block_span=1, |
| 1666 | num_decode_tokens=1, |
| 1667 | window_size=compile_key.window_size, |
| 1668 | swa_index_width=compile_key.swa_index_width, |
| 1669 | image_width=compile_key.image_width, |
| 1670 | compress_ratio=compile_key.compress_ratio, |
| 1671 | topk=compile_key.top_k, |
| 1672 | padded_topk=compile_key.padded_top_k, |
| 1673 | prefill_topk_stride=1, |
| 1674 | decode_compressed_topk=compile_key.decode_compressed_topk, |
| 1675 | decode_compressed_indices_are_local=( |
| 1676 | compile_key.decode_compressed_indices_are_local |
| 1677 | ), |
| 1678 | has_decode_compressed_lens=compile_key.has_decode_compressed_lens, |
| 1679 | window_block_size=compile_key.window_block_size, |
| 1680 | topk_block_size=compile_key.topk_block_size, |
| 1681 | num_warps=num_warps, |
| 1682 | ) |
| 1683 | |
| 1684 | @kernel_launcher |
| 1685 | def __call__( |
| 1686 | self, |
| 1687 | sparse_indices: torch.Tensor, |
| 1688 | sparse_topk_lens: torch.Tensor, |
| 1689 | decode_swa_indices: torch.Tensor, |
| 1690 | decode_compressed_indices: torch.Tensor, |
| 1691 | decode_compressed_topk_lens: torch.Tensor, |
| 1692 | decode_is_valid_token: torch.Tensor, |
| 1693 | prefill_topk_indices: torch.Tensor, |
| 1694 | prefill_left_visible: torch.Tensor, |
| 1695 | prefill_right_visible: torch.Tensor, |
| 1696 | query_start_loc: torch.Tensor, |
| 1697 | seq_lens: torch.Tensor, |
| 1698 | token_to_req_indices: torch.Tensor, |
| 1699 | swa_block_table: torch.Tensor, |
| 1700 | swa_block_size: int, |
| 1701 | swa_block_span: int, |
| 1702 | compressed_block_table: torch.Tensor, |
| 1703 | compressed_block_size: int, |
| 1704 | compressed_block_span: int, |
| 1705 | *, |
| 1706 | num_decode_tokens: int, |
| 1707 | window_size: int, |
| 1708 | swa_index_width: int, |
| 1709 | image_width: int, |
| 1710 | compress_ratio: int, |
| 1711 | topk: int, |
| 1712 | padded_topk: int, |
| 1713 | prefill_topk_stride: int, |
| 1714 | decode_compressed_topk: int, |
| 1715 | decode_compressed_indices_are_local: bool, |
| 1716 | has_decode_compressed_lens: bool, |
| 1717 | window_block_size: int, |
| 1718 | topk_block_size: int, |
| 1719 | num_warps: int, |
| 1720 | ) -> LaunchSpec: |
| 1721 | return (sparse_indices.shape[0],), dict( |
| 1722 | sparse_indices_stride=sparse_indices.stride(0), |
| 1723 | decode_swa_stride=decode_swa_indices.stride(0), |
| 1724 | decode_compressed_stride=decode_compressed_indices.stride(0), |
| 1725 | prefill_topk_stride=prefill_topk_indices.stride(0), |
| 1726 | left_visible_ptr=prefill_left_visible, |
| 1727 | right_visible_ptr=prefill_right_visible, |
| 1728 | swa_block_table_stride=swa_block_table.stride(0), |
| 1729 | compressed_block_table_stride=compressed_block_table.stride(0), |
| 1730 | NUM_DECODE_TOKENS=num_decode_tokens, |
| 1731 | WINDOW_SIZE=window_size, |
| 1732 | SWA_INDEX_WIDTH=swa_index_width, |
| 1733 | IMAGE_WIDTH=image_width, |
| 1734 | COMPRESS_RATIO=compress_ratio, |
| 1735 | TOP_K=topk, |
| 1736 | PADDED_TOP_K=padded_topk, |
| 1737 | PREFILL_TOPK_STRIDE=prefill_topk_stride, |
| 1738 | DECODE_COMPRESSED_TOPK=decode_compressed_topk, |
| 1739 | DECODE_COMPRESSED_INDICES_ARE_LOCAL=decode_compressed_indices_are_local, |
| 1740 | HAS_DECODE_COMPRESSED_LENS=has_decode_compressed_lens, |
| 1741 | WINDOW_BLOCK_SIZE=window_block_size, |
| 1742 | TOPK_BLOCK_SIZE=topk_block_size, |
| 1743 | num_warps=num_warps, |
| 1744 | ) |
| 1745 | |
| 1746 | |
| 1747 | _BUILD_FLASHINFER_MIXED_SPARSE_INDICES_KERNEL = ( |
| 1748 | BuildFlashinferMixedSparseIndicesKernel() |
| 1749 | ) |