@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 659 lines 22.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Triton kernels for Qwen4Exp QSA index selection."""
4
5 import torch
6
7 import vllm.envs as envs
8 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
9 TritonWarmupTensor,
10 )
11 from vllm.platforms import current_platform
12 from vllm.triton_utils import tl, triton
13
14 _TOPK_WORKSPACE_BYTES = 1024 * 1024
15 _DECODE_BLOCK_N = 64
16
17
18 @triton.jit
19 def _qsa_mqa_paged_uniform_kernel(
20 q_ptr,
21 k_cache_ptr,
22 page_table_ptr,
23 visible_blocks_ptr,
24 logits_ptr,
25 stride_q_row,
26 stride_q_head,
27 stride_cache_block,
28 stride_cache_token,
29 stride_table_req,
30 stride_logits_row,
31 PAGE_SIZE: tl.constexpr,
32 PAGE_TABLE_WIDTH: tl.constexpr,
33 NUM_HEADS: tl.constexpr,
34 HEAD_DIM: tl.constexpr,
35 DECODE_QUERY_LEN: tl.constexpr,
36 BLOCK_N: tl.constexpr,
37 TILES_PER_PROG: tl.constexpr,
38 STAGES: tl.constexpr,
39 ) -> None:
40 NUM_COLUMNS: tl.constexpr = PAGE_TABLE_WIDTH * PAGE_SIZE
41 DECODE_QUERY_LEN_PADDED: tl.constexpr = triton.next_power_of_2(DECODE_QUERY_LEN)
42 NUM_HEADS_PADDED: tl.constexpr = triton.next_power_of_2(NUM_HEADS)
43 # tl.dot requires a reduction dimension of at least 16.
44 BLOCK_D: tl.constexpr = max(16, triton.next_power_of_2(HEAD_DIM))
45 request = tl.program_id(0)
46 tile_start = tl.program_id(1) * TILES_PER_PROG
47 query_offsets = tl.arange(0, DECODE_QUERY_LEN_PADDED)
48 valid_query_offsets = query_offsets < DECODE_QUERY_LEN
49 rows = request * DECODE_QUERY_LEN + query_offsets
50 visible = tl.load(
51 visible_blocks_ptr + rows,
52 mask=valid_query_offsets,
53 other=0,
54 )
55 max_visible = tl.max(visible, axis=0)
56 if tile_start * BLOCK_N >= max_visible:
57 return
58 tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(max_visible, BLOCK_N))
59 tile_end = tl.minimum(tile_end, tl.cdiv(NUM_COLUMNS, BLOCK_N))
60
61 dims = tl.arange(0, BLOCK_D)
62 n = tl.arange(0, DECODE_QUERY_LEN_PADDED * NUM_HEADS_PADDED)
63 query_offset = n // NUM_HEADS_PADDED
64 head = n % NUM_HEADS_PADDED
65 valid_query = (query_offset < DECODE_QUERY_LEN) & (head < NUM_HEADS)
66 query = tl.load(
67 q_ptr
68 + (request * DECODE_QUERY_LEN + query_offset)[None, :] * stride_q_row
69 + head[None, :] * stride_q_head
70 + dims[:, None],
71 mask=valid_query[None, :] & (dims[:, None] < HEAD_DIM),
72 other=0.0,
73 )
74 column_offsets = tl.arange(0, BLOCK_N)
75 for tile in tl.range(tile_start, tile_end, num_stages=STAGES):
76 columns = tile * BLOCK_N + column_offsets
77 live = columns < max_visible
78 logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1)
79 page_offset = columns % PAGE_SIZE
80 physical_page = tl.load(
81 page_table_ptr + request * stride_table_req + logical_page,
82 mask=live,
83 other=0,
84 )
85 keys = tl.load(
86 k_cache_ptr
87 + physical_page[:, None].to(tl.int64) * stride_cache_block
88 + page_offset[:, None] * stride_cache_token
89 + dims[None, :],
90 mask=live[:, None] & (dims[None, :] < HEAD_DIM),
91 other=0.0,
92 eviction_policy="evict_first",
93 )
94 scores = tl.dot(keys, query, out_dtype=tl.float32)
95 scores = tl.where(valid_query[None, :], tl.maximum(scores, 0.0), 0.0)
96 scores = tl.reshape(
97 scores,
98 (BLOCK_N, DECODE_QUERY_LEN_PADDED, NUM_HEADS_PADDED),
99 )
100 score = tl.sum(scores, axis=2)
101 tl.store(
102 logits_ptr + rows[None, :] * stride_logits_row + columns[:, None],
103 score,
104 mask=valid_query_offsets[None, :]
105 & (columns[:, None] < NUM_COLUMNS)
106 & (columns[:, None] < visible[None, :]),
107 )
108
109
110 @triton.jit(do_not_specialize=["num_rows", "query_offset"])
111 def _qsa_mqa_paged_prefill_kernel(
112 q_ptr,
113 k_cache_ptr,
114 page_table_ptr,
115 query_start_loc_ptr,
116 visible_blocks_ptr,
117 logits_ptr,
118 stride_q_row,
119 stride_q_head,
120 stride_cache_block,
121 stride_cache_token,
122 stride_table_req,
123 stride_logits_row,
124 num_rows,
125 query_offset,
126 page_table_width,
127 PAGE_SIZE: tl.constexpr,
128 NUM_HEADS: tl.constexpr,
129 HEAD_DIM: tl.constexpr,
130 TILE_R: tl.constexpr,
131 BLOCK_N: tl.constexpr,
132 K_TILES: tl.constexpr,
133 STAGES: tl.constexpr,
134 ) -> None:
135 num_columns = page_table_width * PAGE_SIZE
136 NUM_HEADS_PADDED: tl.constexpr = triton.next_power_of_2(NUM_HEADS)
137 # tl.dot requires a reduction dimension of at least 16.
138 BLOCK_D: tl.constexpr = max(16, triton.next_power_of_2(HEAD_DIM))
139 request = tl.program_id(0)
140 query_base = tl.load(query_start_loc_ptr)
141 query_end = query_offset + num_rows
142 request_start = tl.maximum(
143 tl.load(query_start_loc_ptr + request) - query_base, query_offset
144 )
145 request_end = tl.minimum(
146 tl.load(query_start_loc_ptr + request + 1) - query_base, query_end
147 )
148 absolute_row_start = request_start + tl.program_id(1) * TILE_R
149 if absolute_row_start >= request_end:
150 return
151
152 lanes = tl.arange(0, TILE_R)
153 absolute_rows = absolute_row_start + lanes
154 rows = absolute_rows - query_offset
155 valid_rows = absolute_rows < request_end
156 visible = tl.load(
157 visible_blocks_ptr + absolute_rows,
158 mask=valid_rows,
159 other=0,
160 )
161 max_visible = tl.max(visible, axis=0)
162 k_tile_start = tl.program_id(2) * K_TILES
163 if k_tile_start * BLOCK_N >= max_visible:
164 return
165 k_tile_end = tl.minimum(k_tile_start + K_TILES, tl.cdiv(max_visible, BLOCK_N))
166 k_tile_end = tl.minimum(k_tile_end, tl.cdiv(num_columns, BLOCK_N))
167
168 dims = tl.arange(0, BLOCK_D)
169 m = tl.arange(0, TILE_R * NUM_HEADS_PADDED)
170 q_row_offsets = m // NUM_HEADS_PADDED
171 q_rows = absolute_row_start + q_row_offsets
172 heads = m % NUM_HEADS_PADDED
173 query = tl.load(
174 q_ptr
175 + q_rows[None, :] * stride_q_row
176 + heads[None, :] * stride_q_head
177 + dims[:, None],
178 mask=(heads[None, :] < NUM_HEADS)
179 & (absolute_row_start + q_row_offsets[None, :] < request_end)
180 & (dims[:, None] < HEAD_DIM),
181 other=0.0,
182 )
183 column_offsets = tl.arange(0, BLOCK_N)
184 for tile in tl.range(k_tile_start, k_tile_end, num_stages=STAGES):
185 columns = tile * BLOCK_N + column_offsets
186 live = columns < max_visible
187 logical_page = tl.minimum(columns // PAGE_SIZE, page_table_width - 1)
188 page_offset = columns % PAGE_SIZE
189 physical_page = tl.load(
190 page_table_ptr + request * stride_table_req + logical_page,
191 mask=live,
192 other=0,
193 )
194 keys = tl.load(
195 k_cache_ptr
196 + physical_page[:, None].to(tl.int64) * stride_cache_block
197 + page_offset[:, None] * stride_cache_token
198 + dims[None, :],
199 mask=live[:, None] & (dims[None, :] < HEAD_DIM),
200 other=0.0,
201 eviction_policy="evict_first",
202 )
203 # With an fp8 cache, SM90 lowers this dot to reduced-precision wgmma
204 # accumulation; SM100 tcgen05 accumulates in fp32. The SM90 error
205 # (~1e-4 relative) is far below the e4m3 quantization noise floor.
206 scores = tl.dot(keys, query, out_dtype=tl.float32)
207 scores = tl.reshape(scores, (BLOCK_N, TILE_R, NUM_HEADS_PADDED))
208 score = tl.sum(tl.maximum(scores, 0.0), axis=2)
209 store_mask = (
210 valid_rows[None, :]
211 & (columns[:, None] < visible[None, :])
212 & (columns[:, None] < num_columns)
213 )
214 tl.store(
215 logits_ptr + rows[None, :] * stride_logits_row + columns[:, None],
216 score,
217 mask=store_mask,
218 )
219
220
221 @triton.jit
222 def _expand_qsa_indices_kernel(
223 block_indices_ptr,
224 query_positions_ptr,
225 visible_blocks_ptr,
226 output_ptr,
227 stride_blocks_row,
228 stride_blocks_column,
229 stride_output_row,
230 stride_output_column,
231 BLOCK_TOPK: tl.constexpr,
232 COMPRESS_RATIO: tl.constexpr,
233 COLUMN_BLOCK: tl.constexpr,
234 ) -> None:
235 OUTPUT_WIDTH: tl.constexpr = BLOCK_TOPK * COMPRESS_RATIO + COMPRESS_RATIO - 1
236 row = tl.program_id(0)
237 columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK)
238 query_position = tl.load(query_positions_ptr + row)
239 visible_blocks = tl.load(visible_blocks_ptr + row)
240 complete_blocks = tl.minimum(visible_blocks, BLOCK_TOPK)
241 expanded_count = complete_blocks * COMPRESS_RATIO
242 tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO
243 tail_count = (query_position + 1) - tail_start
244
245 is_expanded = columns < expanded_count
246 block_rank = columns // COMPRESS_RATIO
247 offset = columns % COMPRESS_RATIO
248 safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1)
249 block = tl.load(
250 block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column,
251 mask=is_expanded,
252 other=-1,
253 )
254 expanded = block * COMPRESS_RATIO + offset
255 tail_offset = columns - expanded_count
256 is_tail = (
257 (columns >= expanded_count)
258 & (tail_offset < tail_count)
259 & (tail_offset < COMPRESS_RATIO - 1)
260 )
261 token = tl.where(is_expanded, expanded, tail_start + tail_offset)
262 valid = (columns < OUTPUT_WIDTH) & (is_expanded | is_tail) & (token >= 0)
263 tl.store(
264 output_ptr + row * stride_output_row + columns * stride_output_column,
265 tl.where(valid, token, -1),
266 mask=columns < OUTPUT_WIDTH,
267 )
268 # TRAILING COUNT COLUMN: the packed buffer is OUTPUT_WIDTH+1 wide; the
269 # last column holds this row's valid-entry count (expanded blocks plus
270 # causal tail). It is NOT a token index — the sparse attention kernel
271 # reads it as the row's tile-loop bound.
272 if tl.program_id(1) == 0:
273 tl.store(
274 output_ptr + row * stride_output_row + OUTPUT_WIDTH * stride_output_column,
275 expanded_count + tail_count,
276 )
277
278
279 def _decode_tiles_per_program(num_requests: int, columns: int) -> int:
280 programs = num_requests * triton.cdiv(columns, _DECODE_BLOCK_N)
281 if programs < 16384:
282 return 1
283 if programs < 32768:
284 return 2
285 if programs < 131072:
286 return 4
287 return 8
288
289
290 def _qsa_decode_warmup_profiles(
291 max_dql: int,
292 max_num_reqs: int,
293 max_num_batched_tokens: int,
294 columns: int,
295 ) -> tuple[tuple[int, int], ...]:
296 profiles: list[tuple[int, int]] = []
297 for dql in range(1, max_dql + 1):
298 max_reqs = min(max_num_reqs, max_num_batched_tokens // dql)
299 requests_by_grouping: dict[int, int] = {}
300 for num_requests in range(1, max_reqs + 1):
301 tiles_per_program = _decode_tiles_per_program(num_requests, columns)
302 requests_by_grouping.setdefault(tiles_per_program, num_requests)
303 profiles.extend(
304 (dql, num_requests) for num_requests in requests_by_grouping.values()
305 )
306 return tuple(profiles)
307
308
309 def warmup_qsa_mqa_paged_decode(
310 k_cache: torch.Tensor,
311 page_table: torch.Tensor,
312 *,
313 num_heads: int,
314 head_dim: int,
315 max_decode_query_len: int,
316 max_num_reqs: int,
317 max_num_batched_tokens: int,
318 ) -> tuple[tuple[int, int], ...]:
319 """Compile every reachable decode specialization without launching it."""
320 page_size = k_cache.shape[1]
321 page_table_width = page_table.shape[1]
322 columns = page_table_width * page_size
323 profiles = _qsa_decode_warmup_profiles(
324 max_decode_query_len,
325 max_num_reqs,
326 max_num_batched_tokens,
327 columns,
328 )
329 if not profiles:
330 return ()
331
332 k_cache_ptr = TritonWarmupTensor(k_cache.dtype, shape=tuple(k_cache.shape))
333 page_table_ptr = TritonWarmupTensor(
334 page_table.dtype,
335 shape=(max_num_reqs, page_table_width),
336 )
337 visible_blocks_ptr = TritonWarmupTensor(torch.int32)
338
339 for decode_query_len, num_requests in profiles:
340 num_rows = decode_query_len * num_requests
341 q_ptr = TritonWarmupTensor(
342 k_cache.dtype,
343 shape=(num_rows, num_heads, head_dim),
344 )
345 logits_ptr = TritonWarmupTensor(
346 torch.float32,
347 shape=(num_rows, columns),
348 )
349 tiles_per_program = _decode_tiles_per_program(num_requests, columns)
350 _qsa_mqa_paged_uniform_kernel.warmup(
351 q_ptr,
352 k_cache_ptr,
353 page_table_ptr,
354 visible_blocks_ptr,
355 logits_ptr,
356 num_heads * head_dim,
357 head_dim,
358 k_cache.stride(0),
359 k_cache.stride(1),
360 page_table.stride(0),
361 columns,
362 PAGE_SIZE=page_size,
363 PAGE_TABLE_WIDTH=page_table_width,
364 NUM_HEADS=num_heads,
365 HEAD_DIM=head_dim,
366 DECODE_QUERY_LEN=decode_query_len,
367 BLOCK_N=_DECODE_BLOCK_N,
368 TILES_PER_PROG=tiles_per_program,
369 STAGES=2,
370 # tuned on GB300
371 num_warps=1 if k_cache.dtype == torch.float8_e4m3fn else 2,
372 grid=(
373 num_requests,
374 triton.cdiv(columns, _DECODE_BLOCK_N * tiles_per_program),
375 ),
376 )
377 return profiles
378
379
380 def _prefill_logits(
381 q: torch.Tensor,
382 k_cache: torch.Tensor,
383 page_table: torch.Tensor,
384 query_start_loc: torch.Tensor,
385 visible_blocks: torch.Tensor,
386 logits_workspace: torch.Tensor,
387 max_query_len: int,
388 logits_width: int,
389 query_offset: int,
390 num_queries: int,
391 ) -> torch.Tensor:
392 assert query_start_loc.shape == (page_table.shape[0] + 1,)
393 assert visible_blocks.shape == (q.shape[0],)
394 assert 0 <= query_offset <= query_offset + num_queries <= q.shape[0]
395 assert 0 < logits_width <= page_table.shape[1] * k_cache.shape[1]
396 assert logits_workspace.is_contiguous()
397 assert logits_workspace.numel() >= num_queries * logits_width
398
399 # slice from pre-allocated workspace
400 logits = logits_workspace[: num_queries * logits_width].view(
401 num_queries, logits_width
402 )
403
404 # tuned on GB300
405 if k_cache.dtype == torch.float8_e4m3fn:
406 TILE_R, STAGES, num_warps = 32, 2, 8
407 else:
408 TILE_R, STAGES, num_warps = 64, 2, 4
409 BLOCK_N = 64
410 K_TILES = 16
411 grid = (
412 page_table.shape[0],
413 triton.cdiv(min(num_queries, max_query_len), TILE_R),
414 triton.cdiv(logits_width, BLOCK_N * K_TILES),
415 )
416 _qsa_mqa_paged_prefill_kernel[grid](
417 q,
418 k_cache,
419 page_table,
420 query_start_loc,
421 visible_blocks,
422 logits,
423 *q.stride()[:-1],
424 *k_cache.stride()[:2],
425 *page_table.stride()[:-1],
426 *logits.stride()[:-1],
427 num_queries,
428 query_offset,
429 page_table.shape[1],
430 PAGE_SIZE=k_cache.shape[1],
431 NUM_HEADS=q.shape[1],
432 HEAD_DIM=q.shape[2],
433 TILE_R=TILE_R,
434 BLOCK_N=BLOCK_N,
435 K_TILES=K_TILES,
436 STAGES=STAGES,
437 num_warps=num_warps,
438 )
439 return logits
440
441
442 def expand_qsa_block_indices(
443 block_indices: torch.Tensor,
444 query_positions: torch.Tensor,
445 visible_blocks: torch.Tensor,
446 compress_ratio: int,
447 token_topk: int,
448 out: torch.Tensor,
449 ) -> None:
450 """Expand compressed blocks and compact the causal tail of the open group."""
451 assert token_topk % compress_ratio == 0
452 block_topk = token_topk // compress_ratio
453 output_width = token_topk + compress_ratio - 1
454 assert block_indices.shape == (query_positions.numel(), block_topk)
455 assert visible_blocks.shape == query_positions.shape
456 # +1: the packed buffer's trailing column holds each row's valid-entry
457 # count (never a token index); the index region below only writes
458 # columns [0, output_width).
459 assert out.shape == (block_indices.shape[0], output_width + 1)
460 column_block = 256
461 grid = (block_indices.shape[0], triton.cdiv(output_width, column_block))
462 _expand_qsa_indices_kernel[grid](
463 block_indices,
464 query_positions,
465 visible_blocks,
466 out,
467 *block_indices.stride(),
468 *out.stride(),
469 BLOCK_TOPK=block_topk,
470 COMPRESS_RATIO=compress_ratio,
471 COLUMN_BLOCK=column_block,
472 num_warps=4,
473 )
474
475
476 def _topk(
477 logits: torch.Tensor,
478 visible_blocks: torch.Tensor,
479 token_topk: int,
480 compress_ratio: int,
481 block_indices: torch.Tensor,
482 topk_workspace: torch.Tensor,
483 ) -> None:
484 # similar dispatch logic as DeepSeek indexer
485 block_topk = token_topk // compress_ratio
486 use_cooperative_topk = (
487 logits.shape[0] <= 64
488 and logits.stride(0) % 4 == 0
489 and current_platform.has_device_capability(90)
490 and not current_platform.is_device_capability_family(120)
491 )
492 topk_op = (
493 torch.ops._C.cooperative_topk
494 if use_cooperative_topk
495 else torch.ops._C.persistent_topk
496 )
497 topk_op(
498 logits,
499 visible_blocks,
500 block_indices,
501 topk_workspace,
502 block_topk,
503 logits.shape[1],
504 )
505
506
507 def qsa_select_paged_decode(
508 q: torch.Tensor,
509 k_cache: torch.Tensor,
510 page_table: torch.Tensor,
511 visible_blocks: torch.Tensor,
512 token_topk: int,
513 compress_ratio: int,
514 decode_query_len: int,
515 block_indices: torch.Tensor,
516 ) -> None:
517 """Score and select compressed blocks for a request-major decode batch.
518
519 Args:
520 q: Query tensor shaped ``[num_requests * decode_query_len, heads,
521 head_dim]``.
522 k_cache: Compressed key cache shaped ``[blocks, page_size, 1,
523 head_dim]``.
524 page_table: Request block table shaped ``[num_requests, max_pages]``.
525 visible_blocks: Number of visible compressed blocks per query.
526 token_topk: Number of logical tokens selected per query.
527 compress_ratio: Number of logical tokens represented by a cache row.
528 decode_query_len: Number of query tokens per request.
529 block_indices: Compressed-index output buffer.
530
531 """
532 assert token_topk % compress_ratio == 0
533 assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)
534 assert decode_query_len > 0 and q.shape[0] % decode_query_len == 0
535 assert q.dtype == k_cache.dtype, "Q and the compressed K cache must match"
536 num_requests = q.shape[0] // decode_query_len
537 assert page_table.shape[0] == num_requests
538 assert visible_blocks.shape == (q.shape[0],)
539
540 columns = page_table.shape[1] * k_cache.shape[1]
541 logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device)
542 tiles_per_program = _decode_tiles_per_program(num_requests, columns)
543 grid = (
544 num_requests,
545 triton.cdiv(columns, _DECODE_BLOCK_N * tiles_per_program),
546 )
547 _qsa_mqa_paged_uniform_kernel[grid](
548 q,
549 k_cache,
550 page_table,
551 visible_blocks,
552 logits,
553 *q.stride()[:-1],
554 *k_cache.stride()[:2],
555 *page_table.stride()[:-1],
556 *logits.stride()[:-1],
557 PAGE_SIZE=k_cache.shape[1],
558 PAGE_TABLE_WIDTH=page_table.shape[1],
559 NUM_HEADS=q.shape[1],
560 HEAD_DIM=q.shape[2],
561 DECODE_QUERY_LEN=decode_query_len,
562 BLOCK_N=_DECODE_BLOCK_N,
563 TILES_PER_PROG=tiles_per_program,
564 STAGES=2,
565 # tuned on GB300
566 num_warps=1 if k_cache.dtype == torch.float8_e4m3fn else 2,
567 )
568 _topk(
569 logits,
570 visible_blocks,
571 token_topk,
572 compress_ratio,
573 block_indices,
574 torch.empty((_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device),
575 )
576
577
578 def qsa_select_paged_prefill(
579 q: torch.Tensor,
580 k_cache: torch.Tensor,
581 page_table: torch.Tensor,
582 query_start_loc: torch.Tensor,
583 visible_blocks: torch.Tensor,
584 token_topk: int,
585 compress_ratio: int,
586 max_query_len: int,
587 block_indices: torch.Tensor,
588 max_seq_len: int,
589 ) -> None:
590 """Score and select compressed prefill blocks in bounded chunks.
591
592 Args:
593 q: Packed prefill query tensor shaped ``[num_tokens, heads, head_dim]``.
594 k_cache: Compressed key cache shaped ``[blocks, page_size, 1,
595 head_dim]``.
596 page_table: Block table shaped ``[num_requests, max_pages]``.
597 query_start_loc: Packed prefill query offsets with a terminal offset.
598 Offsets may share the base of a larger backing tensor.
599 visible_blocks: Number of visible compressed blocks per query.
600 token_topk: Number of logical tokens selected per query.
601 compress_ratio: Number of logical tokens represented by a cache row.
602 max_query_len: Maximum number of query tokens in one request.
603 block_indices: Compressed-index output buffer.
604 max_seq_len: Longest context length in the batch this step.
605
606 """
607 assert token_topk % compress_ratio == 0
608 assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)
609 assert q.dtype == k_cache.dtype, "Q and the compressed K cache must match"
610 rows = q.shape[0]
611 # No row scores beyond cdiv(max_seq_len, compress_ratio) compressed
612 # columns. Round up to 64 to keep the logits row stride
613 # cooperative_topk-compatible.
614 logits_width = triton.cdiv(triton.cdiv(max_seq_len, compress_ratio), 64) * 64
615 logits_width = min(max(64, logits_width), page_table.shape[1] * k_cache.shape[1])
616
617 # chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB
618 # always allocate the worst case to avoid memory fragmentation.
619 max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
620 rows_per_chunk = max(1, max_logits_bytes // (logits_width * 4))
621
622 budget_bytes = max(max_logits_bytes, page_table.shape[1] * k_cache.shape[1] * 4)
623 logits_workspace = q.new_empty(budget_bytes // 4, dtype=torch.float32)
624
625 topk_workspace = torch.empty(
626 (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device
627 )
628
629 for query_start in range(0, rows, rows_per_chunk):
630 query_end = min(query_start + rows_per_chunk, rows)
631 query_slice = slice(query_start, query_end)
632 logits = _prefill_logits(
633 q,
634 k_cache,
635 page_table,
636 query_start_loc,
637 visible_blocks,
638 logits_workspace,
639 max_query_len,
640 logits_width,
641 query_offset=query_start,
642 num_queries=query_end - query_start,
643 )
644 _topk(
645 logits,
646 visible_blocks[query_slice],
647 token_topk,
648 compress_ratio,
649 block_indices[query_slice],
650 topk_workspace,
651 )
652
653
654 __all__ = [
655 "expand_qsa_block_indices",
656 "qsa_select_paged_decode",
657 "qsa_select_paged_prefill",
658 "warmup_qsa_mqa_paged_decode",
659 ]