@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 4,341 lines 150 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import functools
4 import importlib
5 import math
6 from collections.abc import Callable
7 from importlib.util import find_spec
8
9 import torch
10 import torch.nn.functional as F
11
12 import vllm.envs as envs
13 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
14 from vllm.config import CUDAGraphMode, get_current_vllm_config
15 from vllm.forward_context import get_forward_context
16 from vllm.logger import init_logger
17 from vllm.platforms import current_platform
18 from vllm.triton_utils import tl, triton
19 from vllm.utils.torch_utils import LayerNameType, direct_register_custom_op
20 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
21 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
22 from vllm.v1.worker.workspace import current_workspace_manager
23
24 if current_platform.is_rocm():
25 from vllm.platforms.rocm import _ON_GFX942, _ON_GFX950
26 else:
27 _ON_GFX942 = False
28 _ON_GFX950 = False
29
30 logger = init_logger(__name__)
31
32 FP8_DTYPE = current_platform.fp8_dtype()
33
34
35 @functools.cache
36 def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None:
37 from vllm._aiter_ops import rocm_aiter_ops
38
39 if not rocm_aiter_ops.is_mla_enabled():
40 return None
41 try:
42 from aiter.ops.pa_sparse_prefill_opus import pa_sparse_prefill_opus
43 except ImportError:
44 return None
45 logger.info_once("Using AITER OPUS for large sparse MLA prefill on gfx950")
46 return pa_sparse_prefill_opus
47
48
49 # Conservative perf gate, not a correctness bound: OPUS is correct for any query
50 # count, but Triton stays faster below this measured crossover.
51 _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024
52
53
54 def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool:
55 """V4.0 C4A is block-flat (NORMAL). Ratio 1 and 2 are 16×16 SHUFFLE."""
56 return compress_ratio == 4
57
58
59 @triton.jit
60 def _localize_aiter_prefill_topk_kernel(
61 indices_ptr,
62 row_starts_ptr,
63 indices_stride,
64 num_topk,
65 BLOCK_SIZE: tl.constexpr,
66 ):
67 row_idx = tl.program_id(0)
68 offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
69 mask = offsets < num_topk
70 indices = tl.load(indices_ptr + row_idx * indices_stride + offsets, mask=mask)
71 row_start = tl.load(row_starts_ptr + row_idx)
72 localized = tl.where(indices < 0, indices, indices - row_start)
73 tl.store(
74 indices_ptr + row_idx * indices_stride + offsets,
75 localized,
76 mask=mask,
77 )
78
79
80 def _localize_aiter_prefill_topk(
81 indices: torch.Tensor,
82 row_starts: torch.Tensor,
83 ) -> None:
84 num_rows, num_topk = indices.shape
85 block_size = 256
86 _localize_aiter_prefill_topk_kernel[(num_rows, triton.cdiv(num_topk, block_size))](
87 indices,
88 row_starts,
89 indices.stride(0),
90 num_topk,
91 BLOCK_SIZE=block_size,
92 )
93
94
95 def _launch_aiter_top_k_per_row_prefill(
96 logits: torch.Tensor,
97 row_starts: torch.Tensor,
98 row_ends: torch.Tensor,
99 indices: torch.Tensor,
100 topk_tokens: int,
101 ) -> None:
102 from vllm._aiter_ops import rocm_aiter_ops
103
104 rocm_aiter_ops.indexer_top_k_prefill(
105 logits, row_starts, row_ends, indices, topk_tokens
106 )
107 _localize_aiter_prefill_topk(indices, row_starts)
108
109
110 @triton.jit
111 def _indexer_k_quant_and_cache_kernel(
112 k_ptr, # [num_tokens, head_dim]
113 kv_cache_ptr, # [n_blks, blk_size//tile_block, head_dim // 16B, tile_block, 16B]
114 # [n_blocks, blk_size, head_dim]
115 kv_cache_scale_ptr, # [n_blks, blk_size]
116 slot_mapping_ptr, # [num_tokens]
117 kv_cache_scale_stride,
118 kv_cache_value_stride,
119 block_size,
120 num_tokens,
121 head_dim: tl.constexpr,
122 LAYOUT: tl.constexpr,
123 BLOCK_TILE_SIZE: tl.constexpr,
124 HEAD_TILE_SIZE: tl.constexpr,
125 IS_FNUZ: tl.constexpr,
126 USE_UE8M0: tl.constexpr,
127 ):
128 tid = tl.program_id(0)
129 offset = tl.arange(0, head_dim)
130 if LAYOUT == "SHUFFLE":
131 tile_offset = (
132 offset // HEAD_TILE_SIZE * BLOCK_TILE_SIZE * HEAD_TILE_SIZE
133 + offset % HEAD_TILE_SIZE
134 )
135 else:
136 tile_offset = offset
137 tile_store_offset = tile_offset
138 # for idx in tl.range(tid, num_tokens, n_program):
139 src_ptr = k_ptr + tid * head_dim
140 slot_id = tl.load(slot_mapping_ptr + tid)
141 if slot_id < 0:
142 return
143 # The packed KV layout makes per-block strides large
144 # enough that block_id * stride can exceed 32-bit range.
145 block_id = (slot_id // block_size).to(tl.int64)
146 block_offset = slot_id % block_size
147 tile_block_id = block_offset // BLOCK_TILE_SIZE
148 tile_block_offset = block_offset % BLOCK_TILE_SIZE
149 val = tl.load(src_ptr + offset)
150 amax = tl.max(val.abs(), axis=-1).to(tl.float32)
151 if IS_FNUZ:
152 scale = tl.maximum(1e-4, amax) / 224.0
153 else:
154 scale = tl.maximum(1e-4, amax) / 448.0
155
156 if USE_UE8M0:
157 scale = tl.exp2(tl.ceil(tl.log2(scale)))
158
159 fp8_val = (val.to(tl.float32) / scale).to(kv_cache_ptr.type.element_ty)
160 if LAYOUT == "SHUFFLE":
161 dst_ptr = (
162 kv_cache_ptr
163 + block_id * kv_cache_value_stride
164 + tile_block_id * BLOCK_TILE_SIZE * head_dim
165 + tile_block_offset * HEAD_TILE_SIZE
166 )
167 else:
168 dst_ptr = (
169 kv_cache_ptr + block_id * kv_cache_value_stride + block_offset * head_dim
170 )
171 tl.store(dst_ptr + tile_store_offset, fp8_val)
172 dst_scale_ptr = kv_cache_scale_ptr + block_id * kv_cache_scale_stride + block_offset
173 tl.store(dst_scale_ptr, scale)
174
175
176 def indexer_k_quant_and_cache_triton(
177 k: torch.Tensor,
178 kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4]
179 slot_mapping: torch.Tensor,
180 quant_block_size,
181 scale_fmt,
182 block_tile_size=16,
183 head_tile_size=16,
184 ):
185 num_blocks = kv_cache.shape[0]
186 head_dim = k.shape[-1]
187 num_tokens = slot_mapping.shape[0]
188 block_size = kv_cache.shape[1]
189 # In real layout, we store the first portion as kv cache value
190 # and second portion as kv cache scale
191 kv_cache = kv_cache.view(num_blocks, -1)
192 kv_cache_value = kv_cache[:, : block_size * head_dim].view(FP8_DTYPE)
193 kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32)
194 head_tile_size = head_tile_size // kv_cache.element_size()
195 layout = "NORMAL" if block_size == 1 else "SHUFFLE"
196 grid = (num_tokens,)
197 _indexer_k_quant_and_cache_kernel[grid](
198 k,
199 kv_cache_value,
200 kv_cache_scale,
201 slot_mapping,
202 kv_cache_scale.stride(0),
203 kv_cache_value.stride(0),
204 block_size,
205 num_tokens,
206 head_dim,
207 layout,
208 block_tile_size,
209 head_tile_size,
210 IS_FNUZ=torch.float8_e4m3fnuz == FP8_DTYPE,
211 USE_UE8M0=scale_fmt == "ue8m0",
212 )
213
214
215 @triton.jit
216 def _cp_gather_indexer_quant_cache_kernel(
217 kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B]
218 # [n_blks, blk_size, head_dim]
219 kv_cache_scale_ptr, # [n_blks, blk_size]
220 k_fp8_ptr, # [num_tokens, head_dim]
221 k_scale_ptr, # [num_tokens]
222 block_table_ptr, # [batch_size, block_table_stride]
223 cu_seqlen_ptr, # [batch_size + 1]
224 token_to_seq_ptr, # [num_tokens]
225 block_size,
226 block_table_stride,
227 kv_cache_stride,
228 kv_cache_scale_stride,
229 LAYOUT: tl.constexpr,
230 HEAD_DIM: tl.constexpr,
231 BLOCK_TILE_SIZE: tl.constexpr,
232 HEAD_TILE_SIZE: tl.constexpr,
233 num_tokens,
234 num_batches,
235 block_table_width,
236 num_blocks,
237 ):
238 tid = tl.program_id(0)
239 offset = tl.arange(0, HEAD_DIM)
240 valid_tid = tid < num_tokens
241 batch_id = tl.load(token_to_seq_ptr + tid, mask=valid_tid, other=-1)
242 valid_batch = (batch_id >= 0) & (batch_id < num_batches)
243 safe_batch_id = tl.where(valid_batch, batch_id, 0)
244 batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0)
245 batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0)
246 batch_offset = tid - batch_start
247 valid_token = valid_tid & valid_batch & (tid >= batch_start) & (tid < batch_end)
248 if not valid_token:
249 return
250 block_table_id = batch_offset // block_size
251 block_offset = batch_offset % block_size
252 valid_block_table = (
253 valid_token
254 & (block_table_id >= 0)
255 & (block_table_id < block_table_width)
256 & (block_offset >= 0)
257 & (block_offset < block_size)
258 )
259 safe_block_table_id = tl.where(valid_block_table, block_table_id, 0)
260 block_table_offset = safe_batch_id * block_table_stride + safe_block_table_id
261 block_id = tl.load(
262 block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
263 )
264 valid_block = valid_block_table & (block_id >= 0) & (block_id < num_blocks)
265 # The packed KV layout makes per-block strides large
266 # enough that block_id * stride can exceed 32-bit range.
267 safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64)
268 safe_block_offset = tl.where(valid_block, block_offset, 0)
269 tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE
270 if LAYOUT == "SHUFFLE":
271 src_cache_offset = (
272 safe_block_id * kv_cache_stride
273 + (safe_block_offset // BLOCK_TILE_SIZE) * HEAD_DIM * BLOCK_TILE_SIZE
274 + tiled_block_offset * HEAD_TILE_SIZE
275 )
276 else:
277 src_cache_offset = (
278 safe_block_id * kv_cache_stride + safe_block_offset * HEAD_DIM
279 )
280 src_scale_offset = safe_block_id * kv_cache_scale_stride + safe_block_offset
281 dst_offset = tid * HEAD_DIM
282 src_scale_ptr = kv_cache_scale_ptr + src_scale_offset
283 src_cache_ptr = kv_cache_ptr + src_cache_offset
284 dst_k_ptr = k_fp8_ptr + dst_offset
285 scale_val = tl.load(src_scale_ptr, mask=valid_block, other=0.0)
286 tl.store(k_scale_ptr + tid, scale_val)
287 if LAYOUT == "SHUFFLE":
288 tiled_src_offset = (
289 offset // HEAD_TILE_SIZE * HEAD_TILE_SIZE * BLOCK_TILE_SIZE
290 + offset % HEAD_TILE_SIZE
291 )
292 else:
293 tiled_src_offset = offset
294 val = tl.load(src_cache_ptr + tiled_src_offset)
295 tl.store(dst_k_ptr + offset, val, mask=valid_block)
296
297
298 @triton.jit(do_not_specialize=["num_batches"])
299 def _cp_gather_indexer_quant_cache_gfx950_kernel(
300 kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B]
301 # [n_blks, blk_size, head_dim]
302 kv_cache_scale_ptr, # [n_blks, blk_size]
303 k_fp8_ptr, # [num_tokens, head_dim]
304 k_scale_ptr, # [num_tokens]
305 block_table_ptr, # [batch_size, block_table_stride]
306 cu_seqlen_ptr, # [batch_size + 1]
307 token_to_seq_ptr, # [num_tokens]
308 block_size,
309 block_table_stride,
310 kv_cache_stride,
311 kv_cache_scale_stride,
312 LAYOUT: tl.constexpr,
313 HEAD_DIM: tl.constexpr,
314 BLOCK_TILE_SIZE: tl.constexpr,
315 HEAD_TILE_SIZE: tl.constexpr,
316 num_batches,
317 BLOCK_TABLE_WIDTH: tl.constexpr,
318 NUM_BLOCKS: tl.constexpr,
319 ):
320 tid = tl.program_id(0)
321 offset = tl.arange(0, HEAD_DIM)
322 batch_id = tl.load(token_to_seq_ptr + tid)
323 valid_batch = (batch_id >= 0) & (batch_id < num_batches)
324 safe_batch_id = tl.where(valid_batch, batch_id, 0)
325 batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0)
326 batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0)
327 batch_offset = tid - batch_start
328 valid_token = valid_batch & (tid >= batch_start) & (tid < batch_end)
329 if not valid_token:
330 return
331 block_table_id = batch_offset // block_size
332 block_offset = batch_offset % block_size
333 valid_block_table = (
334 valid_token
335 & (block_table_id >= 0)
336 & (block_table_id < BLOCK_TABLE_WIDTH)
337 & (block_offset >= 0)
338 & (block_offset < block_size)
339 )
340 safe_block_table_id = tl.where(valid_block_table, block_table_id, 0)
341 block_table_offset = safe_batch_id * block_table_stride + safe_block_table_id
342 block_id = tl.load(
343 block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
344 )
345 valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS)
346 # The packed KV layout makes per-block strides large
347 # enough that block_id * stride can exceed 32-bit range.
348 safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64)
349 safe_block_offset = tl.where(valid_block, block_offset, 0)
350 tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE
351 if LAYOUT == "SHUFFLE":
352 src_cache_offset = (
353 safe_block_id * kv_cache_stride
354 + (safe_block_offset // BLOCK_TILE_SIZE) * HEAD_DIM * BLOCK_TILE_SIZE
355 + tiled_block_offset * HEAD_TILE_SIZE
356 )
357 else:
358 src_cache_offset = (
359 safe_block_id * kv_cache_stride + safe_block_offset * HEAD_DIM
360 )
361 src_scale_offset = safe_block_id * kv_cache_scale_stride + safe_block_offset
362 dst_offset = tid * HEAD_DIM
363 src_scale_ptr = kv_cache_scale_ptr + src_scale_offset
364 src_cache_ptr = kv_cache_ptr + src_cache_offset
365 dst_k_ptr = k_fp8_ptr + dst_offset
366 scale_val = tl.load(src_scale_ptr, mask=valid_block, other=0.0)
367 tl.store(k_scale_ptr + tid, scale_val)
368 if LAYOUT == "SHUFFLE":
369 tiled_src_offset = (
370 offset // HEAD_TILE_SIZE * HEAD_TILE_SIZE * BLOCK_TILE_SIZE
371 + offset % HEAD_TILE_SIZE
372 )
373 else:
374 tiled_src_offset = offset
375 val = tl.load(src_cache_ptr + tiled_src_offset)
376 tl.store(dst_k_ptr + offset, val, mask=valid_block)
377
378
379 def cp_gather_indexer_k_quant_cache_triton(
380 k_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4]
381 k_fp8: torch.Tensor,
382 k_fp8_scale: torch.Tensor,
383 block_table: torch.Tensor,
384 cu_seqlen: torch.Tensor,
385 token_to_seq: torch.Tensor,
386 block_tile_size: int = 16,
387 head_tile_size: int = 16,
388 cache_layout: str | None = None,
389 ):
390 num_tokens = k_fp8.size(0)
391 block_size = k_cache.size(1)
392 block_table_stride = block_table.stride(0)
393 head_dim = k_fp8.shape[-1]
394 num_blocks = k_cache.shape[0]
395 # we assume the kv cache already been split to 2 portion
396 k_cache = k_cache.view(num_blocks, -1)
397 k_cache_value = k_cache[:, : block_size * head_dim].view(FP8_DTYPE)
398 k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)
399 grid = (num_tokens,)
400 k_fp8_scale = k_fp8_scale.view(torch.float32)
401 if cache_layout is None:
402 layout = "NORMAL" if block_size == 1 else "SHUFFLE"
403 else:
404 layout = cache_layout
405 kernel_args = (
406 k_cache_value,
407 k_cache_scale,
408 k_fp8,
409 k_fp8_scale,
410 block_table,
411 cu_seqlen,
412 token_to_seq,
413 block_size,
414 block_table_stride,
415 k_cache_value.stride(0),
416 k_cache_scale.stride(0),
417 layout,
418 head_dim,
419 block_tile_size,
420 head_tile_size,
421 )
422 if _ON_GFX950:
423 _cp_gather_indexer_quant_cache_gfx950_kernel[grid](
424 *kernel_args,
425 cu_seqlen.shape[0] - 1,
426 block_table.shape[1],
427 num_blocks,
428 )
429 else:
430 _cp_gather_indexer_quant_cache_kernel[grid](
431 *kernel_args,
432 num_tokens,
433 cu_seqlen.shape[0] - 1,
434 block_table.shape[1],
435 num_blocks,
436 )
437
438
439 # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156
440 def fp8_paged_mqa_logits_torch(
441 q: torch.Tensor,
442 kv_cache: torch.Tensor,
443 weights: torch.Tensor,
444 context_lens: torch.Tensor,
445 block_tables: torch.Tensor,
446 max_model_len: int,
447 ):
448 from vllm.utils.math_utils import cdiv
449
450 batch_size, next_n, _, dim = q.size()
451 if next_n == 1:
452 block_size = kv_cache.shape[1]
453 logits = torch.full(
454 [batch_size, max_model_len],
455 float("-inf"),
456 device=q.device,
457 dtype=torch.float32,
458 )
459 if context_lens.dim() > 1:
460 context_lens = context_lens.squeeze(-1)
461 kv_cache_flat = kv_cache.view(-1, block_size * (dim + 4))
462 for i in range(batch_size):
463 q_i = q[i, 0].to(torch.float32)
464 q_scale = weights[i]
465 seq_len = int(context_lens[i].item())
466 assert seq_len <= max_model_len
467 num_pages = cdiv(seq_len, block_size)
468 padded_seq_len = num_pages * block_size
469 pages = block_tables[i, :num_pages]
470 cache = kv_cache_flat[pages]
471 scale_offset = block_size * dim
472 cache_value = (
473 cache[..., :scale_offset].view(dtype=FP8_DTYPE).to(torch.float32)
474 )
475 cache_scale = (
476 cache[..., scale_offset:].view(dtype=torch.float32).contiguous()
477 )
478 cache_value = cache_value.view(padded_seq_len, dim)
479 cache_scale = cache_scale.view(padded_seq_len)
480 score = F.linear(cache_value, q_i)
481 score = F.relu(score)
482 score *= q_scale[None, :]
483 score = score.sum(dim=1)
484 score *= cache_scale
485 logits[i, :seq_len] = score[:seq_len]
486 return logits
487
488 kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
489 scale = scale.contiguous().view(torch.float)
490 q = q.float()
491 kv_cache = kv_cache.view(FP8_DTYPE).float() * scale
492 num_block, block_size, _, dim = kv_cache.size()
493 logits = torch.full(
494 [batch_size * next_n, max_model_len],
495 float("-inf"),
496 device=q.device,
497 dtype=torch.float32,
498 )
499 for i in range(batch_size):
500 context_len = context_lens[i]
501 if context_len.ndim == 0:
502 context_len_i = int(context_len.item())
503 q_offsets = torch.arange(
504 context_len_i - next_n, context_len_i, device=q.device
505 )
506 context_limit = torch.full(
507 (next_n,), context_len_i, dtype=torch.int32, device=q.device
508 )
509 else:
510 context_limit = context_len.to(device=q.device, dtype=torch.int32)
511 q_offsets = context_limit - 1
512 weight_slice = (
513 weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
514 )
515 max_context_len = int(context_limit.max().item())
516 for block_rk in range(cdiv(max_context_len, block_size)):
517 block_idx = block_tables[i][block_rk]
518 qx, kx = q[i], kv_cache[block_idx]
519 k_offsets = torch.arange(
520 block_rk * block_size, (block_rk + 1) * block_size, device=q.device
521 )
522 mask = (k_offsets[None, :] < context_limit[:, None]) & (
523 k_offsets[None, :] <= q_offsets[:, None]
524 )
525 s = torch.where(
526 mask[None, :, :],
527 (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(
528 logits.dtype
529 ),
530 float("-inf"),
531 )
532 s = torch.relu(s) * weight_slice[..., None]
533 s = s.sum(dim=0)
534 logits[
535 i * next_n : (i + 1) * next_n,
536 block_rk * block_size : (block_rk + 1) * block_size,
537 ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf"))
538 return logits
539
540
541 @triton.jit
542 def _fp8_paged_mqa_logits_decode_kernel(
543 q_ptr, # fp8 [B, NEXT_N, H, D]
544 kv_val_ptr, # fp8, block-flat: [num_blocks, block_size*(D+4)]
545 kv_scale_ptr, # fp32, block-flat: [num_blocks, block_size*(D+4)//4]
546 weights_ptr, # fp32 [B*NEXT_N, H]
547 ctx_lens_ptr, # int32 [B*NEXT_N] if CTX_PER_ROW else [B]
548 block_tables_ptr, # int32 [B, max_blocks]
549 logits_ptr, # fp32 [B*NEXT_N, max_model_len]
550 stride_q_b,
551 stride_q_n,
552 stride_q_h,
553 stride_w_row,
554 stride_kvblk_fp8,
555 stride_kvblk_f32,
556 scale_region_off, # block_size*D // 4 (fp32 offset of scale region within a block)
557 stride_bt_b,
558 stride_logits_row,
559 max_blocks, # block_tables width; guards the block-table gather
560 max_model_len,
561 NUM_HEADS: tl.constexpr,
562 HEAD_SIZE: tl.constexpr,
563 BLOCK_SIZE: tl.constexpr, # paged-cache page size
564 BLOCK_KV: tl.constexpr, # positions per tile; must divide BLOCK_SIZE
565 N_SPLITS: tl.constexpr, # KV-tile parallelism factor (grid dim 1)
566 NEXT_N: tl.constexpr, # query positions per batch (1 = decode; >1 = MTP verify)
567 CTX_PER_ROW: tl.constexpr, # ctx_lens is already per-(b, n)
568 ):
569 # scale_region_off = block_size*D//4 needs D % 4 == 0 (D=128).
570 tl.static_assert(HEAD_SIZE % 4 == 0)
571 # Grid (B*NEXT_N, N_SPLITS): disjoint KV tiles per program, no host sync.
572 row = tl.program_id(0)
573 split = tl.program_id(1)
574 b = row // NEXT_N
575 n = row % NEXT_N
576 if CTX_PER_ROW:
577 seq_len = tl.load(ctx_lens_ptr + row)
578 else:
579 seq_len = tl.load(ctx_lens_ptr + b) - NEXT_N + n + 1
580 seq_len = tl.minimum(tl.maximum(seq_len, 0), max_model_len)
581
582 h = tl.arange(0, NUM_HEADS)
583 d = tl.arange(0, HEAD_SIZE)
584 # keep q/kv fp8 -> fp8 MFMA (f32 accum), not the slow f32 path
585 q = tl.load(
586 q_ptr + b * stride_q_b + n * stride_q_n + h[:, None] * stride_q_h + d[None, :]
587 )
588 w = tl.load(weights_ptr + row * stride_w_row + h).to(tl.float32) # [H]
589
590 kv_col = tl.arange(0, BLOCK_KV)
591 for kv_start in tl.range(split * BLOCK_KV, seq_len, N_SPLITS * BLOCK_KV):
592 pos = kv_start + kv_col
593 # Tiles are page-aligned (BLOCK_KV | BLOCK_SIZE), so a tile is one page.
594 logical_blk = kv_start // BLOCK_SIZE
595 blk_ok = logical_blk < max_blocks # guard against a wild page -> fault
596 mask_pos = (pos < seq_len) & blk_ok
597 page = tl.load(
598 block_tables_ptr + b * stride_bt_b + logical_blk, mask=blk_ok, other=0
599 ).to(tl.int64)
600 pos_in_blk = pos - logical_blk * BLOCK_SIZE # [BLOCK_KV] == pos % BLOCK_SIZE
601
602 # block-flat layout: values region, then scales region.
603 val_off = page * stride_kvblk_fp8 + pos_in_blk[:, None] * HEAD_SIZE + d[None, :]
604 kv = tl.load(
605 kv_val_ptr + val_off, mask=mask_pos[:, None], other=0.0
606 ) # [BLOCK_KV, D] fp8
607 sc = tl.load(
608 kv_scale_ptr + page * stride_kvblk_f32 + scale_region_off + pos_in_blk,
609 mask=mask_pos,
610 other=0.0,
611 ) # [BLOCK_KV]
612
613 # fp8 upcasts exactly; bit-identical on gfx942, ~1e-5 rel on gfx950
614 # (accum order), top-k unchanged.
615 s = tl.dot(kv, tl.trans(q)) # [BLOCK_KV, H]
616 s = tl.maximum(s, 0.0)
617 s = s * w[None, :]
618 s = tl.sum(s, axis=1) # [BLOCK_KV]
619 s = s * sc
620 tl.store(logits_ptr + row * stride_logits_row + pos, s, mask=mask_pos)
621
622
623 def rocm_fp8_paged_mqa_logits_triton(
624 q_fp8: torch.Tensor,
625 kv_cache_fp8: torch.Tensor,
626 weights: torch.Tensor,
627 context_lens: torch.Tensor,
628 block_tables: torch.Tensor,
629 max_model_len: int,
630 ) -> torch.Tensor:
631 """Triton paged MQA-logits for decode and MTP; matches the torch ref but
632 has no host sync, so it is safe to capture under a full CUDA graph."""
633 batch_size, next_n, num_heads, head_size = q_fp8.shape
634 block_size = kv_cache_fp8.shape[1]
635 BLOCK_KV = 64
636 assert block_size % BLOCK_KV == 0
637
638 fp8_dtype = current_platform.fp8_dtype()
639 num_blocks = kv_cache_fp8.shape[0]
640 kv_flat = kv_cache_fp8.reshape(
641 num_blocks, -1
642 ) # uint8 [num_blocks, block_size*(D+4)]
643 kv_val = kv_flat.view(fp8_dtype) # [num_blocks, block_size*(D+4)] fp8
644 kv_scale = kv_flat.view(torch.float32) # [num_blocks, block_size*(D+4)//4] fp32
645
646 cl = context_lens.reshape(-1)
647 ctx_per_row = not (next_n > 1 and cl.numel() == batch_size)
648
649 max_blocks = block_tables.shape[1]
650 (out_logits,) = current_workspace_manager().get_simultaneous(
651 ((batch_size * next_n, max_model_len), torch.float32),
652 )
653
654 # Memory-bound over the KV range: split each row's keys across programs so
655 # few-row / long-context launches still fill the GPU. Cap splits at the
656 # device CU count (304 on gfx942, 256 on gfx950) rather than a gfx950-sized
657 # constant. All terms are static at launch, so the grid stays CUDA-graph-safe.
658 rows = batch_size * next_n
659 tiles_cap = (max_model_len + BLOCK_KV - 1) // BLOCK_KV
660 N_SPLITS = max(1, min(max(1, _decode_cu_count()), tiles_cap, 1024 // rows))
661 _fp8_paged_mqa_logits_decode_kernel[(rows, N_SPLITS)](
662 q_fp8,
663 kv_val,
664 kv_scale,
665 weights,
666 cl,
667 block_tables,
668 out_logits,
669 q_fp8.stride(0),
670 q_fp8.stride(1),
671 q_fp8.stride(2),
672 weights.stride(0),
673 kv_val.stride(0),
674 kv_scale.stride(0),
675 (block_size * head_size) // 4,
676 block_tables.stride(0),
677 out_logits.stride(0),
678 max_blocks,
679 max_model_len,
680 NUM_HEADS=num_heads,
681 HEAD_SIZE=head_size,
682 BLOCK_SIZE=block_size,
683 BLOCK_KV=BLOCK_KV,
684 N_SPLITS=N_SPLITS,
685 NEXT_N=next_n,
686 CTX_PER_ROW=ctx_per_row,
687 num_warps=4,
688 num_stages=2,
689 )
690 return out_logits
691
692
693 @functools.lru_cache
694 def paged_mqa_logits_module():
695 paged_mqa_logits_module_path = None
696 if find_spec("aiter.ops.triton.pa_mqa_logits") is not None:
697 paged_mqa_logits_module_path = "aiter.ops.triton.pa_mqa_logits"
698 elif find_spec("aiter.ops.triton.attention.pa_mqa_logits") is not None:
699 paged_mqa_logits_module_path = "aiter.ops.triton.attention.pa_mqa_logits"
700
701 if paged_mqa_logits_module_path is not None:
702 try:
703 module = importlib.import_module(paged_mqa_logits_module_path)
704 return module
705 except ImportError:
706 return None
707 return None
708
709
710 def rocm_fp8_paged_mqa_logits(
711 q_fp8: torch.Tensor,
712 kv_cache_fp8: torch.Tensor,
713 weights: torch.Tensor,
714 context_lens: torch.Tensor,
715 block_tables: torch.Tensor,
716 schedule_metadata: torch.Tensor,
717 max_model_len: int,
718 *,
719 compress_ratio: int = 1,
720 ) -> torch.Tensor:
721 """Compute FP8 MQA logits using paged KV-cache.
722
723 Args:
724 q_fp8: Query tensor of shape [B, next_n, H, D]. Casted to
725 `torch.float8_e4m3fn` by caller.
726 kv_cache_fp8: Paged KV-cache in packed FP8+scale layout with shape
727 [num_blocks, block_size, 1, D+4], dtype `torch.uint8`.
728 weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
729 context_lens: Tensor of shape [B], dtype int32; effective context length
730 for each batch element.
731 block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
732 block indices to physical blocks in the paged cache.
733 schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
734 used to distribute work across SMs.
735 max_model_len: Maximum sequence length used to size the logits output.
736 compress_ratio: C4A (4) takes block-flat Triton; 1 and 2 stay on AITER.
737
738 Returns:
739 Logits tensor of shape [B * next_n, max_model_len], dtype
740 `torch.float32`.
741
742 """
743 from vllm._aiter_ops import rocm_aiter_ops
744
745 batch_size, next_n = q_fp8.shape[:2]
746 block_size = kv_cache_fp8.shape[1]
747
748 # C4A only: Flash/DSv3.2 also skip insert but still write SHUFFLE.
749 if (
750 (_ON_GFX950 or _ON_GFX942)
751 and _indexer_k_is_c4a_block_flat(compress_ratio)
752 and block_size > 1
753 ):
754 if block_size % 64 == 0:
755 return rocm_fp8_paged_mqa_logits_triton(
756 q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
757 )
758 # Non 64-aligned page size (not used in prod): eager torch ref.
759 return fp8_paged_mqa_logits_torch(
760 q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
761 )
762
763 aiter_paged_mqa_logits_module = None
764
765 if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():
766 aiter_paged_mqa_logits_module = paged_mqa_logits_module()
767
768 if aiter_paged_mqa_logits_module is not None:
769 if _ON_GFX942 or _ON_GFX950:
770 deepgemm_fp8_paged_mqa_logits = (
771 aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
772 )
773 batch_size, next_n, heads, _ = q_fp8.shape
774 (out_logits,) = current_workspace_manager().get_simultaneous(
775 ((batch_size * next_n, max_model_len), torch.float32),
776 )
777 deepgemm_fp8_paged_mqa_logits(
778 q_fp8,
779 kv_cache_fp8,
780 weights,
781 out_logits,
782 context_lens,
783 block_tables,
784 max_model_len,
785 ChunkK=256,
786 Preshuffle=block_size > 1,
787 KVBlockSize=block_size,
788 WavePerEU=2,
789 )
790 return out_logits
791 deepgemm_fp8_paged_mqa_logits_stage1 = (
792 aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1
793 )
794 batch_size, next_n, heads, _ = q_fp8.shape
795 (out_qk,) = current_workspace_manager().get_simultaneous(
796 ((heads, batch_size * next_n, max_model_len), torch.float32),
797 )
798 out_qk.fill_(float("-inf"))
799 deepgemm_fp8_paged_mqa_logits_stage1(
800 q_fp8,
801 kv_cache_fp8,
802 weights,
803 out_qk,
804 context_lens,
805 block_tables,
806 max_model_len,
807 ChunkQ=heads,
808 )
809 return out_qk.sum(dim=0)
810 else:
811 return fp8_paged_mqa_logits_torch(
812 q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
813 )
814
815
816 # Take from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L84
817 def fp8_mqa_logits_torch(
818 q: torch.Tensor,
819 kv: tuple[torch.Tensor, torch.Tensor],
820 weights: torch.Tensor,
821 cu_seqlen_ks: torch.Tensor,
822 cu_seqlen_ke: torch.Tensor,
823 ) -> torch.Tensor:
824 """Compute FP8 MQA logits for a single sequence without KV paging.
825
826 Args:
827 q: Query tensor of shape [M, H, D]. Casted to
828 `torch.float8_e4m3fn` by caller.
829 kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
830 dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
831 [N, 1]) with dtype `torch.float32`.
832 weights: weights of shape [M, H], dtype `torch.float32`.
833 cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
834 shape [M], dtype int32.
835 cu_seqlen_ke: End indices (exclusive) for valid K per query position,
836 shape [M], dtype int32.
837
838 Returns:
839 Logits tensor of shape [M, N], dtype `torch.float32`.
840
841 """
842 k_fp8, scale = kv
843 seq_len_kv = k_fp8.shape[0]
844 k = k_fp8.to(torch.bfloat16)
845 q = q.to(torch.bfloat16)
846 device = q.device
847
848 mask_lo = (
849 torch.arange(0, seq_len_kv, device=device)[None, :] >= cu_seqlen_ks[:, None]
850 )
851 mask_hi = (
852 torch.arange(0, seq_len_kv, device=device)[None, :] < cu_seqlen_ke[:, None]
853 )
854 mask = mask_lo & mask_hi
855
856 # ``score`` is [H, M, N]; ``scale`` is the per-KV-token scale, which
857 # vLLM callers hand us as ``[N, 1]`` (a ``[N, 4]`` uint8 buffer cast
858 # to fp32). PyTorch right-aligns dimensions for broadcasting, so a
859 # naked ``score * scale`` would align ``scale``'s leading dim with
860 # ``score``'s M dim and raise a shape mismatch. Flatten to ``[N]`` so
861 # broadcasting lines up with the last dim of ``score``.
862 score = torch.einsum("mhd,nd->hmn", q, k).float() * scale.reshape(-1)
863 logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
864 logits = logits.masked_fill(~mask, float("-inf"))
865
866 return logits
867
868
869 @functools.lru_cache
870 def mqa_logits_module():
871 mqa_logits_module_path = None
872 if find_spec("aiter.ops.triton.fp8_mqa_logits") is not None:
873 mqa_logits_module_path = "aiter.ops.triton.fp8_mqa_logits"
874 elif find_spec("aiter.ops.triton.attention.fp8_mqa_logits") is not None:
875 mqa_logits_module_path = "aiter.ops.triton.attention.fp8_mqa_logits"
876
877 if mqa_logits_module_path is not None:
878 try:
879 module = importlib.import_module(mqa_logits_module_path)
880 return module
881 except ImportError:
882 return None
883 return None
884
885
886 def rocm_fp8_mqa_logits(
887 q: torch.Tensor,
888 kv: tuple[torch.Tensor, torch.Tensor],
889 weights: torch.Tensor,
890 cu_seqlen_ks: torch.Tensor,
891 cu_seqlen_ke: torch.Tensor,
892 ) -> torch.Tensor:
893 """Compute FP8 MQA logits for a single sequence without KV paging.
894
895 Args:
896 q: Query tensor of shape [M, H, D]. Casted to
897 `torch.float8_e4m3fn` by caller.
898 kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
899 dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
900 [N, 1]) with dtype `torch.float32`.
901 weights: weights of shape [M, H], dtype `torch.float32`.
902 cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
903 shape [M], dtype int32.
904 cu_seqlen_ke: End indices (exclusive) for valid K per query position,
905 shape [M], dtype int32.
906
907 Returns:
908 Logits tensor of shape [M, N], dtype `torch.float32`.
909
910 """
911 from vllm._aiter_ops import rocm_aiter_ops
912
913 k_fp8, scale = kv
914
915 if _ON_GFX942 and rocm_aiter_ops.is_enabled():
916 from aiter.ops.flydsl import flydsl_fp8_mqa_logits
917
918 return flydsl_fp8_mqa_logits(
919 q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke
920 )
921
922 aiter_mqa_logits_module = None
923 if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():
924 aiter_mqa_logits_module = mqa_logits_module()
925
926 if aiter_mqa_logits_module is not None:
927 fp8_mqa_logits = aiter_mqa_logits_module.fp8_mqa_logits
928 return fp8_mqa_logits(q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke)
929 else:
930 return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)
931
932
933 # Programs along the column axis of the decode mask grid. The shared kernel in
934 # model_executor/kernels/attention/dsa/candidate_blocks.py launches one program
935 # per 1024-column tile, which is 1024 of them per row at max_model_len 1048576
936 # with only the first few doing work. This is a fixed count that strides
937 # instead, measured on gfx950; it is not a portable choice, which is why this
938 # variant lives here rather than replacing the shared one.
939 _MASK_GRID_COLS = 128
940 # Columns each program handles per iteration. Work stops at the tile boundary
941 # containing a row's end rather than at the end itself; the overshoot is
942 # masked the same way the shared kernel masks it, so it costs a tile and
943 # changes nothing.
944 _MASK_TILE = 1024
945
946
947 @triton.jit(do_not_specialize=["width", "nblocks"])
948 def _mask_candidates_strided_kernel(
949 logits,
950 starts,
951 ends,
952 flags,
953 stride_row,
954 stride_col,
955 stride_start,
956 stride_end,
957 width,
958 nblocks,
959 BLOCK_SIZE: tl.constexpr,
960 HAS_STARTS: tl.constexpr,
961 ROW_REPEAT: tl.constexpr,
962 TILE: tl.constexpr,
963 ):
964 row = tl.program_id(0).to(tl.int64)
965 start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0
966 end = tl.load(ends + row // ROW_REPEAT * stride_end)
967 edge = tl.load(flags + row * (nblocks + 1) + nblocks)
968 step = tl.num_programs(1) * TILE
969 tile_start = tl.program_id(1) * TILE
970 # The shared kernel also sanitizes columns at or past `end`. Nothing reads
971 # them: top_k_per_row_decode bounds its scan by the same row ends passed
972 # here, and persistent_topk/cooperative_topk clamp by the same lengths.
973 # Stopping at `end` is what makes the cost track the live context rather
974 # than max_model_len.
975 while tile_start < end:
976 cols = tile_start + tl.arange(0, TILE)
977 valid = (cols >= start) & (cols < end) & (cols < width)
978 block = (cols - start) // BLOCK_SIZE
979 keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0)
980 keep = (keep != 0) | ((cols == width - 1) & (edge != 0))
981 tl.store(
982 logits + row * stride_row + cols * stride_col,
983 -float("inf"),
984 (cols < width) & ~(valid & keep),
985 )
986 tile_start += step
987
988
989 def _apply_candidate_mask_strided(
990 logits: torch.Tensor,
991 row_ks: torch.Tensor | None,
992 row_ke: torch.Tensor,
993 candidate_blocks: torch.Tensor,
994 block_size: int,
995 row_repeat: int = 1,
996 ) -> None:
997 """ROCm decode variant of ``apply_candidate_mask``.
998
999 Same masking semantics over ``[0, end)``, but the grid is sized by a fixed
1000 program count rather than by the logits width. Only worth using where the
1001 width is the ``max_model_len`` workspace and the live context is far
1002 shorter, i.e. the paged decode path below; the prefill chunks pass
1003 chunk-sized logits and stay on the shared kernel.
1004 """
1005 from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
1006 _candidate_flags_kernel,
1007 )
1008
1009 rows, width = logits.shape
1010 if not rows or not width:
1011 return
1012 nblocks = triton.cdiv(width, block_size)
1013 flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8)
1014 start_stride = row_ks.stride(0) if row_ks is not None else 0
1015 _candidate_flags_kernel[(rows,)](
1016 candidate_blocks,
1017 row_ks,
1018 flags,
1019 *candidate_blocks.stride(),
1020 start_stride,
1021 width,
1022 nblocks,
1023 block_size,
1024 candidate_blocks.shape[1],
1025 row_ks is not None,
1026 row_repeat,
1027 )
1028 # Derived from width, which is a tensor shape, so the grid stays static and
1029 # a FULL cudagraph capture remains valid across replays; only the loop trip
1030 # count inside the kernel is data-dependent. The min keeps narrow widths
1031 # from launching programs that would only fall through.
1032 grid_cols = min(_MASK_GRID_COLS, triton.cdiv(width, _MASK_TILE))
1033 _mask_candidates_strided_kernel[(rows, grid_cols)](
1034 logits,
1035 row_ks,
1036 row_ke,
1037 flags,
1038 *logits.stride(),
1039 start_stride,
1040 row_ke.stride(0),
1041 width,
1042 nblocks,
1043 block_size,
1044 row_ks is not None,
1045 row_repeat,
1046 _MASK_TILE,
1047 )
1048
1049
1050 def _max_decode_logits_rows(num_batched_tokens: int) -> int:
1051 """Upper bound on decode rows the paged-MQA logits buffer can ever hold.
1052
1053 ``rocm_fp8_paged_mqa_logits`` sizes its workspace as
1054 ``(batch_size * next_n, max_model_len)``. ``batch_size`` is bounded by
1055 ``max_num_seqs`` and ``next_n`` by ``1 + num_speculative_tokens``, which is
1056 far tighter than ``max_num_batched_tokens`` -- 192 vs 16384 for a typical
1057 32-seq DSpark-5 deployment. The loose bound is harmless at short contexts
1058 but scales with ``max_model_len``, so at the model's full context it asks
1059 for tens of TiB and the engine cannot start. Take whichever valid bound is
1060 smaller; the workspace is locked after profiling, so it must not be under-
1061 estimated.
1062 """
1063 try:
1064 vllm_config = get_current_vllm_config()
1065 except Exception:
1066 return num_batched_tokens
1067 scheduler_config = getattr(vllm_config, "scheduler_config", None)
1068 max_num_seqs = getattr(scheduler_config, "max_num_seqs", None)
1069 if not max_num_seqs:
1070 return num_batched_tokens
1071 speculative_config = getattr(vllm_config, "speculative_config", None)
1072 num_spec = getattr(speculative_config, "num_speculative_tokens", 0) or 0
1073 return min(num_batched_tokens, max_num_seqs * (1 + num_spec))
1074
1075
1076 def rocm_aiter_sparse_attn_indexer_fake(
1077 hidden_states: torch.Tensor,
1078 k_cache_prefix: LayerNameType,
1079 kv_cache: torch.Tensor,
1080 q_fp8: torch.Tensor,
1081 k: torch.Tensor | None,
1082 weights: torch.Tensor,
1083 quant_block_size: int,
1084 scale_fmt: str | None,
1085 topk_tokens: int,
1086 head_dim: int,
1087 max_model_len: int,
1088 total_seq_lens: int,
1089 topk_indices_buffer: torch.Tensor | None,
1090 skip_k_cache_insert: bool = False,
1091 compress_ratio: int = 1,
1092 candidate_blocks: torch.Tensor | None = None,
1093 candidate_block_size: int = 0,
1094 candidate_write: bool = False,
1095 ) -> torch.Tensor:
1096 return topk_indices_buffer
1097
1098
1099 @eager_break_during_capture
1100 def rocm_aiter_sparse_attn_indexer(
1101 hidden_states: torch.Tensor,
1102 k_cache_prefix: LayerNameType,
1103 kv_cache: torch.Tensor,
1104 q_fp8: torch.Tensor,
1105 k: torch.Tensor | None,
1106 weights: torch.Tensor,
1107 quant_block_size: int,
1108 scale_fmt: str | None,
1109 topk_tokens: int,
1110 head_dim: int,
1111 max_model_len: int,
1112 total_seq_lens: int,
1113 topk_indices_buffer: torch.Tensor | None,
1114 skip_k_cache_insert: bool = False,
1115 compress_ratio: int = 1,
1116 candidate_blocks: torch.Tensor | None = None,
1117 candidate_block_size: int = 0,
1118 candidate_write: bool = False,
1119 ) -> torch.Tensor:
1120 from vllm._aiter_ops import rocm_aiter_ops
1121
1122 # careful! this will be None in dummy run
1123 forward_context = get_forward_context()
1124 attn_metadata = forward_context.attn_metadata
1125 from vllm.utils.torch_utils import _resolve_layer_name
1126
1127 k_cache_prefix = _resolve_layer_name(k_cache_prefix)
1128 # assert isinstance(attn_metadata, dict)
1129 if not isinstance(attn_metadata, dict):
1130 # Profiling early-exit: reserve memory to account for runtime
1131 # allocations. Must be in the real impl, not the fake impl —
1132 # torch.compile calls the fake impl under FakeTensor mode where
1133 # workspace manager operations on the locked real workspace
1134 # would corrupt PyTorch's dispatch state.
1135 workspace_manager = current_workspace_manager()
1136
1137 # Prefill k_fp8 and k_scale buffers, used by
1138 # rocm_aiter_sparse_attn_indexer's prefill path
1139 workspace_manager.get_simultaneous(
1140 ((total_seq_lens, head_dim), FP8_DTYPE),
1141 ((total_seq_lens, 4), torch.uint8),
1142 )
1143
1144 # Decode logits buffer, used by rocm_fp8_paged_mqa_logits.
1145 decode_rows = _max_decode_logits_rows(hidden_states.shape[0])
1146 if _ON_GFX942 or _ON_GFX950:
1147 workspace_manager.get_simultaneous(
1148 ((decode_rows, max_model_len), torch.float32),
1149 )
1150 else:
1151 workspace_manager.get_simultaneous(
1152 (
1153 (q_fp8.shape[1], decode_rows, max_model_len),
1154 torch.float32,
1155 ),
1156 )
1157 # Transient logits tensor peak memory, produced by
1158 # rocm_fp8_mqa_logits (prefill) and rocm_fp8_paged_mqa_logits
1159 # (decode). Prefill logits are bounded by
1160 # VLLM_SPARSE_INDEXER_MAX_LOGITS_MB via chunking in
1161 # split_indexer_prefill_chunks; decode logits are smaller.
1162 max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
1163 _ = torch.empty(
1164 max_logits_elems, dtype=torch.uint8, device=hidden_states.device
1165 )
1166
1167 return rocm_aiter_sparse_attn_indexer_fake(
1168 hidden_states,
1169 k_cache_prefix,
1170 kv_cache,
1171 q_fp8,
1172 k,
1173 weights,
1174 quant_block_size,
1175 scale_fmt,
1176 topk_tokens,
1177 head_dim,
1178 max_model_len,
1179 total_seq_lens,
1180 topk_indices_buffer,
1181 skip_k_cache_insert,
1182 compress_ratio,
1183 candidate_blocks,
1184 candidate_block_size,
1185 candidate_write,
1186 )
1187 layer_attn_metadata = attn_metadata[k_cache_prefix]
1188 assert isinstance(layer_attn_metadata, DeepseekV32IndexerMetadata)
1189 assert topk_indices_buffer is not None
1190 assert scale_fmt is not None
1191 slot_mapping = layer_attn_metadata.slot_mapping
1192 has_decode = layer_attn_metadata.num_decodes > 0
1193 has_prefill = layer_attn_metadata.num_prefills > 0
1194 num_decode_tokens = layer_attn_metadata.num_decode_tokens
1195 topk_indices_buffer[: hidden_states.shape[0]] = -1
1196
1197 # during speculative decoding, k may be padded to the CUDA graph batch
1198 # size while slot_mapping only covers actual tokens.
1199 num_tokens = slot_mapping.shape[0]
1200 if k is not None:
1201 k = k[:num_tokens]
1202 elif not skip_k_cache_insert:
1203 raise ValueError("k must be provided when skip_k_cache_insert is False")
1204
1205 if not skip_k_cache_insert:
1206 indexer_k_quant_and_cache_triton(
1207 k,
1208 kv_cache,
1209 slot_mapping,
1210 quant_block_size,
1211 scale_fmt,
1212 )
1213
1214 if has_prefill:
1215 prefill_metadata = layer_attn_metadata.prefill
1216 assert prefill_metadata is not None
1217
1218 workspace_manager = current_workspace_manager()
1219 k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
1220 ((total_seq_lens, head_dim), FP8_DTYPE),
1221 ((total_seq_lens, 4), torch.uint8),
1222 )
1223 for chunk in prefill_metadata.chunks:
1224 k_fp8 = k_fp8_full[: chunk.total_seq_lens]
1225 k_scale = k_scale_full[: chunk.total_seq_lens]
1226 cp_gather_indexer_k_quant_cache_triton(
1227 kv_cache,
1228 k_fp8,
1229 k_scale,
1230 chunk.block_table,
1231 chunk.cu_seq_lens,
1232 token_to_seq=chunk.token_to_seq,
1233 cache_layout=(
1234 "NORMAL" if _indexer_k_is_c4a_block_flat(compress_ratio) else None
1235 ),
1236 )
1237 logits = rocm_fp8_mqa_logits(
1238 q_fp8[chunk.token_start : chunk.token_end],
1239 (k_fp8, k_scale.view(torch.float32)),
1240 weights[chunk.token_start : chunk.token_end],
1241 chunk.cu_seqlen_ks,
1242 chunk.cu_seqlen_ke,
1243 )
1244 if candidate_blocks is not None:
1245 from vllm.model_executor.layers.sparse_attn_indexer import (
1246 _apply_candidate_mask,
1247 _select_candidate_blocks,
1248 )
1249
1250 chunk_candidates = candidate_blocks[chunk.token_start : chunk.token_end]
1251 if candidate_write:
1252 _select_candidate_blocks(
1253 logits,
1254 chunk.cu_seqlen_ks,
1255 chunk.cu_seqlen_ke,
1256 chunk_candidates.shape[1],
1257 candidate_block_size,
1258 chunk_candidates,
1259 )
1260 else:
1261 _apply_candidate_mask(
1262 logits,
1263 chunk.cu_seqlen_ks,
1264 chunk.cu_seqlen_ke,
1265 chunk_candidates,
1266 candidate_block_size,
1267 )
1268 topk_indices = topk_indices_buffer[
1269 chunk.token_start : chunk.token_end, :topk_tokens
1270 ]
1271
1272 num_rows = logits.shape[0]
1273
1274 if rocm_aiter_ops.is_indexer_top_k_supported(
1275 is_prefill=True,
1276 compress_ratio=compress_ratio,
1277 num_rows=num_rows,
1278 ):
1279 _launch_aiter_top_k_per_row_prefill(
1280 logits,
1281 chunk.cu_seqlen_ks,
1282 chunk.cu_seqlen_ke,
1283 topk_indices,
1284 topk_tokens,
1285 )
1286 else:
1287 torch.ops._C.top_k_per_row_prefill(
1288 logits,
1289 chunk.cu_seqlen_ks,
1290 chunk.cu_seqlen_ke,
1291 topk_indices,
1292 num_rows,
1293 logits.stride(0),
1294 logits.stride(1),
1295 topk_tokens,
1296 )
1297
1298 if has_decode:
1299 decode_metadata = layer_attn_metadata.decode
1300 assert decode_metadata is not None
1301 # kv_cache size requirement [num_block, block_size, n_head, head_dim],
1302 # we only have [num_block, block_size, head_dim],
1303 kv_cache = kv_cache.unsqueeze(-2)
1304 decode_lens = decode_metadata.decode_lens
1305 if decode_metadata.requires_padding:
1306 # pad in edge case where we have short chunked prefill length <
1307 # decode_threshold since we unstrictly split
1308 # prefill and decode by decode_threshold
1309 # (currently set to 1 + speculative tokens)
1310 padded_q_fp8_decode_tokens = pack_seq_triton(
1311 q_fp8[:num_decode_tokens], decode_lens
1312 )
1313 else:
1314 padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
1315 decode_lens.shape[0], -1, *q_fp8.shape[1:]
1316 )
1317 # TODO: move and optimize below logic with triton kernels
1318 batch_size = padded_q_fp8_decode_tokens.shape[0]
1319 next_n = padded_q_fp8_decode_tokens.shape[1]
1320 assert batch_size == decode_metadata.seq_lens.shape[0]
1321 num_padded_tokens = batch_size * next_n
1322
1323 logits = rocm_fp8_paged_mqa_logits(
1324 padded_q_fp8_decode_tokens,
1325 kv_cache,
1326 weights[:num_padded_tokens],
1327 decode_metadata.seq_lens,
1328 decode_metadata.block_table,
1329 decode_metadata.schedule_metadata,
1330 max_model_len=max_model_len,
1331 compress_ratio=compress_ratio,
1332 )
1333
1334 if candidate_blocks is not None:
1335 from vllm.model_executor.layers.sparse_attn_indexer import (
1336 _select_candidate_blocks,
1337 )
1338
1339 num_rows = logits.shape[0]
1340 visible = decode_metadata.seq_lens.reshape(-1)
1341 if visible.numel() != num_rows:
1342 visible = visible.repeat_interleave(next_n)
1343 visible = visible[:num_rows].to(torch.int64)
1344 row_starts = torch.zeros_like(visible)
1345 decode_candidates = candidate_blocks[:num_rows]
1346 if candidate_write:
1347 _select_candidate_blocks(
1348 logits,
1349 row_starts,
1350 visible,
1351 decode_candidates.shape[1],
1352 candidate_block_size,
1353 decode_candidates,
1354 )
1355 else:
1356 _apply_candidate_mask_strided(
1357 logits,
1358 row_starts,
1359 visible,
1360 decode_candidates,
1361 candidate_block_size,
1362 )
1363
1364 topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
1365 num_rows = logits.shape[0]
1366
1367 # FULL graphs are not keyed by context length, so use a replay-safe
1368 # upper bound when choosing the captured kernel.
1369 if forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL:
1370 max_compressed_seq_len = max_model_len
1371 else:
1372 max_compressed_seq_len = layer_attn_metadata.max_seq_len // compress_ratio
1373 if rocm_aiter_ops.is_indexer_top_k_supported(
1374 is_prefill=False,
1375 compress_ratio=compress_ratio,
1376 num_rows=num_rows,
1377 max_valid_seq_len=max_compressed_seq_len,
1378 ) and not rocm_aiter_ops.dsv4_indexer_prefers_native_top_k(
1379 num_rows=num_rows,
1380 num_columns=logits.shape[1],
1381 topk_tokens=topk_tokens,
1382 ):
1383 rocm_aiter_ops.indexer_top_k_decode(
1384 logits,
1385 1,
1386 decode_metadata.seq_lens.reshape(-1),
1387 topk_indices,
1388 topk_tokens,
1389 )
1390 else:
1391 torch.ops._C.top_k_per_row_decode(
1392 logits,
1393 next_n,
1394 decode_metadata.seq_lens,
1395 topk_indices,
1396 num_rows,
1397 logits.stride(0),
1398 logits.stride(1),
1399 topk_tokens,
1400 )
1401
1402 if decode_metadata.requires_padding:
1403 # if padded, we need to unpack
1404 # the topk indices removing padded tokens
1405 topk_indices = unpack_seq_triton(
1406 topk_indices.reshape(batch_size, next_n, topk_indices.shape[-1]),
1407 decode_lens,
1408 )
1409 topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (
1410 topk_indices
1411 )
1412
1413 return topk_indices_buffer
1414
1415
1416 def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor:
1417 if scale.dtype == torch.float8_e8m0fnu:
1418 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
1419 _upcast_e8m0_to_fp32,
1420 )
1421
1422 return _upcast_e8m0_to_fp32(scale).contiguous()
1423 if scale.dtype == torch.uint8:
1424 # MXFP8 parameters preserve E8M0 scales as their raw exponent bytes.
1425 # They are biased exponents, not numeric uint8 scale values.
1426 return torch.exp2(scale.to(torch.int16).to(torch.float32) - 127.0)
1427 return scale.to(torch.float32)
1428
1429
1430 def _expand_2d_block_scales(
1431 scale: torch.Tensor,
1432 rows: int,
1433 cols: int,
1434 ) -> torch.Tensor:
1435 scale = _decode_e8m0_scales(scale)
1436 row_blocks, col_blocks = scale.shape[-2:]
1437 row_block = math.ceil(rows / row_blocks)
1438 col_block = math.ceil(cols / col_blocks)
1439 scale = torch.repeat_interleave(scale, row_block, dim=-2)[..., :rows, :]
1440 scale = torch.repeat_interleave(scale, col_block, dim=-1)[..., :, :cols]
1441 return scale
1442
1443
1444 @triton.jit
1445 def _inverse_rope_gptj_kernel(
1446 o_ptr, # [T, H, D] input
1447 out_ptr, # [T, H, D] bf16 output
1448 pos_ptr, # [T] positions
1449 cos_sin_ptr, # [P, rope_dim] fp32 (cos[:half] | sin[half:])
1450 s_t,
1451 s_h, # input row strides (last dim contiguous)
1452 os_t,
1453 os_h, # output row strides
1454 cs_stride, # cos_sin_cache row stride
1455 NOPE: tl.constexpr, # non-rope head dims (passed through)
1456 HALF: tl.constexpr, # rope_dim // 2
1457 BLOCK_NOPE: tl.constexpr,
1458 BLOCK_HALF: tl.constexpr,
1459 ):
1460 """Fused inverse GPT-J RoPE on the trailing rope_dim of each (token, head).
1461
1462 Mirrors ``DeepseekV4ScalingRotaryEmbedding.forward_native(inverse=True)``
1463 for the GPT-J (non-neox) layout, writing bf16 directly. Replaces the
1464 clone + index_select + repeat_interleave + neg + stack + cat + cast chain
1465 (~10 small kernels) with a single launch.
1466 """
1467 t = tl.program_id(0)
1468 h = tl.program_id(1)
1469 in_base = t * s_t + h * s_h
1470 out_base = t * os_t + h * os_h
1471
1472 # NoPE lanes pass through unchanged (only cast to bf16).
1473 n = tl.arange(0, BLOCK_NOPE)
1474 nmask = n < NOPE
1475 vals = tl.load(o_ptr + in_base + n, mask=nmask)
1476 tl.store(out_ptr + out_base + n, vals.to(tl.bfloat16), mask=nmask)
1477
1478 # RoPE lanes: out_even = a*cos + b*sin, out_odd = b*cos - a*sin
1479 # (a = even lane, b = odd lane; sin negated for the inverse rotation).
1480 pos = tl.load(pos_ptr + t).to(tl.int64)
1481 k = tl.arange(0, BLOCK_HALF)
1482 kmask = k < HALF
1483 a = tl.load(o_ptr + in_base + NOPE + 2 * k, mask=kmask).to(tl.float32)
1484 b = tl.load(o_ptr + in_base + NOPE + 2 * k + 1, mask=kmask).to(tl.float32)
1485 cos = tl.load(cos_sin_ptr + pos * cs_stride + k, mask=kmask)
1486 sin = tl.load(cos_sin_ptr + pos * cs_stride + HALF + k, mask=kmask)
1487 out_even = a * cos + b * sin
1488 out_odd = b * cos - a * sin
1489 tl.store(out_ptr + out_base + NOPE + 2 * k, out_even.to(tl.bfloat16), mask=kmask)
1490 tl.store(out_ptr + out_base + NOPE + 2 * k + 1, out_odd.to(tl.bfloat16), mask=kmask)
1491
1492
1493 def _fused_inverse_rope_gptj(
1494 o: torch.Tensor,
1495 positions: torch.Tensor,
1496 cos_sin_cache: torch.Tensor,
1497 rope_head_dim: int,
1498 out: torch.Tensor | None = None,
1499 ) -> torch.Tensor:
1500 """bf16 inverse GPT-J RoPE via a single fused Triton kernel.
1501
1502 ``out`` may alias ``o``: the rotation is a per-row bijection whose kernel
1503 reads both lanes of a pair before storing either.
1504 """
1505 assert o.dim() == 3 and o.stride(-1) == 1, (
1506 "_fused_inverse_rope_gptj expects a [T, H, D] input with a contiguous last dim"
1507 )
1508 assert rope_head_dim > 0 and rope_head_dim % 2 == 0, (
1509 f"_fused_inverse_rope_gptj expects an even rope_head_dim, got {rope_head_dim}"
1510 )
1511 assert cos_sin_cache.shape[-1] == rope_head_dim, (
1512 "_fused_inverse_rope_gptj expects cos_sin_cache laid out as "
1513 f"[P, {rope_head_dim}] = cos | sin, got {tuple(cos_sin_cache.shape)}"
1514 )
1515 num_tokens, num_heads, head_dim = o.shape
1516 if out is None:
1517 out = torch.empty(
1518 (num_tokens, num_heads, head_dim), dtype=torch.bfloat16, device=o.device
1519 )
1520 else:
1521 assert out.dtype == torch.bfloat16, (
1522 f"inverse RoPE writes bf16, got an output buffer of {out.dtype}"
1523 )
1524 if num_tokens == 0:
1525 return out
1526 _inverse_rope_gptj_kernel[(num_tokens, num_heads)](
1527 o,
1528 out,
1529 positions,
1530 cos_sin_cache,
1531 o.stride(0),
1532 o.stride(1),
1533 out.stride(0),
1534 out.stride(1),
1535 cos_sin_cache.stride(0),
1536 NOPE=head_dim - rope_head_dim,
1537 HALF=rope_head_dim // 2,
1538 BLOCK_NOPE=triton.next_power_of_2(head_dim - rope_head_dim),
1539 BLOCK_HALF=triton.next_power_of_2(rope_head_dim // 2),
1540 )
1541 return out
1542
1543
1544 def rocm_inverse_rope_rows_(
1545 o: torch.Tensor,
1546 positions: torch.Tensor,
1547 cos_sin_cache: torch.Tensor,
1548 rope_head_dim: int,
1549 ) -> None:
1550 """Inverse-RoPE attention output rows in place.
1551
1552 For rows no attention kernel rotated in its epilogue. Call it from the
1553 eager attention segment: which rows still owe a rotation depends on the
1554 prefill/decode split, and the o_proj that used to do this runs inside the
1555 compiled region, where a batch-dependent Python value would be frozen at
1556 trace time.
1557 """
1558 if o.shape[0] == 0:
1559 return
1560 _fused_inverse_rope_gptj(o, positions, cos_sin_cache, rope_head_dim, out=o)
1561
1562
1563 @triton.jit
1564 def _inverse_rope_mxfp8_quant_kernel(
1565 o_ptr, # [T, H, D] bf16
1566 q_ptr, # [T, H * D] e4m3
1567 s_ptr, # [T, H * D // 32] uint8 (E8M0)
1568 pos_ptr,
1569 cos_sin_ptr,
1570 s_t,
1571 s_h,
1572 qs_t,
1573 ss_t,
1574 cs_stride,
1575 HEAD_DIM: tl.constexpr,
1576 NOPE: tl.constexpr,
1577 HALF: tl.constexpr,
1578 ):
1579 t = tl.program_id(0)
1580 h = tl.program_id(1)
1581 lanes = tl.arange(0, HEAD_DIM)
1582 x = tl.load(o_ptr + t * s_t + h * s_h + lanes).to(tl.float32)
1583 # Same lane-pair formulation as the reduce epilogue: NoPE pairs take
1584 # cos=1 / sin=0, so one expression rotates the whole row.
1585 pos = tl.load(pos_ptr + t).to(tl.int64)
1586 pair_idx = tl.arange(0, HEAD_DIM // 2) - (NOPE // 2)
1587 is_rope = pair_idx >= 0
1588 k = tl.where(is_rope, pair_idx, 0)
1589 cos = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + k), 1.0)
1590 sin = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + HALF + k), 0.0)
1591 even, odd = tl.split(tl.reshape(x, (HEAD_DIM // 2, 2)))
1592 x = tl.reshape(
1593 tl.join(even * cos + odd * sin, odd * cos - even * sin), (1, HEAD_DIM)
1594 )
1595 xq, bits = _mxfp8_quantize_rows(x, 1, HEAD_DIM)
1596 tl.store(
1597 q_ptr + t * qs_t + h * HEAD_DIM + lanes,
1598 tl.reshape(xq, (HEAD_DIM,)).to(q_ptr.dtype.element_ty),
1599 )
1600 scale_offsets = tl.arange(0, HEAD_DIM // 32)
1601 tl.store(
1602 s_ptr + t * ss_t + h * (HEAD_DIM // 32) + scale_offsets,
1603 tl.reshape(bits, (HEAD_DIM // 32,)).to(tl.uint8),
1604 )
1605
1606
1607 def rocm_inverse_rope_mxfp8_rows(
1608 o: torch.Tensor,
1609 positions: torch.Tensor,
1610 cos_sin_cache: torch.Tensor,
1611 rope_head_dim: int,
1612 out_data: torch.Tensor,
1613 out_scale: torch.Tensor,
1614 ) -> None:
1615 """Inverse-RoPE bf16 attention rows and MXFP8-quantize them for wo_a.
1616
1617 The counterpart of ``rocm_inverse_rope_rows_`` for layers whose attention
1618 output is MXFP8: rows the decode reduce did not emit (prefill) go through
1619 here. ``o`` is [T, H, D]; ``out_data`` [T, H * D] e4m3 and ``out_scale``
1620 [T, H * D // 32] E8M0, the layout the reduce epilogue writes.
1621 """
1622 num_tokens, num_heads, head_dim = o.shape
1623 if num_tokens == 0:
1624 return
1625 assert o.stride(-1) == 1 and out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
1626 assert out_data.shape == (num_tokens, num_heads * head_dim)
1627 assert out_scale.shape == (num_tokens, num_heads * head_dim // 32)
1628 assert cos_sin_cache.shape[-1] == rope_head_dim
1629 _inverse_rope_mxfp8_quant_kernel[(num_tokens, num_heads)](
1630 o,
1631 out_data,
1632 out_scale,
1633 positions,
1634 cos_sin_cache,
1635 o.stride(0),
1636 o.stride(1),
1637 out_data.stride(0),
1638 out_scale.stride(0),
1639 cos_sin_cache.stride(0),
1640 HEAD_DIM=head_dim,
1641 NOPE=head_dim - rope_head_dim,
1642 HALF=rope_head_dim // 2,
1643 num_warps=4,
1644 )
1645
1646
1647 def _get_cached_wo_a_bf16(
1648 wo_a: torch.nn.Module,
1649 n_local_groups: int,
1650 o_lora_rank: int,
1651 hidden_dim: int,
1652 ) -> torch.Tensor:
1653 """Dequantize wo_a to bf16 once and cache it on the module.
1654
1655 wo_a weights are static, so the fp8 -> fp32 -> (* block scale) -> bf16
1656 dequant only needs to run once. Recomputing it every decode step shows up
1657 in the profile as the largest copy/mul kernels (``direct_copy float`` ~55us
1658 and ``MulFunctor float`` ~31us per two layers). SGLang / ATOM keep wo_a in
1659 bf16 and feed a plain bf16 GEMM; this mirrors that.
1660 """
1661 cached = getattr(wo_a, "_dsv4_wo_a_bf16", None)
1662 if cached is not None:
1663 return cached
1664 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
1665 get_fp8_block_weight_scale,
1666 )
1667
1668 wo_a_scale_param = get_fp8_block_weight_scale(wo_a)
1669 if wo_a_scale_param is None:
1670 # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the
1671 # historical ``_inv`` suffix.
1672 wo_a_scale_param = getattr(wo_a, "weight_scale", None)
1673 # Emulated MXFP8 kernels can replace the original one-byte weight with an
1674 # already-dequantized BF16 tensor while retaining the scale attribute for
1675 # metadata. Applying that retained scale again would double-dequantize the
1676 # weight. Block scaling is only valid while the one-byte FP8 storage remains.
1677 if wo_a_scale_param is not None and wo_a.weight.element_size() == 1:
1678 wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
1679 torch.float32
1680 )
1681 wo_a_scale = _expand_2d_block_scales(
1682 wo_a_scale_param.view(n_local_groups, -1, wo_a_scale_param.shape[-1]),
1683 o_lora_rank,
1684 hidden_dim,
1685 )
1686 cached = (wo_a_weight * wo_a_scale).to(torch.bfloat16)
1687 else:
1688 cached = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
1689 torch.bfloat16
1690 )
1691 wo_a._dsv4_wo_a_bf16 = cached
1692 return cached
1693
1694
1695 def rocm_inv_rope_einsum(
1696 rotary_emb: torch.nn.Module,
1697 o: torch.Tensor,
1698 positions: torch.Tensor,
1699 rope_head_dim: int,
1700 n_local_groups: int,
1701 o_lora_rank: int,
1702 wo_a: torch.nn.Module,
1703 inverse_rope: bool = True,
1704 ) -> torch.Tensor:
1705 """Inverse-RoPE + WO_A bmm path used on ROCm.
1706
1707 Fuses the inverse GPT-J RoPE into one Triton kernel and caches the bf16
1708 wo_a weight so the per-step dequant disappears. Callers whose attention
1709 already rotated every row pass ``inverse_rope=False``; that is a property
1710 of the attention backend, not of the batch, so it stays constant across
1711 steps and is safe to read from compiled code.
1712 """
1713 if inverse_rope:
1714 o_ref = _fused_inverse_rope_gptj(
1715 o, positions, rotary_emb.cos_sin_cache, rope_head_dim
1716 )
1717 else:
1718 assert o.dtype == torch.bfloat16, (
1719 "a pre-rotated attention output feeds the wo_a bmm directly, so it "
1720 f"must already be bf16, got {o.dtype}"
1721 )
1722 o_ref = o
1723 o_ref = o_ref.reshape(o.shape[0], n_local_groups, -1)
1724
1725 wo_a_weight = _get_cached_wo_a_bf16(
1726 wo_a, n_local_groups, o_lora_rank, o_ref.shape[-1]
1727 )
1728
1729 return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
1730
1731
1732 @triton.jit
1733 def _mxfp8_wo_a_bmm_kernel(
1734 a_ptr, # [T, G * K] e4m3
1735 as_ptr, # [T, G * K // 32] E8M0
1736 w_ptr, # [G * R, K] e4m3
1737 ws_ptr, # [G * R // SCALE_ROWS, K // 32] E8M0
1738 out_ptr, # [T, G * R]
1739 num_tokens,
1740 stride_at,
1741 stride_ast,
1742 stride_wn,
1743 stride_wsn,
1744 stride_out,
1745 R: tl.constexpr,
1746 K: tl.constexpr,
1747 SCALE_ROWS: tl.constexpr,
1748 BLOCK_M: tl.constexpr,
1749 BLOCK_N: tl.constexpr,
1750 BLOCK_K: tl.constexpr,
1751 ):
1752 # Output columns are group-major, so a BLOCK_N tile never straddles two
1753 # groups (R % BLOCK_N == 0) and picks its group's K slice of A.
1754 pid_n = tl.program_id(0)
1755 pid_m = tl.program_id(1)
1756 group = (pid_n * BLOCK_N) // R
1757 offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
1758 offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
1759 m_mask = offs_m < num_tokens
1760 offs_k = tl.arange(0, BLOCK_K)
1761 offs_sk = tl.arange(0, BLOCK_K // 32)
1762
1763 a_ptrs = a_ptr + offs_m[:, None] * stride_at + (group * K + offs_k)[None, :]
1764 as_ptrs = (
1765 as_ptr + offs_m[:, None] * stride_ast + (group * (K // 32) + offs_sk)[None, :]
1766 )
1767 w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :]
1768 ws_ptrs = ws_ptr + (offs_n // SCALE_ROWS)[:, None] * stride_wsn + offs_sk[None, :]
1769
1770 acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
1771 for _ in range(K // BLOCK_K):
1772 a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
1773 a_s = tl.load(as_ptrs, mask=m_mask[:, None], other=127)
1774 w = tl.load(w_ptrs)
1775 w_s = tl.load(ws_ptrs)
1776 acc += tl.dot_scaled(a, a_s, "e4m3", w.T, w_s, "e4m3")
1777 a_ptrs += BLOCK_K
1778 as_ptrs += BLOCK_K // 32
1779 w_ptrs += BLOCK_K
1780 ws_ptrs += BLOCK_K // 32
1781
1782 tl.store(
1783 out_ptr + offs_m[:, None] * stride_out + offs_n[None, :],
1784 acc.to(out_ptr.dtype.element_ty),
1785 mask=m_mask[:, None],
1786 )
1787
1788
1789 def _mxfp8_wo_a_bmm_config(num_tokens: int, n_groups: int) -> tuple[int, ...]:
1790 """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for gfx950.
1791
1792 Tuned under HIP graphs with a cold weight at G = 4 and 2, over every
1793 decode shape of conc 1-128 x 0-5 spec tokens plus prefill chunks up to
1794 8K tokens. The best tile tracks the total work T * G, so the tiers are
1795 keyed on it.
1796
1797 This will be replaced after new GEMM kernel from AITER with proper 32x32 scale
1798 shape GEMM fp8 enabled.
1799 """
1800 work = num_tokens * n_groups
1801 if work <= 64:
1802 return 16, 16, 1024, 2, 3
1803 if work <= 128:
1804 return 32, 16, 1024, 2, 3
1805 if work <= 256:
1806 return 32, 32, 512, 2, 3
1807 if work <= 512:
1808 return 64, 32, 512, 2, 3
1809 if work <= 1024:
1810 return 64, 64, 512, 4, 2
1811 if work <= 2048:
1812 return 64, 64, 256, 4, 2
1813 if work <= 3072:
1814 return 64, 64, 256, 2, 1
1815 if work <= 4096:
1816 return 128, 128, 256, 8, 2
1817 return 128, 128, 128, 4, 2
1818
1819
1820 def _rocm_mxfp8_wo_a_bmm_impl(
1821 a: torch.Tensor,
1822 a_scale: torch.Tensor,
1823 weight: torch.Tensor,
1824 weight_scale: torch.Tensor,
1825 n_groups: int,
1826 o_lora_rank: int,
1827 ) -> torch.Tensor:
1828 num_tokens = a.shape[0]
1829 group_dim = a.shape[1] // n_groups
1830 out = torch.empty(
1831 (num_tokens, n_groups * o_lora_rank), dtype=torch.bfloat16, device=a.device
1832 )
1833 if num_tokens == 0:
1834 return out
1835 block_m, block_n, block_k, num_warps, num_stages = _mxfp8_wo_a_bmm_config(
1836 num_tokens, n_groups
1837 )
1838 grid = (n_groups * o_lora_rank // block_n, triton.cdiv(num_tokens, block_m))
1839 _mxfp8_wo_a_bmm_kernel[grid](
1840 a,
1841 a_scale,
1842 weight,
1843 weight_scale,
1844 out,
1845 num_tokens,
1846 a.stride(0),
1847 a_scale.stride(0),
1848 weight.stride(0),
1849 weight_scale.stride(0),
1850 out.stride(0),
1851 R=o_lora_rank,
1852 K=group_dim,
1853 SCALE_ROWS=weight.shape[0] // weight_scale.shape[0],
1854 BLOCK_M=block_m,
1855 BLOCK_N=block_n,
1856 BLOCK_K=block_k,
1857 num_warps=num_warps,
1858 num_stages=num_stages,
1859 )
1860 return out
1861
1862
1863 def _rocm_mxfp8_wo_a_bmm_fake(
1864 a: torch.Tensor,
1865 a_scale: torch.Tensor,
1866 weight: torch.Tensor,
1867 weight_scale: torch.Tensor,
1868 n_groups: int,
1869 o_lora_rank: int,
1870 ) -> torch.Tensor:
1871 return a.new_empty((a.shape[0], n_groups * o_lora_rank), dtype=torch.bfloat16)
1872
1873
1874 # An opaque op: the tile choice branches on the token count, which the
1875 # compiled o_proj must not freeze at its trace-time value.
1876 direct_register_custom_op(
1877 op_name="rocm_dsv41_mxfp8_wo_a_bmm",
1878 op_func=_rocm_mxfp8_wo_a_bmm_impl,
1879 fake_impl=_rocm_mxfp8_wo_a_bmm_fake,
1880 )
1881
1882
1883 def rocm_mxfp8_wo_a_bmm(
1884 a: torch.Tensor,
1885 a_scale: torch.Tensor,
1886 wo_a: torch.nn.Module,
1887 n_groups: int,
1888 o_lora_rank: int,
1889 ) -> torch.Tensor:
1890 """Grouped MXFP8 wo_a: ``out[t, g, :] = a[t, g, :] @ W[g].T``, bf16 out.
1891
1892 ``a`` is the [T, G * K] e4m3 attention output and ``a_scale`` its
1893 [T, G * K // 32] E8M0 scales, as the sparse decode reduce writes them.
1894 The weight is the checkpoint's MXFP8 ``wo_a`` as loaded, [G * R, K] with
1895 either [G * R // 32, K // 32] block scales or [G * R, K // 32] per-row
1896 scales, so there is no dequantized copy to keep.
1897 Returns [T, G * R].
1898 """
1899 return torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(
1900 a, a_scale, wo_a.weight, wo_a.weight_scale, n_groups, o_lora_rank
1901 )
1902
1903
1904 _DSV4_SPARSE_NOPE_DIM = 448
1905 _DSV4_SPARSE_ROPE_DIM = 64
1906
1907
1908 def _validate_sparse_dims(
1909 head_dim: int,
1910 nope_head_dim: int,
1911 rope_head_dim: int,
1912 op_name: str,
1913 ) -> None:
1914 assert head_dim > 0, f"{op_name} expected a positive head_dim, got {head_dim}"
1915 assert nope_head_dim > 0, (
1916 f"{op_name} expected a positive NoPE dimension, got {nope_head_dim}"
1917 )
1918 assert rope_head_dim >= 0, (
1919 f"{op_name} expected a non-negative RoPE dimension, got {rope_head_dim}"
1920 )
1921 assert head_dim == nope_head_dim + rope_head_dim, (
1922 f"{op_name} expected head_dim={nope_head_dim + rope_head_dim}, got {head_dim}"
1923 )
1924
1925
1926 def _validate_dsv4_sparse_dims(
1927 head_dim: int,
1928 nope_head_dim: int,
1929 rope_head_dim: int,
1930 op_name: str,
1931 ) -> None:
1932 _validate_sparse_dims(head_dim, nope_head_dim, rope_head_dim, op_name)
1933 assert (
1934 nope_head_dim == _DSV4_SPARSE_NOPE_DIM
1935 and rope_head_dim == _DSV4_SPARSE_ROPE_DIM
1936 ), (
1937 f"{op_name} expects {_DSV4_SPARSE_NOPE_DIM} NoPE dims and "
1938 f"{_DSV4_SPARSE_ROPE_DIM} RoPE dims"
1939 )
1940
1941
1942 @triton.jit
1943 def _count_valid_row_entries_kernel(
1944 indices_ptr,
1945 lengths_ptr,
1946 counts_ptr,
1947 indptr_ptr,
1948 indices_stride0,
1949 num_rows_limit,
1950 num_queries,
1951 row_width,
1952 BLOCK_R: tl.constexpr,
1953 BLOCK_W: tl.constexpr,
1954 ):
1955 pid = tl.program_id(0)
1956 # The scan fills indptr[1:], so write indptr[0] here and skip a zero fill.
1957 if pid == 0:
1958 tl.store(indptr_ptr, 0)
1959
1960 rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)
1961 row_ok = rows < num_queries
1962 row_len = tl.load(lengths_ptr + rows, mask=row_ok, other=0)
1963
1964 acc = tl.zeros((BLOCK_R,), dtype=tl.int32)
1965 for start in range(0, row_width, BLOCK_W):
1966 cols = start + tl.arange(0, BLOCK_W)
1967 in_row = cols[None, :] < row_width
1968 vals = tl.load(
1969 indices_ptr + rows[:, None] * indices_stride0 + cols[None, :],
1970 mask=row_ok[:, None] & in_row,
1971 other=-1,
1972 ).to(tl.int32)
1973 keep = (
1974 row_ok[:, None] & in_row & (cols[None, :] < row_len[:, None]) & (vals >= 0)
1975 )
1976 if num_rows_limit >= 0:
1977 keep = keep & (vals < num_rows_limit)
1978 acc += tl.sum(keep.to(tl.int32), axis=1)
1979
1980 tl.store(counts_ptr + rows, acc, mask=row_ok)
1981
1982
1983 @triton.jit
1984 def _compact_dense_row_to_ragged_kernel(
1985 indices_ptr,
1986 lengths_ptr,
1987 indptr_ptr,
1988 out_ptr,
1989 indices_stride0,
1990 num_rows_limit,
1991 row_width,
1992 BLOCK_W: tl.constexpr,
1993 ):
1994 row_idx = tl.program_id(0)
1995 row_len = tl.load(lengths_ptr + row_idx)
1996 out_start = tl.load(indptr_ptr + row_idx)
1997
1998 written = 0
1999 for start in range(0, row_width, BLOCK_W):
2000 cols = start + tl.arange(0, BLOCK_W)
2001 in_row = cols < row_width
2002 vals = tl.load(
2003 indices_ptr + row_idx * indices_stride0 + cols, mask=in_row, other=-1
2004 ).to(tl.int32)
2005 keep = in_row & (cols < row_len) & (vals >= 0)
2006 if num_rows_limit >= 0:
2007 keep = keep & (vals < num_rows_limit)
2008 rank = tl.cumsum(keep.to(tl.int32), axis=0) - 1
2009 tl.store(out_ptr + out_start + written + rank, vals, mask=keep)
2010 written += tl.sum(keep.to(tl.int32), axis=0)
2011
2012
2013 def build_ragged_indices_from_dense(
2014 indices: torch.Tensor,
2015 lengths: torch.Tensor,
2016 num_rows: int = -1,
2017 ) -> tuple[torch.Tensor, torch.Tensor]:
2018 indices = indices.reshape(indices.shape[0], -1)
2019 lengths = lengths.to(device=indices.device, dtype=torch.int32).reshape(-1)
2020 assert lengths.numel() == indices.shape[0], (
2021 f"Expected one length per row, got {lengths.shape} for indices {indices.shape}"
2022 )
2023
2024 max_width = indices.shape[1] if indices.ndim == 2 else 0
2025 # Both kernels bound the row by row_width and row_len, so no clamp needed.
2026 lengths = lengths.contiguous()
2027
2028 # Dense rows reserve slots they may not fill, leaving -1 anywhere in the
2029 # first `lengths` entries. Keep only the valid ones, since the consumers
2030 # index the KV pool without checking the sign.
2031 num_queries = indices.shape[0]
2032 # max_width is whatever the caller packed (align(topk + window, 128) for the
2033 # V4.1 indexer) and nothing caps it, so loop the row rather than size the
2034 # block by it.
2035 block = min(triton.next_power_of_2(max_width), 1024) if max_width > 0 else 1
2036 counts = torch.empty(num_queries, dtype=torch.int32, device=indices.device)
2037 indptr = torch.empty(num_queries + 1, dtype=torch.int32, device=indices.device)
2038 if num_queries > 0 and max_width > 0:
2039 # A program per row only gets row_width elements, which is launch
2040 # bound for narrow rows. Tile rows to fill the block.
2041 block_r = max(1, 1024 // block)
2042 _count_valid_row_entries_kernel[(triton.cdiv(num_queries, block_r),)](
2043 indices,
2044 lengths,
2045 counts,
2046 indptr,
2047 indices.stride(0),
2048 int(num_rows),
2049 num_queries,
2050 max_width,
2051 BLOCK_R=block_r,
2052 BLOCK_W=block,
2053 )
2054 else:
2055 counts.zero_()
2056 indptr.zero_()
2057
2058 torch.cumsum(counts, dim=0, out=indptr[1:])
2059
2060 if indices.numel() == 0:
2061 flat = torch.empty(0, dtype=torch.int32, device=indices.device)
2062 else:
2063 flat = torch.empty(
2064 indices.shape[0] * max_width,
2065 dtype=torch.int32,
2066 device=indices.device,
2067 )
2068 if flat.numel() > 0:
2069 _compact_dense_row_to_ragged_kernel[(indices.shape[0],)](
2070 indices,
2071 lengths,
2072 indptr,
2073 flat,
2074 indices.stride(0),
2075 int(num_rows),
2076 max_width,
2077 BLOCK_W=block,
2078 )
2079
2080 return flat, indptr
2081
2082
2083 def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor:
2084 if x.dtype == torch.int32 and x.ndim == 1 and x.is_contiguous():
2085 return x
2086 return x.to(torch.int32).contiguous()
2087
2088
2089 @triton.jit
2090 def _sparse_kv_row_offset(slot, stride):
2091 # A global token slot fits in int32, but its byte/element offset may not.
2092 return slot.to(tl.int64) * stride
2093
2094
2095 @triton.jit
2096 def _sparse_attn_prefill_ragged_kernel(
2097 q_ptr,
2098 kv_ptr,
2099 kv_indices_ptr,
2100 kv_indptr_ptr,
2101 attn_sink_ptr,
2102 out_ptr,
2103 q_stride_t,
2104 q_stride_h,
2105 q_stride_d,
2106 kv_stride_n,
2107 kv_stride_d,
2108 out_stride_t,
2109 out_stride_h,
2110 out_stride_d,
2111 num_heads,
2112 head_dim,
2113 num_kv,
2114 scale,
2115 HAS_ATTN_SINK: tl.constexpr,
2116 OUT_DV: tl.constexpr,
2117 BLOCK_H: tl.constexpr,
2118 BLOCK_D: tl.constexpr,
2119 BLOCK_K: tl.constexpr,
2120 ):
2121 query_idx = tl.program_id(0)
2122 pid_h = tl.program_id(1)
2123
2124 head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
2125 dim_offsets = tl.arange(0, BLOCK_D)
2126 head_mask = head_offsets < num_heads
2127 dim_mask = dim_offsets < head_dim
2128
2129 q = tl.load(
2130 q_ptr
2131 + query_idx * q_stride_t
2132 + head_offsets[:, None] * q_stride_h
2133 + dim_offsets[None, :] * q_stride_d,
2134 mask=head_mask[:, None] & dim_mask[None, :],
2135 other=0.0,
2136 )
2137
2138 neg_large = -3.4028234663852886e38
2139 m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
2140 l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
2141 acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
2142
2143 kv_start = tl.load(kv_indptr_ptr + query_idx)
2144 kv_end = tl.load(kv_indptr_ptr + query_idx + 1)
2145 kv_len = kv_end - kv_start
2146
2147 k_offsets = tl.arange(0, BLOCK_K)
2148 slot = tl.load(
2149 kv_indices_ptr + kv_start + k_offsets, mask=k_offsets < kv_len, other=-1
2150 )
2151 for k_start in tl.range(0, kv_len, BLOCK_K):
2152 k_pos = k_start + k_offsets
2153 in_range = k_pos < kv_len
2154 valid = in_range & (slot >= 0) & (slot < num_kv)
2155 safe_slot = tl.where(valid, slot, 0)
2156
2157 kv = tl.load(
2158 kv_ptr
2159 + _sparse_kv_row_offset(safe_slot[:, None], kv_stride_n)
2160 + dim_offsets[None, :] * kv_stride_d,
2161 mask=valid[:, None] & dim_mask[None, :],
2162 other=0.0,
2163 )
2164
2165 next_k_pos = k_start + BLOCK_K + k_offsets
2166 slot = tl.load(
2167 kv_indices_ptr + kv_start + next_k_pos, mask=next_k_pos < kv_len, other=-1
2168 )
2169
2170 scores = tl.dot(q, tl.trans(kv)) * scale
2171 scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
2172
2173 m_block = tl.max(scores, axis=1)
2174 m_new = tl.maximum(m_i, m_block)
2175 alpha = tl.exp(m_i - m_new)
2176 p = tl.exp(scores - m_new[:, None])
2177 p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
2178 l_new = l_i * alpha + tl.sum(p, axis=1)
2179
2180 acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
2181 m_i = m_new
2182 l_i = l_new
2183
2184 if HAS_ATTN_SINK:
2185 sink = tl.load(
2186 attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large
2187 ).to(tl.float32)
2188 m_final = tl.maximum(m_i, sink)
2189 alpha = tl.exp(m_i - m_final)
2190 l_final = l_i * alpha + tl.exp(sink - m_final)
2191 denom = tl.maximum(l_final, 1.0e-30)
2192 out = tl.where(
2193 l_final[:, None] > 0.0,
2194 (acc * alpha[:, None]) / denom[:, None],
2195 0.0,
2196 )
2197 else:
2198 denom = tl.maximum(l_i, 1.0e-30)
2199 out = tl.where(l_i[:, None] > 0.0, acc / denom[:, None], 0.0)
2200
2201 tl.store(
2202 out_ptr
2203 + query_idx * out_stride_t
2204 + head_offsets[:, None] * out_stride_h
2205 + dim_offsets[None, :] * out_stride_d,
2206 out,
2207 mask=head_mask[:, None] & (dim_offsets[None, :] < OUT_DV),
2208 )
2209
2210
2211 @triton.jit
2212 def _decode_e8m0_scales_triton(encoded_scales):
2213 scale_bits = encoded_scales.to(tl.int32) << 23
2214 scale_bits = tl.where(encoded_scales == 0, 1 << 22, scale_bits)
2215 return scale_bits.to(tl.float32, bitcast=True)
2216
2217
2218 @triton.jit
2219 def _load_fp8_ds_mla_gfx950_nope_exact_chunk(
2220 token_data_ptr,
2221 token_scale_ptr,
2222 valid,
2223 CHUNK_START: tl.constexpr,
2224 CHUNK_SIZE: tl.constexpr,
2225 BLOCK_K: tl.constexpr,
2226 IS_FNUZ: tl.constexpr,
2227 ):
2228 offsets = CHUNK_START + tl.arange(0, CHUNK_SIZE)
2229 x_uint8 = tl.load(
2230 token_data_ptr[:, None] + offsets[None, :],
2231 mask=valid[:, None],
2232 other=0,
2233 )
2234 scale_offsets = CHUNK_START // 64 + tl.arange(0, CHUNK_SIZE // 64)
2235 encoded_scales = tl.load(
2236 token_scale_ptr[:, None] + scale_offsets[None, :],
2237 mask=valid[:, None],
2238 other=127,
2239 )
2240 scales = _decode_e8m0_scales_triton(encoded_scales)
2241 scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, CHUNK_SIZE // 64, 64))
2242 scales = tl.reshape(scales, (BLOCK_K, CHUNK_SIZE))
2243 if IS_FNUZ:
2244 x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32)
2245 else:
2246 x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32)
2247 value = (x_f32 * scales).to(tl.bfloat16)
2248 zero = tl.zeros((BLOCK_K, CHUNK_SIZE), dtype=tl.bfloat16)
2249 return tl.where(valid[:, None], value, zero)
2250
2251
2252 @triton.jit
2253 def _load_fp8_ds_mla_gfx950_tail128(
2254 token_data_ptr,
2255 token_scale_ptr,
2256 valid,
2257 NOPE_DIM: tl.constexpr,
2258 BLOCK_K: tl.constexpr,
2259 IS_FNUZ: tl.constexpr,
2260 ):
2261 tail_offsets = 384 + tl.arange(0, 128)
2262 nope_mask = tail_offsets < NOPE_DIM
2263 x_uint8 = tl.load(
2264 token_data_ptr[:, None] + tail_offsets[None, :],
2265 mask=valid[:, None] & nope_mask[None, :],
2266 other=0,
2267 )
2268 scale_offsets = 6 + tl.arange(0, 2)
2269 scale_mask = scale_offsets < NOPE_DIM // 64
2270 encoded_scales = tl.load(
2271 token_scale_ptr[:, None] + scale_offsets[None, :],
2272 mask=valid[:, None] & scale_mask[None, :],
2273 other=127,
2274 )
2275 scales = _decode_e8m0_scales_triton(encoded_scales)
2276 scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, 2, 64))
2277 scales = tl.reshape(scales, (BLOCK_K, 128))
2278 if IS_FNUZ:
2279 x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32)
2280 else:
2281 x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32)
2282 nope = (x_f32 * scales).to(tl.bfloat16)
2283
2284 rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
2285 rope = tl.load(
2286 rope_ptr[:, None] + (tail_offsets[None, :] - NOPE_DIM),
2287 mask=valid[:, None] & ~nope_mask[None, :],
2288 other=0.0,
2289 )
2290 value = tl.where(nope_mask[None, :], nope, rope)
2291 zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16)
2292 return tl.where(valid[:, None], value, zero)
2293
2294
2295 @triton.jit
2296 def _sparse_attn_decode_ragged_kernel(
2297 q_ptr,
2298 main_cache_ptr,
2299 main_indices_ptr,
2300 main_indptr_ptr,
2301 extra_cache_ptr,
2302 extra_indices_ptr,
2303 extra_indptr_ptr,
2304 attn_sink_ptr,
2305 out_ptr,
2306 q_stride0,
2307 q_stride1,
2308 out_stride0,
2309 out_stride1,
2310 main_cache_stride0,
2311 extra_cache_stride0,
2312 main_num_rows,
2313 extra_num_rows,
2314 main_block_size,
2315 extra_block_size,
2316 scale,
2317 num_heads,
2318 HAS_ATTN_SINK: tl.constexpr,
2319 HAS_EXTRA: tl.constexpr,
2320 NOPE_DIM: tl.constexpr,
2321 NOPE_BLOCK: tl.constexpr,
2322 ROPE_DIM: tl.constexpr,
2323 # SWA K-cache (main): C++ encoder writes FNUZ on gfx942, OCP on gfx950.
2324 # Compressed K-cache (extra): Triton encoder writes OCP everywhere.
2325 IS_FNUZ_MAIN: tl.constexpr,
2326 IS_FNUZ_EXTRA: tl.constexpr,
2327 BLOCK_H: tl.constexpr,
2328 BLOCK_K: tl.constexpr,
2329 ):
2330 query_idx = tl.program_id(0)
2331 pid_h = tl.program_id(1)
2332
2333 head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
2334 head_mask = head_offsets < num_heads
2335 nope_offsets = tl.arange(0, NOPE_BLOCK)
2336 nope_mask = nope_offsets < NOPE_DIM
2337 rope_offsets = tl.arange(0, ROPE_DIM)
2338
2339 q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1
2340 q_nope = tl.load(
2341 q_row_ptr + nope_offsets[None, :],
2342 mask=head_mask[:, None] & nope_mask[None, :],
2343 other=0.0,
2344 )
2345 q_rope = tl.load(
2346 q_row_ptr + NOPE_DIM + rope_offsets[None, :],
2347 mask=head_mask[:, None],
2348 other=0.0,
2349 )
2350
2351 neg_large = -3.4028234663852886e38
2352 m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
2353 l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
2354 acc_nope = tl.zeros((BLOCK_H, NOPE_BLOCK), dtype=tl.float32)
2355 acc_rope = tl.zeros((BLOCK_H, ROPE_DIM), dtype=tl.float32)
2356 k_offsets = tl.arange(0, BLOCK_K)
2357
2358 main_start = tl.load(main_indptr_ptr + query_idx)
2359 main_end = tl.load(main_indptr_ptr + query_idx + 1)
2360 main_len = main_end - main_start
2361
2362 zero_nope = tl.zeros((BLOCK_K, NOPE_BLOCK), dtype=tl.bfloat16)
2363 zero_rope = tl.zeros((BLOCK_K, ROPE_DIM), dtype=tl.bfloat16)
2364
2365 for k_start in tl.range(0, main_len, BLOCK_K):
2366 k_pos = k_start + k_offsets
2367 in_range = k_pos < main_len
2368 slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1)
2369 valid = in_range & (slot >= 0) & (slot < main_num_rows)
2370 safe_slot = tl.where(valid, slot, 0)
2371
2372 block_idx = safe_slot // main_block_size
2373 pos_in_block = safe_slot % main_block_size
2374 cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0
2375 token_data_ptr = cache_block_ptr + pos_in_block * 576
2376 token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8
2377
2378 x_uint8 = tl.load(
2379 token_data_ptr[:, None] + nope_offsets[None, :],
2380 mask=valid[:, None] & nope_mask[None, :],
2381 other=0,
2382 )
2383 if IS_FNUZ_MAIN:
2384 x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
2385 else:
2386 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
2387 encoded_scales = tl.load(
2388 token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
2389 mask=valid[:, None] & nope_mask[None, :],
2390 other=127,
2391 )
2392 scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
2393 k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
2394 k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
2395 k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
2396
2397 rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
2398 k_rope = tl.load(
2399 rope_ptr[:, None] + rope_offsets[None, :],
2400 mask=valid[:, None],
2401 other=0.0,
2402 )
2403 k_rope = tl.where(valid[:, None], k_rope, zero_rope)
2404 k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
2405
2406 scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(q_rope, tl.trans(k_rope))
2407 scores *= scale
2408 scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
2409
2410 m_block = tl.max(scores, axis=1)
2411 m_new = tl.maximum(m_i, m_block)
2412 alpha = tl.exp(m_i - m_new)
2413 p = tl.exp(scores - m_new[:, None])
2414 p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
2415 l_new = l_i * alpha + tl.sum(p, axis=1)
2416
2417 acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
2418 acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
2419 m_i = m_new
2420 l_i = l_new
2421
2422 if HAS_EXTRA:
2423 extra_start = tl.load(extra_indptr_ptr + query_idx)
2424 extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
2425 extra_len = extra_end - extra_start
2426
2427 for k_start in tl.range(0, extra_len, BLOCK_K):
2428 k_pos = k_start + k_offsets
2429 in_range = k_pos < extra_len
2430 slot = tl.load(
2431 extra_indices_ptr + extra_start + k_pos, mask=in_range, other=-1
2432 )
2433 valid = in_range & (slot >= 0) & (slot < extra_num_rows)
2434 safe_slot = tl.where(valid, slot, 0)
2435
2436 block_idx = safe_slot // extra_block_size
2437 pos_in_block = safe_slot % extra_block_size
2438 cache_block_ptr = (
2439 extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0
2440 )
2441 token_data_ptr = cache_block_ptr + pos_in_block * 576
2442 token_scale_ptr = (
2443 cache_block_ptr + extra_block_size * 576 + pos_in_block * 8
2444 )
2445
2446 x_uint8 = tl.load(
2447 token_data_ptr[:, None] + nope_offsets[None, :],
2448 mask=valid[:, None] & nope_mask[None, :],
2449 other=0,
2450 )
2451 if IS_FNUZ_EXTRA:
2452 x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
2453 else:
2454 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
2455 encoded_scales = tl.load(
2456 token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
2457 mask=valid[:, None] & nope_mask[None, :],
2458 other=127,
2459 )
2460 scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
2461 k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
2462 k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
2463 k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
2464
2465 rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
2466 k_rope = tl.load(
2467 rope_ptr[:, None] + rope_offsets[None, :],
2468 mask=valid[:, None],
2469 other=0.0,
2470 )
2471 k_rope = tl.where(valid[:, None], k_rope, zero_rope)
2472 k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
2473
2474 scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(
2475 q_rope,
2476 tl.trans(k_rope),
2477 )
2478 scores *= scale
2479 scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
2480
2481 m_block = tl.max(scores, axis=1)
2482 m_new = tl.maximum(m_i, m_block)
2483 alpha = tl.exp(m_i - m_new)
2484 p = tl.exp(scores - m_new[:, None])
2485 p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
2486 l_new = l_i * alpha + tl.sum(p, axis=1)
2487
2488 acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
2489 acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
2490 m_i = m_new
2491 l_i = l_new
2492
2493 if HAS_ATTN_SINK:
2494 sink = tl.load(
2495 attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large
2496 ).to(tl.float32)
2497 m_final = tl.maximum(m_i, sink)
2498 alpha = tl.exp(m_i - m_final)
2499 l_final = l_i * alpha + tl.exp(sink - m_final)
2500 denom = tl.maximum(l_final, 1.0e-30)
2501 out_nope = tl.where(
2502 l_final[:, None] > 0.0,
2503 (acc_nope * alpha[:, None]) / denom[:, None],
2504 0.0,
2505 )
2506 out_rope = tl.where(
2507 l_final[:, None] > 0.0,
2508 (acc_rope * alpha[:, None]) / denom[:, None],
2509 0.0,
2510 )
2511 else:
2512 denom = tl.maximum(l_i, 1.0e-30)
2513 out_nope = tl.where(l_i[:, None] > 0.0, acc_nope / denom[:, None], 0.0)
2514 out_rope = tl.where(l_i[:, None] > 0.0, acc_rope / denom[:, None], 0.0)
2515
2516 out_row_ptr = (
2517 out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1
2518 )
2519 tl.store(
2520 out_row_ptr + nope_offsets[None, :],
2521 out_nope,
2522 mask=head_mask[:, None] & nope_mask[None, :],
2523 )
2524 tl.store(
2525 out_row_ptr + NOPE_DIM + rope_offsets[None, :],
2526 out_rope,
2527 mask=head_mask[:, None],
2528 )
2529
2530
2531 @triton.jit
2532 def _sparse_attn_decode_partial_kernel(
2533 q_ptr,
2534 main_cache_ptr,
2535 main_indices_ptr,
2536 main_indptr_ptr,
2537 extra_cache_ptr,
2538 extra_indices_ptr,
2539 extra_indptr_ptr,
2540 part_m_ptr,
2541 part_l_ptr,
2542 part_acc_ptr,
2543 q_stride0,
2544 q_stride1,
2545 main_cache_stride0,
2546 extra_cache_stride0,
2547 pm_stride0,
2548 pm_stride_s,
2549 pa_stride0,
2550 pa_stride_s,
2551 pa_stride_h,
2552 main_num_rows,
2553 extra_num_rows,
2554 main_block_size,
2555 extra_block_size,
2556 scale,
2557 num_heads,
2558 HAS_EXTRA: tl.constexpr,
2559 NOPE_DIM: tl.constexpr,
2560 NOPE_BLOCK: tl.constexpr,
2561 ROPE_DIM: tl.constexpr,
2562 # `main_cache` is the SWA K-cache (written by the C++ encoder, FNUZ on
2563 # gfx942 / OCP on gfx950). `extra_cache` is the compressed K-cache
2564 # (Triton encoder, OCP on every platform). Reading both with the same
2565 # `IS_FNUZ` would decode one of them with the wrong FNUZ/OCP scale ratio.
2566 IS_FNUZ_MAIN: tl.constexpr,
2567 IS_FNUZ_EXTRA: tl.constexpr,
2568 BLOCK_H: tl.constexpr,
2569 BLOCK_K: tl.constexpr,
2570 NUM_SPLITS: tl.constexpr,
2571 NUM_STAGES: tl.constexpr,
2572 ):
2573 query_idx = tl.program_id(0)
2574 split_id = tl.program_id(1)
2575 pid_h = tl.program_id(2)
2576
2577 head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
2578 head_mask = head_offsets < num_heads
2579 nope_offsets = tl.arange(0, NOPE_BLOCK)
2580 nope_mask = nope_offsets < NOPE_DIM
2581 rope_offsets = tl.arange(0, ROPE_DIM)
2582
2583 q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1
2584 q_nope = tl.load(
2585 q_row_ptr + nope_offsets[None, :],
2586 mask=head_mask[:, None] & nope_mask[None, :],
2587 other=0.0,
2588 )
2589 q_rope = tl.load(
2590 q_row_ptr + NOPE_DIM + rope_offsets[None, :],
2591 mask=head_mask[:, None],
2592 other=0.0,
2593 )
2594
2595 neg_large = -3.4028234663852886e38
2596 m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
2597 l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
2598 acc_nope = tl.zeros((BLOCK_H, NOPE_BLOCK), dtype=tl.float32)
2599 acc_rope = tl.zeros((BLOCK_H, ROPE_DIM), dtype=tl.float32)
2600 k_offsets = tl.arange(0, BLOCK_K)
2601
2602 zero_nope = tl.zeros((BLOCK_K, NOPE_BLOCK), dtype=tl.bfloat16)
2603 zero_rope = tl.zeros((BLOCK_K, ROPE_DIM), dtype=tl.bfloat16)
2604
2605 # Each split processes a contiguous slice of this query's main (SWA) and
2606 # extra (topk) segments. Slices are handled independently so a block never
2607 # straddles the main/extra boundary.
2608 main_start = tl.load(main_indptr_ptr + query_idx)
2609 main_end = tl.load(main_indptr_ptr + query_idx + 1)
2610 main_len = main_end - main_start
2611 main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS
2612 main_lo = split_id * main_chunk
2613 main_hi = tl.minimum(main_lo + main_chunk, main_len)
2614
2615 for k_start in tl.range(main_lo, main_hi, BLOCK_K, num_stages=NUM_STAGES):
2616 k_pos = k_start + k_offsets
2617 in_range = k_pos < main_hi
2618 slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1)
2619 valid = in_range & (slot >= 0) & (slot < main_num_rows)
2620 safe_slot = tl.where(valid, slot, 0)
2621
2622 block_idx = safe_slot // main_block_size
2623 pos_in_block = safe_slot % main_block_size
2624 cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0
2625 token_data_ptr = cache_block_ptr + pos_in_block * 576
2626 token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8
2627
2628 x_uint8 = tl.load(
2629 token_data_ptr[:, None] + nope_offsets[None, :],
2630 mask=valid[:, None] & nope_mask[None, :],
2631 other=0,
2632 )
2633 if IS_FNUZ_MAIN:
2634 x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
2635 else:
2636 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
2637 encoded_scales = tl.load(
2638 token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
2639 mask=valid[:, None] & nope_mask[None, :],
2640 other=127,
2641 )
2642 scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
2643 k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
2644 k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
2645 k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
2646
2647 rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
2648 k_rope = tl.load(
2649 rope_ptr[:, None] + rope_offsets[None, :],
2650 mask=valid[:, None],
2651 other=0.0,
2652 )
2653 k_rope = tl.where(valid[:, None], k_rope, zero_rope)
2654 k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
2655
2656 scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(q_rope, tl.trans(k_rope))
2657 scores *= scale
2658 scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
2659
2660 m_block = tl.max(scores, axis=1)
2661 m_new = tl.maximum(m_i, m_block)
2662 alpha = tl.exp(m_i - m_new)
2663 p = tl.exp(scores - m_new[:, None])
2664 p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
2665 l_new = l_i * alpha + tl.sum(p, axis=1)
2666
2667 acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
2668 acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
2669 m_i = m_new
2670 l_i = l_new
2671
2672 if HAS_EXTRA:
2673 extra_start = tl.load(extra_indptr_ptr + query_idx)
2674 extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
2675 extra_len = extra_end - extra_start
2676 extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS
2677 extra_lo = split_id * extra_chunk
2678 extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len)
2679
2680 for k_start in tl.range(extra_lo, extra_hi, BLOCK_K, num_stages=NUM_STAGES):
2681 k_pos = k_start + k_offsets
2682 in_range = k_pos < extra_hi
2683 slot = tl.load(
2684 extra_indices_ptr + extra_start + k_pos, mask=in_range, other=-1
2685 )
2686 valid = in_range & (slot >= 0) & (slot < extra_num_rows)
2687 safe_slot = tl.where(valid, slot, 0)
2688
2689 block_idx = safe_slot // extra_block_size
2690 pos_in_block = safe_slot % extra_block_size
2691 cache_block_ptr = (
2692 extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0
2693 )
2694 token_data_ptr = cache_block_ptr + pos_in_block * 576
2695 token_scale_ptr = (
2696 cache_block_ptr + extra_block_size * 576 + pos_in_block * 8
2697 )
2698
2699 x_uint8 = tl.load(
2700 token_data_ptr[:, None] + nope_offsets[None, :],
2701 mask=valid[:, None] & nope_mask[None, :],
2702 other=0,
2703 )
2704 if IS_FNUZ_EXTRA:
2705 x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True)
2706 else:
2707 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
2708 encoded_scales = tl.load(
2709 token_scale_ptr[:, None] + nope_offsets[None, :] // 64,
2710 mask=valid[:, None] & nope_mask[None, :],
2711 other=127,
2712 )
2713 scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0)
2714 k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16)
2715 k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope)
2716 k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope)
2717
2718 rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16))
2719 k_rope = tl.load(
2720 rope_ptr[:, None] + rope_offsets[None, :],
2721 mask=valid[:, None],
2722 other=0.0,
2723 )
2724 k_rope = tl.where(valid[:, None], k_rope, zero_rope)
2725 k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope)
2726
2727 scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(
2728 q_rope,
2729 tl.trans(k_rope),
2730 )
2731 scores *= scale
2732 scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large)
2733
2734 m_block = tl.max(scores, axis=1)
2735 m_new = tl.maximum(m_i, m_block)
2736 alpha = tl.exp(m_i - m_new)
2737 p = tl.exp(scores - m_new[:, None])
2738 p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
2739 l_new = l_i * alpha + tl.sum(p, axis=1)
2740
2741 acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope)
2742 acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope)
2743 m_i = m_new
2744 l_i = l_new
2745
2746 # Store raw (un-normalized) partial state for this split. Softmax sink and
2747 # final normalization happen in the reduce kernel.
2748 pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + head_offsets
2749 tl.store(part_m_ptr + pm_base, m_i, mask=head_mask)
2750 tl.store(part_l_ptr + pm_base, l_i, mask=head_mask)
2751 acc_base = (
2752 part_acc_ptr
2753 + query_idx * pa_stride0
2754 + split_id * pa_stride_s
2755 + head_offsets[:, None] * pa_stride_h
2756 )
2757 tl.store(
2758 acc_base + nope_offsets[None, :],
2759 acc_nope,
2760 mask=head_mask[:, None] & nope_mask[None, :],
2761 )
2762 tl.store(
2763 acc_base + NOPE_DIM + rope_offsets[None, :],
2764 acc_rope,
2765 mask=head_mask[:, None],
2766 )
2767
2768
2769 @triton.jit
2770 def _sparse_attn_decode_gfx950_partial_loaded_tile(
2771 q_combined,
2772 cache_ptr,
2773 slot,
2774 valid,
2775 cache_stride0,
2776 scale: tl.constexpr,
2777 head_mask,
2778 m_i,
2779 l_i,
2780 acc_nope_0a,
2781 acc_nope_0b,
2782 acc_nope_1,
2783 acc_tail,
2784 BLOCK_SIZE: tl.constexpr,
2785 NOPE_DIM: tl.constexpr,
2786 BLOCK_K: tl.constexpr,
2787 IS_FNUZ: tl.constexpr,
2788 TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr,
2789 ):
2790 safe_slot = tl.where(valid, slot, 0)
2791 block_idx = safe_slot // BLOCK_SIZE
2792 pos_in_block = safe_slot % BLOCK_SIZE
2793 cache_block_ptr = cache_ptr + block_idx.to(tl.int64) * cache_stride0
2794 token_data_ptr = cache_block_ptr + pos_in_block * 576
2795 token_scale_ptr = cache_block_ptr + BLOCK_SIZE * 576 + pos_in_block * 8
2796 k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk(
2797 token_data_ptr,
2798 token_scale_ptr,
2799 valid,
2800 0,
2801 128,
2802 BLOCK_K,
2803 IS_FNUZ,
2804 )
2805 k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk(
2806 token_data_ptr,
2807 token_scale_ptr,
2808 valid,
2809 128,
2810 128,
2811 BLOCK_K,
2812 IS_FNUZ,
2813 )
2814 k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk(
2815 token_data_ptr,
2816 token_scale_ptr,
2817 valid,
2818 256,
2819 128,
2820 BLOCK_K,
2821 IS_FNUZ,
2822 )
2823 k_tail = _load_fp8_ds_mla_gfx950_tail128(
2824 token_data_ptr,
2825 token_scale_ptr,
2826 valid,
2827 NOPE_DIM,
2828 BLOCK_K,
2829 IS_FNUZ,
2830 )
2831 if not TRUST_EXTRA_CACHE_NAN_FREE:
2832 zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16)
2833 k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero)
2834 k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero)
2835 k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero)
2836 k_tail = tl.where(k_tail == k_tail, k_tail, zero)
2837 k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1)
2838 k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1)
2839 k_combined = tl.cat(k_nope_0, k_tail_256, dim=1)
2840
2841 scores = tl.dot(q_combined, tl.trans(k_combined))
2842 scores *= scale * 1.4426950408889634
2843 scores = tl.where(
2844 head_mask[:, None] & valid[None, :],
2845 scores,
2846 -3.4028234663852886e38,
2847 )
2848 m_block = tl.max(scores, axis=1)
2849 m_new = tl.maximum(m_i, m_block)
2850 alpha = tl.exp2(m_i - m_new)
2851 p = tl.exp2(scores - m_new[:, None])
2852 p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
2853 l_new = l_i * alpha + tl.sum(p, axis=1)
2854 p_bf16 = p.to(k_nope_0a.dtype)
2855 acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a)
2856 acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b)
2857 acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1)
2858 acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail)
2859 return (
2860 m_new,
2861 l_new,
2862 acc_nope_0a,
2863 acc_nope_0b,
2864 acc_nope_1,
2865 acc_tail,
2866 )
2867
2868
2869 @triton.jit
2870 def _sparse_attn_decode_gfx950_partial_kernel(
2871 q_ptr,
2872 main_cache_ptr,
2873 main_indices_ptr,
2874 main_indptr_ptr,
2875 extra_cache_ptr,
2876 extra_indices_ptr,
2877 extra_indptr_ptr,
2878 part_m_ptr,
2879 part_l_ptr,
2880 part_acc_ptr,
2881 q_stride0: tl.constexpr,
2882 q_stride1: tl.constexpr,
2883 main_cache_stride0: tl.constexpr,
2884 extra_cache_stride0: tl.constexpr,
2885 main_num_rows,
2886 extra_num_rows,
2887 MAIN_BLOCK_SIZE: tl.constexpr,
2888 EXTRA_BLOCK_SIZE: tl.constexpr,
2889 scale: tl.constexpr,
2890 num_heads: tl.constexpr,
2891 HAS_EXTRA: tl.constexpr,
2892 NOPE_DIM: tl.constexpr,
2893 ROPE_DIM: tl.constexpr,
2894 IS_FNUZ_MAIN: tl.constexpr,
2895 IS_FNUZ_EXTRA: tl.constexpr,
2896 TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr,
2897 ADAPTIVE_SPLITS: tl.constexpr,
2898 ONE_WAVE_SPLITS: tl.constexpr,
2899 BLOCK_H: tl.constexpr,
2900 BLOCK_K: tl.constexpr,
2901 NUM_SPLITS: tl.constexpr,
2902 NUM_STAGES: tl.constexpr,
2903 ):
2904 query_idx = tl.program_id(0)
2905 split_id = tl.program_id(1)
2906 pid_h = tl.program_id(2)
2907
2908 tl.static_assert(NOPE_DIM == 448)
2909 tl.static_assert(ROPE_DIM == 64)
2910
2911 head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
2912 if num_heads % BLOCK_H == 0:
2913 head_mask = tl.full((BLOCK_H,), True, tl.int1)
2914 else:
2915 head_mask = head_offsets < num_heads
2916 neg_large = -3.4028234663852886e38
2917
2918 if ADAPTIVE_SPLITS:
2919 main_start = tl.load(main_indptr_ptr + query_idx)
2920 main_end = tl.load(main_indptr_ptr + query_idx + 1)
2921 main_len = main_end - main_start
2922 if HAS_EXTRA:
2923 extra_start = tl.load(extra_indptr_ptr + query_idx)
2924 extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
2925 extra_len = extra_end - extra_start
2926 else:
2927 extra_start = 0
2928 extra_len = 0
2929 split4_span: tl.constexpr = 4 * BLOCK_K
2930 split4_iters = (main_len + split4_span - 1) // split4_span
2931 split4_iters += (extra_len + split4_span - 1) // split4_span
2932 use_four_splits = split4_iters <= 3
2933 work_splits = NUM_SPLITS
2934 if ONE_WAVE_SPLITS > 4 and ONE_WAVE_SPLITS < NUM_SPLITS:
2935 one_wave_span: tl.constexpr = ONE_WAVE_SPLITS * BLOCK_K
2936 one_wave_iters = (main_len + one_wave_span - 1) // one_wave_span
2937 one_wave_iters += (extra_len + one_wave_span - 1) // one_wave_span
2938 work_splits = tl.where(one_wave_iters <= 3, ONE_WAVE_SPLITS, work_splits)
2939 work_splits = tl.where(use_four_splits, 4, work_splits)
2940 if split_id >= work_splits:
2941 pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets
2942 tl.store(part_m_ptr + pm_base, neg_large, mask=head_mask)
2943 tl.store(part_l_ptr + pm_base, 0.0, mask=head_mask)
2944 return
2945 else:
2946 work_splits = NUM_SPLITS
2947
2948 nope_offsets_0a = tl.arange(0, 128)
2949 nope_offsets_0b = 128 + tl.arange(0, 128)
2950 nope_offsets_0 = tl.arange(0, 256)
2951 tail_offsets = 256 + tl.arange(0, 256)
2952 nope_offsets_1 = 256 + tl.arange(0, 128)
2953 tail_offsets_128 = 384 + tl.arange(0, 128)
2954
2955 q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1
2956 q_nope_0 = tl.load(
2957 q_row_ptr + nope_offsets_0[None, :],
2958 mask=head_mask[:, None],
2959 other=0.0,
2960 )
2961 q_tail = tl.load(
2962 q_row_ptr + tail_offsets[None, :],
2963 mask=head_mask[:, None],
2964 other=0.0,
2965 )
2966 q_combined = tl.cat(q_nope_0, q_tail, dim=1)
2967
2968 m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
2969 l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
2970 acc_nope_0a = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2971 acc_nope_0b = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2972 acc_nope_1 = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2973 acc_tail = tl.zeros((BLOCK_H, 128), dtype=tl.float32)
2974 k_offsets = tl.arange(0, BLOCK_K)
2975
2976 if not ADAPTIVE_SPLITS:
2977 main_start = tl.load(main_indptr_ptr + query_idx)
2978 main_end = tl.load(main_indptr_ptr + query_idx + 1)
2979 main_len = main_end - main_start
2980 main_chunk = (main_len + work_splits - 1) // work_splits
2981 main_lo = split_id * main_chunk
2982 main_hi = tl.minimum(main_lo + main_chunk, main_len)
2983
2984 for k_start in tl.range(
2985 main_lo,
2986 main_hi,
2987 BLOCK_K,
2988 num_stages=NUM_STAGES,
2989 ):
2990 k_pos = k_start + k_offsets
2991 in_range = k_pos < main_hi
2992 slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1)
2993 valid = in_range & (slot >= 0) & (slot < main_num_rows)
2994 (
2995 m_i,
2996 l_i,
2997 acc_nope_0a,
2998 acc_nope_0b,
2999 acc_nope_1,
3000 acc_tail,
3001 ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
3002 q_combined,
3003 main_cache_ptr,
3004 slot,
3005 valid,
3006 main_cache_stride0,
3007 scale,
3008 head_mask,
3009 m_i,
3010 l_i,
3011 acc_nope_0a,
3012 acc_nope_0b,
3013 acc_nope_1,
3014 acc_tail,
3015 MAIN_BLOCK_SIZE,
3016 NOPE_DIM,
3017 BLOCK_K,
3018 IS_FNUZ_MAIN,
3019 False,
3020 )
3021
3022 if HAS_EXTRA:
3023 if not ADAPTIVE_SPLITS:
3024 extra_start = tl.load(extra_indptr_ptr + query_idx)
3025 extra_end = tl.load(extra_indptr_ptr + query_idx + 1)
3026 extra_len = extra_end - extra_start
3027 extra_chunk = (extra_len + work_splits - 1) // work_splits
3028 extra_lo = split_id * extra_chunk
3029 extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len)
3030
3031 outer_block_k: tl.constexpr = 2 * BLOCK_K
3032 outer_k_offsets = tl.arange(0, outer_block_k)
3033 extra_hi_full = (
3034 extra_lo + ((extra_hi - extra_lo) // outer_block_k) * outer_block_k
3035 )
3036 for k_start in tl.range(
3037 extra_lo,
3038 extra_hi_full,
3039 outer_block_k,
3040 num_stages=NUM_STAGES,
3041 ):
3042 slot = tl.load(extra_indices_ptr + extra_start + k_start + outer_k_offsets)
3043 valid = (slot >= 0) & (slot < extra_num_rows)
3044 slot_pairs = tl.trans(tl.reshape(slot, (2, BLOCK_K)))
3045 valid_pairs = tl.trans(tl.reshape(valid, (2, BLOCK_K)))
3046 slot_lo, slot_hi = tl.split(slot_pairs)
3047 valid_lo, valid_hi = tl.split(valid_pairs)
3048 (
3049 m_i,
3050 l_i,
3051 acc_nope_0a,
3052 acc_nope_0b,
3053 acc_nope_1,
3054 acc_tail,
3055 ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
3056 q_combined,
3057 extra_cache_ptr,
3058 slot_lo,
3059 valid_lo,
3060 extra_cache_stride0,
3061 scale,
3062 head_mask,
3063 m_i,
3064 l_i,
3065 acc_nope_0a,
3066 acc_nope_0b,
3067 acc_nope_1,
3068 acc_tail,
3069 EXTRA_BLOCK_SIZE,
3070 NOPE_DIM,
3071 BLOCK_K,
3072 IS_FNUZ_EXTRA,
3073 TRUST_EXTRA_CACHE_NAN_FREE,
3074 )
3075 (
3076 m_i,
3077 l_i,
3078 acc_nope_0a,
3079 acc_nope_0b,
3080 acc_nope_1,
3081 acc_tail,
3082 ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
3083 q_combined,
3084 extra_cache_ptr,
3085 slot_hi,
3086 valid_hi,
3087 extra_cache_stride0,
3088 scale,
3089 head_mask,
3090 m_i,
3091 l_i,
3092 acc_nope_0a,
3093 acc_nope_0b,
3094 acc_nope_1,
3095 acc_tail,
3096 EXTRA_BLOCK_SIZE,
3097 NOPE_DIM,
3098 BLOCK_K,
3099 IS_FNUZ_EXTRA,
3100 TRUST_EXTRA_CACHE_NAN_FREE,
3101 )
3102 for tail_idx in tl.static_range(2):
3103 tail_start = extra_hi_full + tail_idx * BLOCK_K
3104 if tail_start < extra_hi:
3105 k_pos = tail_start + k_offsets
3106 in_range = k_pos < extra_hi
3107 slot = tl.load(
3108 extra_indices_ptr + extra_start + k_pos,
3109 mask=in_range,
3110 other=-1,
3111 )
3112 valid = in_range & (slot >= 0) & (slot < extra_num_rows)
3113 (
3114 m_i,
3115 l_i,
3116 acc_nope_0a,
3117 acc_nope_0b,
3118 acc_nope_1,
3119 acc_tail,
3120 ) = _sparse_attn_decode_gfx950_partial_loaded_tile(
3121 q_combined,
3122 extra_cache_ptr,
3123 slot,
3124 valid,
3125 extra_cache_stride0,
3126 scale,
3127 head_mask,
3128 m_i,
3129 l_i,
3130 acc_nope_0a,
3131 acc_nope_0b,
3132 acc_nope_1,
3133 acc_tail,
3134 EXTRA_BLOCK_SIZE,
3135 NOPE_DIM,
3136 BLOCK_K,
3137 IS_FNUZ_EXTRA,
3138 TRUST_EXTRA_CACHE_NAN_FREE,
3139 )
3140
3141 pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets
3142 m_store = tl.where(l_i > 0.0, m_i * 0.6931471805599453, neg_large)
3143 tl.store(part_m_ptr + pm_base, m_store, mask=head_mask)
3144 tl.store(part_l_ptr + pm_base, l_i, mask=head_mask)
3145 acc_base = part_acc_ptr + (
3146 (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets[:, None]
3147 ) * (NOPE_DIM + ROPE_DIM)
3148 tl.store(
3149 acc_base + nope_offsets_0a[None, :],
3150 acc_nope_0a,
3151 mask=head_mask[:, None],
3152 )
3153 tl.store(
3154 acc_base + nope_offsets_0b[None, :],
3155 acc_nope_0b,
3156 mask=head_mask[:, None],
3157 )
3158 tl.store(
3159 acc_base + nope_offsets_1[None, :],
3160 acc_nope_1,
3161 mask=head_mask[:, None],
3162 )
3163 tl.store(
3164 acc_base + tail_offsets_128[None, :],
3165 acc_tail,
3166 mask=head_mask[:, None],
3167 )
3168
3169
3170 @triton.jit
3171 def _mxfp8_scale_bits(amax):
3172 """Biased E8M0 exponent that puts ``amax`` at the top of the e4m3 range.
3173
3174 Same rounding as ``mxfp8_e4m3_quantize``, so the output is bit-identical to
3175 quantizing the tensor there.
3176 """
3177 amax = tl.maximum(amax, 1.1754943508222875e-38)
3178 bits = tl.ceil(tl.log2(amax / 448.0)) + 127.0
3179 return tl.minimum(tl.maximum(bits, 0.0), 254.0)
3180
3181
3182 @triton.jit
3183 def _mxfp8_quantize_rows(x, ROWS: tl.constexpr, COLS: tl.constexpr):
3184 """MXFP8-quantize ``x`` [ROWS, COLS] in registers, one scale per 32 lanes.
3185
3186 Returns the rescaled fp32 values (to be cast to e4m3 on store) and the
3187 [ROWS, COLS // 32] biased E8M0 exponents.
3188 """
3189 blocks = tl.reshape(x, (ROWS, COLS // 32, 32))
3190 bits = _mxfp8_scale_bits(tl.max(tl.abs(blocks), axis=2))
3191 # Multiply by the reciprocal: a divisor of 2**-127 would be subnormal and
3192 # flush to zero, turning an all-zero block into NaN.
3193 q = blocks * tl.exp2(127.0 - bits)[:, :, None]
3194 return tl.reshape(q, (ROWS, COLS)), bits
3195
3196
3197 @triton.jit
3198 def _sparse_attn_decode_reduce_kernel(
3199 part_m_ptr,
3200 part_l_ptr,
3201 part_acc_ptr,
3202 attn_sink_ptr,
3203 out_ptr,
3204 pos_ptr,
3205 cos_sin_ptr,
3206 out_scale_ptr,
3207 out_stride0,
3208 out_stride1,
3209 os_stride0,
3210 os_stride1,
3211 pm_stride0,
3212 pm_stride_s,
3213 pa_stride0,
3214 pa_stride_s,
3215 pa_stride_h,
3216 cs_stride,
3217 num_heads,
3218 HAS_ATTN_SINK: tl.constexpr,
3219 ADAPTIVE_SPLITS: tl.constexpr,
3220 COMB_DIM: tl.constexpr,
3221 BLOCK_H: tl.constexpr,
3222 NUM_SPLITS: tl.constexpr,
3223 SPLITS_PAD: tl.constexpr,
3224 FUSE_INV_ROPE: tl.constexpr,
3225 NOPE: tl.constexpr,
3226 HALF: tl.constexpr,
3227 QUANT_OUT: tl.constexpr,
3228 ):
3229 query_idx = tl.program_id(0)
3230 pid_h = tl.program_id(1)
3231
3232 head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
3233 head_mask = head_offsets < num_heads
3234 comb_offsets = tl.arange(0, COMB_DIM)
3235 # SPLITS_PAD is NUM_SPLITS rounded up to a power of two so the parallel
3236 # split-axis load is a legal arange for any split count; padding lanes are
3237 # masked off.
3238 split_offsets = tl.arange(0, SPLITS_PAD)
3239 split_mask = split_offsets < NUM_SPLITS
3240
3241 neg_large = -3.4028234663852886e38
3242
3243 # Phase 1: load every split's running max/sum at once and reduce the max
3244 # in parallel (tl.max over the split axis) instead of walking the splits
3245 # serially. This breaks the long online-softmax dependency chain that made
3246 # the reduce latency-bound.
3247 load_mask = split_mask[:, None] & head_mask[None, :]
3248 pm_split = (
3249 part_m_ptr
3250 + query_idx * pm_stride0
3251 + split_offsets[:, None] * pm_stride_s
3252 + head_offsets[None, :]
3253 )
3254 m_all = tl.load(pm_split, mask=load_mask, other=neg_large) # [S, H]
3255 l_all = tl.load(
3256 part_l_ptr
3257 + query_idx * pm_stride0
3258 + split_offsets[:, None] * pm_stride_s
3259 + head_offsets[None, :],
3260 mask=load_mask,
3261 other=0.0,
3262 )
3263
3264 m_comb = tl.max(m_all, axis=0) # [H]
3265 if HAS_ATTN_SINK:
3266 sink = tl.load(
3267 attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large
3268 ).to(tl.float32)
3269 m_final = tl.maximum(m_comb, sink)
3270 else:
3271 m_final = m_comb
3272
3273 w_all = tl.exp(m_all - m_final[None, :]) # [S, H]
3274 w_all = tl.where(load_mask, w_all, 0.0)
3275 l_final = tl.sum(w_all * l_all, axis=0) # [H]
3276 if HAS_ATTN_SINK:
3277 l_final = l_final + tl.exp(sink - m_final)
3278 denom = tl.maximum(l_final, 1.0e-30)
3279
3280 # Phase 2: weighted sum of the per-split accumulators. The combine weight
3281 # for each split only depends on the (already known) global max, so the
3282 # acc loads carry no cross-split dependency and the compiler can pipeline
3283 # them; only the cheap FMA into `acc` is loop-carried.
3284 acc = tl.zeros((BLOCK_H, COMB_DIM), dtype=tl.float32)
3285 for s in tl.static_range(NUM_SPLITS):
3286 m_s = tl.load(
3287 part_m_ptr + query_idx * pm_stride0 + s * pm_stride_s + head_offsets,
3288 mask=head_mask,
3289 other=neg_large,
3290 )
3291 w_s = tl.exp(m_s - m_final)
3292 if ADAPTIVE_SPLITS:
3293 active_split = m_s > neg_large
3294 w_s = tl.where(head_mask & active_split, w_s, 0.0)
3295 acc_base = (
3296 part_acc_ptr
3297 + query_idx * pa_stride0
3298 + s * pa_stride_s
3299 + head_offsets[:, None] * pa_stride_h
3300 )
3301 if ADAPTIVE_SPLITS:
3302 acc_s = tl.load(
3303 acc_base + comb_offsets[None, :],
3304 mask=head_mask[:, None] & active_split[:, None],
3305 other=0.0,
3306 )
3307 else:
3308 acc_s = tl.load(
3309 acc_base + comb_offsets[None, :],
3310 mask=head_mask[:, None],
3311 other=0.0,
3312 )
3313 acc += w_s[:, None] * acc_s
3314
3315 out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0)
3316
3317 if FUSE_INV_ROPE:
3318 # Inverse GPT-J RoPE on the trailing rope lanes, straight out of the
3319 # combine registers: out_even = a*cos + b*sin, out_odd = b*cos - a*sin.
3320 # NoPE lanes take cos=1/sin=0 so one expression covers the whole row
3321 # and the o_proj rotation pass disappears for these tokens.
3322 pos = tl.load(pos_ptr + query_idx).to(tl.int64)
3323 pair_idx = tl.arange(0, COMB_DIM // 2) - (NOPE // 2)
3324 is_rope = pair_idx >= 0
3325 k = tl.where(is_rope, pair_idx, 0)
3326 cos = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + k), 1.0)
3327 sin = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + HALF + k), 0.0)
3328 even, odd = tl.split(tl.reshape(out, (BLOCK_H, COMB_DIM // 2, 2)))
3329 out = tl.reshape(
3330 tl.join(
3331 even * cos[None, :] + odd * sin[None, :],
3332 odd * cos[None, :] - even * sin[None, :],
3333 ),
3334 (BLOCK_H, COMB_DIM),
3335 )
3336
3337 out_row_ptr = (
3338 out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1
3339 )
3340 if QUANT_OUT:
3341 # Emit wo_a's MXFP8 input directly instead of a bf16 row that a
3342 # separate pass would read back to quantize.
3343 out, scale_bits = _mxfp8_quantize_rows(out, BLOCK_H, COMB_DIM)
3344 scale_offsets = tl.arange(0, COMB_DIM // 32)
3345 tl.store(
3346 out_scale_ptr
3347 + query_idx * os_stride0
3348 + head_offsets[:, None] * os_stride1
3349 + scale_offsets[None, :],
3350 scale_bits.to(tl.uint8),
3351 mask=head_mask[:, None],
3352 )
3353 tl.store(
3354 out_row_ptr + comb_offsets[None, :],
3355 out.to(out_ptr.dtype.element_ty),
3356 mask=head_mask[:, None],
3357 )
3358
3359
3360 def _rocm_sparse_attn_prefill_ragged_triton(
3361 q: torch.Tensor,
3362 kv: torch.Tensor,
3363 indices: torch.Tensor,
3364 indptr: torch.Tensor,
3365 scale: float,
3366 attn_sink: torch.Tensor | None,
3367 nope_head_dim: int,
3368 rope_head_dim: int,
3369 out: torch.Tensor | None = None,
3370 ) -> torch.Tensor:
3371 assert q.ndim == 3, f"expected q=[sq,h,d], got {q.shape}"
3372 assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}"
3373 assert indices.ndim == 1, f"expected indices=[nnz], got {indices.shape}"
3374 assert indptr.ndim == 1, f"expected indptr=[sq+1], got {indptr.shape}"
3375 assert not q.is_cpu and not kv.is_cpu and not indices.is_cpu and not indptr.is_cpu
3376
3377 indices = _as_int32_contiguous_1d(indices)
3378 indptr = _as_int32_contiguous_1d(indptr)
3379 has_attn_sink = attn_sink is not None
3380 if attn_sink is None:
3381 attn_sink = torch.empty(1, device=q.device, dtype=torch.float32)
3382 else:
3383 attn_sink = attn_sink.contiguous()
3384
3385 num_queries, num_heads, head_dim = q.shape
3386 assert indptr.numel() == num_queries + 1, (
3387 f"expected indptr shape [{num_queries + 1}], got {indptr.shape}"
3388 )
3389 _validate_sparse_dims(
3390 head_dim,
3391 nope_head_dim,
3392 rope_head_dim,
3393 "_rocm_sparse_attn_prefill_ragged_triton",
3394 )
3395
3396 block_h = 16
3397 block_d = triton.next_power_of_2(head_dim)
3398 block_k = 16 if head_dim >= 256 else 32
3399 num_warps = 4
3400 if out is None:
3401 out = torch.empty_like(q)
3402 else:
3403 assert out.shape[:2] == q.shape[:2], (
3404 f"expected out=[{num_queries},{num_heads},*], got {out.shape}"
3405 )
3406 assert out.shape[-1] in (nope_head_dim, head_dim), (
3407 f"expected out width {nope_head_dim} or {head_dim}, got {out.shape[-1]}"
3408 )
3409 _sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))](
3410 q,
3411 kv,
3412 indices,
3413 indptr,
3414 attn_sink,
3415 out,
3416 q.stride(0),
3417 q.stride(1),
3418 q.stride(2),
3419 kv.stride(0),
3420 kv.stride(1),
3421 out.stride(0),
3422 out.stride(1),
3423 out.stride(2),
3424 num_heads,
3425 head_dim,
3426 kv.shape[0],
3427 float(scale),
3428 HAS_ATTN_SINK=has_attn_sink,
3429 OUT_DV=out.shape[-1],
3430 BLOCK_H=block_h,
3431 BLOCK_D=block_d,
3432 BLOCK_K=block_k,
3433 num_warps=num_warps,
3434 )
3435 return out
3436
3437
3438 def _rocm_sparse_attn_prefill_triton(
3439 q: torch.Tensor,
3440 kv: torch.Tensor,
3441 indices: torch.Tensor,
3442 scale: float,
3443 attn_sink: torch.Tensor | None,
3444 nope_head_dim: int,
3445 rope_head_dim: int,
3446 topk_length: torch.Tensor | None = None,
3447 out: torch.Tensor | None = None,
3448 ) -> torch.Tensor:
3449 ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
3450 indices,
3451 topk_length
3452 if topk_length is not None
3453 else (indices >= 0).sum(dim=-1, dtype=torch.int32),
3454 num_rows=kv.shape[0],
3455 )
3456 return _rocm_sparse_attn_prefill_ragged_triton(
3457 q=q,
3458 kv=kv,
3459 indices=ragged_indices,
3460 indptr=ragged_indptr,
3461 scale=scale,
3462 attn_sink=attn_sink,
3463 nope_head_dim=nope_head_dim,
3464 rope_head_dim=rope_head_dim,
3465 out=out,
3466 )
3467
3468
3469 def _can_use_aiter_sparse_prefill_opus(
3470 q: torch.Tensor,
3471 kv: torch.Tensor,
3472 attn_sink: torch.Tensor | None,
3473 output: torch.Tensor,
3474 on_gfx950: bool = _ON_GFX950,
3475 ) -> bool:
3476 return (
3477 on_gfx950
3478 and q.shape[0] >= _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES
3479 and q.is_cuda
3480 and output.shape == q.shape
3481 and kv.shape[-1] == q.shape[-1]
3482 and q.dtype in (torch.bfloat16, torch.float16)
3483 and kv.dtype == q.dtype
3484 and output.dtype == q.dtype
3485 and kv.device == q.device
3486 and output.device == q.device
3487 and q.stride(-1) == 1
3488 and kv.stride(-1) == 1
3489 and output.stride() == q.stride()
3490 and attn_sink is not None
3491 and attn_sink.shape == (q.shape[1],)
3492 and attn_sink.dtype == torch.float32
3493 and attn_sink.device == q.device
3494 )
3495
3496
3497 def _rocm_sparse_attn_prefill_ragged_aiter_opus(
3498 q: torch.Tensor,
3499 kv: torch.Tensor,
3500 indices: torch.Tensor,
3501 indptr: torch.Tensor,
3502 scale: float,
3503 attn_sink: torch.Tensor,
3504 output: torch.Tensor,
3505 ) -> bool:
3506 pa_sparse_prefill_opus = _get_aiter_sparse_prefill_opus()
3507 if pa_sparse_prefill_opus is None:
3508 return False
3509
3510 indices = _as_int32_contiguous_1d(indices)
3511 indptr = _as_int32_contiguous_1d(indptr)
3512 empty_indices = indices[:0]
3513 empty_indptr = torch.zeros_like(indptr)
3514 pa_sparse_prefill_opus(
3515 q,
3516 kv,
3517 indices,
3518 indptr,
3519 kv[:1],
3520 empty_indices,
3521 empty_indptr,
3522 attn_sink.contiguous(),
3523 float(scale),
3524 out=output,
3525 )
3526 return True
3527
3528
3529 @functools.lru_cache
3530 def _decode_cu_count() -> int:
3531 try:
3532 return torch.cuda.get_device_properties(0).multi_processor_count
3533 except Exception:
3534 return 256 # For gfx950 arch, gated behind a fallback path for other archs.
3535
3536
3537 def _decode_partial_iters(
3538 avg_main_len: float, avg_extra_len: float, splits: int, block_k: int
3539 ) -> int:
3540 """BLOCK_K iterations one partial workgroup walks for ``splits`` splits.
3541
3542 Each split processes ``ceil(seg_len / splits)`` tokens of a segment, walked
3543 ``BLOCK_K`` at a time, and the main/extra segments are handled separately.
3544 """
3545 main_iters = (
3546 math.ceil(math.ceil(avg_main_len / splits) / block_k) if avg_main_len > 0 else 0
3547 )
3548 extra_iters = (
3549 math.ceil(math.ceil(avg_extra_len / splits) / block_k)
3550 if avg_extra_len > 0
3551 else 0
3552 )
3553 return main_iters + extra_iters
3554
3555
3556 def _decode_num_splits(
3557 num_queries: int,
3558 heads_blocks: int,
3559 avg_main_len: float = 0.0,
3560 avg_extra_len: float = 0.0,
3561 block_k: int = 32,
3562 ) -> int:
3563 """Pick a flash-decode split count to keep the GPU busy across batch sizes.
3564
3565 Decode launches only ``num_queries * heads_blocks`` workgroups otherwise,
3566 which severely under-fills the device for the low-concurrency regime that
3567 dominates latency. Splitting the KV sequence adds parallelism.
3568
3569 We model the relative partial-kernel latency for a given split count ``s``
3570 as ``waves * (1/s + mu)`` where ``waves = ceil(base * s / CU)`` and ``mu``
3571 is a small per-wave overhead penalty:
3572
3573 - ``waves / s`` captures the partial compute: each wave walks roughly
3574 ``total_tokens / s`` tokens and there are ``waves`` of them, so dividing
3575 by ``s`` makes more splits cheaper *until* they spill into extra waves.
3576 - ``mu * waves`` charges per-wave launch/tail overhead so we do not
3577 over-split into many mostly-idle waves (e.g. batch 224 on 256 CUs is
3578 best left at 1 split rather than 8 splits across 7 waves).
3579
3580 The minimiser naturally prefers split counts that pack the device into full
3581 waves (``base * s`` near a multiple of ``CU``) and falls back to 1 split
3582 once the batch already fills the device. Ties favour the smaller split
3583 count (less reduce work).
3584
3585 Finally we "snap down" the chosen split count to the smallest value that
3586 yields the same wave count *and* the same per-workgroup BLOCK_K iteration
3587 count. Because latency tracks iteration count (not raw token count), extra
3588 splits that do not lower the iteration count add only reduce/HBM overhead
3589 for no parallelism gain (e.g. batch 24: s8 and s10 both walk 4 extra iters
3590 in one wave, so s8 is strictly better). Snapping needs the average segment
3591 lengths, which the caller derives sync-free from the ragged index sizes.
3592 """
3593 base = max(1, num_queries * heads_blocks)
3594 # Target ~1 workgroup per CU: enough to fill the device while keeping the
3595 # reduce cost (which grows with split count) small. Tuned on gfx950.
3596 cu = max(1, _decode_cu_count())
3597 # Per-wave overhead penalty: higher values discourage split counts that
3598 # spill into extra GPU waves. Tuned on gfx950.
3599 mu = 0.04
3600 best_splits = 1
3601 best_cost = None
3602 # Search up to 16 splits; beyond that the reduce/HBM overhead dominates.
3603 for splits in range(1, 17):
3604 waves = (base * splits + cu - 1) // cu
3605 cost = waves * (1.0 / splits + mu)
3606 if best_cost is None or cost < best_cost - 1e-9:
3607 best_splits = splits
3608 best_cost = cost
3609
3610 if best_splits > 1 and (avg_main_len > 0 or avg_extra_len > 0):
3611 target_waves = (base * best_splits + cu - 1) // cu
3612 target_iters = _decode_partial_iters(
3613 avg_main_len, avg_extra_len, best_splits, block_k
3614 )
3615 for splits in range(1, best_splits):
3616 waves = (base * splits + cu - 1) // cu
3617 iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k)
3618 if waves == target_waves and iters == target_iters:
3619 best_splits = splits
3620 break
3621 return best_splits
3622
3623
3624 def _decode_gfx950_num_splits(
3625 num_queries: int,
3626 heads_blocks: int,
3627 avg_main_len: float = 0.0,
3628 avg_extra_len: float = 0.0,
3629 block_k: int = 32,
3630 ) -> int:
3631 base = max(1, num_queries * heads_blocks)
3632 cu = max(1, _decode_cu_count())
3633 target_workgroups = 2 * cu
3634 num_splits = min(
3635 32,
3636 max(
3637 1,
3638 math.ceil(target_workgroups / base),
3639 ),
3640 )
3641 if (
3642 base >= 16
3643 and num_splits > 4
3644 and _decode_partial_iters(avg_main_len, avg_extra_len, 4, block_k) <= 3
3645 ):
3646 return 4
3647 if 16 <= base < 64:
3648 one_wave_splits = max(1, cu // base)
3649 one_wave_iters = _decode_partial_iters(
3650 avg_main_len, avg_extra_len, one_wave_splits, block_k
3651 )
3652 target_waves = 1 if one_wave_iters <= 9 else 2
3653 num_splits = min(num_splits, max(1, target_waves * cu // base))
3654 if base >= 16 and num_splits > 1:
3655 target_waves = (base * num_splits + cu - 1) // cu
3656 target_iters = _decode_partial_iters(
3657 avg_main_len, avg_extra_len, num_splits, block_k
3658 )
3659 for splits in range(1, num_splits):
3660 waves = (base * splits + cu - 1) // cu
3661 iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k)
3662 if waves == target_waves and iters == target_iters:
3663 return splits
3664 return num_splits
3665
3666
3667 def _rocm_sparse_attn_decode_ragged_triton(
3668 q: torch.Tensor,
3669 main_cache: torch.Tensor,
3670 main_indices: torch.Tensor,
3671 main_indptr: torch.Tensor,
3672 scale: float,
3673 attn_sink: torch.Tensor | None,
3674 nope_head_dim: int,
3675 rope_head_dim: int,
3676 extra_cache: torch.Tensor | None = None,
3677 extra_indices: torch.Tensor | None = None,
3678 extra_indptr: torch.Tensor | None = None,
3679 out: torch.Tensor | None = None,
3680 extra_cache_nan_free: bool = False,
3681 adaptive_splits: bool = False,
3682 inv_rope_positions: torch.Tensor | None = None,
3683 inv_rope_cos_sin_cache: torch.Tensor | None = None,
3684 out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
3685 ) -> torch.Tensor:
3686 """Split-K sparse decode; returns the attention output.
3687
3688 With ``out_mxfp8 = (data, scale)`` the reduce writes MXFP8 instead of
3689 bf16: ``data`` is [b, h * d] e4m3 and ``scale`` [b, h * d // 32] E8M0, and
3690 ``data`` viewed as [b, h, d] is returned.
3691 """
3692 assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}"
3693 assert main_cache.ndim == 3, (
3694 f"expected main_cache=[blocks,block,bytes], got {main_cache.shape}"
3695 )
3696 assert main_indices.ndim == 1, (
3697 f"expected main_indices=[nnz], got {main_indices.shape}"
3698 )
3699 assert main_indptr.ndim == 1, f"expected main_indptr=[b+1], got {main_indptr.shape}"
3700 assert (
3701 not q.is_cpu
3702 and not main_cache.is_cpu
3703 and not main_indices.is_cpu
3704 and not main_indptr.is_cpu
3705 )
3706
3707 main_indices = _as_int32_contiguous_1d(main_indices)
3708 main_indptr = _as_int32_contiguous_1d(main_indptr)
3709 has_attn_sink = attn_sink is not None
3710 if attn_sink is None:
3711 attn_sink = torch.empty(1, device=q.device, dtype=torch.float32)
3712 else:
3713 attn_sink = attn_sink.contiguous()
3714
3715 num_queries, num_heads, head_dim = q.shape
3716 assert main_indptr.numel() == num_queries + 1, (
3717 f"expected main_indptr shape [{num_queries + 1}], got {main_indptr.shape}"
3718 )
3719 _validate_dsv4_sparse_dims(
3720 head_dim,
3721 nope_head_dim,
3722 rope_head_dim,
3723 "_rocm_sparse_attn_decode_ragged_triton",
3724 )
3725
3726 has_extra = (
3727 extra_cache is not None
3728 and extra_indices is not None
3729 and extra_indptr is not None
3730 )
3731 assert not extra_cache_nan_free or (_ON_GFX950 and has_extra), (
3732 "extra_cache_nan_free requires a gfx950 compressed cache with trusted "
3733 "canonical-writer provenance"
3734 )
3735 if has_extra:
3736 assert extra_cache is not None
3737 assert extra_indices is not None
3738 assert extra_indptr is not None
3739 assert extra_indices.ndim == 1, (
3740 f"expected extra_indices=[nnz], got {extra_indices.shape}"
3741 )
3742 assert extra_indptr.ndim == 1, (
3743 f"expected extra_indptr=[b+1], got {extra_indptr.shape}"
3744 )
3745 extra_indices = _as_int32_contiguous_1d(extra_indices)
3746 extra_indptr = _as_int32_contiguous_1d(extra_indptr)
3747 assert extra_indptr.numel() == num_queries + 1, (
3748 f"expected extra_indptr shape [{num_queries + 1}], got {extra_indptr.shape}"
3749 )
3750 else:
3751 extra_cache = main_cache
3752 extra_indices = torch.empty(0, device=q.device, dtype=torch.int32)
3753 extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32)
3754
3755 block_h = 16
3756 out_scale = None
3757 if out_mxfp8 is not None:
3758 assert out is None, "out and out_mxfp8 are mutually exclusive"
3759 assert _ON_GFX950, "the MXFP8 reduce epilogue is gfx950-only"
3760 assert inv_rope_positions is not None, (
3761 "the MXFP8 output feeds wo_a, so it must be inverse-RoPE'd first"
3762 )
3763 out_data, out_scale = out_mxfp8
3764 assert out_data.dtype == torch.float8_e4m3fn and out_scale.dtype == (
3765 torch.uint8
3766 ), f"expected e4m3/uint8 MXFP8 buffers, got {out_data.dtype}/{out_scale.dtype}"
3767 assert out_data.shape == (num_queries, num_heads * head_dim), (
3768 f"expected MXFP8 data [{num_queries}, {num_heads * head_dim}], "
3769 f"got {tuple(out_data.shape)}"
3770 )
3771 assert out_scale.shape == (num_queries, num_heads * head_dim // 32), (
3772 f"expected MXFP8 scale [{num_queries}, {num_heads * head_dim // 32}], "
3773 f"got {tuple(out_scale.shape)}"
3774 )
3775 assert out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
3776 out = out_data.view(num_queries, num_heads, head_dim)
3777 elif out is None:
3778 out = torch.empty_like(q, dtype=torch.bfloat16)
3779 else:
3780 assert out.shape == q.shape, f"expected out shape {q.shape}, got {out.shape}"
3781 assert out.device == q.device, (
3782 f"expected out on device {q.device}, got {out.device}"
3783 )
3784 assert out.dtype == torch.bfloat16, (
3785 f"expected out dtype {torch.bfloat16}, got {out.dtype}"
3786 )
3787 heads_blocks = triton.cdiv(num_heads, block_h)
3788 nope_block = triton.next_power_of_2(nope_head_dim)
3789 comb_dim = nope_head_dim + rope_head_dim
3790 is_fnuz = current_platform.is_fp8_fnuz()
3791
3792 if not (_ON_GFX942 or _ON_GFX950): # Fallback path for un-tuned architectures.
3793 block_k = 16 if head_dim >= 256 else 32
3794 _sparse_attn_decode_ragged_kernel[(num_queries, heads_blocks)](
3795 q,
3796 main_cache,
3797 main_indices,
3798 main_indptr,
3799 extra_cache,
3800 extra_indices,
3801 extra_indptr,
3802 attn_sink,
3803 out,
3804 q.stride(0),
3805 q.stride(1),
3806 out.stride(0),
3807 out.stride(1),
3808 main_cache.stride(0),
3809 extra_cache.stride(0),
3810 main_cache.shape[0] * main_cache.shape[1],
3811 extra_cache.shape[0] * extra_cache.shape[1],
3812 main_cache.shape[1],
3813 extra_cache.shape[1],
3814 scale,
3815 num_heads,
3816 HAS_ATTN_SINK=has_attn_sink,
3817 HAS_EXTRA=has_extra,
3818 NOPE_DIM=nope_head_dim,
3819 NOPE_BLOCK=nope_block,
3820 ROPE_DIM=rope_head_dim,
3821 IS_FNUZ_MAIN=is_fnuz,
3822 IS_FNUZ_EXTRA=False,
3823 BLOCK_H=block_h,
3824 BLOCK_K=block_k,
3825 num_warps=8,
3826 )
3827 return out
3828
3829 block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950.
3830 if _ON_GFX950:
3831 inv_q = 1.0 / max(1, num_queries)
3832 avg_main_len = main_indices.numel() * inv_q
3833 avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0
3834 num_splits = _decode_gfx950_num_splits(
3835 num_queries,
3836 heads_blocks,
3837 avg_main_len,
3838 avg_extra_len,
3839 block_k,
3840 )
3841 else:
3842 # Average per-query segment lengths, read sync-free from the ragged
3843 # index sizes, let the split heuristic avoid over-splitting.
3844 inv_q = 1.0 / max(1, num_queries)
3845 avg_main_len = main_indices.numel() * inv_q
3846 avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0
3847 num_splits = _decode_num_splits(
3848 num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k
3849 )
3850
3851 base_workgroups = num_queries * heads_blocks
3852 adaptive_splits = (
3853 _ON_GFX950 and adaptive_splits and base_workgroups >= 16 and num_splits > 4
3854 )
3855 one_wave_splits = (
3856 max(1, _decode_cu_count() // base_workgroups)
3857 if adaptive_splits and 16 <= base_workgroups < 64
3858 else num_splits
3859 )
3860
3861 part_m = torch.empty(
3862 (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device
3863 )
3864 part_l = torch.empty_like(part_m)
3865 part_acc = torch.empty(
3866 (num_queries, num_splits, num_heads, comb_dim),
3867 dtype=torch.float32,
3868 device=q.device,
3869 )
3870
3871 if _ON_GFX950:
3872 _sparse_attn_decode_gfx950_partial_kernel[
3873 (num_queries, num_splits, heads_blocks)
3874 ](
3875 q,
3876 main_cache,
3877 main_indices,
3878 main_indptr,
3879 extra_cache,
3880 extra_indices,
3881 extra_indptr,
3882 part_m,
3883 part_l,
3884 part_acc,
3885 q.stride(0),
3886 q.stride(1),
3887 main_cache.stride(0),
3888 extra_cache.stride(0),
3889 main_cache.shape[0] * main_cache.shape[1],
3890 extra_cache.shape[0] * extra_cache.shape[1],
3891 main_cache.shape[1],
3892 extra_cache.shape[1],
3893 scale,
3894 num_heads,
3895 HAS_EXTRA=has_extra,
3896 NOPE_DIM=nope_head_dim,
3897 ROPE_DIM=rope_head_dim,
3898 IS_FNUZ_MAIN=is_fnuz,
3899 IS_FNUZ_EXTRA=False,
3900 TRUST_EXTRA_CACHE_NAN_FREE=extra_cache_nan_free,
3901 ADAPTIVE_SPLITS=adaptive_splits,
3902 ONE_WAVE_SPLITS=one_wave_splits,
3903 BLOCK_H=block_h,
3904 BLOCK_K=block_k,
3905 NUM_SPLITS=num_splits,
3906 NUM_STAGES=1,
3907 num_warps=4,
3908 waves_per_eu=0,
3909 )
3910 else:
3911 _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)](
3912 q,
3913 main_cache,
3914 main_indices,
3915 main_indptr,
3916 extra_cache,
3917 extra_indices,
3918 extra_indptr,
3919 part_m,
3920 part_l,
3921 part_acc,
3922 q.stride(0),
3923 q.stride(1),
3924 main_cache.stride(0),
3925 extra_cache.stride(0),
3926 part_m.stride(0),
3927 part_m.stride(1),
3928 part_acc.stride(0),
3929 part_acc.stride(1),
3930 part_acc.stride(2),
3931 main_cache.shape[0] * main_cache.shape[1],
3932 extra_cache.shape[0] * extra_cache.shape[1],
3933 main_cache.shape[1],
3934 extra_cache.shape[1],
3935 scale,
3936 num_heads,
3937 HAS_EXTRA=has_extra,
3938 NOPE_DIM=nope_head_dim,
3939 NOPE_BLOCK=nope_block,
3940 ROPE_DIM=rope_head_dim,
3941 # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950).
3942 # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere).
3943 # Reading both with a single IS_FNUZ would decode one of them with the
3944 # wrong FNUZ/OCP scale ratio (~1.87×).
3945 IS_FNUZ_MAIN=is_fnuz,
3946 IS_FNUZ_EXTRA=False,
3947 BLOCK_H=block_h,
3948 BLOCK_K=block_k,
3949 NUM_SPLITS=num_splits,
3950 NUM_STAGES=1,
3951 num_warps=4,
3952 )
3953
3954 fuse_inv_rope = inv_rope_positions is not None
3955 if fuse_inv_rope:
3956 assert inv_rope_cos_sin_cache is not None
3957 assert inv_rope_cos_sin_cache.shape[-1] == rope_head_dim, (
3958 "fused inverse RoPE expects cos_sin_cache laid out as "
3959 f"[P, {rope_head_dim}] = cos | sin, got "
3960 f"{tuple(inv_rope_cos_sin_cache.shape)}"
3961 )
3962 assert nope_head_dim % 2 == 0 and rope_head_dim % 2 == 0, (
3963 "fused inverse RoPE pairs adjacent lanes, so both head dims must "
3964 f"be even, got nope={nope_head_dim} rope={rope_head_dim}"
3965 )
3966
3967 _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)](
3968 part_m,
3969 part_l,
3970 part_acc,
3971 attn_sink,
3972 out,
3973 inv_rope_positions,
3974 inv_rope_cos_sin_cache,
3975 out_scale,
3976 out.stride(0),
3977 out.stride(1),
3978 out_scale.stride(0) if out_scale is not None else 0,
3979 head_dim // 32,
3980 part_m.stride(0),
3981 part_m.stride(1),
3982 part_acc.stride(0),
3983 part_acc.stride(1),
3984 part_acc.stride(2),
3985 inv_rope_cos_sin_cache.stride(0) if inv_rope_cos_sin_cache is not None else 0,
3986 num_heads,
3987 HAS_ATTN_SINK=has_attn_sink,
3988 ADAPTIVE_SPLITS=adaptive_splits,
3989 COMB_DIM=comb_dim,
3990 BLOCK_H=1,
3991 NUM_SPLITS=num_splits,
3992 SPLITS_PAD=triton.next_power_of_2(num_splits),
3993 FUSE_INV_ROPE=fuse_inv_rope,
3994 NOPE=nope_head_dim,
3995 HALF=rope_head_dim // 2,
3996 QUANT_OUT=out_scale is not None,
3997 num_warps=4,
3998 )
3999 return out
4000
4001
4002 def _rocm_sparse_attn_decode_triton(
4003 q: torch.Tensor,
4004 main_cache: torch.Tensor,
4005 main_indices: torch.Tensor,
4006 scale: float,
4007 attn_sink: torch.Tensor | None,
4008 nope_head_dim: int,
4009 rope_head_dim: int,
4010 extra_cache: torch.Tensor | None = None,
4011 extra_indices: torch.Tensor | None = None,
4012 main_lengths: torch.Tensor | None = None,
4013 extra_lengths: torch.Tensor | None = None,
4014 main_ragged_indices: torch.Tensor | None = None,
4015 main_ragged_indptr: torch.Tensor | None = None,
4016 extra_ragged_indices: torch.Tensor | None = None,
4017 extra_ragged_indptr: torch.Tensor | None = None,
4018 out: torch.Tensor | None = None,
4019 extra_cache_nan_free: bool = False,
4020 adaptive_splits: bool = False,
4021 inv_rope_positions: torch.Tensor | None = None,
4022 inv_rope_cos_sin_cache: torch.Tensor | None = None,
4023 out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
4024 ) -> torch.Tensor:
4025 if main_ragged_indices is None or main_ragged_indptr is None:
4026 main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense(
4027 main_indices,
4028 main_lengths
4029 if main_lengths is not None
4030 else (main_indices >= 0).sum(dim=-1, dtype=torch.int32),
4031 num_rows=main_cache.shape[0] * main_cache.shape[1],
4032 )
4033
4034 if (
4035 (extra_ragged_indices is None or extra_ragged_indptr is None)
4036 and extra_cache is not None
4037 and extra_indices is not None
4038 ):
4039 extra_ragged_indices, extra_ragged_indptr = build_ragged_indices_from_dense(
4040 extra_indices,
4041 extra_lengths
4042 if extra_lengths is not None
4043 else (extra_indices >= 0).sum(dim=-1, dtype=torch.int32),
4044 num_rows=extra_cache.shape[0] * extra_cache.shape[1],
4045 )
4046
4047 return _rocm_sparse_attn_decode_ragged_triton(
4048 q=q,
4049 main_cache=main_cache,
4050 main_indices=main_ragged_indices,
4051 main_indptr=main_ragged_indptr,
4052 scale=scale,
4053 attn_sink=attn_sink,
4054 nope_head_dim=nope_head_dim,
4055 rope_head_dim=rope_head_dim,
4056 extra_cache=extra_cache,
4057 extra_indices=extra_ragged_indices,
4058 extra_indptr=extra_ragged_indptr,
4059 out=out,
4060 extra_cache_nan_free=extra_cache_nan_free,
4061 adaptive_splits=adaptive_splits,
4062 inv_rope_positions=inv_rope_positions,
4063 inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
4064 out_mxfp8=out_mxfp8,
4065 )
4066
4067
4068 def rocm_sparse_attn_prefill(
4069 q: torch.Tensor,
4070 kv: torch.Tensor,
4071 indices: torch.Tensor | None,
4072 topk_length: torch.Tensor | None,
4073 scale: float,
4074 head_dim: int,
4075 nope_head_dim: int,
4076 rope_head_dim: int,
4077 attn_sink: torch.Tensor | None,
4078 output: torch.Tensor,
4079 ragged_indices: torch.Tensor | None = None,
4080 ragged_indptr: torch.Tensor | None = None,
4081 ) -> None:
4082 assert kv.ndim == 3 and kv.shape[1] == 1, (
4083 f"ROCm Triton sparse prefill expects kv=[skv,1,d], got {kv.shape}"
4084 )
4085 _validate_sparse_dims(
4086 head_dim,
4087 nope_head_dim,
4088 rope_head_dim,
4089 "rocm_sparse_attn_prefill",
4090 )
4091 opus_attn_sink = None if attn_sink is None else attn_sink[: q.shape[1]]
4092 if (
4093 _can_use_aiter_sparse_prefill_opus(q, kv.squeeze(1), opus_attn_sink, output)
4094 and _get_aiter_sparse_prefill_opus() is not None
4095 ):
4096 if ragged_indices is None or ragged_indptr is None:
4097 assert indices is not None
4098 indices_2d = indices.reshape(indices.shape[0], -1)
4099 ragged_indices, ragged_indptr = build_ragged_indices_from_dense(
4100 indices_2d,
4101 topk_length
4102 if topk_length is not None
4103 else (indices_2d >= 0).sum(dim=-1, dtype=torch.int32),
4104 num_rows=kv.shape[0],
4105 )
4106 assert opus_attn_sink is not None
4107 if _rocm_sparse_attn_prefill_ragged_aiter_opus(
4108 q=q,
4109 kv=kv.squeeze(1),
4110 indices=ragged_indices,
4111 indptr=ragged_indptr,
4112 scale=scale,
4113 attn_sink=opus_attn_sink,
4114 output=output,
4115 ):
4116 return
4117
4118 if ragged_indices is not None and ragged_indptr is not None:
4119 _rocm_sparse_attn_prefill_ragged_triton(
4120 q=q,
4121 kv=kv.squeeze(1),
4122 indices=ragged_indices,
4123 indptr=ragged_indptr,
4124 scale=scale,
4125 attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
4126 nope_head_dim=nope_head_dim,
4127 rope_head_dim=rope_head_dim,
4128 out=output,
4129 )
4130 else:
4131 assert indices is not None
4132 indices_2d = indices.reshape(indices.shape[0], -1)
4133 _rocm_sparse_attn_prefill_triton(
4134 q=q,
4135 kv=kv.squeeze(1),
4136 indices=indices_2d,
4137 scale=scale,
4138 attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
4139 nope_head_dim=nope_head_dim,
4140 rope_head_dim=rope_head_dim,
4141 topk_length=topk_length,
4142 out=output,
4143 )
4144
4145
4146 def rocm_sparse_attn_decode(
4147 q: torch.Tensor,
4148 kv_cache: torch.Tensor | None,
4149 swa_k_cache: torch.Tensor,
4150 swa_only: bool,
4151 topk_indices: torch.Tensor | None,
4152 topk_lens: torch.Tensor | None,
4153 swa_indices: torch.Tensor,
4154 swa_lens: torch.Tensor,
4155 swa_ragged_indices: torch.Tensor | None,
4156 swa_ragged_indptr: torch.Tensor | None,
4157 topk_ragged_indices: torch.Tensor | None,
4158 topk_ragged_indptr: torch.Tensor | None,
4159 attn_sink: torch.Tensor | None,
4160 scale: float,
4161 head_dim: int,
4162 nope_head_dim: int,
4163 rope_head_dim: int,
4164 output: torch.Tensor | None,
4165 extra_cache_nan_free: bool = False,
4166 adaptive_splits: bool = False,
4167 inv_rope_positions: torch.Tensor | None = None,
4168 inv_rope_cos_sin_cache: torch.Tensor | None = None,
4169 output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
4170 ) -> int:
4171 """Run sparse MLA decode into ``output``.
4172
4173 Passing ``inv_rope_positions`` folds the inverse RoPE into the reduce
4174 epilogue. Returns how many leading rows of ``output`` came back rotated,
4175 so a caller mixing in a decode path that does not fuse still knows what it
4176 owes the standalone pass. Read it from the eager attention segment only.
4177
4178 ``output_mxfp8 = (data, scale)`` replaces ``output``: the reduce also
4179 MXFP8-quantizes the rotated rows for the FP8 wo_a (see
4180 ``_rocm_sparse_attn_decode_ragged_triton``). It needs gfx950 and the fused
4181 inverse RoPE, and always covers every row.
4182 """
4183 assert swa_k_cache.dtype == torch.uint8, (
4184 "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, "
4185 f"got {swa_k_cache.dtype}"
4186 )
4187 _validate_dsv4_sparse_dims(
4188 head_dim,
4189 nope_head_dim,
4190 rope_head_dim,
4191 "rocm_sparse_attn_decode",
4192 )
4193
4194 main_indices = swa_indices.reshape(swa_indices.shape[0], -1)
4195
4196 extra_cache = None
4197 extra_indices = None
4198 if not swa_only:
4199 assert kv_cache is not None
4200 assert topk_indices is not None or (
4201 topk_ragged_indices is not None and topk_ragged_indptr is not None
4202 )
4203 assert kv_cache.dtype == torch.uint8, (
4204 "ROCm Triton sparse decode expects uint8 fp8_ds_mla extra cache, "
4205 f"got {kv_cache.dtype}"
4206 )
4207 extra_cache = kv_cache
4208 if topk_indices is not None:
4209 extra_indices = topk_indices.reshape(topk_indices.shape[0], -1)
4210
4211 if output_mxfp8 is not None:
4212 assert output is None, "output and output_mxfp8 are mutually exclusive"
4213 direct_out = None
4214 else:
4215 assert output is not None
4216 direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
4217 attn_out = _rocm_sparse_attn_decode_triton(
4218 q=q,
4219 main_cache=swa_k_cache,
4220 main_indices=main_indices,
4221 scale=scale,
4222 attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]],
4223 nope_head_dim=nope_head_dim,
4224 rope_head_dim=rope_head_dim,
4225 extra_cache=extra_cache,
4226 extra_indices=extra_indices,
4227 main_lengths=swa_lens,
4228 extra_lengths=topk_lens,
4229 main_ragged_indices=swa_ragged_indices,
4230 main_ragged_indptr=swa_ragged_indptr,
4231 extra_ragged_indices=topk_ragged_indices,
4232 extra_ragged_indptr=topk_ragged_indptr,
4233 out=direct_out,
4234 extra_cache_nan_free=extra_cache_nan_free,
4235 adaptive_splits=adaptive_splits,
4236 inv_rope_positions=inv_rope_positions,
4237 inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
4238 out_mxfp8=output_mxfp8,
4239 )
4240 if output_mxfp8 is not None:
4241 return q.shape[0]
4242 assert output is not None
4243 if direct_out is None:
4244 output.copy_(attn_out.to(output.dtype))
4245 return output.shape[0] if inv_rope_positions is not None else 0
4246
4247
4248 @triton.jit(do_not_specialize=["token_offset"])
4249 def _prefill_topk_global_slots_kernel(
4250 out_ptr,
4251 lens_ptr,
4252 topk_ptr,
4253 topk_stride,
4254 token_to_req_ptr,
4255 query_start_loc_ptr,
4256 seq_lens_ptr,
4257 is_valid_token_ptr,
4258 block_table_ptr,
4259 block_table_stride,
4260 token_offset,
4261 num_compressed,
4262 TOPK: tl.constexpr,
4263 COMPRESS_RATIO: tl.constexpr,
4264 BLOCK_SIZE: tl.constexpr,
4265 BLOCK_W: tl.constexpr,
4266 ):
4267 row = tl.program_id(0)
4268 token = row + token_offset
4269 req = tl.load(token_to_req_ptr + token)
4270 query_start = tl.load(query_start_loc_ptr + req)
4271 query_len = tl.load(query_start_loc_ptr + req + 1) - query_start
4272 pos = tl.load(seq_lens_ptr + req) - query_len + token - query_start
4273 # Only (pos + 1) // ratio compressed tokens exist yet; the rest of the row
4274 # can hold stale indices.
4275 row_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOPK)
4276 row_len = tl.where(tl.load(is_valid_token_ptr + token), row_len, 0)
4277
4278 for start in range(0, TOPK, BLOCK_W):
4279 cols = start + tl.arange(0, BLOCK_W)
4280 live = cols < row_len
4281 local = tl.load(topk_ptr + row * topk_stride + cols, mask=live, other=-1)
4282 valid = live & (local >= 0) & (local < num_compressed)
4283 local = tl.where(valid, local, 0)
4284 block = tl.load(
4285 block_table_ptr + req * block_table_stride + local // BLOCK_SIZE,
4286 mask=valid,
4287 other=0,
4288 )
4289 slot = tl.where(valid, block * BLOCK_SIZE + local % BLOCK_SIZE, -1)
4290 tl.store(out_ptr + row * TOPK + cols, slot, mask=cols < TOPK)
4291 tl.store(lens_ptr + row, row_len)
4292
4293
4294 def build_prefill_topk_ragged_indices(
4295 topk_indices: torch.Tensor,
4296 token_to_req_indices: torch.Tensor,
4297 query_start_loc: torch.Tensor,
4298 seq_lens: torch.Tensor,
4299 is_valid_token: torch.Tensor,
4300 block_table: torch.Tensor,
4301 block_size: int,
4302 compress_ratio: int,
4303 num_compressed: int,
4304 token_offset: int,
4305 num_rows: int = -1,
4306 ) -> tuple[torch.Tensor, torch.Tensor]:
4307 """Map prefill top-k rows to a ragged stream of compressed-cache slots.
4308
4309 ``topk_indices`` holds local compressed positions for the prefill tokens,
4310 which sit at ``token_offset`` in the batch; ``token_to_req_indices``,
4311 ``query_start_loc``, ``seq_lens`` and ``block_table`` are batch-wide.
4312 ``block_size`` is the compressed cache's, i.e. already divided by the ratio.
4313 """
4314 topk_indices = topk_indices.reshape(topk_indices.shape[0], -1)
4315 num_tokens, width = topk_indices.shape
4316 dense = torch.empty(
4317 (num_tokens, width), dtype=torch.int32, device=topk_indices.device
4318 )
4319 lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
4320 if num_tokens > 0 and width > 0:
4321 _prefill_topk_global_slots_kernel[(num_tokens,)](
4322 dense,
4323 lens,
4324 topk_indices,
4325 topk_indices.stride(0),
4326 token_to_req_indices,
4327 query_start_loc,
4328 seq_lens,
4329 is_valid_token,
4330 block_table,
4331 block_table.stride(0),
4332 token_offset,
4333 num_compressed,
4334 TOPK=width,
4335 COMPRESS_RATIO=compress_ratio,
4336 BLOCK_SIZE=block_size,
4337 BLOCK_W=min(triton.next_power_of_2(width), 1024),
4338 )
4339 else:
4340 lens.zero_()
4341 return build_ragged_indices_from_dense(dense, lens, num_rows=num_rows)