@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 612 lines 20.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """AITER page-16 sparse paged-attention helpers for MiniMax-M3 on ROCm."""
4
5 import torch
6
7 try:
8 from vllm.triton_utils import tl, triton
9 except ModuleNotFoundError:
10 import triton
11 import triton.language as tl
12
13 from vllm.models.minimax_m3.common.ops.sparse_attn import SPARSE_BLOCK_SIZE
14
15 ASM_PAGE_SIZE = 16
16 PAGES_PER_SPARSE_BLOCK = SPARSE_BLOCK_SIZE // ASM_PAGE_SIZE
17
18 _FP8_DTYPES = {
19 dtype
20 for dtype in (
21 getattr(torch, "float8_e4m3fn", None),
22 getattr(torch, "float8_e4m3fnuz", None),
23 getattr(torch, "float8_e5m2", None),
24 getattr(torch, "float8_e5m2fnuz", None),
25 )
26 if dtype is not None
27 }
28
29
30 def _is_fp8_kv_cache_tensor(kv_cache: torch.Tensor) -> bool:
31 return kv_cache.dtype in _FP8_DTYPES
32
33
34 @triton.jit
35 def _write_sparse_block_table_row_from_values(
36 blk, # [BLOCK_SIZE_T] int32 selected logical block ids for this query
37 bt_row, # [max_blocks] int32, this request's logical page table
38 sbt_row, # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
39 ctx_ptr, # int32, this query's attended context length
40 abs_pos, # absolute position of this query token, may be negative
41 max_topk,
42 SPARSE_BLOCK_SIZE_C: tl.constexpr,
43 PAGES_PER_BLOCK: tl.constexpr,
44 BLOCK_PAGE_STRIDE: tl.constexpr,
45 BLOCK_SIZE_T: tl.constexpr,
46 ):
47 """Compact one query's selected logical blocks into physical page-16 ids.
48
49 ``BLOCK_PAGE_STRIDE`` is how many page ids a block spans, which is
50 ``PAGES_PER_BLOCK`` when each K/V side is its own dense plane and twice
51 that when both sides share a block. Padded rows carry a negative
52 ``abs_pos``, which clamps the causal range to empty.
53 """
54 causal_len = tl.maximum(abs_pos + 1, 0)
55 self_blk = abs_pos // SPARSE_BLOCK_SIZE_C
56
57 off_t = tl.arange(0, BLOCK_SIZE_T)
58 valid = (off_t < max_topk) & (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
59 is_tail = valid & (blk == self_blk)
60 is_full = valid & (blk < self_blk)
61
62 n_full = tl.sum(is_full.to(tl.int32), axis=0)
63 n_valid = tl.sum(valid.to(tl.int32), axis=0)
64 earlier_full = tl.cumsum(is_full.to(tl.int32), axis=0) - is_full.to(tl.int32)
65 slot = tl.where(is_full, earlier_full, n_full)
66
67 logical_page = tl.load(bt_row + blk, mask=valid, other=0).to(tl.int32)
68 base_phys = logical_page * BLOCK_PAGE_STRIDE
69 dst_base = slot * PAGES_PER_BLOCK
70
71 for j in tl.static_range(PAGES_PER_BLOCK):
72 tl.store(sbt_row + dst_base + j, base_phys + j, mask=valid)
73
74 n_used = n_valid * PAGES_PER_BLOCK
75 off_w = tl.arange(0, BLOCK_SIZE_T * PAGES_PER_BLOCK)
76 row_width = max_topk * PAGES_PER_BLOCK
77 tl.store(
78 sbt_row + off_w,
79 tl.zeros_like(off_w),
80 mask=(off_w >= n_used) & (off_w < row_width),
81 )
82
83 tail_tokens = causal_len - self_blk * SPARSE_BLOCK_SIZE_C
84 has_tail = tl.sum(is_tail.to(tl.int32), axis=0) > 0
85 ctx = n_full * SPARSE_BLOCK_SIZE_C + tl.where(has_tail, tail_tokens, 0)
86 ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, causal_len))
87 tl.store(ctx_ptr, ctx)
88
89
90 @triton.jit
91 def _write_sparse_block_table_row(
92 topk_row, # [topk] int32, selected logical block ids for this query
93 bt_row, # [max_blocks] int32, this request's logical page table
94 sbt_row, # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
95 ctx_ptr, # int32, this query's attended context length
96 abs_pos, # absolute position of this query token, may be negative
97 max_topk,
98 stride_topk_k,
99 SPARSE_BLOCK_SIZE_C: tl.constexpr,
100 PAGES_PER_BLOCK: tl.constexpr,
101 BLOCK_PAGE_STRIDE: tl.constexpr,
102 BLOCK_SIZE_T: tl.constexpr,
103 ):
104 off_t = tl.arange(0, BLOCK_SIZE_T)
105 blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
106 _write_sparse_block_table_row_from_values(
107 blk,
108 bt_row,
109 sbt_row,
110 ctx_ptr,
111 abs_pos,
112 max_topk,
113 SPARSE_BLOCK_SIZE_C,
114 PAGES_PER_BLOCK,
115 BLOCK_PAGE_STRIDE,
116 BLOCK_SIZE_T,
117 )
118
119
120 @triton.jit
121 def _build_sparse_block_table_kernel(
122 topk_ptr, # [1, batch * DECODE_QUERY_LEN, topk] int32, logical block ids
123 block_table_ptr, # [batch, max_blocks] int32, logical 128-page table
124 seq_lens_ptr, # [batch] int32
125 sparse_bt_ptr, # [batch * DECODE_QUERY_LEN, topk * 8] int32
126 sparse_ctx_ptr, # [batch * DECODE_QUERY_LEN] int32
127 max_topk,
128 stride_topk_n,
129 stride_topk_k,
130 stride_bt_b,
131 stride_sbt_b,
132 SPARSE_BLOCK_SIZE_C: tl.constexpr,
133 PAGES_PER_BLOCK: tl.constexpr,
134 BLOCK_PAGE_STRIDE: tl.constexpr,
135 BLOCK_SIZE_T: tl.constexpr,
136 DECODE_QUERY_LEN: tl.constexpr,
137 ):
138 # One program per query token. The decode indexer flattens speculative rows
139 # as ``req_id * DECODE_QUERY_LEN + local_q``, and a request's query tokens
140 # are the last DECODE_QUERY_LEN positions of its sequence, so both the
141 # request id and the absolute position follow from the program id.
142 pid_n = tl.program_id(0)
143 req_id = pid_n // DECODE_QUERY_LEN
144 local_q = pid_n % DECODE_QUERY_LEN
145 seq_len = tl.load(seq_lens_ptr + req_id)
146 abs_pos = seq_len - DECODE_QUERY_LEN + local_q
147
148 _write_sparse_block_table_row(
149 topk_ptr + pid_n * stride_topk_n,
150 block_table_ptr + req_id * stride_bt_b,
151 sparse_bt_ptr + pid_n * stride_sbt_b,
152 sparse_ctx_ptr + pid_n,
153 abs_pos,
154 max_topk,
155 stride_topk_k,
156 SPARSE_BLOCK_SIZE_C,
157 PAGES_PER_BLOCK,
158 BLOCK_PAGE_STRIDE,
159 BLOCK_SIZE_T,
160 )
161
162
163 @triton.jit
164 def _build_sparse_block_table_prefill_kernel(
165 topk_ptr, # [1, total_q, topk]
166 block_table_ptr, # [batch, max_blocks]
167 req_id_ptr, # [total_q]
168 abs_pos_ptr, # [total_q]
169 sparse_bt_ptr, # [total_q, topk * 8]
170 sparse_ctx_ptr, # [total_q]
171 max_topk,
172 stride_topk_n,
173 stride_topk_k,
174 stride_bt_b,
175 stride_sbt_n,
176 SPARSE_BLOCK_SIZE_C: tl.constexpr,
177 PAGES_PER_BLOCK: tl.constexpr,
178 BLOCK_PAGE_STRIDE: tl.constexpr,
179 BLOCK_SIZE_T: tl.constexpr,
180 ):
181 pid_n = tl.program_id(0)
182 req_id = tl.load(req_id_ptr + pid_n)
183 abs_pos = tl.load(abs_pos_ptr + pid_n)
184
185 _write_sparse_block_table_row(
186 topk_ptr + pid_n * stride_topk_n,
187 block_table_ptr + req_id * stride_bt_b,
188 sparse_bt_ptr + pid_n * stride_sbt_n,
189 sparse_ctx_ptr + pid_n,
190 abs_pos,
191 max_topk,
192 stride_topk_k,
193 SPARSE_BLOCK_SIZE_C,
194 PAGES_PER_BLOCK,
195 BLOCK_PAGE_STRIDE,
196 BLOCK_SIZE_T,
197 )
198
199
200 def minimax_m3_alloc_sparse_block_table(
201 topk_idx: torch.Tensor,
202 ) -> tuple[torch.Tensor, torch.Tensor]:
203 assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
204 total_q = topk_idx.shape[1]
205 width = topk_idx.shape[-1] * PAGES_PER_SPARSE_BLOCK
206 return (
207 torch.empty((total_q, width), dtype=torch.int32, device=topk_idx.device),
208 torch.empty((total_q,), dtype=torch.int32, device=topk_idx.device),
209 )
210
211
212 @torch.no_grad()
213 def minimax_m3_build_sparse_block_table_prefill(
214 topk_idx: torch.Tensor, # [1, total_q, topk]
215 block_table: torch.Tensor, # [batch, max_blocks]
216 query_req_id: torch.Tensor, # [total_q]
217 query_abs_pos: torch.Tensor, # [total_q]
218 block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
219 ) -> tuple[torch.Tensor, torch.Tensor]:
220 """Build one page-16 sparse block table row per prefill query token."""
221 sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
222 topk = topk_idx.shape[-1]
223 _build_sparse_block_table_prefill_kernel[(topk_idx.shape[1],)](
224 topk_idx,
225 block_table,
226 query_req_id,
227 query_abs_pos,
228 sparse_bt,
229 sparse_ctx,
230 topk,
231 topk_idx.stride(1),
232 topk_idx.stride(2),
233 block_table.stride(0),
234 sparse_bt.stride(0),
235 SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
236 PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
237 BLOCK_PAGE_STRIDE=block_page_stride,
238 BLOCK_SIZE_T=triton.next_power_of_2(topk),
239 )
240 return sparse_bt, sparse_ctx
241
242
243 @torch.no_grad()
244 def minimax_m3_build_sparse_block_table_decode(
245 topk_idx: torch.Tensor, # [1, batch * decode_query_len, topk]
246 block_table: torch.Tensor, # [batch, max_blocks]
247 seq_lens: torch.Tensor, # [batch]
248 decode_query_len: int = 1,
249 block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
250 ) -> tuple[torch.Tensor, torch.Tensor]:
251 """Build one page-16 sparse block table row per decode query token."""
252 total_q = topk_idx.shape[1]
253 expected_q = block_table.shape[0] * decode_query_len
254 assert total_q == expected_q, (
255 "MiniMax-M3 decode top-k rows must equal batch * decode_query_len: "
256 f"{total_q} != {expected_q}"
257 )
258 sparse_bt, sparse_ctx = minimax_m3_alloc_sparse_block_table(topk_idx)
259 topk = topk_idx.shape[-1]
260 _build_sparse_block_table_kernel[(total_q,)](
261 topk_idx,
262 block_table,
263 seq_lens,
264 sparse_bt,
265 sparse_ctx,
266 topk,
267 topk_idx.stride(1),
268 topk_idx.stride(2),
269 block_table.stride(0),
270 sparse_bt.stride(0),
271 SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
272 PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
273 BLOCK_PAGE_STRIDE=block_page_stride,
274 BLOCK_SIZE_T=triton.next_power_of_2(topk),
275 DECODE_QUERY_LEN=decode_query_len,
276 )
277 return sparse_bt, sparse_ctx
278
279
280 # Retained for the decode_query_len == 1 callers (and tests) that predate the
281 # speculative-decode rows.
282 minimax_m3_build_sparse_block_table = minimax_m3_build_sparse_block_table_decode
283
284
285 @triton.jit
286 def _insert_index_cache_kernel(
287 index_k_ptr, # [num_tokens, head_dim]
288 index_cache_ptr, # [num_blocks, block_size, head_dim]
289 slot_mapping_ptr, # [num_tokens]
290 stride_k_t,
291 stride_k_d,
292 stride_c_b,
293 stride_c_t,
294 stride_c_d,
295 stride_slot_t,
296 CACHE_BLOCK_SIZE: tl.constexpr,
297 HEAD_DIM: tl.constexpr,
298 BLOCK_D: tl.constexpr,
299 ):
300 pid_t = tl.program_id(0)
301 offs_d = tl.arange(0, BLOCK_D)
302 slot = tl.load(slot_mapping_ptr + pid_t * stride_slot_t)
303 block_id = slot // CACHE_BLOCK_SIZE
304 block_off = slot - block_id * CACHE_BLOCK_SIZE
305
306 src = index_k_ptr + pid_t * stride_k_t + offs_d * stride_k_d
307 dst = (
308 index_cache_ptr
309 + block_id * stride_c_b
310 + block_off * stride_c_t
311 + offs_d * stride_c_d
312 )
313 mask = (slot >= 0) & (offs_d < HEAD_DIM)
314 value = tl.load(src, mask=offs_d < HEAD_DIM, other=0.0)
315 # The side cache carries its own dtype (bf16, or e4m3 for the fp8 indexer),
316 # so convert rather than leaning on the store's implicit cast.
317 tl.store(dst, value.to(dst.dtype.element_ty), mask=mask)
318
319
320 @torch.no_grad()
321 def minimax_m3_insert_index_cache(
322 index_k: torch.Tensor,
323 index_cache: torch.Tensor,
324 index_slot_mapping: torch.Tensor,
325 ) -> None:
326 """Scatter index keys into MiniMax-M3's key-only side cache."""
327 if index_k.numel() == 0 or index_cache.numel() == 0:
328 return
329 if index_k.dim() != 2 or index_cache.dim() != 3:
330 raise ValueError("MiniMax-M3 index cache insert expects [N,D] and [B,T,D]")
331 if index_k.shape[1] != index_cache.shape[2]:
332 raise ValueError("MiniMax-M3 index key dim must match index cache head dim")
333 if index_slot_mapping.dim() != 1 or index_slot_mapping.shape[0] != index_k.shape[0]:
334 raise ValueError("MiniMax-M3 index slot mapping must be a length-N vector")
335 if index_cache.stride(2) != 1:
336 raise ValueError("MiniMax-M3 index cache requires contiguous head dimension")
337
338 head_dim = index_k.shape[1]
339 _insert_index_cache_kernel[(index_k.shape[0],)](
340 index_k,
341 index_cache,
342 index_slot_mapping,
343 index_k.stride(0),
344 index_k.stride(1),
345 index_cache.stride(0),
346 index_cache.stride(1),
347 index_cache.stride(2),
348 index_slot_mapping.stride(0),
349 CACHE_BLOCK_SIZE=index_cache.shape[1],
350 HEAD_DIM=head_dim,
351 BLOCK_D=triton.next_power_of_2(head_dim),
352 num_warps=4,
353 )
354
355
356 def _sides_are_packed(k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool:
357 """Whether a block holds both its K and V pages instead of one side."""
358 return v_cache.shape[0] != k_cache.shape[0]
359
360
361 def minimax_m3_sparse_block_page_stride(
362 k_cache: torch.Tensor, v_cache: torch.Tensor
363 ) -> int:
364 """How many page ids one sparse block spans.
365
366 Each side owns ``PAGES_PER_SPARSE_BLOCK`` pages. When the sides are dense
367 planes a block is exactly that wide; when they interleave, a block covers
368 both sides' pages and V is reached from the same page id at a fixed offset.
369 """
370 return PAGES_PER_SPARSE_BLOCK * (2 if _sides_are_packed(k_cache, v_cache) else 1)
371
372
373 # K/V sides sharing one block, which is how many sides' pages a block spans.
374 # The attend publishes two head slots and the indexer's side cache one, and
375 # specs that mix HNC shapes narrow the resolved layout to a block-compact one
376 # (LHBNC is not), so this path always lands on an interleaved block.
377 PAGE16_SIDES_PER_BLOCK = 2
378
379
380 def minimax_m3_rebase_block_table_to_page16(
381 block_table: torch.Tensor,
382 out: torch.Tensor | None = None,
383 ) -> torch.Tensor:
384 """Rebase a logical page table onto AITER's page-16 page numbering.
385
386 The page ids AITER's top-k emits for the attend are
387 ``block_table[blk] * pages_per_block + j``, and ``pages_per_block`` is fixed
388 at AITER build time to one side's pages. An interleaved block spans both
389 sides, so the only way to reach its pages through that kernel is to hand it
390 a table already scaled to the wider stride. The caller does this once per
391 step, since every sparse layer resolves its selection through the same
392 table.
393 """
394 if out is None:
395 out = torch.empty_like(block_table)
396 return torch.mul(block_table, PAGE16_SIDES_PER_BLOCK, out=out)
397
398
399 def _gluon_scale_arg(
400 scale: torch.Tensor | None,
401 *,
402 num_phys_pages: int,
403 num_kv_heads: int,
404 ) -> torch.Tensor | None:
405 if scale is None:
406 return None
407 if scale.numel() == 1:
408 return scale
409 if scale.dim() != 2 or scale.shape[0] != num_kv_heads:
410 raise ValueError(
411 "MiniMax-M3 AITER sparse PA supports scalar KV scales or "
412 "[num_kv_heads, max_kv_tokens] scales"
413 )
414 max_tokens = num_phys_pages * ASM_PAGE_SIZE
415 if scale.shape[1] < max_tokens:
416 raise ValueError(
417 "MiniMax-M3 AITER sparse PA KV scale tensor is smaller than the "
418 f"cache token capacity ({scale.shape[1]} < {max_tokens})"
419 )
420 scale = scale[:, :max_tokens]
421 return (
422 scale.transpose(0, 1)
423 .contiguous()
424 .view(num_phys_pages, ASM_PAGE_SIZE, num_kv_heads)
425 .permute(0, 2, 1)
426 .contiguous()
427 .unsqueeze(-1)
428 )
429
430
431 @torch.no_grad()
432 def minimax_m3_sparse_attn_decode_aiter(
433 q: torch.Tensor, # [batch * decode_query_len, num_heads, head_dim]
434 k_cache: torch.Tensor, # [phys16, num_kv_heads, head_dim // x, 16, x]
435 v_cache: torch.Tensor, # [phys16, num_kv_heads, 16 // x, head_dim, x]
436 topk_idx: torch.Tensor, # [1, batch * decode_query_len, topk]
437 block_table: torch.Tensor, # [batch, max_blocks]
438 seq_lens: torch.Tensor, # [batch]
439 num_kv_heads: int,
440 sm_scale: float,
441 output: torch.Tensor, # [batch * decode_query_len, num_heads, head_dim]
442 k_scale: torch.Tensor | None = None,
443 v_scale: torch.Tensor | None = None,
444 decode_query_len: int = 1,
445 sparse_block_table: torch.Tensor | None = None,
446 sparse_context_lens: torch.Tensor | None = None,
447 ) -> None:
448 if (sparse_block_table is None) != (sparse_context_lens is None):
449 raise ValueError(
450 "MiniMax-M3 prepared sparse block table and context lengths "
451 "must be provided together"
452 )
453 if sparse_block_table is None:
454 sparse_block_table, sparse_context_lens = (
455 minimax_m3_build_sparse_block_table_decode(
456 topk_idx,
457 block_table,
458 seq_lens,
459 decode_query_len,
460 minimax_m3_sparse_block_page_stride(k_cache, v_cache),
461 )
462 )
463 assert sparse_context_lens is not None
464 _run_gluon_decode(
465 q,
466 k_cache,
467 v_cache,
468 sparse_block_table,
469 sparse_context_lens,
470 num_kv_heads,
471 sm_scale,
472 output,
473 k_scale,
474 v_scale,
475 )
476
477
478 @torch.no_grad()
479 def minimax_m3_sparse_attn_prefill_aiter(
480 q: torch.Tensor, # [total_q, num_heads, head_dim]
481 k_cache: torch.Tensor,
482 v_cache: torch.Tensor,
483 topk_idx: torch.Tensor, # [1, total_q, topk]
484 block_table: torch.Tensor,
485 query_req_id: torch.Tensor, # [total_q] int32, owning request per token
486 query_abs_pos: torch.Tensor, # [total_q] int32, absolute position per token
487 num_kv_heads: int,
488 sm_scale: float,
489 output: torch.Tensor,
490 k_scale: torch.Tensor | None = None,
491 v_scale: torch.Tensor | None = None,
492 sparse_bt: torch.Tensor | None = None,
493 sparse_ctx: torch.Tensor | None = None,
494 ) -> None:
495 if (sparse_bt is None) != (sparse_ctx is None):
496 raise ValueError(
497 "MiniMax-M3 prepared sparse block table and context lengths "
498 "must be provided together"
499 )
500 if sparse_bt is None:
501 sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
502 topk_idx,
503 block_table,
504 query_req_id,
505 query_abs_pos,
506 minimax_m3_sparse_block_page_stride(k_cache, v_cache),
507 )
508 assert sparse_ctx is not None
509 _run_gluon_decode(
510 q,
511 k_cache,
512 v_cache,
513 sparse_bt,
514 sparse_ctx,
515 num_kv_heads,
516 sm_scale,
517 output,
518 k_scale,
519 v_scale,
520 )
521
522
523 def _run_gluon_decode(
524 q: torch.Tensor,
525 k_cache: torch.Tensor,
526 v_cache: torch.Tensor,
527 sparse_bt: torch.Tensor,
528 sparse_ctx: torch.Tensor,
529 num_kv_heads: int,
530 sm_scale: float,
531 output: torch.Tensor,
532 k_scale: torch.Tensor | None,
533 v_scale: torch.Tensor | None,
534 ) -> None:
535 from aiter import dtypes as aiter_dtypes
536 from aiter.ops.triton.gluon.pa_decode_gluon import (
537 get_recommended_splits,
538 pa_decode_gluon,
539 )
540
541 if not q.is_contiguous():
542 q = q.contiguous()
543 if not output.is_contiguous():
544 raise ValueError("MiniMax-M3 AITER sparse PA output must be contiguous")
545
546 total_q, num_q_heads, head_size = q.shape
547 if head_size != 128:
548 raise ValueError("MiniMax-M3 AITER sparse PA requires head_dim == 128")
549 group_size = num_q_heads // num_kv_heads
550
551 nphys16, hkv = k_cache.shape[0], k_cache.shape[1]
552 # V is offset by half a block when the sides interleave, so it holds fewer
553 # pages than K even though both are addressed by the same page id.
554 v_nphys16 = v_cache.shape[0]
555 k_cache_view = k_cache.view(nphys16 * hkv, 1, *k_cache.shape[2:])
556 v_cache_view = v_cache.view(v_nphys16 * hkv, 1, *v_cache.shape[2:])
557 q_view = q.view(total_q * num_kv_heads, group_size, head_size)
558 out_view = output.view(total_q * num_kv_heads, group_size, head_size)
559
560 num_seqs = total_q * num_kv_heads
561 max_context_partition_num = get_recommended_splits(num_seqs, 1)
562 context_partition_size = 256
563 intermediate_shape = (num_seqs, 1, max_context_partition_num, group_size)
564 exp_sums = torch.empty(intermediate_shape, dtype=torch.float32, device=q.device)
565 max_logits = torch.empty_like(exp_sums)
566 temporary_output = torch.empty(
567 *intermediate_shape, head_size, dtype=q.dtype, device=q.device
568 )
569
570 is_fp8 = _is_fp8_kv_cache_tensor(k_cache)
571 compute_type = aiter_dtypes.fp8 if is_fp8 else q.dtype
572 if is_fp8:
573 if _sides_are_packed(k_cache, v_cache) and not all(
574 scale is None or scale.numel() == 1 for scale in (k_scale, v_scale)
575 ):
576 raise NotImplementedError(
577 "MiniMax-M3 AITER sparse PA supports only scalar KV scales "
578 "when K and V pages share a block; per-token scales would "
579 "need the same page renumbering."
580 )
581 k_scale_arg = _gluon_scale_arg(
582 k_scale, num_phys_pages=nphys16, num_kv_heads=hkv
583 )
584 v_scale_arg = _gluon_scale_arg(
585 v_scale, num_phys_pages=v_nphys16, num_kv_heads=hkv
586 )
587 else:
588 k_scale_arg = v_scale_arg = None
589
590 pa_decode_gluon(
591 output=out_view,
592 query=q_view,
593 key_cache=k_cache_view,
594 value_cache=v_cache_view,
595 context_lengths=sparse_ctx,
596 block_tables=sparse_bt,
597 softmax_scale=sm_scale,
598 query_length=1,
599 max_context_partition_num=max_context_partition_num,
600 context_partition_size=context_partition_size,
601 compute_type=compute_type,
602 query_scale=None,
603 key_scale=k_scale_arg,
604 value_scale=v_scale_arg,
605 exp_sums=exp_sums,
606 max_logits=max_logits,
607 temporary_output=temporary_output,
608 alibi_slopes=None,
609 sinks=None,
610 sliding_window=-1,
611 ps=True,
612 )