| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton kernels for the Qwen4Exp weight-free QSA path.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import math |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm import _custom_ops as ops |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.triton_utils import HAS_TRITON, tl, triton |
| 14 | |
| 15 | _LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024 |
| 16 | _TOPK_WORKSPACE_BYTES = 1024 * 1024 |
| 17 | |
| 18 | |
| 19 | @triton.jit |
| 20 | def _qsa_mqa_paged_kernel( |
| 21 | q_ptr, |
| 22 | k_cache_ptr, |
| 23 | page_table_ptr, |
| 24 | token_to_req_ptr, |
| 25 | query_positions_ptr, |
| 26 | sequence_lengths_ptr, |
| 27 | visible_blocks_ptr, |
| 28 | logits_ptr, |
| 29 | stride_q_row, |
| 30 | stride_q_head, |
| 31 | stride_q_dim, |
| 32 | stride_cache_block, |
| 33 | stride_cache_token, |
| 34 | stride_cache_dim, |
| 35 | stride_table_req, |
| 36 | stride_table_page, |
| 37 | stride_logits_row, |
| 38 | num_rows, |
| 39 | num_columns, |
| 40 | num_pages, |
| 41 | num_requests, |
| 42 | score_divisor, |
| 43 | PAGE_SIZE: tl.constexpr, |
| 44 | PAGE_TABLE_WIDTH: tl.constexpr, |
| 45 | NUM_HEADS: tl.constexpr, |
| 46 | HEAD_DIM: tl.constexpr, |
| 47 | BLOCK_N: tl.constexpr, |
| 48 | BLOCK_D: tl.constexpr, |
| 49 | COMPRESS_RATIO: tl.constexpr, |
| 50 | ) -> None: |
| 51 | row = tl.program_id(0) |
| 52 | columns = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) |
| 53 | dims = tl.arange(0, BLOCK_D) |
| 54 | request = tl.load(token_to_req_ptr + row) |
| 55 | safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) |
| 56 | query_position = tl.load(query_positions_ptr + row) |
| 57 | sequence_length = tl.load( |
| 58 | sequence_lengths_ptr + safe_request, |
| 59 | mask=(request >= 0) & (request < num_requests), |
| 60 | other=0, |
| 61 | ) |
| 62 | visible = tl.minimum( |
| 63 | (query_position + 1) // COMPRESS_RATIO, |
| 64 | sequence_length // COMPRESS_RATIO, |
| 65 | ) |
| 66 | if tl.program_id(1) == 0: |
| 67 | tl.store(visible_blocks_ptr + row, visible) |
| 68 | logical_page = columns // PAGE_SIZE |
| 69 | page_offset = columns % PAGE_SIZE |
| 70 | valid = ( |
| 71 | (row < num_rows) |
| 72 | & (columns < num_columns) |
| 73 | & (columns < visible) |
| 74 | & (request >= 0) |
| 75 | & (request < num_requests) |
| 76 | & (logical_page < PAGE_TABLE_WIDTH) |
| 77 | ) |
| 78 | safe_logical_page = tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1) |
| 79 | physical_page = tl.load( |
| 80 | page_table_ptr |
| 81 | + safe_request * stride_table_req |
| 82 | + safe_logical_page * stride_table_page, |
| 83 | mask=valid, |
| 84 | other=-1, |
| 85 | ) |
| 86 | valid &= (physical_page >= 0) & (physical_page < num_pages) |
| 87 | # physical_page * block stride can overflow int32 for large caches. |
| 88 | safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64) |
| 89 | score = tl.zeros((BLOCK_N,), dtype=tl.float32) |
| 90 | |
| 91 | for head in tl.static_range(0, NUM_HEADS): |
| 92 | query = tl.load( |
| 93 | q_ptr + row * stride_q_row + head * stride_q_head + dims * stride_q_dim, |
| 94 | mask=dims < HEAD_DIM, |
| 95 | other=0.0, |
| 96 | ).to(tl.float32) |
| 97 | keys = tl.load( |
| 98 | k_cache_ptr |
| 99 | + safe_physical_page[:, None] * stride_cache_block |
| 100 | + page_offset[:, None] * stride_cache_token |
| 101 | + dims[None, :] * stride_cache_dim, |
| 102 | mask=valid[:, None] & (dims[None, :] < HEAD_DIM), |
| 103 | other=0.0, |
| 104 | ).to(tl.float32) |
| 105 | dot = tl.sum(keys * query[None, :], axis=1) |
| 106 | score += tl.maximum(dot, 0.0) |
| 107 | |
| 108 | score /= score_divisor |
| 109 | tl.store( |
| 110 | logits_ptr + row * stride_logits_row + columns, |
| 111 | tl.where(valid, score, -float("inf")), |
| 112 | mask=(row < num_rows) & (columns < num_columns), |
| 113 | ) |
| 114 | |
| 115 | |
| 116 | @triton.jit |
| 117 | def _expand_qsa_indices_kernel( |
| 118 | block_indices_ptr, |
| 119 | query_positions_ptr, |
| 120 | sequence_lengths_ptr, |
| 121 | token_to_req_ptr, |
| 122 | output_ptr, |
| 123 | stride_blocks_row, |
| 124 | stride_blocks_column, |
| 125 | stride_output_row, |
| 126 | stride_output_column, |
| 127 | rows, |
| 128 | num_requests, |
| 129 | BLOCK_TOPK: tl.constexpr, |
| 130 | COMPRESS_RATIO: tl.constexpr, |
| 131 | TOKEN_TOPK: tl.constexpr, |
| 132 | OUTPUT_WIDTH: tl.constexpr, |
| 133 | COLUMN_BLOCK: tl.constexpr, |
| 134 | ) -> None: |
| 135 | row = tl.program_id(0) |
| 136 | columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK) |
| 137 | query_position = tl.load(query_positions_ptr + row) |
| 138 | request = tl.load(token_to_req_ptr + row) |
| 139 | safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) |
| 140 | sequence_length = tl.load( |
| 141 | sequence_lengths_ptr + safe_request, |
| 142 | mask=(request >= 0) & (request < num_requests), |
| 143 | other=0, |
| 144 | ) |
| 145 | complete_blocks = tl.minimum( |
| 146 | tl.minimum( |
| 147 | (query_position + 1) // COMPRESS_RATIO, |
| 148 | sequence_length // COMPRESS_RATIO, |
| 149 | ), |
| 150 | BLOCK_TOPK, |
| 151 | ) |
| 152 | expanded_count = complete_blocks * COMPRESS_RATIO |
| 153 | tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO |
| 154 | tail_count = (query_position + 1) - tail_start |
| 155 | |
| 156 | is_expanded = columns < expanded_count |
| 157 | block_rank = columns // COMPRESS_RATIO |
| 158 | offset = columns % COMPRESS_RATIO |
| 159 | safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1) |
| 160 | block = tl.load( |
| 161 | block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column, |
| 162 | mask=(row < rows) & is_expanded, |
| 163 | other=-1, |
| 164 | ) |
| 165 | expanded = block * COMPRESS_RATIO + offset |
| 166 | tail_offset = columns - expanded_count |
| 167 | is_tail = ( |
| 168 | (columns >= expanded_count) |
| 169 | & (tail_offset < tail_count) |
| 170 | & (tail_offset < COMPRESS_RATIO - 1) |
| 171 | ) |
| 172 | token = tl.where(is_expanded, expanded, tail_start + tail_offset) |
| 173 | valid = ( |
| 174 | (row < rows) |
| 175 | & (columns < OUTPUT_WIDTH) |
| 176 | & (is_expanded | is_tail) |
| 177 | & (token >= 0) |
| 178 | & (token < sequence_length) |
| 179 | ) |
| 180 | tl.store( |
| 181 | output_ptr + row * stride_output_row + columns * stride_output_column, |
| 182 | tl.where(valid, token, -1), |
| 183 | mask=(row < rows) & (columns < OUTPUT_WIDTH), |
| 184 | ) |
| 185 | |
| 186 | |
| 187 | @triton.jit |
| 188 | def _qsa_sparse_paged_gqa_splitk_kernel( |
| 189 | q_ptr, |
| 190 | k_cache_ptr, |
| 191 | v_cache_ptr, |
| 192 | indices_ptr, |
| 193 | block_table_ptr, |
| 194 | token_to_req_ptr, |
| 195 | partial_output_ptr, |
| 196 | partial_lse_ptr, |
| 197 | output_ptr, |
| 198 | stride_q_row, |
| 199 | stride_q_head, |
| 200 | stride_k_block, |
| 201 | stride_k_token, |
| 202 | stride_k_head, |
| 203 | stride_v_block, |
| 204 | stride_v_token, |
| 205 | stride_v_head, |
| 206 | stride_indices_row, |
| 207 | stride_table_req, |
| 208 | stride_output_row, |
| 209 | stride_output_head, |
| 210 | num_rows, |
| 211 | num_cache_blocks, |
| 212 | num_requests, |
| 213 | TOPK: tl.constexpr, |
| 214 | PAGE_SIZE: tl.constexpr, |
| 215 | PAGE_TABLE_WIDTH: tl.constexpr, |
| 216 | GROUP_SIZE: tl.constexpr, |
| 217 | HEAD_DIM: tl.constexpr, |
| 218 | NUM_QUERY_HEADS: tl.constexpr, |
| 219 | NUM_SPLITS: tl.constexpr, |
| 220 | NUM_TILES: tl.constexpr, |
| 221 | BLOCK_M: tl.constexpr, |
| 222 | BLOCK_N: tl.constexpr, |
| 223 | ) -> None: |
| 224 | row = tl.program_id(0) |
| 225 | kv_head = tl.program_id(1) |
| 226 | split_id = tl.program_id(2) |
| 227 | request = tl.load(token_to_req_ptr + row) |
| 228 | safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) |
| 229 | |
| 230 | head_offsets = tl.arange(0, BLOCK_M) |
| 231 | dim_offsets = tl.arange(0, HEAD_DIM) |
| 232 | column_offsets = tl.arange(0, BLOCK_N) |
| 233 | first_head = kv_head * GROUP_SIZE |
| 234 | query = tl.load( |
| 235 | q_ptr |
| 236 | + row * stride_q_row |
| 237 | + (first_head + head_offsets[:, None]) * stride_q_head |
| 238 | + dim_offsets[None, :], |
| 239 | mask=head_offsets[:, None] < GROUP_SIZE, |
| 240 | other=0.0, |
| 241 | ) |
| 242 | |
| 243 | max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32) |
| 244 | normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32) |
| 245 | accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32) |
| 246 | softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634 |
| 247 | |
| 248 | # Dynamic bounds avoid padded main-loop iterations for uneven splits. |
| 249 | split_tile_start = split_id * NUM_TILES // NUM_SPLITS |
| 250 | split_tile_end = (split_id + 1) * NUM_TILES // NUM_SPLITS |
| 251 | for tile in range(split_tile_start, split_tile_end): |
| 252 | columns = tile * BLOCK_N + column_offsets |
| 253 | logical_token = tl.load( |
| 254 | indices_ptr + row * stride_indices_row + columns, |
| 255 | mask=columns < TOPK, |
| 256 | other=-1, |
| 257 | ) |
| 258 | safe_token = tl.maximum(logical_token, 0) |
| 259 | logical_page = safe_token // PAGE_SIZE |
| 260 | page_offset = safe_token % PAGE_SIZE |
| 261 | valid = ( |
| 262 | (request >= 0) |
| 263 | & (request < num_requests) |
| 264 | & (logical_token >= 0) |
| 265 | & (logical_page < PAGE_TABLE_WIDTH) |
| 266 | ) |
| 267 | physical_page = tl.load( |
| 268 | block_table_ptr |
| 269 | + safe_request * stride_table_req |
| 270 | + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1), |
| 271 | mask=valid, |
| 272 | other=-1, |
| 273 | ) |
| 274 | valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) |
| 275 | # physical_page * block stride can overflow int32 for large caches. |
| 276 | safe_page = tl.maximum(physical_page, 0).to(tl.int64) |
| 277 | keys = tl.load( |
| 278 | k_cache_ptr |
| 279 | + safe_page[None, :] * stride_k_block |
| 280 | + page_offset[None, :] * stride_k_token |
| 281 | + kv_head * stride_k_head |
| 282 | + dim_offsets[:, None], |
| 283 | mask=valid[None, :], |
| 284 | other=0.0, |
| 285 | ) |
| 286 | values = tl.load( |
| 287 | v_cache_ptr |
| 288 | + safe_page[:, None] * stride_v_block |
| 289 | + page_offset[:, None] * stride_v_token |
| 290 | + kv_head * stride_v_head |
| 291 | + dim_offsets[None, :], |
| 292 | mask=valid[:, None], |
| 293 | other=0.0, |
| 294 | ) |
| 295 | scores = tl.dot(query, keys) |
| 296 | # Scaling scores avoids re-quantizing a scaled query to BF16. |
| 297 | scores *= softmax_scale_log2 |
| 298 | scores = tl.where(valid[None, :], scores, -1.0e20) |
| 299 | next_max = tl.maximum(max_value, tl.max(scores, axis=1)) |
| 300 | alpha = tl.math.exp2(max_value - next_max) |
| 301 | probabilities = tl.where( |
| 302 | valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 |
| 303 | ) |
| 304 | accumulator = tl.dot( |
| 305 | probabilities.to(values.dtype), |
| 306 | values, |
| 307 | acc=accumulator * alpha[:, None], |
| 308 | ) |
| 309 | normalizer = normalizer * alpha + tl.sum(probabilities, axis=1) |
| 310 | max_value = next_max |
| 311 | |
| 312 | has_values = normalizer > 0 |
| 313 | normalized_output = tl.where( |
| 314 | has_values[:, None], |
| 315 | accumulator / tl.maximum(normalizer[:, None], 1.0e-20), |
| 316 | 0.0, |
| 317 | ) |
| 318 | output_mask = head_offsets[:, None] < GROUP_SIZE |
| 319 | if NUM_SPLITS == 1: |
| 320 | tl.store( |
| 321 | output_ptr |
| 322 | + row * stride_output_row |
| 323 | + (first_head + head_offsets[:, None]) * stride_output_head |
| 324 | + dim_offsets[None, :], |
| 325 | normalized_output, |
| 326 | mask=output_mask, |
| 327 | ) |
| 328 | else: |
| 329 | partial_lse = tl.where( |
| 330 | has_values, |
| 331 | max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)), |
| 332 | -float("inf"), |
| 333 | ) |
| 334 | tl.store( |
| 335 | partial_output_ptr |
| 336 | + ( |
| 337 | (split_id * num_rows + row) * NUM_QUERY_HEADS |
| 338 | + first_head |
| 339 | + head_offsets[:, None] |
| 340 | ) |
| 341 | * HEAD_DIM |
| 342 | + dim_offsets[None, :], |
| 343 | normalized_output, |
| 344 | mask=output_mask, |
| 345 | ) |
| 346 | tl.store( |
| 347 | partial_lse_ptr |
| 348 | + (split_id * num_rows + row) * NUM_QUERY_HEADS |
| 349 | + first_head |
| 350 | + head_offsets, |
| 351 | partial_lse, |
| 352 | mask=head_offsets < GROUP_SIZE, |
| 353 | ) |
| 354 | |
| 355 | |
| 356 | @triton.jit |
| 357 | def _qsa_merge_splitk_kernel( |
| 358 | partial_output_ptr, |
| 359 | partial_lse_ptr, |
| 360 | output_ptr, |
| 361 | stride_output_row, |
| 362 | stride_output_head, |
| 363 | num_rows, |
| 364 | HEAD_DIM: tl.constexpr, |
| 365 | NUM_QUERY_HEADS: tl.constexpr, |
| 366 | NUM_SPLITS: tl.constexpr, |
| 367 | BLOCK_SPLITS: tl.constexpr, |
| 368 | ) -> None: |
| 369 | row = tl.program_id(0) |
| 370 | head = tl.program_id(1) |
| 371 | split_offsets = tl.arange(0, BLOCK_SPLITS) |
| 372 | dim_offsets = tl.arange(0, HEAD_DIM) |
| 373 | split_mask = split_offsets < NUM_SPLITS |
| 374 | lse = tl.load( |
| 375 | partial_lse_ptr + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head, |
| 376 | mask=split_mask, |
| 377 | other=-float("inf"), |
| 378 | ) |
| 379 | lse_max = tl.max(lse, axis=0) |
| 380 | has_values = lse_max > -float("inf") |
| 381 | shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf")) |
| 382 | weights = tl.math.exp2(shifted) |
| 383 | denominator = tl.sum(weights, axis=0) |
| 384 | partial_output = tl.load( |
| 385 | partial_output_ptr |
| 386 | + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head) |
| 387 | * HEAD_DIM |
| 388 | + dim_offsets[None, :], |
| 389 | mask=split_mask[:, None], |
| 390 | other=0.0, |
| 391 | ) |
| 392 | merged = tl.sum(partial_output * weights[:, None], axis=0) |
| 393 | merged = tl.where(denominator > 0, merged / denominator, 0.0) |
| 394 | tl.store( |
| 395 | output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, |
| 396 | merged, |
| 397 | ) |
| 398 | |
| 399 | |
| 400 | @triton.jit |
| 401 | def _store_qsa_rows_kernel( |
| 402 | cache_ptr, |
| 403 | slots_ptr, |
| 404 | rows_ptr, |
| 405 | stride_cache_block, |
| 406 | stride_cache_token, |
| 407 | stride_cache_dim, |
| 408 | stride_rows_row, |
| 409 | stride_rows_dim, |
| 410 | num_rows, |
| 411 | num_blocks, |
| 412 | PAGE_SIZE: tl.constexpr, |
| 413 | WIDTH: tl.constexpr, |
| 414 | BLOCK_D: tl.constexpr, |
| 415 | ) -> None: |
| 416 | row = tl.program_id(0) |
| 417 | dims = tl.arange(0, BLOCK_D) |
| 418 | slot = tl.load(slots_ptr + row) |
| 419 | valid = (row < num_rows) & (slot >= 0) & (slot < num_blocks * PAGE_SIZE) |
| 420 | block = tl.maximum(slot, 0) // PAGE_SIZE |
| 421 | token = tl.maximum(slot, 0) % PAGE_SIZE |
| 422 | values = tl.load( |
| 423 | rows_ptr + row * stride_rows_row + dims * stride_rows_dim, |
| 424 | mask=valid & (dims < WIDTH), |
| 425 | other=0, |
| 426 | ) |
| 427 | tl.store( |
| 428 | cache_ptr |
| 429 | + block * stride_cache_block |
| 430 | + token * stride_cache_token |
| 431 | + dims * stride_cache_dim, |
| 432 | values, |
| 433 | mask=valid & (dims < WIDTH), |
| 434 | ) |
| 435 | |
| 436 | |
| 437 | @triton.jit |
| 438 | def _compress_qsa_groups_kernel( |
| 439 | raw_keys_ptr, # this step's raw key rows, straight from activations |
| 440 | raw_positions_ptr, # this step's per-token positions |
| 441 | compressor_state_cache_ptr, # per-request ring of previous raw keys |
| 442 | rope_cache_ptr, # packed RoPE position tail of the ring |
| 443 | compressor_state_table_ptr, |
| 444 | token_to_req_ptr, |
| 445 | query_start_loc_ptr, |
| 446 | logical_positions_ptr, |
| 447 | compressed_slots_ptr, |
| 448 | pooled_ptr, |
| 449 | first_positions_ptr, |
| 450 | stride_raw_row, |
| 451 | stride_raw_dim, |
| 452 | stride_raw_positions_row, |
| 453 | stride_raw_positions_dim, |
| 454 | stride_compressor_state_block, |
| 455 | stride_compressor_state_token, |
| 456 | stride_compressor_state_dim, |
| 457 | stride_rope_block, |
| 458 | stride_rope_token, |
| 459 | stride_rope_dim, |
| 460 | stride_compressor_state_table_req, |
| 461 | stride_pooled_row, |
| 462 | stride_pooled_dim, |
| 463 | stride_positions_row, |
| 464 | stride_positions_dim, |
| 465 | num_rows, |
| 466 | num_compressor_state_blocks, |
| 467 | num_requests, |
| 468 | COMPRESSOR_STATE_SIZE: tl.constexpr, |
| 469 | COMPRESS_RATIO: tl.constexpr, |
| 470 | HEAD_DIM: tl.constexpr, |
| 471 | BLOCK_D: tl.constexpr, |
| 472 | LOAD_ROPE_POSITIONS: tl.constexpr, |
| 473 | ) -> None: |
| 474 | row = tl.program_id(0) |
| 475 | dims = tl.arange(0, BLOCK_D) |
| 476 | request = tl.load(token_to_req_ptr + row) |
| 477 | end_position = tl.load(logical_positions_ptr + row) |
| 478 | compressed_slot = tl.load(compressed_slots_ptr + row) |
| 479 | valid_request = (request >= 0) & (request < num_requests) |
| 480 | safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) |
| 481 | query_row_start = tl.load( |
| 482 | query_start_loc_ptr + safe_request, mask=valid_request, other=0 |
| 483 | ) |
| 484 | query_row_end = tl.load( |
| 485 | query_start_loc_ptr + safe_request + 1, mask=valid_request, other=0 |
| 486 | ) |
| 487 | chunk_start_position = end_position - (row - query_row_start) |
| 488 | compressor_state_block = tl.load( |
| 489 | compressor_state_table_ptr + safe_request * stride_compressor_state_table_req, |
| 490 | mask=valid_request, |
| 491 | other=-1, |
| 492 | ) |
| 493 | valid_compressor_state_block = (compressor_state_block >= 0) & ( |
| 494 | compressor_state_block < num_compressor_state_blocks |
| 495 | ) |
| 496 | valid_row = ( |
| 497 | (row < num_rows) |
| 498 | & valid_request |
| 499 | & (row >= query_row_start) |
| 500 | & (row < query_row_end) |
| 501 | & (end_position >= COMPRESS_RATIO - 1) |
| 502 | & (compressed_slot >= 0) |
| 503 | ) |
| 504 | accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) |
| 505 | |
| 506 | # A group can span the compressor-state ring (older members) and this |
| 507 | # step's raw rows (members at positions >= chunk_start_position). |
| 508 | for group_offset in tl.range(0, COMPRESS_RATIO): |
| 509 | position = end_position - (COMPRESS_RATIO - 1 - group_offset) |
| 510 | use_raw = position >= chunk_start_position |
| 511 | raw_row = query_row_start + position - chunk_start_position |
| 512 | raw_values = tl.load( |
| 513 | raw_keys_ptr + raw_row * stride_raw_row + dims * stride_raw_dim, |
| 514 | mask=valid_row |
| 515 | & use_raw |
| 516 | & (raw_row >= query_row_start) |
| 517 | & (raw_row < query_row_end) |
| 518 | & (raw_row < num_rows) |
| 519 | & (dims < HEAD_DIM), |
| 520 | other=0.0, |
| 521 | ).to(tl.float32) |
| 522 | compressor_state_values = tl.load( |
| 523 | compressor_state_cache_ptr |
| 524 | + tl.maximum(compressor_state_block, 0).to(tl.int64) |
| 525 | * stride_compressor_state_block |
| 526 | + (position % COMPRESSOR_STATE_SIZE) * stride_compressor_state_token |
| 527 | + dims * stride_compressor_state_dim, |
| 528 | mask=valid_row |
| 529 | & ~use_raw |
| 530 | & valid_compressor_state_block |
| 531 | & (dims < HEAD_DIM), |
| 532 | other=0.0, |
| 533 | ).to(tl.float32) |
| 534 | accumulator += tl.where(use_raw, raw_values, compressor_state_values) |
| 535 | |
| 536 | tl.store( |
| 537 | pooled_ptr + row * stride_pooled_row + dims * stride_pooled_dim, |
| 538 | accumulator / COMPRESS_RATIO, |
| 539 | mask=(row < num_rows) & (dims < HEAD_DIM), |
| 540 | ) |
| 541 | |
| 542 | position_dims = tl.arange(0, 4) |
| 543 | first_position = end_position - COMPRESS_RATIO + 1 |
| 544 | if LOAD_ROPE_POSITIONS: |
| 545 | first_from_raw = first_position >= chunk_start_position |
| 546 | raw_first_row = query_row_start + first_position - chunk_start_position |
| 547 | raw_position_values = tl.load( |
| 548 | raw_positions_ptr |
| 549 | + raw_first_row * stride_raw_positions_row |
| 550 | + position_dims * stride_raw_positions_dim, |
| 551 | mask=valid_row |
| 552 | & first_from_raw |
| 553 | & (raw_first_row >= query_row_start) |
| 554 | & (raw_first_row < query_row_end) |
| 555 | & (raw_first_row < num_rows) |
| 556 | & (position_dims < 3), |
| 557 | other=0, |
| 558 | ) |
| 559 | compressor_state_position_values = tl.load( |
| 560 | rope_cache_ptr |
| 561 | + tl.maximum(compressor_state_block, 0).to(tl.int64) * stride_rope_block |
| 562 | + (first_position % COMPRESSOR_STATE_SIZE) * stride_rope_token |
| 563 | + position_dims * stride_rope_dim, |
| 564 | mask=valid_row |
| 565 | & ~first_from_raw |
| 566 | & valid_compressor_state_block |
| 567 | & (position_dims < 3), |
| 568 | other=0, |
| 569 | ) |
| 570 | position_values = tl.where( |
| 571 | first_from_raw, |
| 572 | raw_position_values, |
| 573 | compressor_state_position_values, |
| 574 | ) |
| 575 | else: |
| 576 | position_values = tl.where(valid_row, first_position, 0) |
| 577 | tl.store( |
| 578 | first_positions_ptr |
| 579 | + row * stride_positions_row |
| 580 | + position_dims * stride_positions_dim, |
| 581 | position_values, |
| 582 | mask=(row < num_rows) & (position_dims < 3), |
| 583 | ) |
| 584 | |
| 585 | |
| 586 | def _validate_mqa(q: torch.Tensor) -> None: |
| 587 | if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0: |
| 588 | raise ValueError("QSA query must be [rows, heads, head_dim]") |
| 589 | |
| 590 | |
| 591 | def qsa_mqa_paged( |
| 592 | q: torch.Tensor, |
| 593 | k_cache: torch.Tensor, |
| 594 | page_table: torch.Tensor, |
| 595 | token_to_req: torch.Tensor, |
| 596 | query_positions: torch.Tensor, |
| 597 | sequence_lengths: torch.Tensor, |
| 598 | compress_ratio: int, |
| 599 | num_columns: int | None = None, |
| 600 | score_scale: float | None = None, |
| 601 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 602 | """Compute QSA scores directly from a paged compressed-key cache.""" |
| 603 | _validate_mqa(q) |
| 604 | if not q.is_cuda or not HAS_TRITON: |
| 605 | raise RuntimeError("paged QSA scoring requires a GPU and Triton") |
| 606 | if k_cache.ndim != 4 or k_cache.shape[2] != 1: |
| 607 | raise ValueError("QSA cache must be [pages, page_size, 1, head_dim]") |
| 608 | if k_cache.shape[3] != q.shape[2]: |
| 609 | raise ValueError("QSA query and cache dimensions must match") |
| 610 | if page_table.ndim != 2: |
| 611 | raise ValueError("QSA page table must be two-dimensional") |
| 612 | if q.shape[0] and (not all(k_cache.shape[:2]) or not all(page_table.shape)): |
| 613 | raise ValueError("QSA paged scoring cache and page table must be nonempty") |
| 614 | if token_to_req.shape != (q.shape[0],): |
| 615 | raise ValueError("QSA request mapping must match query rows") |
| 616 | if query_positions.shape != (q.shape[0],): |
| 617 | raise ValueError("QSA query positions must match query rows") |
| 618 | if sequence_lengths.shape != (page_table.shape[0],): |
| 619 | raise ValueError("QSA sequence lengths must match page-table requests") |
| 620 | if compress_ratio <= 0: |
| 621 | raise ValueError("QSA compression ratio must be positive") |
| 622 | score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale |
| 623 | if score_divisor <= 0: |
| 624 | raise ValueError("QSA score scale must be positive") |
| 625 | |
| 626 | capacity = page_table.shape[1] * k_cache.shape[1] |
| 627 | columns = capacity if num_columns is None else num_columns |
| 628 | if columns < 0: |
| 629 | raise ValueError("QSA score width must be non-negative") |
| 630 | logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) |
| 631 | visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) |
| 632 | if not q.shape[0] or not columns: |
| 633 | return logits, visible_blocks |
| 634 | block_n = 32 |
| 635 | _qsa_mqa_paged_kernel[(q.shape[0], triton.cdiv(columns, block_n))]( |
| 636 | q, |
| 637 | k_cache, |
| 638 | page_table, |
| 639 | token_to_req, |
| 640 | query_positions, |
| 641 | sequence_lengths, |
| 642 | visible_blocks, |
| 643 | logits, |
| 644 | q.stride(0), |
| 645 | q.stride(1), |
| 646 | q.stride(2), |
| 647 | k_cache.stride(0), |
| 648 | k_cache.stride(1), |
| 649 | k_cache.stride(3), |
| 650 | page_table.stride(0), |
| 651 | page_table.stride(1), |
| 652 | logits.stride(0), |
| 653 | q.shape[0], |
| 654 | columns, |
| 655 | k_cache.shape[0], |
| 656 | page_table.shape[0], |
| 657 | float(score_divisor), |
| 658 | PAGE_SIZE=k_cache.shape[1], |
| 659 | PAGE_TABLE_WIDTH=page_table.shape[1], |
| 660 | NUM_HEADS=q.shape[1], |
| 661 | HEAD_DIM=q.shape[2], |
| 662 | BLOCK_N=block_n, |
| 663 | BLOCK_D=triton.next_power_of_2(q.shape[2]), |
| 664 | COMPRESS_RATIO=compress_ratio, |
| 665 | num_warps=4, |
| 666 | ) |
| 667 | return logits, visible_blocks |
| 668 | |
| 669 | |
| 670 | def expand_qsa_block_indices_cuda( |
| 671 | block_indices: torch.Tensor, |
| 672 | query_positions: torch.Tensor, |
| 673 | sequence_lengths: torch.Tensor, |
| 674 | token_to_req: torch.Tensor, |
| 675 | compress_ratio: int, |
| 676 | token_topk: int, |
| 677 | out: torch.Tensor | None = None, |
| 678 | ) -> torch.Tensor: |
| 679 | """Expand compressed blocks and compact the causal tail of the open group.""" |
| 680 | if not block_indices.is_cuda or not HAS_TRITON: |
| 681 | raise RuntimeError("QSA index expansion requires a GPU and Triton") |
| 682 | if token_topk % compress_ratio: |
| 683 | raise ValueError("QSA token top-k must be divisible by compression ratio") |
| 684 | block_topk = token_topk // compress_ratio |
| 685 | output_width = token_topk + compress_ratio - 1 |
| 686 | if block_indices.shape != (query_positions.numel(), block_topk): |
| 687 | raise ValueError("QSA compressed top-k has an invalid shape") |
| 688 | if token_to_req.shape != query_positions.shape: |
| 689 | raise ValueError("QSA request mapping must match query positions") |
| 690 | if sequence_lengths.ndim != 1 or not sequence_lengths.shape[0]: |
| 691 | raise ValueError("QSA request sequence lengths must be nonempty") |
| 692 | if out is None: |
| 693 | out = torch.empty( |
| 694 | (block_indices.shape[0], output_width), |
| 695 | dtype=torch.int32, |
| 696 | device=block_indices.device, |
| 697 | ) |
| 698 | elif out.shape != (block_indices.shape[0], output_width): |
| 699 | raise ValueError("QSA expansion output has an invalid shape") |
| 700 | if not block_indices.shape[0]: |
| 701 | return out |
| 702 | column_block = 256 |
| 703 | _expand_qsa_indices_kernel[ |
| 704 | (block_indices.shape[0], triton.cdiv(output_width, column_block)) |
| 705 | ]( |
| 706 | block_indices, |
| 707 | query_positions, |
| 708 | sequence_lengths, |
| 709 | token_to_req, |
| 710 | out, |
| 711 | block_indices.stride(0), |
| 712 | block_indices.stride(1), |
| 713 | out.stride(0), |
| 714 | out.stride(1), |
| 715 | block_indices.shape[0], |
| 716 | sequence_lengths.shape[0], |
| 717 | BLOCK_TOPK=block_topk, |
| 718 | COMPRESS_RATIO=compress_ratio, |
| 719 | TOKEN_TOPK=token_topk, |
| 720 | OUTPUT_WIDTH=output_width, |
| 721 | COLUMN_BLOCK=column_block, |
| 722 | num_warps=4, |
| 723 | ) |
| 724 | return out |
| 725 | |
| 726 | |
| 727 | def qsa_select_paged_tokens( |
| 728 | q: torch.Tensor, |
| 729 | k_cache: torch.Tensor, |
| 730 | page_table: torch.Tensor, |
| 731 | token_to_req: torch.Tensor, |
| 732 | query_positions: torch.Tensor, |
| 733 | sequence_lengths: torch.Tensor, |
| 734 | token_topk: int, |
| 735 | compress_ratio: int, |
| 736 | out: torch.Tensor | None = None, |
| 737 | ) -> torch.Tensor: |
| 738 | """Score, select, and expand QSA indices without host synchronization.""" |
| 739 | rows = q.shape[0] |
| 740 | output_width = token_topk + compress_ratio - 1 |
| 741 | if out is None: |
| 742 | out = torch.empty((rows, output_width), dtype=torch.int32, device=q.device) |
| 743 | if out.shape != (rows, output_width): |
| 744 | raise ValueError("QSA selection output has an invalid shape") |
| 745 | if not rows: |
| 746 | return out |
| 747 | |
| 748 | columns = page_table.shape[1] * k_cache.shape[1] |
| 749 | block_topk = token_topk // compress_ratio |
| 750 | rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)) |
| 751 | chunk_rows = min(rows, rows_per_chunk) |
| 752 | blocks_buffer = torch.empty( |
| 753 | (chunk_rows, block_topk), dtype=torch.int32, device=q.device |
| 754 | ) |
| 755 | topk_workspace = torch.empty( |
| 756 | (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device |
| 757 | ) |
| 758 | for row_start in range(0, rows, rows_per_chunk): |
| 759 | row_end = min(row_start + rows_per_chunk, rows) |
| 760 | row_slice = slice(row_start, row_end) |
| 761 | logits, visible_blocks = qsa_mqa_paged( |
| 762 | q[row_slice], |
| 763 | k_cache, |
| 764 | page_table, |
| 765 | token_to_req[row_slice], |
| 766 | query_positions[row_slice], |
| 767 | sequence_lengths, |
| 768 | compress_ratio, |
| 769 | ) |
| 770 | blocks = blocks_buffer[: row_end - row_start] |
| 771 | use_cooperative_topk = ( |
| 772 | current_platform.is_cuda() |
| 773 | and blocks.shape[0] <= 32 |
| 774 | and logits.stride(0) % 4 == 0 |
| 775 | and current_platform.has_device_capability(90) |
| 776 | and not current_platform.is_device_capability_family(120) |
| 777 | ) |
| 778 | if use_cooperative_topk: |
| 779 | torch.ops._C.cooperative_topk( |
| 780 | logits, |
| 781 | visible_blocks, |
| 782 | blocks, |
| 783 | topk_workspace, |
| 784 | block_topk, |
| 785 | columns, |
| 786 | ) |
| 787 | elif current_platform.is_cuda(): |
| 788 | torch.ops._C.persistent_topk( |
| 789 | logits, |
| 790 | visible_blocks, |
| 791 | blocks, |
| 792 | topk_workspace, |
| 793 | block_topk, |
| 794 | columns, |
| 795 | ) |
| 796 | else: |
| 797 | ops.top_k_per_row_decode( |
| 798 | logits, |
| 799 | 1, |
| 800 | visible_blocks, |
| 801 | blocks, |
| 802 | blocks.shape[0], |
| 803 | logits.stride(0), |
| 804 | logits.stride(1), |
| 805 | block_topk, |
| 806 | ) |
| 807 | expand_qsa_block_indices_cuda( |
| 808 | blocks, |
| 809 | query_positions[row_slice], |
| 810 | sequence_lengths, |
| 811 | token_to_req[row_slice], |
| 812 | compress_ratio, |
| 813 | token_topk, |
| 814 | out[row_slice], |
| 815 | ) |
| 816 | return out |
| 817 | |
| 818 | |
| 819 | def qsa_sparse_paged_attention( |
| 820 | q: torch.Tensor, |
| 821 | k_cache: torch.Tensor, |
| 822 | v_cache: torch.Tensor, |
| 823 | logical_indices: torch.Tensor, |
| 824 | block_table: torch.Tensor, |
| 825 | token_to_req: torch.Tensor, |
| 826 | out: torch.Tensor | None = None, |
| 827 | ) -> torch.Tensor: |
| 828 | """Run sparse GQA directly over paged BF16 K/V caches.""" |
| 829 | if not q.is_cuda or not HAS_TRITON: |
| 830 | raise RuntimeError("paged QSA sparse attention requires a GPU and Triton") |
| 831 | if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: |
| 832 | raise ValueError("QSA sparse attention received invalid Q/K/V shapes") |
| 833 | if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]: |
| 834 | raise ValueError("QSA indices must have one row per query") |
| 835 | if token_to_req.shape != (q.shape[0],) or block_table.ndim != 2: |
| 836 | raise ValueError("QSA sparse attention metadata has invalid shapes") |
| 837 | if not all(k_cache.shape[:3]) or not all(block_table.shape): |
| 838 | raise ValueError("QSA sparse attention cache and block table must be nonempty") |
| 839 | if logical_indices.shape[1] <= 0: |
| 840 | raise ValueError("QSA sparse attention requires a positive selection width") |
| 841 | if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: |
| 842 | raise ValueError("QSA sparse attention requires valid grouped-query heads") |
| 843 | head_dim = q.shape[2] |
| 844 | assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 |
| 845 | assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16 |
| 846 | assert logical_indices.dtype == block_table.dtype == torch.int32 |
| 847 | assert token_to_req.dtype == torch.int32 |
| 848 | assert q.device == k_cache.device == v_cache.device |
| 849 | assert q.device == logical_indices.device == block_table.device |
| 850 | assert q.device == token_to_req.device |
| 851 | assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 |
| 852 | assert logical_indices.stride(1) == block_table.stride(1) == 1 |
| 853 | assert token_to_req.stride(0) == 1 |
| 854 | if out is None: |
| 855 | out = torch.empty_like(q) |
| 856 | if out.shape != q.shape: |
| 857 | raise ValueError("QSA sparse output must match its query") |
| 858 | assert out.dtype == q.dtype and out.device == q.device |
| 859 | assert out.stride(2) == 1 |
| 860 | if not q.shape[0]: |
| 861 | return out |
| 862 | |
| 863 | group_size = q.shape[1] // k_cache.shape[2] |
| 864 | block_m = triton.next_power_of_2(group_size) |
| 865 | base_programs = q.shape[0] * k_cache.shape[2] |
| 866 | small_profile_limit = 8 if block_m <= 8 else 4 |
| 867 | |
| 868 | # Tuned on GB300 for the Qwen-Air TP1, TP2, and TP4 attention shapes. |
| 869 | # Narrow tiles favor decode; wide tiles improve throughput for prefill. |
| 870 | if base_programs <= small_profile_limit: |
| 871 | block_n, target_splits, partial_warps = 16, 64, 4 |
| 872 | elif base_programs < 32: |
| 873 | block_n, target_splits, partial_warps = 16, 32, 4 |
| 874 | elif base_programs <= 256: |
| 875 | block_n, target_splits, partial_warps = 64, 8, 2 |
| 876 | elif base_programs <= 512: |
| 877 | block_n, target_splits, partial_warps = 64, 4, 2 |
| 878 | else: |
| 879 | block_n, target_splits, partial_warps = 64, 1, 2 |
| 880 | # gfx942 and gfx950 have a 64 KiB LDS limit. One software-pipelining |
| 881 | # stage keeps the wide TP4 tile within that shared-memory budget. |
| 882 | partial_stages = 1 if current_platform.is_rocm() else 2 |
| 883 | |
| 884 | num_tiles = triton.cdiv(logical_indices.shape[1], block_n) |
| 885 | # Avoid empty splits when the selection width is smaller than the profile. |
| 886 | max_useful_splits = 1 << (num_tiles.bit_length() - 1) |
| 887 | num_splits = min(max_useful_splits, target_splits) |
| 888 | |
| 889 | # Split=1 writes output directly and compiles out all workspace accesses. |
| 890 | if num_splits == 1: |
| 891 | partial_output = out |
| 892 | partial_lse = out |
| 893 | else: |
| 894 | # FP32 partials preserve accuracy when merging independently normalized |
| 895 | # splits. |
| 896 | partial_output = torch.empty( |
| 897 | (num_splits, *q.shape), dtype=torch.float32, device=q.device |
| 898 | ) |
| 899 | partial_lse = torch.empty( |
| 900 | (num_splits, q.shape[0], q.shape[1]), |
| 901 | dtype=torch.float32, |
| 902 | device=q.device, |
| 903 | ) |
| 904 | |
| 905 | partial_grid = (q.shape[0], k_cache.shape[2], num_splits) |
| 906 | _qsa_sparse_paged_gqa_splitk_kernel[partial_grid]( |
| 907 | q, |
| 908 | k_cache, |
| 909 | v_cache, |
| 910 | logical_indices, |
| 911 | block_table, |
| 912 | token_to_req, |
| 913 | partial_output, |
| 914 | partial_lse, |
| 915 | out, |
| 916 | q.stride(0), |
| 917 | q.stride(1), |
| 918 | k_cache.stride(0), |
| 919 | k_cache.stride(1), |
| 920 | k_cache.stride(2), |
| 921 | v_cache.stride(0), |
| 922 | v_cache.stride(1), |
| 923 | v_cache.stride(2), |
| 924 | logical_indices.stride(0), |
| 925 | block_table.stride(0), |
| 926 | out.stride(0), |
| 927 | out.stride(1), |
| 928 | q.shape[0], |
| 929 | k_cache.shape[0], |
| 930 | block_table.shape[0], |
| 931 | TOPK=logical_indices.shape[1], |
| 932 | PAGE_SIZE=k_cache.shape[1], |
| 933 | PAGE_TABLE_WIDTH=block_table.shape[1], |
| 934 | GROUP_SIZE=group_size, |
| 935 | HEAD_DIM=q.shape[2], |
| 936 | NUM_QUERY_HEADS=q.shape[1], |
| 937 | NUM_SPLITS=num_splits, |
| 938 | NUM_TILES=num_tiles, |
| 939 | BLOCK_M=block_m, |
| 940 | BLOCK_N=block_n, |
| 941 | num_warps=partial_warps, |
| 942 | num_stages=partial_stages, |
| 943 | ) |
| 944 | if num_splits == 1: |
| 945 | return out |
| 946 | |
| 947 | _qsa_merge_splitk_kernel[(q.shape[0], q.shape[1])]( |
| 948 | partial_output, |
| 949 | partial_lse, |
| 950 | out, |
| 951 | out.stride(0), |
| 952 | out.stride(1), |
| 953 | q.shape[0], |
| 954 | HEAD_DIM=q.shape[2], |
| 955 | NUM_QUERY_HEADS=q.shape[1], |
| 956 | NUM_SPLITS=num_splits, |
| 957 | BLOCK_SPLITS=triton.next_power_of_2(num_splits), |
| 958 | num_warps=2, |
| 959 | num_stages=1, |
| 960 | ) |
| 961 | return out |
| 962 | |
| 963 | |
| 964 | def qsa_store_cache_rows( |
| 965 | cache: torch.Tensor, |
| 966 | slot_mapping: torch.Tensor, |
| 967 | rows: torch.Tensor, |
| 968 | ) -> None: |
| 969 | """Store fixed-width rows in a QSA cache without boolean indexing.""" |
| 970 | if not cache.is_cuda or not HAS_TRITON: |
| 971 | raise RuntimeError("QSA cache stores require a GPU and Triton") |
| 972 | if cache.ndim != 4 or cache.shape[2] != 1: |
| 973 | raise ValueError("QSA cache must be [pages, page_size, 1, width]") |
| 974 | if not all(cache.shape): |
| 975 | raise ValueError("QSA cache dimensions must be nonzero") |
| 976 | if rows.ndim == 3: |
| 977 | if rows.shape[1] != 1: |
| 978 | raise ValueError("QSA cache rows must have one head") |
| 979 | rows = rows[:, 0] |
| 980 | if rows.shape != (slot_mapping.numel(), cache.shape[3]): |
| 981 | raise ValueError("QSA cache rows and slots have incompatible shapes") |
| 982 | if not rows.shape[0]: |
| 983 | return |
| 984 | _store_qsa_rows_kernel[(rows.shape[0],)]( |
| 985 | cache, |
| 986 | slot_mapping, |
| 987 | rows, |
| 988 | cache.stride(0), |
| 989 | cache.stride(1), |
| 990 | cache.stride(3), |
| 991 | rows.stride(0), |
| 992 | rows.stride(1), |
| 993 | rows.shape[0], |
| 994 | cache.shape[0], |
| 995 | PAGE_SIZE=cache.shape[1], |
| 996 | WIDTH=cache.shape[3], |
| 997 | BLOCK_D=triton.next_power_of_2(cache.shape[3]), |
| 998 | num_warps=4, |
| 999 | ) |
| 1000 | |
| 1001 | |
| 1002 | def qsa_compress_groups_with_ratio( |
| 1003 | raw_keys: torch.Tensor, # this step's raw key rows [rows, 1, head_size] |
| 1004 | raw_positions: torch.Tensor, # this step's positions [rows, 1, 3] int64 |
| 1005 | compressor_state_cache: torch.Tensor, |
| 1006 | compressor_state_block_table: torch.Tensor, |
| 1007 | token_to_req: torch.Tensor, |
| 1008 | query_start_loc: torch.Tensor, |
| 1009 | logical_positions: torch.Tensor, |
| 1010 | compressed_slots: torch.Tensor, |
| 1011 | compress_ratio: int, |
| 1012 | rope_cache: torch.Tensor | None = None, |
| 1013 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1014 | """Pool completed groups from the compressor-state ring and raw token rows.""" |
| 1015 | if not raw_keys.is_cuda or not HAS_TRITON: |
| 1016 | raise RuntimeError("QSA compression requires a GPU and Triton") |
| 1017 | rows = token_to_req.numel() |
| 1018 | if compress_ratio <= 0: |
| 1019 | raise ValueError("QSA compression ratio must be positive") |
| 1020 | if raw_keys.ndim != 3 or raw_keys.shape[:2] != (rows, 1): |
| 1021 | raise ValueError("QSA raw keys must be [rows, 1, head_size]") |
| 1022 | if raw_positions.shape != (rows, 1, 3) or raw_positions.dtype != torch.int64: |
| 1023 | raise ValueError("QSA raw positions must be [rows, 1, 3] int64") |
| 1024 | if logical_positions.shape != (rows,) or compressed_slots.shape != (rows,): |
| 1025 | raise ValueError("QSA compression metadata must match token rows") |
| 1026 | if compressor_state_cache.ndim != 4 or compressor_state_cache.shape[2] != 1: |
| 1027 | raise ValueError("QSA compressor-state cache has an invalid shape") |
| 1028 | if ( |
| 1029 | # The ring is wider than one group so speculative rows cannot alias |
| 1030 | # onto the committed keys of the group still being collected. |
| 1031 | compressor_state_cache.shape[1] < compress_ratio |
| 1032 | or compressor_state_cache.shape[3] != raw_keys.shape[2] |
| 1033 | or compressor_state_cache.dtype != raw_keys.dtype |
| 1034 | ): |
| 1035 | raise ValueError( |
| 1036 | "QSA compressor-state cache does not match the compression layout" |
| 1037 | ) |
| 1038 | if ( |
| 1039 | compressor_state_block_table.ndim != 2 |
| 1040 | or compressor_state_block_table.shape[1] < 1 |
| 1041 | ): |
| 1042 | raise ValueError( |
| 1043 | "QSA compressor-state block table must contain one block per request" |
| 1044 | ) |
| 1045 | if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2: |
| 1046 | raise ValueError("QSA query starts must contain a terminal offset") |
| 1047 | num_requests = query_start_loc.shape[0] - 1 |
| 1048 | if compressor_state_block_table.shape[0] < num_requests: |
| 1049 | raise ValueError("QSA compressor-state block table has too few request rows") |
| 1050 | if rope_cache is not None and ( |
| 1051 | rope_cache.ndim != 4 |
| 1052 | or rope_cache.shape[:3] != compressor_state_cache.shape[:3] |
| 1053 | or rope_cache.shape[3] != 3 |
| 1054 | or rope_cache.dtype != torch.int64 |
| 1055 | ): |
| 1056 | raise ValueError("QSA packed position view has an invalid shape or dtype") |
| 1057 | if rows and ( |
| 1058 | not all(compressor_state_cache.shape) |
| 1059 | or not all(compressor_state_block_table.shape) |
| 1060 | ): |
| 1061 | raise ValueError("QSA compressor-state cache and block table must be nonempty") |
| 1062 | pooled = torch.empty( |
| 1063 | (rows, 1, raw_keys.shape[2]), |
| 1064 | dtype=raw_keys.dtype, |
| 1065 | device=raw_keys.device, |
| 1066 | ) |
| 1067 | first_positions = torch.empty((rows, 3), dtype=torch.int64, device=raw_keys.device) |
| 1068 | if not rows: |
| 1069 | return pooled, first_positions |
| 1070 | if rope_cache is None: |
| 1071 | rope_cache = compressor_state_cache |
| 1072 | load_rope_positions = False |
| 1073 | else: |
| 1074 | load_rope_positions = True |
| 1075 | _compress_qsa_groups_kernel[(rows,)]( |
| 1076 | raw_keys, |
| 1077 | raw_positions, |
| 1078 | compressor_state_cache, |
| 1079 | rope_cache, |
| 1080 | compressor_state_block_table, |
| 1081 | token_to_req, |
| 1082 | query_start_loc, |
| 1083 | logical_positions, |
| 1084 | compressed_slots, |
| 1085 | pooled, |
| 1086 | first_positions, |
| 1087 | raw_keys.stride(0), |
| 1088 | raw_keys.stride(2), |
| 1089 | raw_positions.stride(0), |
| 1090 | raw_positions.stride(2), |
| 1091 | compressor_state_cache.stride(0), |
| 1092 | compressor_state_cache.stride(1), |
| 1093 | compressor_state_cache.stride(3), |
| 1094 | rope_cache.stride(0), |
| 1095 | rope_cache.stride(1), |
| 1096 | rope_cache.stride(3), |
| 1097 | compressor_state_block_table.stride(0), |
| 1098 | pooled.stride(0), |
| 1099 | pooled.stride(2), |
| 1100 | first_positions.stride(0), |
| 1101 | first_positions.stride(1), |
| 1102 | rows, |
| 1103 | compressor_state_cache.shape[0], |
| 1104 | num_requests, |
| 1105 | COMPRESSOR_STATE_SIZE=compressor_state_cache.shape[1], |
| 1106 | COMPRESS_RATIO=compress_ratio, |
| 1107 | HEAD_DIM=raw_keys.shape[2], |
| 1108 | LOAD_ROPE_POSITIONS=load_rope_positions, |
| 1109 | BLOCK_D=triton.next_power_of_2(raw_keys.shape[2]), |
| 1110 | num_warps=4, |
| 1111 | ) |
| 1112 | return pooled, first_positions |
| 1113 | |
| 1114 | |
| 1115 | __all__ = [ |
| 1116 | "expand_qsa_block_indices_cuda", |
| 1117 | "qsa_compress_groups_with_ratio", |
| 1118 | "qsa_mqa_paged", |
| 1119 | "qsa_select_paged_tokens", |
| 1120 | "qsa_sparse_paged_attention", |
| 1121 | "qsa_store_cache_rows", |
| 1122 | ] |