| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton kernels for Qwen4Exp QSA sparse attention and cache updates.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from functools import lru_cache |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 12 | TritonWarmupTensor, |
| 13 | triton_scalar_specialization_rep, |
| 14 | ) |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.triton_utils import HAS_TRITON, tl, triton |
| 17 | |
| 18 | |
| 19 | @lru_cache(maxsize=1) |
| 20 | def _is_sm120() -> bool: |
| 21 | """True on sm_120 (RTX PRO 6000 Blackwell): selects the sm_120 tuning table.""" |
| 22 | return current_platform.get_device_capability() == (12, 0) |
| 23 | |
| 24 | |
| 25 | @lru_cache(maxsize=1) |
| 26 | def _is_sm90() -> bool: |
| 27 | """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table.""" |
| 28 | return current_platform.get_device_capability() == (9, 0) |
| 29 | |
| 30 | |
| 31 | @triton.jit(do_not_specialize=["num_rows", "num_requests"]) |
| 32 | def _qsa_sparse_paged_gqa_splitk_kernel( |
| 33 | q_ptr, |
| 34 | k_cache_ptr, |
| 35 | v_cache_ptr, |
| 36 | indices_ptr, |
| 37 | block_table_ptr, |
| 38 | token_to_req_ptr, |
| 39 | partial_output_ptr, |
| 40 | partial_lse_ptr, |
| 41 | output_ptr, |
| 42 | softmax_scale, |
| 43 | output_scale, |
| 44 | output_gate_ptr, |
| 45 | stride_q_row, |
| 46 | stride_q_head, |
| 47 | stride_k_block, |
| 48 | stride_k_token, |
| 49 | stride_k_head, |
| 50 | stride_v_block, |
| 51 | stride_v_token, |
| 52 | stride_v_head, |
| 53 | stride_indices_row, |
| 54 | stride_table_req, |
| 55 | stride_output_row, |
| 56 | stride_output_head, |
| 57 | stride_output_gate_row, |
| 58 | stride_output_gate_head, |
| 59 | num_rows, |
| 60 | num_cache_blocks, |
| 61 | num_requests, |
| 62 | TOPK: tl.constexpr, |
| 63 | PAGE_SIZE: tl.constexpr, |
| 64 | PAGE_TABLE_WIDTH: tl.constexpr, |
| 65 | GROUP_SIZE: tl.constexpr, |
| 66 | HEAD_DIM: tl.constexpr, |
| 67 | NUM_QUERY_HEADS: tl.constexpr, |
| 68 | NUM_SPLITS: tl.constexpr, |
| 69 | NUM_TILES: tl.constexpr, |
| 70 | BLOCK_M: tl.constexpr, |
| 71 | BLOCK_N: tl.constexpr, |
| 72 | IS_FP8: tl.constexpr, |
| 73 | ) -> None: |
| 74 | row = tl.program_id(0) |
| 75 | kv_head = tl.program_id(1) |
| 76 | split_id = tl.program_id(2) |
| 77 | request = tl.load(token_to_req_ptr + row) |
| 78 | safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) |
| 79 | |
| 80 | # The packed selection buffer carries one TRAILING COUNT COLUMN per row |
| 81 | # (column TOPK of a TOPK+1-wide buffer): the row's valid-entry count, |
| 82 | # written by the expand kernel. It is never a token index — the tile loop |
| 83 | # and the index load below only ever cover columns [0, TOPK). |
| 84 | valid_count = tl.load(indices_ptr + row * stride_indices_row + TOPK) |
| 85 | |
| 86 | head_offsets = tl.arange(0, BLOCK_M) |
| 87 | dim_offsets = tl.arange(0, HEAD_DIM) |
| 88 | column_offsets = tl.arange(0, BLOCK_N) |
| 89 | first_head = kv_head * GROUP_SIZE |
| 90 | query = tl.load( |
| 91 | q_ptr |
| 92 | + row * stride_q_row |
| 93 | + (first_head + head_offsets[:, None]) * stride_q_head |
| 94 | + dim_offsets[None, :], |
| 95 | mask=head_offsets[:, None] < GROUP_SIZE, |
| 96 | other=0.0, |
| 97 | ) |
| 98 | |
| 99 | max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32) |
| 100 | normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32) |
| 101 | accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32) |
| 102 | # softmax_scale is the host-side attention scale (1/sqrt(head_dim), with the |
| 103 | # fp8 K dequant scale already pre-multiplied in); convert to log2 units once |
| 104 | # here for the exp2-based online softmax. |
| 105 | score_scale = softmax_scale * 1.4426950408889634 |
| 106 | |
| 107 | tile_end = tl.minimum(NUM_TILES, tl.cdiv(tl.minimum(valid_count, TOPK), BLOCK_N)) |
| 108 | |
| 109 | for tile in range(split_id, tile_end, NUM_SPLITS): |
| 110 | columns = tile * BLOCK_N + column_offsets |
| 111 | logical_token = tl.load( |
| 112 | indices_ptr + row * stride_indices_row + columns, |
| 113 | mask=columns < TOPK, |
| 114 | other=-1, |
| 115 | ) |
| 116 | safe_token = tl.maximum(logical_token, 0) |
| 117 | logical_page = safe_token // PAGE_SIZE |
| 118 | page_offset = safe_token % PAGE_SIZE |
| 119 | valid = ( |
| 120 | (request >= 0) |
| 121 | & (request < num_requests) |
| 122 | & (logical_token >= 0) |
| 123 | & (logical_page < PAGE_TABLE_WIDTH) |
| 124 | ) |
| 125 | physical_page = tl.load( |
| 126 | block_table_ptr |
| 127 | + safe_request * stride_table_req |
| 128 | + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1), |
| 129 | mask=valid, |
| 130 | other=-1, |
| 131 | ) |
| 132 | valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) |
| 133 | # physical_page * block stride can overflow int32 for large caches. |
| 134 | safe_page = tl.maximum(physical_page, 0).to(tl.int64) |
| 135 | keys = tl.load( |
| 136 | k_cache_ptr |
| 137 | + safe_page[None, :] * stride_k_block |
| 138 | + page_offset[None, :] * stride_k_token |
| 139 | + kv_head * stride_k_head |
| 140 | + dim_offsets[:, None], |
| 141 | mask=valid[None, :], |
| 142 | other=0.0, |
| 143 | ) |
| 144 | values = tl.load( |
| 145 | v_cache_ptr |
| 146 | + safe_page[:, None] * stride_v_block |
| 147 | + page_offset[:, None] * stride_v_token |
| 148 | + kv_head * stride_v_head |
| 149 | + dim_offsets[None, :], |
| 150 | mask=valid[:, None], |
| 151 | other=0.0, |
| 152 | ) |
| 153 | if IS_FP8: |
| 154 | # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK |
| 155 | # measured slower here and less accurate). |
| 156 | keys = keys.to(query.dtype) |
| 157 | scores = tl.dot(query, keys) |
| 158 | # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8 |
| 159 | # caches the K dequant scale is already folded into softmax_scale on the |
| 160 | # host. |
| 161 | scores *= score_scale |
| 162 | scores = tl.where(valid[None, :], scores, -1.0e20) |
| 163 | next_max = tl.maximum(max_value, tl.max(scores, axis=1)) |
| 164 | alpha = tl.math.exp2(max_value - next_max) |
| 165 | probabilities = tl.where( |
| 166 | valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 |
| 167 | ) |
| 168 | if IS_FP8: |
| 169 | # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online |
| 170 | # softmax) so fp16 has the range, its wider mantissa is more |
| 171 | # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. |
| 172 | values = values.to(tl.float16) |
| 173 | accumulator = tl.dot( |
| 174 | probabilities.to(values.dtype), |
| 175 | values, |
| 176 | acc=accumulator * alpha[:, None], |
| 177 | ) |
| 178 | normalizer = normalizer * alpha + tl.sum(probabilities, axis=1) |
| 179 | max_value = next_max |
| 180 | |
| 181 | has_values = normalizer > 0 |
| 182 | # Fold the fp8 V dequant scale (output_scale, 1.0 for bf16) into a per-row |
| 183 | # reciprocal normalizer, so the output is a per-row multiply rather than a |
| 184 | # HEAD_DIM-wide scale. The split-K LSE below keeps the unscaled normalizer, |
| 185 | # so the merge stays correct. |
| 186 | inv_normalizer = output_scale / tl.maximum(normalizer, 1.0e-20) |
| 187 | normalized_output = tl.where( |
| 188 | has_values[:, None], |
| 189 | accumulator * inv_normalizer[:, None], |
| 190 | 0.0, |
| 191 | ) |
| 192 | output_mask = head_offsets[:, None] < GROUP_SIZE |
| 193 | if NUM_SPLITS == 1: |
| 194 | # Preserve the unfused path's BF16 attention-output rounding before |
| 195 | # applying the gate in FP32. |
| 196 | normalized_output = normalized_output.to(output_ptr.dtype.element_ty) |
| 197 | output_gate = tl.load( |
| 198 | output_gate_ptr |
| 199 | + row * stride_output_gate_row |
| 200 | + (first_head + head_offsets[:, None]) * stride_output_gate_head |
| 201 | + dim_offsets[None, :], |
| 202 | mask=output_mask, |
| 203 | other=0.0, |
| 204 | ).to(tl.float32) |
| 205 | normalized_output = normalized_output.to(tl.float32) * tl.sigmoid(output_gate) |
| 206 | tl.store( |
| 207 | output_ptr |
| 208 | + row * stride_output_row |
| 209 | + (first_head + head_offsets[:, None]) * stride_output_head |
| 210 | + dim_offsets[None, :], |
| 211 | normalized_output, |
| 212 | mask=output_mask, |
| 213 | ) |
| 214 | else: |
| 215 | partial_lse = tl.where( |
| 216 | has_values, |
| 217 | max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)), |
| 218 | -float("inf"), |
| 219 | ) |
| 220 | partial_row = (split_id * num_rows + row).to(tl.int64) |
| 221 | tl.store( |
| 222 | partial_output_ptr |
| 223 | + (partial_row * NUM_QUERY_HEADS + first_head + head_offsets[:, None]) |
| 224 | * HEAD_DIM |
| 225 | + dim_offsets[None, :], |
| 226 | normalized_output, |
| 227 | mask=output_mask, |
| 228 | ) |
| 229 | tl.store( |
| 230 | partial_lse_ptr + partial_row * NUM_QUERY_HEADS + first_head + head_offsets, |
| 231 | partial_lse, |
| 232 | mask=head_offsets < GROUP_SIZE, |
| 233 | ) |
| 234 | |
| 235 | |
| 236 | @triton.jit(do_not_specialize=["num_rows"]) |
| 237 | def _qsa_merge_splitk_kernel( |
| 238 | partial_output_ptr, |
| 239 | partial_lse_ptr, |
| 240 | output_ptr, |
| 241 | output_gate_ptr, |
| 242 | stride_output_row, |
| 243 | stride_output_head, |
| 244 | stride_output_gate_row, |
| 245 | stride_output_gate_head, |
| 246 | num_rows, |
| 247 | HEAD_DIM: tl.constexpr, |
| 248 | NUM_QUERY_HEADS: tl.constexpr, |
| 249 | NUM_SPLITS: tl.constexpr, |
| 250 | BLOCK_SPLITS: tl.constexpr, |
| 251 | ) -> None: |
| 252 | row = tl.program_id(0) |
| 253 | head = tl.program_id(1) |
| 254 | split_offsets = tl.arange(0, BLOCK_SPLITS) |
| 255 | dim_offsets = tl.arange(0, HEAD_DIM) |
| 256 | split_mask = split_offsets < NUM_SPLITS |
| 257 | lse = tl.load( |
| 258 | partial_lse_ptr + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head, |
| 259 | mask=split_mask, |
| 260 | other=-float("inf"), |
| 261 | ) |
| 262 | lse_max = tl.max(lse, axis=0) |
| 263 | has_values = lse_max > -float("inf") |
| 264 | shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf")) |
| 265 | weights = tl.math.exp2(shifted) |
| 266 | denominator = tl.sum(weights, axis=0) |
| 267 | split_rows = split_offsets.to(tl.int64) * num_rows + row |
| 268 | partial_output = tl.load( |
| 269 | partial_output_ptr |
| 270 | + (split_rows[:, None] * NUM_QUERY_HEADS + head) * HEAD_DIM |
| 271 | + dim_offsets[None, :], |
| 272 | mask=split_mask[:, None], |
| 273 | other=0.0, |
| 274 | ) |
| 275 | merged = tl.sum(partial_output * weights[:, None], axis=0) |
| 276 | merged = tl.where(denominator > 0, merged / denominator, 0.0) |
| 277 | # Preserve the unfused path's BF16 attention-output rounding before |
| 278 | # applying the gate in FP32. |
| 279 | merged = merged.to(output_ptr.dtype.element_ty) |
| 280 | output_gate = tl.load( |
| 281 | output_gate_ptr |
| 282 | + row * stride_output_gate_row |
| 283 | + head * stride_output_gate_head |
| 284 | + dim_offsets |
| 285 | ).to(tl.float32) |
| 286 | merged = merged.to(tl.float32) * tl.sigmoid(output_gate) |
| 287 | tl.store( |
| 288 | output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, |
| 289 | merged, |
| 290 | ) |
| 291 | |
| 292 | |
| 293 | @triton.jit |
| 294 | def _store_qsa_rows_kernel( |
| 295 | cache_ptr, |
| 296 | slots_ptr, |
| 297 | rows_ptr, |
| 298 | stride_cache_block, |
| 299 | stride_cache_token, |
| 300 | stride_cache_dim, |
| 301 | stride_rows_row, |
| 302 | stride_rows_dim, |
| 303 | num_rows, |
| 304 | num_blocks, |
| 305 | PAGE_SIZE: tl.constexpr, |
| 306 | WIDTH: tl.constexpr, |
| 307 | BLOCK_D: tl.constexpr, |
| 308 | ) -> None: |
| 309 | row = tl.program_id(0) |
| 310 | dims = tl.arange(0, BLOCK_D) |
| 311 | slot = tl.load(slots_ptr + row) |
| 312 | valid = (row < num_rows) & (slot >= 0) & (slot < num_blocks * PAGE_SIZE) |
| 313 | block = tl.maximum(slot, 0) // PAGE_SIZE |
| 314 | token = tl.maximum(slot, 0) % PAGE_SIZE |
| 315 | values = tl.load( |
| 316 | rows_ptr + row * stride_rows_row + dims * stride_rows_dim, |
| 317 | mask=valid & (dims < WIDTH), |
| 318 | other=0, |
| 319 | ) |
| 320 | tl.store( |
| 321 | cache_ptr |
| 322 | + block * stride_cache_block |
| 323 | + token * stride_cache_token |
| 324 | + dims * stride_cache_dim, |
| 325 | values, |
| 326 | mask=valid & (dims < WIDTH), |
| 327 | ) |
| 328 | |
| 329 | |
| 330 | @triton.jit |
| 331 | def _compress_qsa_groups_kernel( |
| 332 | raw_keys_ptr, # this step's raw key rows, straight from activations |
| 333 | raw_positions_ptr, # this step's per-token positions |
| 334 | compressor_state_cache_ptr, # per-request ring of previous raw keys |
| 335 | rope_cache_ptr, # packed RoPE position tail of the ring |
| 336 | compressor_state_table_ptr, |
| 337 | token_to_req_ptr, |
| 338 | query_start_loc_ptr, |
| 339 | logical_positions_ptr, |
| 340 | compressed_slots_ptr, |
| 341 | pooled_ptr, |
| 342 | first_positions_ptr, |
| 343 | stride_raw_row, |
| 344 | stride_raw_dim, |
| 345 | stride_raw_positions_row, |
| 346 | stride_raw_positions_dim, |
| 347 | stride_compressor_state_block, |
| 348 | stride_compressor_state_token, |
| 349 | stride_compressor_state_dim, |
| 350 | stride_rope_block, |
| 351 | stride_rope_token, |
| 352 | stride_rope_dim, |
| 353 | stride_compressor_state_table_req, |
| 354 | stride_pooled_row, |
| 355 | stride_pooled_dim, |
| 356 | stride_positions_row, |
| 357 | stride_positions_dim, |
| 358 | num_rows, |
| 359 | num_compressor_state_blocks, |
| 360 | num_requests, |
| 361 | COMPRESSOR_STATE_SIZE: tl.constexpr, |
| 362 | COMPRESS_RATIO: tl.constexpr, |
| 363 | HEAD_DIM: tl.constexpr, |
| 364 | BLOCK_D: tl.constexpr, |
| 365 | LOAD_ROPE_POSITIONS: tl.constexpr, |
| 366 | ) -> None: |
| 367 | row = tl.program_id(0) |
| 368 | dims = tl.arange(0, BLOCK_D) |
| 369 | request = tl.load(token_to_req_ptr + row) |
| 370 | end_position = tl.load(logical_positions_ptr + row) |
| 371 | compressed_slot = tl.load(compressed_slots_ptr + row) |
| 372 | valid_request = (request >= 0) & (request < num_requests) |
| 373 | safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) |
| 374 | query_row_start = tl.load( |
| 375 | query_start_loc_ptr + safe_request, mask=valid_request, other=0 |
| 376 | ) |
| 377 | query_row_end = tl.load( |
| 378 | query_start_loc_ptr + safe_request + 1, mask=valid_request, other=0 |
| 379 | ) |
| 380 | chunk_start_position = end_position - (row - query_row_start) |
| 381 | compressor_state_block = tl.load( |
| 382 | compressor_state_table_ptr + safe_request * stride_compressor_state_table_req, |
| 383 | mask=valid_request, |
| 384 | other=-1, |
| 385 | ) |
| 386 | valid_compressor_state_block = (compressor_state_block >= 0) & ( |
| 387 | compressor_state_block < num_compressor_state_blocks |
| 388 | ) |
| 389 | valid_row = ( |
| 390 | (row < num_rows) |
| 391 | & valid_request |
| 392 | & (row >= query_row_start) |
| 393 | & (row < query_row_end) |
| 394 | & (end_position >= COMPRESS_RATIO - 1) |
| 395 | & (compressed_slot >= 0) |
| 396 | ) |
| 397 | accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) |
| 398 | |
| 399 | # A group can span the compressor-state ring (older members) and this |
| 400 | # step's raw rows (members at positions >= chunk_start_position). |
| 401 | for group_offset in tl.range(0, COMPRESS_RATIO): |
| 402 | position = end_position - (COMPRESS_RATIO - 1 - group_offset) |
| 403 | use_raw = position >= chunk_start_position |
| 404 | raw_row = query_row_start + position - chunk_start_position |
| 405 | raw_values = tl.load( |
| 406 | raw_keys_ptr + raw_row * stride_raw_row + dims * stride_raw_dim, |
| 407 | mask=valid_row |
| 408 | & use_raw |
| 409 | & (raw_row >= query_row_start) |
| 410 | & (raw_row < query_row_end) |
| 411 | & (raw_row < num_rows) |
| 412 | & (dims < HEAD_DIM), |
| 413 | other=0.0, |
| 414 | ).to(tl.float32) |
| 415 | compressor_state_values = tl.load( |
| 416 | compressor_state_cache_ptr |
| 417 | + tl.maximum(compressor_state_block, 0).to(tl.int64) |
| 418 | * stride_compressor_state_block |
| 419 | + (position % COMPRESSOR_STATE_SIZE) * stride_compressor_state_token |
| 420 | + dims * stride_compressor_state_dim, |
| 421 | mask=valid_row |
| 422 | & ~use_raw |
| 423 | & valid_compressor_state_block |
| 424 | & (dims < HEAD_DIM), |
| 425 | other=0.0, |
| 426 | ).to(tl.float32) |
| 427 | accumulator += tl.where(use_raw, raw_values, compressor_state_values) |
| 428 | |
| 429 | tl.store( |
| 430 | pooled_ptr + row * stride_pooled_row + dims * stride_pooled_dim, |
| 431 | accumulator / COMPRESS_RATIO, |
| 432 | mask=(row < num_rows) & (dims < HEAD_DIM), |
| 433 | ) |
| 434 | |
| 435 | position_dims = tl.arange(0, 4) |
| 436 | first_position = end_position - COMPRESS_RATIO + 1 |
| 437 | if LOAD_ROPE_POSITIONS: |
| 438 | first_from_raw = first_position >= chunk_start_position |
| 439 | raw_first_row = query_row_start + first_position - chunk_start_position |
| 440 | raw_position_values = tl.load( |
| 441 | raw_positions_ptr |
| 442 | + raw_first_row * stride_raw_positions_row |
| 443 | + position_dims * stride_raw_positions_dim, |
| 444 | mask=valid_row |
| 445 | & first_from_raw |
| 446 | & (raw_first_row >= query_row_start) |
| 447 | & (raw_first_row < query_row_end) |
| 448 | & (raw_first_row < num_rows) |
| 449 | & (position_dims < 3), |
| 450 | other=0, |
| 451 | ) |
| 452 | compressor_state_position_values = tl.load( |
| 453 | rope_cache_ptr |
| 454 | + tl.maximum(compressor_state_block, 0).to(tl.int64) * stride_rope_block |
| 455 | + (first_position % COMPRESSOR_STATE_SIZE) * stride_rope_token |
| 456 | + position_dims * stride_rope_dim, |
| 457 | mask=valid_row |
| 458 | & ~first_from_raw |
| 459 | & valid_compressor_state_block |
| 460 | & (position_dims < 3), |
| 461 | other=0, |
| 462 | ) |
| 463 | position_values = tl.where( |
| 464 | first_from_raw, |
| 465 | raw_position_values, |
| 466 | compressor_state_position_values, |
| 467 | ) |
| 468 | else: |
| 469 | position_values = tl.where(valid_row, first_position, 0) |
| 470 | tl.store( |
| 471 | first_positions_ptr |
| 472 | + row * stride_positions_row |
| 473 | + position_dims * stride_positions_dim, |
| 474 | position_values, |
| 475 | mask=(row < num_rows) & (position_dims < 3), |
| 476 | ) |
| 477 | |
| 478 | |
| 479 | def _select_sm120_config( |
| 480 | base_programs: int, use_prefill_config: bool, is_fp8: bool |
| 481 | ) -> tuple[int, int, int]: |
| 482 | """(block_n, target_splits, num_warps) retuned on sm_120 (RTX PRO 6000 |
| 483 | Blackwell), split by cache dtype: bf16 and fp8 favour different configs on |
| 484 | sm_120, most visibly on the large-prefill region. Each entry is the fastest |
| 485 | config that stays correct on its own path; the main win is more warps. |
| 486 | """ |
| 487 | if is_fp8: |
| 488 | if base_programs > 2048: |
| 489 | return (32, 2, 1) if use_prefill_config else (64, 1, 2) |
| 490 | if base_programs <= 24: |
| 491 | return 64, 64, 8 |
| 492 | if base_programs <= 32: |
| 493 | return 128, 8, 4 |
| 494 | if base_programs <= 64: |
| 495 | return 64, 8, 4 |
| 496 | if base_programs <= 128: |
| 497 | return 32, 4, 4 |
| 498 | if base_programs <= 256: |
| 499 | return 128, 4, 4 |
| 500 | if base_programs <= 512: |
| 501 | return 64, 4, 4 |
| 502 | return 32, 2, 1 |
| 503 | # bf16 K/V cache. |
| 504 | if base_programs > 2048: |
| 505 | return 64, 2, 2 |
| 506 | if base_programs <= 24: |
| 507 | return 64, 64, 8 |
| 508 | if base_programs <= 64: |
| 509 | return 64, 8, 8 |
| 510 | if base_programs <= 128: |
| 511 | return 32, 8, 8 |
| 512 | if base_programs <= 256: |
| 513 | return 32, 8, 1 |
| 514 | if base_programs <= 512: |
| 515 | return 64, 4, 2 |
| 516 | return 32, 4, 1 |
| 517 | |
| 518 | |
| 519 | def _select_sm90_config( |
| 520 | base_programs: int, use_prefill_config: bool, is_fp8: bool |
| 521 | ) -> tuple[int, int, int]: |
| 522 | """(block_n, target_splits, num_warps) tuned on sm_90 (H20).""" |
| 523 | if base_programs > 2048: |
| 524 | return (32, 1, 1) if use_prefill_config else (64, 1, 2) |
| 525 | if is_fp8: |
| 526 | if base_programs <= 24: |
| 527 | return 64, 64, 2 |
| 528 | if base_programs <= 32: |
| 529 | return 32, 16, 1 |
| 530 | if base_programs <= 64: |
| 531 | return 32, 8, 1 |
| 532 | if base_programs <= 128: |
| 533 | return 32, 4, 1 |
| 534 | if base_programs <= 256: |
| 535 | return 32, 8, 1 |
| 536 | return 32, 4, 1 |
| 537 | if base_programs <= 32: |
| 538 | return 64, 64, 2 |
| 539 | if base_programs <= 64: |
| 540 | return 32, 16, 1 |
| 541 | if base_programs <= 256: |
| 542 | return 32, 8, 1 |
| 543 | return 32, 4, 1 |
| 544 | |
| 545 | |
| 546 | def _select_config( |
| 547 | num_rows: int, |
| 548 | num_kv_heads: int, |
| 549 | use_prefill_config: bool, |
| 550 | num_columns: int, |
| 551 | is_fp8: bool = False, |
| 552 | ) -> tuple[int, int, int, int]: |
| 553 | """Select (block_n, num_warps, num_tiles, num_splits) for the kernel. |
| 554 | |
| 555 | Keyed on base_programs = num_rows * num_kv_heads. The bp > 2048 region splits |
| 556 | on use_prefill_config (capture-stable: at FULL-graph capture max_query_len is |
| 557 | the uniform decode/verify length). This default table was tuned on GB300; |
| 558 | sm_120 (RTX PRO 6000 Blackwell) dispatches to _select_sm120_config instead, |
| 559 | and sm_90 (Hopper) to _select_sm90_config. |
| 560 | """ |
| 561 | base_programs = num_rows * num_kv_heads |
| 562 | if _is_sm120(): |
| 563 | BLOCK_N, target_splits, num_warps = _select_sm120_config( |
| 564 | base_programs, use_prefill_config, is_fp8 |
| 565 | ) |
| 566 | elif _is_sm90(): |
| 567 | BLOCK_N, target_splits, num_warps = _select_sm90_config( |
| 568 | base_programs, use_prefill_config, is_fp8 |
| 569 | ) |
| 570 | elif base_programs > 2048: |
| 571 | BLOCK_N, target_splits, num_warps = ( |
| 572 | (32, 1, 1) if use_prefill_config else (64, 1, 2) |
| 573 | ) |
| 574 | elif base_programs <= 24: |
| 575 | BLOCK_N, target_splits, num_warps = 32, 64, 4 |
| 576 | elif base_programs <= 32: |
| 577 | BLOCK_N, target_splits, num_warps = 32, 16, 1 |
| 578 | elif base_programs <= 64: |
| 579 | BLOCK_N, target_splits, num_warps = 32, 8, 1 |
| 580 | elif base_programs <= 128: |
| 581 | BLOCK_N, target_splits, num_warps = 32, 4, 1 |
| 582 | elif base_programs <= 256: |
| 583 | BLOCK_N, target_splits, num_warps = 32, 8, 1 |
| 584 | elif base_programs <= 512: |
| 585 | BLOCK_N, target_splits, num_warps = 64, 4, 2 |
| 586 | else: |
| 587 | BLOCK_N, target_splits, num_warps = 64, 1, 2 |
| 588 | num_tiles = triton.cdiv(num_columns, BLOCK_N) |
| 589 | # Never more splits than tiles, never empty. |
| 590 | num_splits = min(target_splits, num_tiles) |
| 591 | return BLOCK_N, num_warps, num_tiles, num_splits |
| 592 | |
| 593 | |
| 594 | def qsa_sparse_paged_attention( |
| 595 | q: torch.Tensor, |
| 596 | k_cache: torch.Tensor, |
| 597 | v_cache: torch.Tensor, |
| 598 | logical_indices: torch.Tensor, |
| 599 | block_table: torch.Tensor, |
| 600 | token_to_req: torch.Tensor, |
| 601 | use_prefill_config: bool, |
| 602 | out: torch.Tensor | None = None, |
| 603 | k_scale: float | None = None, |
| 604 | v_scale: float | None = None, |
| 605 | *, |
| 606 | output_gate: torch.Tensor, |
| 607 | ) -> torch.Tensor: |
| 608 | """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches. |
| 609 | |
| 610 | With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales |
| 611 | as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is |
| 612 | pre-multiplied into the softmax scale and v_scale becomes the kernel's |
| 613 | output scale, so the kernel needs no device scale buffers. |
| 614 | |
| 615 | logical_indices is the PACKED selection buffer: [rows, selection_width + 1] |
| 616 | with the trailing column holding each row's valid-entry count (written by |
| 617 | the expand kernel; never a token index). The kernel reads it as the |
| 618 | tile-loop bound. use_prefill_config only steers the top of the config table; see |
| 619 | _select_config. |
| 620 | """ |
| 621 | if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: |
| 622 | raise ValueError("QSA sparse attention received invalid Q/K/V shapes") |
| 623 | if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]: |
| 624 | raise ValueError("QSA indices must have one row per query") |
| 625 | if token_to_req.shape != (q.shape[0],) or block_table.ndim != 2: |
| 626 | raise ValueError("QSA sparse attention metadata has invalid shapes") |
| 627 | if not all(k_cache.shape[:3]) or not all(block_table.shape): |
| 628 | raise ValueError("QSA sparse attention cache and block table must be nonempty") |
| 629 | if logical_indices.shape[1] < 2: |
| 630 | raise ValueError( |
| 631 | "QSA packed indices need selection columns plus the count column" |
| 632 | ) |
| 633 | if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: |
| 634 | raise ValueError("QSA sparse attention requires valid grouped-query heads") |
| 635 | head_dim = q.shape[2] |
| 636 | assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 |
| 637 | assert q.dtype == torch.bfloat16 |
| 638 | assert k_cache.dtype == v_cache.dtype |
| 639 | is_fp8 = k_cache.dtype == torch.float8_e4m3fn |
| 640 | if is_fp8: |
| 641 | assert k_scale is not None and v_scale is not None |
| 642 | # Host pre-multiply: fold the K dequant scale into the attention scale |
| 643 | # and pass V's dequant scale as the kernel's output scale. |
| 644 | softmax_scale = (head_dim**-0.5) * float(k_scale) |
| 645 | output_scale = float(v_scale) |
| 646 | else: |
| 647 | assert k_cache.dtype == torch.bfloat16 |
| 648 | softmax_scale = head_dim**-0.5 |
| 649 | output_scale = 1.0 |
| 650 | assert logical_indices.dtype == block_table.dtype == torch.int32 |
| 651 | assert token_to_req.dtype == torch.int32 |
| 652 | assert q.device == k_cache.device == v_cache.device |
| 653 | assert q.device == logical_indices.device == block_table.device |
| 654 | assert q.device == token_to_req.device |
| 655 | assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 |
| 656 | assert logical_indices.stride(1) == block_table.stride(1) == 1 |
| 657 | assert token_to_req.stride(0) == 1 |
| 658 | |
| 659 | if out is None: |
| 660 | out = torch.empty_like(q) |
| 661 | if out.shape != q.shape: |
| 662 | raise ValueError("QSA sparse output must match its query") |
| 663 | assert out.dtype == q.dtype and out.device == q.device |
| 664 | assert out.stride(2) == 1 |
| 665 | assert output_gate.is_contiguous() |
| 666 | output_gate_view = output_gate.view_as(q) |
| 667 | if not q.shape[0]: |
| 668 | return out |
| 669 | |
| 670 | group_size = q.shape[1] // k_cache.shape[2] |
| 671 | block_m = triton.next_power_of_2(group_size) |
| 672 | selection_width = logical_indices.shape[1] - 1 # trailing column is the count |
| 673 | block_n, partial_warps, num_tiles, num_splits = _select_config( |
| 674 | q.shape[0], k_cache.shape[2], use_prefill_config, selection_width, is_fp8 |
| 675 | ) |
| 676 | |
| 677 | # Split=1 writes output directly and compiles out all workspace accesses. |
| 678 | if num_splits == 1: |
| 679 | partial_output = out |
| 680 | partial_lse = out |
| 681 | else: |
| 682 | # FP32 partials preserve accuracy when merging independently normalized |
| 683 | # splits. |
| 684 | partial_output = torch.empty( |
| 685 | (num_splits, *q.shape), dtype=torch.float32, device=q.device |
| 686 | ) |
| 687 | partial_lse = torch.empty( |
| 688 | (num_splits, q.shape[0], q.shape[1]), |
| 689 | dtype=torch.float32, |
| 690 | device=q.device, |
| 691 | ) |
| 692 | |
| 693 | partial_grid = (q.shape[0], k_cache.shape[2], num_splits) |
| 694 | _qsa_sparse_paged_gqa_splitk_kernel[partial_grid]( |
| 695 | q, |
| 696 | k_cache, |
| 697 | v_cache, |
| 698 | logical_indices, |
| 699 | block_table, |
| 700 | token_to_req, |
| 701 | partial_output, |
| 702 | partial_lse, |
| 703 | out, |
| 704 | softmax_scale, |
| 705 | output_scale, |
| 706 | output_gate_view, |
| 707 | q.stride(0), |
| 708 | q.stride(1), |
| 709 | k_cache.stride(0), |
| 710 | k_cache.stride(1), |
| 711 | k_cache.stride(2), |
| 712 | v_cache.stride(0), |
| 713 | v_cache.stride(1), |
| 714 | v_cache.stride(2), |
| 715 | logical_indices.stride(0), |
| 716 | block_table.stride(0), |
| 717 | out.stride(0), |
| 718 | out.stride(1), |
| 719 | output_gate_view.stride(0), |
| 720 | output_gate_view.stride(1), |
| 721 | q.shape[0], |
| 722 | k_cache.shape[0], |
| 723 | block_table.shape[0], |
| 724 | TOPK=selection_width, |
| 725 | PAGE_SIZE=k_cache.shape[1], |
| 726 | PAGE_TABLE_WIDTH=block_table.shape[1], |
| 727 | GROUP_SIZE=group_size, |
| 728 | HEAD_DIM=q.shape[2], |
| 729 | NUM_QUERY_HEADS=q.shape[1], |
| 730 | NUM_SPLITS=num_splits, |
| 731 | NUM_TILES=num_tiles, |
| 732 | BLOCK_M=block_m, |
| 733 | BLOCK_N=block_n, |
| 734 | IS_FP8=is_fp8, |
| 735 | num_warps=partial_warps, |
| 736 | num_stages=2, |
| 737 | ) |
| 738 | if num_splits == 1: |
| 739 | return out |
| 740 | |
| 741 | _qsa_merge_splitk_kernel[(q.shape[0], q.shape[1])]( |
| 742 | partial_output, |
| 743 | partial_lse, |
| 744 | out, |
| 745 | output_gate_view, |
| 746 | out.stride(0), |
| 747 | out.stride(1), |
| 748 | output_gate_view.stride(0), |
| 749 | output_gate_view.stride(1), |
| 750 | q.shape[0], |
| 751 | HEAD_DIM=q.shape[2], |
| 752 | NUM_QUERY_HEADS=q.shape[1], |
| 753 | NUM_SPLITS=num_splits, |
| 754 | BLOCK_SPLITS=triton.next_power_of_2(num_splits), |
| 755 | num_warps=2, |
| 756 | num_stages=1, |
| 757 | ) |
| 758 | return out |
| 759 | |
| 760 | |
| 761 | def warmup_qsa_sparse_paged_attention( |
| 762 | kv_cache: torch.Tensor, |
| 763 | block_table: torch.Tensor, |
| 764 | *, |
| 765 | num_query_heads: int, |
| 766 | selection_width: int, |
| 767 | ) -> tuple[tuple[int, int, int], ...]: |
| 768 | """Compile every production-reachable split-K/merge specialization.""" |
| 769 | head_dim = kv_cache.shape[-1] // 2 |
| 770 | key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) |
| 771 | # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time. |
| 772 | is_fp8 = kv_cache.dtype == torch.uint8 |
| 773 | cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype |
| 774 | num_kv_heads = key_cache.shape[2] |
| 775 | group_size = num_query_heads // num_kv_heads |
| 776 | block_m = triton.next_power_of_2(group_size) |
| 777 | |
| 778 | # Every config the dispatch can pick for this group size. |
| 779 | profiles = { |
| 780 | _select_config( |
| 781 | num_rows, num_kv_heads, use_prefill_config, selection_width, is_fp8 |
| 782 | ) |
| 783 | for num_rows in range(1, 8193) |
| 784 | for use_prefill_config in (False, True) |
| 785 | } |
| 786 | |
| 787 | # Scalars constant per deployment get their real values (their divisibility |
| 788 | # specialization is wanted); the batch-varying ones are do_not_specialize'd |
| 789 | # on the kernels, so any value here compiles the only variant. |
| 790 | num_rows = 16 |
| 791 | num_requests = 16 |
| 792 | q_ptr = TritonWarmupTensor( |
| 793 | torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) |
| 794 | ) |
| 795 | k_cache_ptr = TritonWarmupTensor( |
| 796 | cache_dtype, |
| 797 | shape=tuple(key_cache.shape), |
| 798 | strides=tuple(key_cache.stride()), |
| 799 | ) |
| 800 | v_cache_ptr = TritonWarmupTensor( |
| 801 | cache_dtype, |
| 802 | shape=tuple(value_cache.shape), |
| 803 | strides=tuple(value_cache.stride()), |
| 804 | ) |
| 805 | # +1: the packed buffer's trailing count column. |
| 806 | indices_ptr = TritonWarmupTensor(torch.int32, shape=(num_rows, selection_width + 1)) |
| 807 | block_table_ptr = TritonWarmupTensor( |
| 808 | block_table.dtype, |
| 809 | shape=tuple(block_table.shape), |
| 810 | strides=tuple(block_table.stride()), |
| 811 | ) |
| 812 | token_to_req_ptr = TritonWarmupTensor(torch.int32) |
| 813 | output_ptr = TritonWarmupTensor( |
| 814 | torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) |
| 815 | ) |
| 816 | # The output gate is mandatory at runtime; warm the gated specialization. |
| 817 | output_gate_ptr = TritonWarmupTensor( |
| 818 | torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) |
| 819 | ) |
| 820 | head_stride = head_dim |
| 821 | row_stride = num_query_heads * head_dim |
| 822 | num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0]) |
| 823 | |
| 824 | warmed = [] |
| 825 | for block_n, warps, num_tiles, num_splits in sorted(profiles): |
| 826 | if num_splits == 1: |
| 827 | partial_output_ptr = output_ptr |
| 828 | partial_lse_ptr = output_ptr |
| 829 | else: |
| 830 | partial_output_ptr = TritonWarmupTensor( |
| 831 | torch.float32, |
| 832 | shape=(num_splits, num_rows, num_query_heads, head_dim), |
| 833 | ) |
| 834 | partial_lse_ptr = TritonWarmupTensor( |
| 835 | torch.float32, shape=(num_splits, num_rows, num_query_heads) |
| 836 | ) |
| 837 | _qsa_sparse_paged_gqa_splitk_kernel.warmup( |
| 838 | q_ptr, |
| 839 | k_cache_ptr, |
| 840 | v_cache_ptr, |
| 841 | indices_ptr, |
| 842 | block_table_ptr, |
| 843 | token_to_req_ptr, |
| 844 | partial_output_ptr, |
| 845 | partial_lse_ptr, |
| 846 | output_ptr, |
| 847 | 1.0, |
| 848 | 1.0, |
| 849 | output_gate_ptr, |
| 850 | row_stride, |
| 851 | head_stride, |
| 852 | key_cache.stride(0), |
| 853 | key_cache.stride(1), |
| 854 | key_cache.stride(2), |
| 855 | value_cache.stride(0), |
| 856 | value_cache.stride(1), |
| 857 | value_cache.stride(2), |
| 858 | selection_width + 1, |
| 859 | block_table.stride(0), |
| 860 | row_stride, |
| 861 | head_stride, |
| 862 | row_stride, |
| 863 | head_stride, |
| 864 | num_rows, |
| 865 | num_cache_blocks, |
| 866 | num_requests, |
| 867 | TOPK=selection_width, |
| 868 | PAGE_SIZE=key_cache.shape[1], |
| 869 | PAGE_TABLE_WIDTH=block_table.shape[1], |
| 870 | GROUP_SIZE=group_size, |
| 871 | HEAD_DIM=head_dim, |
| 872 | NUM_QUERY_HEADS=num_query_heads, |
| 873 | NUM_SPLITS=num_splits, |
| 874 | NUM_TILES=num_tiles, |
| 875 | BLOCK_M=block_m, |
| 876 | BLOCK_N=block_n, |
| 877 | IS_FP8=is_fp8, |
| 878 | num_warps=warps, |
| 879 | num_stages=2, |
| 880 | grid=(num_rows, num_kv_heads, num_splits), |
| 881 | ) |
| 882 | if num_splits > 1: |
| 883 | _qsa_merge_splitk_kernel.warmup( |
| 884 | partial_output_ptr, |
| 885 | partial_lse_ptr, |
| 886 | output_ptr, |
| 887 | output_gate_ptr, |
| 888 | row_stride, |
| 889 | head_stride, |
| 890 | row_stride, |
| 891 | head_stride, |
| 892 | num_rows, |
| 893 | HEAD_DIM=head_dim, |
| 894 | NUM_QUERY_HEADS=num_query_heads, |
| 895 | NUM_SPLITS=num_splits, |
| 896 | BLOCK_SPLITS=triton.next_power_of_2(num_splits), |
| 897 | num_warps=2, |
| 898 | num_stages=1, |
| 899 | grid=(num_rows, num_query_heads), |
| 900 | ) |
| 901 | warmed.append((block_n, num_splits, warps)) |
| 902 | return tuple(warmed) |
| 903 | |
| 904 | |
| 905 | def qsa_store_cache_rows( |
| 906 | cache: torch.Tensor, |
| 907 | slot_mapping: torch.Tensor, |
| 908 | rows: torch.Tensor, |
| 909 | ) -> None: |
| 910 | """Store fixed-width rows in a QSA cache without boolean indexing.""" |
| 911 | if not cache.is_cuda or not HAS_TRITON: |
| 912 | raise RuntimeError("QSA CUDA cache stores require Triton") |
| 913 | if cache.ndim != 4 or cache.shape[2] != 1: |
| 914 | raise ValueError("QSA cache must be [pages, page_size, 1, width]") |
| 915 | if not all(cache.shape): |
| 916 | raise ValueError("QSA cache dimensions must be nonzero") |
| 917 | if rows.ndim == 3: |
| 918 | if rows.shape[1] != 1: |
| 919 | raise ValueError("QSA cache rows must have one head") |
| 920 | rows = rows[:, 0] |
| 921 | if rows.shape != (slot_mapping.numel(), cache.shape[3]): |
| 922 | raise ValueError("QSA cache rows and slots have incompatible shapes") |
| 923 | if not rows.shape[0]: |
| 924 | return |
| 925 | _store_qsa_rows_kernel[(rows.shape[0],)]( |
| 926 | cache, |
| 927 | slot_mapping, |
| 928 | rows, |
| 929 | cache.stride(0), |
| 930 | cache.stride(1), |
| 931 | cache.stride(3), |
| 932 | rows.stride(0), |
| 933 | rows.stride(1), |
| 934 | rows.shape[0], |
| 935 | cache.shape[0], |
| 936 | PAGE_SIZE=cache.shape[1], |
| 937 | WIDTH=cache.shape[3], |
| 938 | BLOCK_D=triton.next_power_of_2(cache.shape[3]), |
| 939 | num_warps=4, |
| 940 | ) |
| 941 | |
| 942 | |
| 943 | def qsa_compress_groups_with_ratio( |
| 944 | raw_keys: torch.Tensor, # this step's raw key rows [rows, 1, head_size] |
| 945 | raw_positions: torch.Tensor, # this step's positions [rows, 1, 3] int64 |
| 946 | compressor_state_cache: torch.Tensor, |
| 947 | compressor_state_block_table: torch.Tensor, |
| 948 | token_to_req: torch.Tensor, |
| 949 | query_start_loc: torch.Tensor, |
| 950 | logical_positions: torch.Tensor, |
| 951 | compressed_slots: torch.Tensor, |
| 952 | compress_ratio: int, |
| 953 | rope_cache: torch.Tensor | None = None, |
| 954 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 955 | """Pool completed groups from the compressor-state ring and raw token rows.""" |
| 956 | if not raw_keys.is_cuda or not HAS_TRITON: |
| 957 | raise RuntimeError("QSA CUDA compression requires Triton") |
| 958 | rows = token_to_req.numel() |
| 959 | if compress_ratio <= 0: |
| 960 | raise ValueError("QSA compression ratio must be positive") |
| 961 | if raw_keys.ndim != 3 or raw_keys.shape[:2] != (rows, 1): |
| 962 | raise ValueError("QSA raw keys must be [rows, 1, head_size]") |
| 963 | if raw_positions.shape != (rows, 1, 3) or raw_positions.dtype != torch.int64: |
| 964 | raise ValueError("QSA raw positions must be [rows, 1, 3] int64") |
| 965 | if logical_positions.shape != (rows,) or compressed_slots.shape != (rows,): |
| 966 | raise ValueError("QSA compression metadata must match token rows") |
| 967 | if compressor_state_cache.ndim != 4 or compressor_state_cache.shape[2] != 1: |
| 968 | raise ValueError("QSA compressor-state cache has an invalid shape") |
| 969 | if ( |
| 970 | # The ring is wider than one group so speculative rows cannot alias |
| 971 | # onto the committed keys of the group still being collected. |
| 972 | compressor_state_cache.shape[1] < compress_ratio |
| 973 | or compressor_state_cache.shape[3] != raw_keys.shape[2] |
| 974 | or compressor_state_cache.dtype != raw_keys.dtype |
| 975 | ): |
| 976 | raise ValueError( |
| 977 | "QSA compressor-state cache does not match the compression layout" |
| 978 | ) |
| 979 | if ( |
| 980 | compressor_state_block_table.ndim != 2 |
| 981 | or compressor_state_block_table.shape[1] < 1 |
| 982 | ): |
| 983 | raise ValueError( |
| 984 | "QSA compressor-state block table must contain one block per request" |
| 985 | ) |
| 986 | if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2: |
| 987 | raise ValueError("QSA query starts must contain a terminal offset") |
| 988 | num_requests = query_start_loc.shape[0] - 1 |
| 989 | if compressor_state_block_table.shape[0] < num_requests: |
| 990 | raise ValueError("QSA compressor-state block table has too few request rows") |
| 991 | if rope_cache is not None and ( |
| 992 | rope_cache.ndim != 4 |
| 993 | or rope_cache.shape[:3] != compressor_state_cache.shape[:3] |
| 994 | or rope_cache.shape[3] != 3 |
| 995 | or rope_cache.dtype != torch.int64 |
| 996 | ): |
| 997 | raise ValueError("QSA packed position view has an invalid shape or dtype") |
| 998 | if rows and ( |
| 999 | not all(compressor_state_cache.shape) |
| 1000 | or not all(compressor_state_block_table.shape) |
| 1001 | ): |
| 1002 | raise ValueError("QSA compressor-state cache and block table must be nonempty") |
| 1003 | pooled = torch.empty( |
| 1004 | (rows, 1, raw_keys.shape[2]), |
| 1005 | dtype=raw_keys.dtype, |
| 1006 | device=raw_keys.device, |
| 1007 | ) |
| 1008 | first_positions = torch.empty((rows, 3), dtype=torch.int64, device=raw_keys.device) |
| 1009 | if not rows: |
| 1010 | return pooled, first_positions |
| 1011 | if rope_cache is None: |
| 1012 | rope_cache = compressor_state_cache |
| 1013 | load_rope_positions = False |
| 1014 | else: |
| 1015 | load_rope_positions = True |
| 1016 | _compress_qsa_groups_kernel[(rows,)]( |
| 1017 | raw_keys, |
| 1018 | raw_positions, |
| 1019 | compressor_state_cache, |
| 1020 | rope_cache, |
| 1021 | compressor_state_block_table, |
| 1022 | token_to_req, |
| 1023 | query_start_loc, |
| 1024 | logical_positions, |
| 1025 | compressed_slots, |
| 1026 | pooled, |
| 1027 | first_positions, |
| 1028 | raw_keys.stride(0), |
| 1029 | raw_keys.stride(2), |
| 1030 | raw_positions.stride(0), |
| 1031 | raw_positions.stride(2), |
| 1032 | compressor_state_cache.stride(0), |
| 1033 | compressor_state_cache.stride(1), |
| 1034 | compressor_state_cache.stride(3), |
| 1035 | rope_cache.stride(0), |
| 1036 | rope_cache.stride(1), |
| 1037 | rope_cache.stride(3), |
| 1038 | compressor_state_block_table.stride(0), |
| 1039 | pooled.stride(0), |
| 1040 | pooled.stride(2), |
| 1041 | first_positions.stride(0), |
| 1042 | first_positions.stride(1), |
| 1043 | rows, |
| 1044 | compressor_state_cache.shape[0], |
| 1045 | num_requests, |
| 1046 | COMPRESSOR_STATE_SIZE=compressor_state_cache.shape[1], |
| 1047 | COMPRESS_RATIO=compress_ratio, |
| 1048 | HEAD_DIM=raw_keys.shape[2], |
| 1049 | LOAD_ROPE_POSITIONS=load_rope_positions, |
| 1050 | BLOCK_D=triton.next_power_of_2(raw_keys.shape[2]), |
| 1051 | num_warps=4, |
| 1052 | ) |
| 1053 | return pooled, first_positions |
| 1054 | |
| 1055 | |
| 1056 | __all__ = [ |
| 1057 | "qsa_compress_groups_with_ratio", |
| 1058 | "qsa_sparse_paged_attention", |
| 1059 | "qsa_store_cache_rows", |
| 1060 | "warmup_qsa_sparse_paged_attention", |
| 1061 | ] |