@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ae5b8e4a8d

[ROCm][PERF] Enable AITER PA gluon decode for MiniMax-M3 MTP and dense layers (#52849)

Signed-off-by: ukannika <uma.kannikanti@amd.com> Signed-off-by: ykamiset <yaminipreethi.kamisetty@amd.com> Co-authored-by: ykamiset <yaminipreethi.kamisetty@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

ukannika committed Aug 28, 2026 at 15:54 UTC ae5b8e4a8d77a5de404a0cf9a91a56e8c7f2a94e
5 files changed +574 -201
vllm/models/minimax_m3/amd/model.py
+39 -12
@@ -101,6 +101,8 @@ from vllm.models.minimax_m3.common.mm_preprocess import (
101 from vllm.models.minimax_m3.common.sparse_attention import (
102 MiniMaxM3SparseBackend,
103 MiniMaxM3SparseImpl,
104 + MiniMaxM3SparseMetadata,
105 + minimax_m3_rebase_slots_to_page16,
106 minimax_m3_use_aiter_sparse_pa,
107 select_main_impl_cls,
108 )
@@ -755,35 +757,45 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
757 if is_quantized_kv_cache(self.kv_cache_dtype):
758 kv_cache = kv_cache.view(self.impl.kv_cache_fp8_dtype)
759 key_cache, value_cache = kv_cache.unbind(1)
758 - if not key_cache.is_contiguous() or not value_cache.is_contiguous():
759 - raise RuntimeError(
760 - "MiniMax-M3 AITER sparse PA requires K/V-separated KV cache "
761 - "storage. Set VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1 before "
762 - "initializing the engine."
763 - )
760
765 - x = 16 // key_cache.element_size()
761 + x = 16 // kv_cache.element_size()
762 if self.head_dim % x != 0:
763 raise RuntimeError(
764 "MiniMax-M3 AITER sparse PA requires head_dim divisible by "
765 f"16 / dtype_size, got head_dim={self.head_dim}, x={x}"
766 )
771 - num_blocks = key_cache.shape[0]
772 - num_phys16 = num_blocks * 8
773 - self.kv_cache_k = key_cache.view(
767 +
768 + num_blocks, _, block_size, _ = kv_cache.shape
769 + pages_in_side = block_size // 16
770 + if key_cache.is_contiguous() and value_cache.is_contiguous():
771 + k_src, v_src = key_cache, value_cache
772 + block_pages, v_page_offset = pages_in_side, 0
773 + elif kv_cache.is_contiguous():
774 + k_src = v_src = kv_cache
775 + block_pages, v_page_offset = 2 * pages_in_side, pages_in_side
776 + else:
777 + raise RuntimeError(
778 + "MiniMax-M3 AITER sparse PA needs each K/V head slot stored as "
779 + "whole 16-token pages, but the resolved KV cache layout gives "
780 + "neither dense planes nor block-contiguous pages."
781 + )
782 +
783 + num_phys16 = num_blocks * block_pages
784 + self.kv_cache_k = k_src.view(
785 num_phys16,
786 self.num_kv_heads,
787 self.head_dim // x,
788 16,
789 x,
790 )
780 - self.kv_cache_v = value_cache.view(
791 + # Offsetting V by half a block lets both sides share one page id.
792 + self.kv_cache_v = v_src.view(
793 num_phys16,
794 self.num_kv_heads,
795 16 // x,
796 self.head_dim,
797 x,
786 - )
798 + )[v_page_offset:]
799 self._aiter_sparse_pa_cache_data_ptr = self.kv_cache.data_ptr()
800
801 def get_aiter_sparse_pa_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]:
@@ -807,6 +819,21 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
819 )
820
821 key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache()
822 + if value_cache.shape[0] != key_cache.shape[0]:
823 + attn_metadata = get_forward_context().attn_metadata
824 + page16_slot_mapping = None
825 + if isinstance(attn_metadata, dict):
826 + main_md = attn_metadata[self.layer_name]
827 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
828 + page16_slot_mapping = main_md.page16_slot_mapping
829 + if (
830 + page16_slot_mapping is None
831 + or page16_slot_mapping.shape != slot_mapping.shape
832 + ):
833 + page16_slot_mapping = minimax_m3_rebase_slots_to_page16(
834 + slot_mapping, self.kv_cache.shape[2]
835 + )
836 + slot_mapping = page16_slot_mapping
837 kv_cache_dtype = (
838 self.kv_cache_dtype
839 if is_quantized_kv_cache(self.kv_cache_dtype)
vllm/models/minimax_m3/amd/ops/sparse_pa.py
+155 -118
@@ -32,35 +32,34 @@ def _is_fp8_kv_cache_tensor(kv_cache: torch.Tensor) -> bool:
32
33
34 @triton.jit
35 -def _build_sparse_block_table_kernel(
36 - topk_ptr, # [1, batch, topk] int32, selected logical 128-block ids
37 - block_table_ptr, # [batch, max_blocks] int32, logical 128-page table
38 - seq_lens_ptr, # [batch] int32
39 - sparse_bt_ptr, # [batch, topk * 8] int32, physical 16-page table
40 - sparse_ctx_ptr, # [batch] int32
35 +def _write_sparse_block_table_row(
36 + topk_row, # [topk] int32, selected logical block ids for this query
37 + bt_row, # [max_blocks] int32, this request's logical page table
38 + sbt_row, # [topk * PAGES_PER_BLOCK] int32, physical 16-page table
39 + ctx_ptr, # int32, this query's attended context length
40 + abs_pos, # absolute position of this query token, may be negative
41 max_topk,
42 - stride_topk_n,
42 stride_topk_k,
44 - stride_bt_b,
45 - stride_sbt_b,
43 SPARSE_BLOCK_SIZE_C: tl.constexpr,
44 PAGES_PER_BLOCK: tl.constexpr,
45 + BLOCK_PAGE_STRIDE: tl.constexpr,
46 BLOCK_SIZE_T: tl.constexpr,
47 ):
50 - pid_b = tl.program_id(0)
51 - seq_len = tl.load(seq_lens_ptr + pid_b)
52 - has_tokens = seq_len > 0
53 - last_blk = tl.maximum((seq_len - 1) // SPARSE_BLOCK_SIZE_C, 0)
48 + """Compact one query's selected logical blocks into physical page-16 ids.
49
55 - topk_row = topk_ptr + pid_b * stride_topk_n
56 - bt_row = block_table_ptr + pid_b * stride_bt_b
57 - sbt_row = sparse_bt_ptr + pid_b * stride_sbt_b
50 + ``BLOCK_PAGE_STRIDE`` is how many page ids a block spans, which is
51 + ``PAGES_PER_BLOCK`` when each K/V side is its own dense plane and twice
52 + that when both sides share a block. Padded rows carry a negative
53 + ``abs_pos``, which clamps the causal range to empty.
54 + """
55 + causal_len = tl.maximum(abs_pos + 1, 0)
56 + self_blk = abs_pos // SPARSE_BLOCK_SIZE_C
57
58 off_t = tl.arange(0, BLOCK_SIZE_T)
59 blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
61 - valid = has_tokens & (blk >= 0)
62 - is_tail = valid & (blk == last_blk)
63 - is_full = valid & (blk != last_blk)
60 + valid = (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
61 + is_tail = valid & (blk == self_blk)
62 + is_full = valid & (blk < self_blk)
63
64 n_full = tl.sum(is_full.to(tl.int32), axis=0)
65 n_valid = tl.sum(valid.to(tl.int32), axis=0)
@@ -68,7 +67,7 @@ def _build_sparse_block_table_kernel(
67 slot = tl.where(is_full, earlier_full, n_full)
68
69 logical_page = tl.load(bt_row + blk, mask=valid, other=0).to(tl.int32)
71 - base_phys = logical_page * PAGES_PER_BLOCK
70 + base_phys = logical_page * BLOCK_PAGE_STRIDE
71 dst_base = slot * PAGES_PER_BLOCK
72
73 for j in tl.static_range(PAGES_PER_BLOCK):
@@ -83,42 +82,54 @@ def _build_sparse_block_table_kernel(
82 mask=(off_w >= n_used) & (off_w < row_width),
83 )
84
86 - tail_tokens = tl.where(has_tokens, seq_len - last_blk * SPARSE_BLOCK_SIZE_C, 0)
85 + tail_tokens = causal_len - self_blk * SPARSE_BLOCK_SIZE_C
86 has_tail = tl.sum(is_tail.to(tl.int32), axis=0) > 0
87 ctx = n_full * SPARSE_BLOCK_SIZE_C + tl.where(has_tail, tail_tokens, 0)
89 - ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, seq_len))
90 - tl.store(sparse_ctx_ptr + pid_b, ctx)
88 + ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, causal_len))
89 + tl.store(ctx_ptr, ctx)
90
91
93 -@torch.no_grad()
94 -def minimax_m3_build_sparse_block_table(
95 - topk_idx: torch.Tensor, # [1, batch, topk]
96 - block_table: torch.Tensor, # [batch, max_blocks]
97 - seq_lens: torch.Tensor, # [batch]
98 -) -> tuple[torch.Tensor, torch.Tensor]:
99 - """Compact selected logical sparse blocks into physical page-16 tables."""
100 - assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
101 - batch = topk_idx.shape[1]
102 - topk = topk_idx.shape[-1]
103 - width = topk * PAGES_PER_SPARSE_BLOCK
104 - sparse_bt = torch.empty((batch, width), dtype=torch.int32, device=topk_idx.device)
105 - sparse_ctx = torch.empty((batch,), dtype=torch.int32, device=topk_idx.device)
106 - _build_sparse_block_table_kernel[(batch,)](
107 - topk_idx,
108 - block_table,
109 - seq_lens,
110 - sparse_bt,
111 - sparse_ctx,
112 - topk,
113 - topk_idx.stride(1),
114 - topk_idx.stride(2),
115 - block_table.stride(0),
116 - sparse_bt.stride(0),
117 - SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
118 - PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
119 - BLOCK_SIZE_T=triton.next_power_of_2(topk),
92 +@triton.jit
93 +def _build_sparse_block_table_kernel(
94 + topk_ptr, # [1, batch * DECODE_QUERY_LEN, topk] int32, logical block ids
95 + block_table_ptr, # [batch, max_blocks] int32, logical 128-page table
96 + seq_lens_ptr, # [batch] int32
97 + sparse_bt_ptr, # [batch * DECODE_QUERY_LEN, topk * 8] int32
98 + sparse_ctx_ptr, # [batch * DECODE_QUERY_LEN] int32
99 + max_topk,
100 + stride_topk_n,
101 + stride_topk_k,
102 + stride_bt_b,
103 + stride_sbt_b,
104 + SPARSE_BLOCK_SIZE_C: tl.constexpr,
105 + PAGES_PER_BLOCK: tl.constexpr,
106 + BLOCK_PAGE_STRIDE: tl.constexpr,
107 + BLOCK_SIZE_T: tl.constexpr,
108 + DECODE_QUERY_LEN: tl.constexpr,
109 +):
110 + # One program per query token. The decode indexer flattens speculative rows
111 + # as ``req_id * DECODE_QUERY_LEN + local_q``, and a request's query tokens
112 + # are the last DECODE_QUERY_LEN positions of its sequence, so both the
113 + # request id and the absolute position follow from the program id.
114 + pid_n = tl.program_id(0)
115 + req_id = pid_n // DECODE_QUERY_LEN
116 + local_q = pid_n % DECODE_QUERY_LEN
117 + seq_len = tl.load(seq_lens_ptr + req_id)
118 + abs_pos = seq_len - DECODE_QUERY_LEN + local_q
119 +
120 + _write_sparse_block_table_row(
121 + topk_ptr + pid_n * stride_topk_n,
122 + block_table_ptr + req_id * stride_bt_b,
123 + sparse_bt_ptr + pid_n * stride_sbt_b,
124 + sparse_ctx_ptr + pid_n,
125 + abs_pos,
126 + max_topk,
127 + stride_topk_k,
128 + SPARSE_BLOCK_SIZE_C,
129 + PAGES_PER_BLOCK,
130 + BLOCK_PAGE_STRIDE,
131 + BLOCK_SIZE_T,
132 )
121 - return sparse_bt, sparse_ctx
133
134
135 @triton.jit
@@ -136,52 +147,39 @@ def _build_sparse_block_table_prefill_kernel(
147 stride_sbt_n,
148 SPARSE_BLOCK_SIZE_C: tl.constexpr,
149 PAGES_PER_BLOCK: tl.constexpr,
150 + BLOCK_PAGE_STRIDE: tl.constexpr,
151 BLOCK_SIZE_T: tl.constexpr,
152 ):
153 pid_n = tl.program_id(0)
154 req_id = tl.load(req_id_ptr + pid_n)
155 abs_pos = tl.load(abs_pos_ptr + pid_n)
144 - # Padded speculative rows can have negative positions; clamp to empty range.
145 - causal_len = tl.maximum(abs_pos + 1, 0)
146 - self_blk = abs_pos // SPARSE_BLOCK_SIZE_C
147 -
148 - topk_row = topk_ptr + pid_n * stride_topk_n
149 - bt_row = block_table_ptr + req_id * stride_bt_b
150 - sbt_row = sparse_bt_ptr + pid_n * stride_sbt_n
151 -
152 - off_t = tl.arange(0, BLOCK_SIZE_T)
153 - blk = tl.load(topk_row + off_t * stride_topk_k, mask=off_t < max_topk, other=-1)
154 - valid = (causal_len > 0) & (blk >= 0) & (blk <= self_blk)
155 - is_tail = valid & (blk == self_blk)
156 - is_full = valid & (blk < self_blk)
156
158 - n_full = tl.sum(is_full.to(tl.int32), axis=0)
159 - n_valid = tl.sum(valid.to(tl.int32), axis=0)
160 - earlier_full = tl.cumsum(is_full.to(tl.int32), axis=0) - is_full.to(tl.int32)
161 - slot = tl.where(is_full, earlier_full, n_full)
162 -
163 - logical_page = tl.load(bt_row + blk, mask=valid, other=0).to(tl.int32)
164 - base_phys = logical_page * PAGES_PER_BLOCK
165 - dst_base = slot * PAGES_PER_BLOCK
157 + _write_sparse_block_table_row(
158 + topk_ptr + pid_n * stride_topk_n,
159 + block_table_ptr + req_id * stride_bt_b,
160 + sparse_bt_ptr + pid_n * stride_sbt_n,
161 + sparse_ctx_ptr + pid_n,
162 + abs_pos,
163 + max_topk,
164 + stride_topk_k,
165 + SPARSE_BLOCK_SIZE_C,
166 + PAGES_PER_BLOCK,
167 + BLOCK_PAGE_STRIDE,
168 + BLOCK_SIZE_T,
169 + )
170
167 - for j in tl.static_range(PAGES_PER_BLOCK):
168 - tl.store(sbt_row + dst_base + j, base_phys + j, mask=valid)
171
170 - n_used = n_valid * PAGES_PER_BLOCK
171 - off_w = tl.arange(0, BLOCK_SIZE_T * PAGES_PER_BLOCK)
172 - row_width = max_topk * PAGES_PER_BLOCK
173 - tl.store(
174 - sbt_row + off_w,
175 - tl.zeros_like(off_w),
176 - mask=(off_w >= n_used) & (off_w < row_width),
172 +def _alloc_sparse_block_table(
173 + topk_idx: torch.Tensor,
174 +) -> tuple[torch.Tensor, torch.Tensor]:
175 + assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
176 + total_q = topk_idx.shape[1]
177 + width = topk_idx.shape[-1] * PAGES_PER_SPARSE_BLOCK
178 + return (
179 + torch.empty((total_q, width), dtype=torch.int32, device=topk_idx.device),
180 + torch.empty((total_q,), dtype=torch.int32, device=topk_idx.device),
181 )
182
179 - tail_tokens = causal_len - self_blk * SPARSE_BLOCK_SIZE_C
180 - has_tail = tl.sum(is_tail.to(tl.int32), axis=0) > 0
181 - ctx = n_full * SPARSE_BLOCK_SIZE_C + tl.where(has_tail, tail_tokens, 0)
182 - ctx = tl.where(has_tail, ctx, tl.minimum(n_valid * SPARSE_BLOCK_SIZE_C, causal_len))
183 - tl.store(sparse_ctx_ptr + pid_n, ctx)
184 -
183
184 @torch.no_grad()
185 def minimax_m3_build_sparse_block_table_prefill(
@@ -189,15 +187,12 @@ def minimax_m3_build_sparse_block_table_prefill(
187 block_table: torch.Tensor, # [batch, max_blocks]
188 query_req_id: torch.Tensor, # [total_q]
189 query_abs_pos: torch.Tensor, # [total_q]
190 + block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
191 ) -> tuple[torch.Tensor, torch.Tensor]:
192 """Build one page-16 sparse block table row per prefill query token."""
194 - assert topk_idx.shape[0] == 1, "AITER sparse PA requires num_kv_heads == 1"
195 - total_q = topk_idx.shape[1]
193 + sparse_bt, sparse_ctx = _alloc_sparse_block_table(topk_idx)
194 topk = topk_idx.shape[-1]
197 - width = topk * PAGES_PER_SPARSE_BLOCK
198 - sparse_bt = torch.empty((total_q, width), dtype=torch.int32, device=topk_idx.device)
199 - sparse_ctx = torch.empty((total_q,), dtype=torch.int32, device=topk_idx.device)
200 - _build_sparse_block_table_prefill_kernel[(total_q,)](
195 + _build_sparse_block_table_prefill_kernel[(topk_idx.shape[1],)](
196 topk_idx,
197 block_table,
198 query_req_id,
@@ -211,6 +206,7 @@ def minimax_m3_build_sparse_block_table_prefill(
206 sparse_bt.stride(0),
207 SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
208 PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
209 + BLOCK_PAGE_STRIDE=block_page_stride,
210 BLOCK_SIZE_T=triton.next_power_of_2(topk),
211 )
212 return sparse_bt, sparse_ctx
@@ -221,27 +217,41 @@ def minimax_m3_build_sparse_block_table_decode(
217 topk_idx: torch.Tensor, # [1, batch * decode_query_len, topk]
218 block_table: torch.Tensor, # [batch, max_blocks]
219 seq_lens: torch.Tensor, # [batch]
224 - decode_query_len: int,
220 + decode_query_len: int = 1,
221 + block_page_stride: int = PAGES_PER_SPARSE_BLOCK,
222 ) -> tuple[torch.Tensor, torch.Tensor]:
223 """Build one page-16 sparse block table row per decode query token."""
227 - if decode_query_len == 1:
228 - return minimax_m3_build_sparse_block_table(topk_idx, block_table, seq_lens)
229 -
224 total_q = topk_idx.shape[1]
225 expected_q = block_table.shape[0] * decode_query_len
226 assert total_q == expected_q, (
227 "MiniMax-M3 decode top-k rows must equal batch * decode_query_len: "
228 f"{total_q} != {expected_q}"
229 )
236 - pos = torch.arange(total_q, dtype=torch.int32, device=topk_idx.device)
237 - query_req_id = torch.div(pos, decode_query_len, rounding_mode="floor")
238 - q_offset = pos - query_req_id * decode_query_len
239 - query_abs_pos = (seq_lens[query_req_id] - decode_query_len + q_offset).to(
240 - torch.int32
241 - )
242 - return minimax_m3_build_sparse_block_table_prefill(
243 - topk_idx, block_table, query_req_id, query_abs_pos
230 + sparse_bt, sparse_ctx = _alloc_sparse_block_table(topk_idx)
231 + topk = topk_idx.shape[-1]
232 + _build_sparse_block_table_kernel[(total_q,)](
233 + topk_idx,
234 + block_table,
235 + seq_lens,
236 + sparse_bt,
237 + sparse_ctx,
238 + topk,
239 + topk_idx.stride(1),
240 + topk_idx.stride(2),
241 + block_table.stride(0),
242 + sparse_bt.stride(0),
243 + SPARSE_BLOCK_SIZE_C=SPARSE_BLOCK_SIZE,
244 + PAGES_PER_BLOCK=PAGES_PER_SPARSE_BLOCK,
245 + BLOCK_PAGE_STRIDE=block_page_stride,
246 + BLOCK_SIZE_T=triton.next_power_of_2(topk),
247 + DECODE_QUERY_LEN=decode_query_len,
248 )
249 + return sparse_bt, sparse_ctx
250 +
251 +
252 +# Retained for the decode_query_len == 1 callers (and tests) that predate the
253 +# speculative-decode rows.
254 +minimax_m3_build_sparse_block_table = minimax_m3_build_sparse_block_table_decode
255
256
257 @triton.jit
@@ -313,6 +323,21 @@ def minimax_m3_insert_index_cache(
323 )
324
325
326 +def _sides_are_packed(k_cache: torch.Tensor, v_cache: torch.Tensor) -> bool:
327 + """Whether a block holds both its K and V pages instead of one side."""
328 + return v_cache.shape[0] != k_cache.shape[0]
329 +
330 +
331 +def _block_page_stride(k_cache: torch.Tensor, v_cache: torch.Tensor) -> int:
332 + """How many page ids one sparse block spans.
333 +
334 + Each side owns ``PAGES_PER_SPARSE_BLOCK`` pages. When the sides are dense
335 + planes a block is exactly that wide; when they interleave, a block covers
336 + both sides' pages and V is reached from the same page id at a fixed offset.
337 + """
338 + return PAGES_PER_SPARSE_BLOCK * (2 if _sides_are_packed(k_cache, v_cache) else 1)
339 +
340 +
341 def _gluon_scale_arg(
342 scale: torch.Tensor | None,
343 *,
@@ -361,7 +386,11 @@ def minimax_m3_sparse_attn_decode_aiter(
386 decode_query_len: int = 1,
387 ) -> None:
388 sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_decode(
364 - topk_idx, block_table, seq_lens, decode_query_len
389 + topk_idx,
390 + block_table,
391 + seq_lens,
392 + decode_query_len,
393 + _block_page_stride(k_cache, v_cache),
394 )
395 _run_gluon_decode(
396 q,
@@ -384,23 +413,20 @@ def minimax_m3_sparse_attn_prefill_aiter(
413 v_cache: torch.Tensor,
414 topk_idx: torch.Tensor, # [1, total_q, topk]
415 block_table: torch.Tensor,
387 - cu_seqlens_q: torch.Tensor,
388 - prefix_lens: torch.Tensor,
416 + query_req_id: torch.Tensor, # [total_q] int32, owning request per token
417 + query_abs_pos: torch.Tensor, # [total_q] int32, absolute position per token
418 num_kv_heads: int,
419 sm_scale: float,
420 output: torch.Tensor,
421 k_scale: torch.Tensor | None = None,
422 v_scale: torch.Tensor | None = None,
423 ) -> None:
395 - total_q = q.shape[0]
396 - pos = torch.arange(total_q, dtype=torch.int32, device=q.device)
397 - query_req_id = torch.searchsorted(cu_seqlens_q[1:].contiguous(), pos, right=True)
398 - query_req_id = query_req_id.to(torch.int32)
399 - query_abs_pos = (prefix_lens[query_req_id] + (pos - cu_seqlens_q[query_req_id])).to(
400 - torch.int32
401 - )
424 sparse_bt, sparse_ctx = minimax_m3_build_sparse_block_table_prefill(
403 - topk_idx, block_table, query_req_id, query_abs_pos
425 + topk_idx,
426 + block_table,
427 + query_req_id,
428 + query_abs_pos,
429 + _block_page_stride(k_cache, v_cache),
430 )
431 _run_gluon_decode(
432 q,
@@ -445,8 +471,11 @@ def _run_gluon_decode(
471 group_size = num_q_heads // num_kv_heads
472
473 nphys16, hkv = k_cache.shape[0], k_cache.shape[1]
474 + # V is offset by half a block when the sides interleave, so it holds fewer
475 + # pages than K even though both are addressed by the same page id.
476 + v_nphys16 = v_cache.shape[0]
477 k_cache_view = k_cache.view(nphys16 * hkv, 1, *k_cache.shape[2:])
449 - v_cache_view = v_cache.view(nphys16 * hkv, 1, *v_cache.shape[2:])
478 + v_cache_view = v_cache.view(v_nphys16 * hkv, 1, *v_cache.shape[2:])
479 q_view = q.view(total_q * num_kv_heads, group_size, head_size)
480 out_view = output.view(total_q * num_kv_heads, group_size, head_size)
481
@@ -463,11 +492,19 @@ def _run_gluon_decode(
492 is_fp8 = _is_fp8_kv_cache_tensor(k_cache)
493 compute_type = aiter_dtypes.fp8 if is_fp8 else q.dtype
494 if is_fp8:
495 + if _sides_are_packed(k_cache, v_cache) and not all(
496 + scale is None or scale.numel() == 1 for scale in (k_scale, v_scale)
497 + ):
498 + raise NotImplementedError(
499 + "MiniMax-M3 AITER sparse PA supports only scalar KV scales "
500 + "when K and V pages share a block; per-token scales would "
501 + "need the same page renumbering."
502 + )
503 k_scale_arg = _gluon_scale_arg(
504 k_scale, num_phys_pages=nphys16, num_kv_heads=hkv
505 )
506 v_scale_arg = _gluon_scale_arg(
470 - v_scale, num_phys_pages=nphys16, num_kv_heads=hkv
507 + v_scale, num_phys_pages=v_nphys16, num_kv_heads=hkv
508 )
509 else:
510 k_scale_arg = v_scale_arg = None
vllm/models/minimax_m3/amd/sparse_attention_msa.py
+3 -2
@@ -75,14 +75,15 @@ class MiniMaxM3SparseAiterPAImpl(MiniMaxM3SparseImpl):
75 if main_md.num_prefills > 0:
76 p = main_md.prefill
77 assert p is not None
78 + assert p.query_req_id is not None and p.query_abs_pos is not None
79 minimax_m3_sparse_attn_prefill_aiter(
80 q[nd:],
81 k_cache,
82 v_cache,
83 topk[:, nd:num_tokens, :],
84 p.block_table,
84 - p.cu_seqlens_q,
85 - p.context_lens,
85 + p.query_req_id,
86 + p.query_abs_pos,
87 self.num_kv_heads,
88 self.scale,
89 out[nd:],
vllm/models/minimax_m3/common/sparse_attention.py
+97 -7
@@ -66,6 +66,49 @@ def _minimax_m3_aiter_sparse_pa_requested() -> bool:
66 return rocm_aiter_ops.is_enabled() and rocm_aiter_ops.is_shuffle_kv_cache_enabled()
67
68
69 +def minimax_m3_rebase_slots_to_page16(
70 + slot_mapping: torch.Tensor,
71 + block_size: int,
72 + out: torch.Tensor | None = None,
73 +) -> torch.Tensor:
74 + """Rebase a token slot mapping onto AITER's page-16 page numbering.
75 +
76 + AITER's KV writer derives the destination page from ``slot // 16`` and
77 + assumes pages are numbered consecutively. When the resolved layout stores
78 + both K/V sides inside a block, a block spans twice as many pages, so only
79 + the block component of the slot doubles -- the offset within the page must
80 + not move. Clamping before the division leaves padding slots at their
81 + negative sentinel, which is what tells the writer to skip them.
82 + """
83 + if out is None:
84 + out = torch.empty_like(slot_mapping)
85 + torch.clamp(slot_mapping, min=0, out=out)
86 + return (
87 + out.div_(block_size, rounding_mode="floor").mul_(block_size).add_(slot_mapping)
88 + )
89 +
90 +
91 +def minimax_m3_query_token_positions(
92 + cu_seqlens_q: torch.Tensor,
93 + prefix_lens: torch.Tensor,
94 + total_q: int,
95 +) -> tuple[torch.Tensor, torch.Tensor]:
96 + """Map each prefill query token to its request id and absolute KV position.
97 +
98 + Both tensors are the same for every sparse layer in a step, so the AITER
99 + sparse PA block-table builder reads them from the metadata instead of
100 + rebuilding them once per layer.
101 + """
102 + pos = torch.arange(total_q, dtype=torch.int32, device=cu_seqlens_q.device)
103 + query_req_id = torch.searchsorted(
104 + cu_seqlens_q[1:].contiguous(), pos, right=True
105 + ).to(torch.int32)
106 + query_abs_pos = (prefix_lens[query_req_id] + (pos - cu_seqlens_q[query_req_id])).to(
107 + torch.int32
108 + )
109 + return query_req_id, query_abs_pos
110 +
111 +
112 def minimax_m3_use_aiter_sparse_pa(num_kv_heads: int) -> bool:
113 """Whether to use the ROCm AITER page-16 sparse PA prototype."""
114 requested = _minimax_m3_aiter_sparse_pa_requested()
@@ -117,10 +160,14 @@ class MiniMaxM3SparseBackend(AttentionBackend):
160
161 @classmethod
162 def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None:
120 - # AITER sparse PA unbinds K and V as separate contiguous regions, which
121 - # requires the head dimension to sit outside the block dimension.
163 + # AITER sparse PA reads K and V as separate page-16 tensors. LHBNC
164 + # gives each side a dense plane; LBHNC interleaves the sides inside a
165 + # block, which the page-id rebasing in ops/sparse_pa.py absorbs. LBHNC
166 + # is listed first because these layers publish two head slots while the
167 + # indexer publishes one, and mixed HNC shapes need a block-compact
168 + # layout to share the block-id pool.
169 if _minimax_m3_aiter_sparse_pa_requested():
123 - return (KVCacheLayout.LHBNC,)
170 + return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
171 return super().supported_kv_cache_layouts()
172
173
@@ -137,6 +184,12 @@ class MiniMaxM3SparsePrefillMetadata:
184 max_seq_len: int
185 total_kv_blocks: int
186
187 + # Per-query-token request id and absolute KV position, built once per step
188 + # for the AITER sparse PA block-table kernel. See
189 + # ``minimax_m3_query_token_positions``.
190 + query_req_id: torch.Tensor | None = None
191 + query_abs_pos: torch.Tensor | None = None
192 +
193
194 @dataclass
195 class MiniMaxM3SparseDecodeMetadata:
@@ -167,6 +220,11 @@ class MiniMaxM3SparseMetadata(AttentionMetadata):
220 prefill: MiniMaxM3SparsePrefillMetadata | None = None
221 decode: MiniMaxM3SparseDecodeMetadata | None = None
222
223 + # ``slot_mapping`` rebased onto the page-16 numbering AITER's KV writer
224 + # uses when both K/V sides share a block. Only set for the AITER sparse PA
225 + # path; see ``MiniMaxM3SparseMetadataBuilder.build``.
226 + page16_slot_mapping: torch.Tensor | None = None
227 +
228
229 class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMetadata]):
230 # Full cudagraphs for uniform decode batches, incl. spec-decode verify
@@ -191,6 +249,19 @@ class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMet
249 dtype=torch.int32,
250 device=device,
251 )
252 + # Every sparse layer shares one slot mapping, so the AITER page-16
253 + # rebase is done here once per step instead of once per layer. Stable
254 + # buffer for the same reason as the context lengths above.
255 + self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa(
256 + kv_cache_spec.num_kv_heads
257 + )
258 + self.page16_slot_mapping_buffer: torch.Tensor | None = None
259 + if self.use_aiter_sparse_pa:
260 + self.page16_slot_mapping_buffer = torch.empty(
261 + vllm_config.scheduler_config.max_num_batched_tokens,
262 + dtype=torch.int64,
263 + device=device,
264 + )
265
266 def build(
267 self,
@@ -236,17 +307,26 @@ class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMet
307 )
308 prefill_cu_seqlens_k[0] = 0
309 torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:])
310 + prefill_cu_seqlens_q = (
311 + query_start_loc[num_decodes:] - num_decode_tokens
312 + ).to(torch.int32)
313 + prefill_context_lens = context_lens[num_decodes:]
314 + query_req_id = query_abs_pos = None
315 + if self.use_aiter_sparse_pa:
316 + query_req_id, query_abs_pos = minimax_m3_query_token_positions(
317 + prefill_cu_seqlens_q, prefill_context_lens, num_prefill_tokens
318 + )
319 prefill_metadata = MiniMaxM3SparsePrefillMetadata(
240 - cu_seqlens_q=(query_start_loc[num_decodes:] - num_decode_tokens).to(
241 - torch.int32
242 - ),
320 + cu_seqlens_q=prefill_cu_seqlens_q,
321 cu_seqlens_k=prefill_cu_seqlens_k,
322 seq_lens=prefill_kv_lens,
245 - context_lens=context_lens[num_decodes:],
323 + context_lens=prefill_context_lens,
324 block_table=block_table[num_decodes:],
325 max_query_len=common_attn_metadata.max_query_len,
326 max_seq_len=common_attn_metadata.max_seq_len,
327 total_kv_blocks=prefill_total_kv_blocks,
328 + query_req_id=query_req_id,
329 + query_abs_pos=query_abs_pos,
330 )
331
332 decode_metadata: MiniMaxM3SparseDecodeMetadata | None = None
@@ -265,10 +345,20 @@ class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMet
345 decode_query_len=decode_query_len,
346 )
347
348 + page16_slot_mapping = None
349 + if self.page16_slot_mapping_buffer is not None:
350 + slot_mapping = common_attn_metadata.slot_mapping
351 + page16_slot_mapping = minimax_m3_rebase_slots_to_page16(
352 + slot_mapping,
353 + self.kv_cache_spec.block_size,
354 + out=self.page16_slot_mapping_buffer[: slot_mapping.shape[0]],
355 + )
356 +
357 return MiniMaxM3SparseMetadata(
358 seq_lens=seq_lens,
359 max_seq_len=common_attn_metadata.max_seq_len,
360 slot_mapping=common_attn_metadata.slot_mapping,
361 + page16_slot_mapping=page16_slot_mapping,
362 num_actual_tokens=num_tokens,
363 num_decodes=num_decodes,
364 num_decode_tokens=num_decode_tokens,
vllm/v1/attention/backends/rocm_aiter_fa.py
+280 -62
@@ -2,13 +2,17 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Attention layer with AiterFlashAttention."""
4
5 -from dataclasses import dataclass
5 +from dataclasses import dataclass, replace
6 from typing import ClassVar
7
8 import torch
9
10 from vllm._aiter_ops import rocm_aiter_ops
11 -from vllm.config import VllmConfig, get_layers_from_vllm_config
11 +from vllm.config import (
12 + VllmConfig,
13 + get_current_vllm_config_or_none,
14 + get_layers_from_vllm_config,
15 +)
16 from vllm.config.cache import CacheDType
17 from vllm.logger import init_logger
18 from vllm.model_executor.layers.attention import Attention
@@ -17,7 +21,7 @@ from vllm.platforms.interface import DeviceCapability
21 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
22 from vllm.utils.math_utils import cdiv
23 from vllm.utils.platform_utils import num_compute_units
20 -from vllm.utils.torch_utils import is_quantized_kv_cache
24 +from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache
25 from vllm.v1.attention.backend import (
26 AttentionBackend,
27 AttentionCGSupport,
@@ -34,9 +38,43 @@ from vllm.v1.attention.backends.utils import (
38 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
39 from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
40
41 +_PA_GLUON_MAX_QUERY_LEN = 4
42 +_PA_GLUON_MAX_QUERY_GROUP_SIZE = 64
43 +# Query group sizes the gluon paged-attention decode kernel is validated for: 8, 16
44 +_PA_GLUON_QUERY_GROUP_SIZES = (8, 16)
45 +
46 +# The kernel is only validated for this head size and kernel block size.
47 +_PA_GLUON_HEAD_SIZE = 128
48 +_PA_GLUON_BLOCK_SIZE = 128
49 +
50 +
51 +def _pa_gluon_supports(num_heads_q: int, num_heads_kv: int, head_size: int) -> bool:
52 + """Whether the head config can use the PA decode gluon kernel.
53 +
54 + Requires the shuffle KV cache layout (the kernel reads K/V in that layout)
55 + plus a head config the kernel is validated for. Both the advertised kernel
56 + block sizes and the decode dispatch go through this, so a config can never
57 + be offered a 128-token page that gluon will then decline to serve.
58 + """
59 + return (
60 + rocm_aiter_ops.is_shuffle_kv_cache_enabled()
61 + and num_heads_kv > 0
62 + and num_heads_q % num_heads_kv == 0
63 + and num_heads_q // num_heads_kv in _PA_GLUON_QUERY_GROUP_SIZES
64 + and head_size == _PA_GLUON_HEAD_SIZE
65 + )
66 +
67 +
68 _PARTITION_SIZE_ROCM = 256
69 _CP_TOKENS_PER_ITER_ROCM = 32 * 1024
70 if current_platform.is_rocm():
71 + from aiter.ops.triton.gluon.pa_decode_gluon import (
72 + get_recommended_splits,
73 + )
74 + from aiter.ops.triton.gluon.pa_decode_gluon import (
75 + pa_decode_gluon as _pa_decode_gluon,
76 + )
77 +
78 from vllm.triton_utils import tl, triton
79
80 def block_size(x, head_dim):
@@ -126,15 +164,17 @@ if current_platform.is_rocm():
164 # for kv cache layout as
165 # K: [num_blocks, num_head, head_dim // x, page_size, x]
166 # V: [num_blocks, num_head, page_size // x, head_dim, x]
167 + # The sides only sit a dense page apart when they are separate
168 + # planes (H outermost), so take the block stride from the tensor.
169 key_cache_ptr_offset = (
170 key_cache_ptr
131 - + block_id * num_heads * head_size * PAGE_SIZE
171 + + block_id * k_cache_stride0
172 + head_id * head_size * PAGE_SIZE
173 + slot_id * x
174 )
175 value_cache_ptr_offset = (
176 value_cache_ptr
137 - + block_id * num_heads * head_size * PAGE_SIZE
177 + + block_id * v_cache_stride0
178 + head_id * head_size * PAGE_SIZE
179 + (slot_id // x) * head_size * x
180 + slot_id % x
@@ -319,6 +359,7 @@ logger = init_logger(__name__)
359 @dataclass
360 class AiterFlashAttentionDecodeMetadata:
361 max_query_len: int
362 + uniform_query_len: int | None
363
364
365 @dataclass
@@ -475,10 +516,11 @@ class AiterFlashAttentionMetadataBuilder(
516 and self.scale.numel() == 1
517 and is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype)
518 ):
478 - layers = get_layers_from_vllm_config(self.vllm_config, Attention)
479 - first_layer_name = [k for k in layers][0]
519 + # Size the scales from a layer this builder owns. The draft model
520 + # runs its own builder over its own KV cache, so the first layer of
521 + # the whole config can carry an unrelated block count.
522 kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[
481 - first_layer_name
523 + self.layer_names[0]
524 ].kv_cache.shape
525 num_blocks = kv_cache_shape[0]
526 self.scale = torch.ones(
@@ -510,8 +552,15 @@ class AiterFlashAttentionMetadataBuilder(
552
553 decode_metadata = None
554 if num_decodes > 0:
555 + decode_max_query_len = query_lens_cpu[:num_decodes].max().item()
556 + uniform_query_len = (
557 + decode_max_query_len
558 + if num_decode_tokens == num_decodes * decode_max_query_len
559 + else None
560 + )
561 decode_metadata = AiterFlashAttentionDecodeMetadata(
514 - max_query_len=query_lens_cpu[:num_decodes].max().item(),
562 + max_query_len=decode_max_query_len,
563 + uniform_query_len=uniform_query_len,
564 )
565
566 prefill_metadata = None
@@ -682,9 +731,27 @@ class AiterFlashAttentionMetadataBuilder(
731 """
732 num_reqs = common_attn_metadata.num_reqs
733 num_tokens = common_attn_metadata.num_actual_tokens
734 + max_query_len = common_attn_metadata.max_query_len
735 +
736 + # Uniform-decode assumption does not hold for the
737 + # drafter's first forward after a target step: it inherits the target's
738 + # per-request query lengths, so rows can be longer than gluon's limit or
739 + # ragged. Those batches need the real split, which costs a sync.
740 + # _PA_GLUON_MAX_QUERY_LEN only binds when gluon is the decode consumer,
741 + # so test that rather than the shuffle layout alone.
742 + if _pa_gluon_supports(self.num_heads_q, self.num_heads_kv, self.headdim) and (
743 + max_query_len > _PA_GLUON_MAX_QUERY_LEN
744 + or num_tokens != num_reqs * max_query_len
745 + ):
746 + return self.build(
747 + common_prefix_len=0, common_attn_metadata=common_attn_metadata
748 + )
749
750 decode_metadata = AiterFlashAttentionDecodeMetadata(
687 - max_query_len=common_attn_metadata.max_query_len,
751 + max_query_len=max_query_len,
752 + uniform_query_len=(
753 + max_query_len if num_tokens == num_reqs * max_query_len else None
754 + ),
755 )
756
757 return AiterFlashAttentionMetadata(
@@ -739,6 +806,21 @@ class AiterFlashAttentionBackend(AttentionBackend):
806
807 @staticmethod
808 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
809 + if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():
810 + return [16, 32]
811 + # Only gluon serves 128-token pages; the pa_fwd_asm/ll4mi decode
812 + # fallback is limited to 16 and 32. Advertise 128 only when gluon can
813 + # run so selection never picks a page we cannot serve.
814 + vllm_config = get_current_vllm_config_or_none()
815 + if vllm_config is not None and vllm_config.model_config is not None:
816 + mc = vllm_config.model_config
817 + pc = vllm_config.parallel_config
818 + if _pa_gluon_supports(
819 + mc.get_num_attention_heads(pc),
820 + mc.get_num_kv_heads(pc),
821 + mc.get_head_size(),
822 + ):
823 + return [16, 32, 128]
824 return [16, 32]
825
826 @classmethod
@@ -765,13 +847,33 @@ class AiterFlashAttentionBackend(AttentionBackend):
847
848 @classmethod
849 def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
768 - """Validate the block size the ROCm gather kernels require."""
850 + """Validate the block size the ROCm gather kernels require, and under
851 + the shuffle layout publish K and V as two head slots.
852 +
853 + The shuffle read/write kernels rearrange a whole ``head_size x
854 + block_size`` tile per side, so each side needs its own contiguous run
855 + in the page. Packed into the content dim they interleave per token
856 + instead, which no layout can undo.
857 + """
858 # block_size == 1 is the per-token page-size probe (see
859 # Platform.get_page_size_bytes); real blocks must be gatherable in
860 # 16-token units by the ROCm kernel.
861 if spec.block_size != 1 and spec.block_size % 16 != 0:
862 raise ValueError("Block size must be a multiple of 16.")
774 - return spec
863 + if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():
864 + return spec
865 + if spec.state_content_bytes is not None:
866 + return spec
867 + assert spec.head_size == spec.head_size_v, (
868 + "Separate K/V head slots require symmetric K/V head sizes."
869 + )
870 + return replace(
871 + spec,
872 + num_head_slots=2,
873 + state_content_bytes=spec.num_kv_heads
874 + * spec.head_size
875 + * get_dtype_size(spec.dtype),
876 + )
877
878 @classmethod
879 def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
@@ -1160,11 +1262,25 @@ class AiterFlashAttentionImpl(AttentionImpl):
1262 if num_decodes > 0:
1263 assert attn_metadata.decode_metadata is not None
1264 decode_max_query_len = attn_metadata.decode_metadata.max_query_len
1265 + decode_query_len = attn_metadata.decode_metadata.uniform_query_len
1266
1164 - # Use unified_attention for speculative decoding (multi-token),
1165 - # sliding window, or sinks
1166 - # (pa_fwd_asm and paged_attention_v1 don't support sinks)
1167 - if (
1267 + # check if we can use the gluon paged-attention decode kernel
1268 + use_gluon = (
1269 + _pa_gluon_supports(
1270 + self.num_heads, self.num_kv_heads, self.head_size
1271 + )
1272 + and key_cache.shape[1] == _PA_GLUON_BLOCK_SIZE
1273 + and decode_query_len is not None
1274 + and decode_query_len <= _PA_GLUON_MAX_QUERY_LEN
1275 + and decode_query_len * (self.num_heads // self.num_kv_heads)
1276 + <= _PA_GLUON_MAX_QUERY_GROUP_SIZE
1277 + and (decode_query_len == 1 or attn_metadata.causal)
1278 + )
1279 + # Use unified_attention for the decodes the paged kernels can't
1280 + # take: sliding window, sinks, or a multi-token batch that gluon
1281 + # declined (pa_fwd_asm and paged_attention_v1 don't support
1282 + # sinks).
1283 + if not use_gluon and (
1284 self.sliding_window[0] != -1
1285 or decode_max_query_len > 1
1286 or self.sinks is not None
@@ -1287,21 +1403,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
1403 )
1404 elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1405 _, num_heads, head_size = query.shape
1290 - num_seqs = attn_metadata.seq_lens.shape[0]
1291 - max_num_partitions = (
1292 - attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
1293 - ) // _PARTITION_SIZE_ROCM
1294 - tmp_out = torch.empty(
1295 - (num_seqs, num_heads, max_num_partitions, head_size),
1296 - dtype=query.dtype,
1297 - device=query.device,
1298 - )
1299 - exp_sums = torch.empty(
1300 - (num_seqs, num_heads, max_num_partitions),
1301 - dtype=torch.float32,
1302 - device=query.device,
1303 - )
1304 - max_logits = torch.empty_like(exp_sums)
1406 num_blocks, block_size, num_kv_heads, _ = key_cache.shape
1407 x = 16 // key_cache.element_size()
1408 new_key_cache = key_cache.reshape(
@@ -1310,37 +1411,146 @@ class AiterFlashAttentionImpl(AttentionImpl):
1411 new_value_cache = value_cache.reshape(
1412 num_blocks, num_kv_heads, block_size // x, head_size, x
1413 )
1313 - k_qscale = (
1314 - layer._k_scale
1315 - if attn_metadata.k_scale is None
1316 - else attn_metadata.k_scale
1317 - )
1318 - v_qscale = (
1319 - layer._v_scale
1320 - if attn_metadata.v_scale is None
1321 - else attn_metadata.v_scale
1322 - )
1323 - rocm_aiter_ops.paged_attention_common(
1324 - Q=query[:num_decode_tokens],
1325 - K=new_key_cache,
1326 - V=new_value_cache,
1327 - tmp_out=tmp_out,
1328 - max_logits=max_logits,
1329 - exp_sums=exp_sums,
1330 - max_seq_len=attn_metadata.max_seq_len,
1331 - block_tables=attn_metadata.block_table[:num_decodes],
1332 - context_lens=attn_metadata.seq_lens[:num_decodes],
1333 - block_tables_stride0=attn_metadata.block_table[
1334 - :num_decodes
1335 - ].stride(0),
1336 - scale=self.scale,
1337 - K_QScale_hip=k_qscale,
1338 - V_QScale_hip=v_qscale,
1339 - K_QScale_asm=k_qscale,
1340 - V_QScale_asm=v_qscale,
1341 - out_=output[:num_decode_tokens],
1342 - kv_cache_dtype=self.kv_cache_dtype,
1343 - )
1414 +
1415 + if use_gluon:
1416 + is_fp8_kv = is_quantized_kv_cache(self.kv_cache_dtype)
1417 + # Per-tensor descale, as a float32 [1] tensor.
1418 + k_scale_gluon = (
1419 + layer._k_scale.reshape(1).to(torch.float32)
1420 + if is_fp8_kv
1421 + else None
1422 + )
1423 + v_scale_gluon = (
1424 + layer._v_scale.reshape(1).to(torch.float32)
1425 + if is_fp8_kv
1426 + else None
1427 + )
1428 + compute_type = (
1429 + current_platform.fp8_dtype() if is_fp8_kv else query.dtype
1430 + )
1431 + # The kernel folds the query positions into the group
1432 + # dim, so the intermediate buffers are sized by the
1433 + # combined extent.
1434 + query_group_size = decode_query_len * (
1435 + num_heads // num_kv_heads
1436 + )
1437 +
1438 + sliding_window_int = (
1439 + self.sliding_window[0] + 1
1440 + if self.sliding_window[0] > 0
1441 + else 0
1442 + )
1443 + if sliding_window_int > 0:
1444 + max_context_partition_num = 1
1445 + context_partition_size = 128
1446 + else:
1447 + max_context_partition_num = get_recommended_splits(
1448 + num_decodes, num_kv_heads
1449 + )
1450 + context_partition_size = _PARTITION_SIZE_ROCM
1451 +
1452 + intermediate_shape = (
1453 + num_decodes,
1454 + num_kv_heads,
1455 + max_context_partition_num,
1456 + query_group_size,
1457 + )
1458 + exp_sums = torch.empty(
1459 + intermediate_shape,
1460 + dtype=torch.float32,
1461 + device=query.device,
1462 + )
1463 + max_logits = torch.empty_like(exp_sums)
1464 + temporary_output = torch.empty(
1465 + (*intermediate_shape, head_size),
1466 + dtype=output.dtype,
1467 + device=query.device,
1468 + )
1469 +
1470 + _pa_decode_gluon(
1471 + output=output[:num_decode_tokens],
1472 + query=query[:num_decode_tokens],
1473 + key_cache=new_key_cache,
1474 + value_cache=new_value_cache,
1475 + context_lengths=attn_metadata.seq_lens[:num_decodes].to(
1476 + torch.int32
1477 + ),
1478 + block_tables=attn_metadata.block_table[:num_decodes].to(
1479 + torch.int32
1480 + ),
1481 + softmax_scale=self.scale,
1482 + query_length=decode_query_len,
1483 + max_context_partition_num=max_context_partition_num,
1484 + context_partition_size=context_partition_size,
1485 + compute_type=compute_type,
1486 + query_scale=None,
1487 + key_scale=k_scale_gluon,
1488 + value_scale=v_scale_gluon,
1489 + exp_sums=exp_sums,
1490 + max_logits=max_logits,
1491 + temporary_output=temporary_output,
1492 + alibi_slopes=self.alibi_slopes,
1493 + sinks=self.sinks,
1494 + sliding_window=sliding_window_int,
1495 + ps=True,
1496 + )
1497 + else:
1498 + # Unlike gluon, this kernel derives block addresses
1499 + # arithmetically, so it only reads the right bytes when
1500 + # each side is a dense plane of pages.
1501 + assert new_key_cache.stride(0) == (
1502 + num_kv_heads * head_size * block_size
1503 + ), (
1504 + "paged_attention_common needs the K/V sides as "
1505 + "dense planes; the resolved KV cache layout "
1506 + "interleaves them within a block."
1507 + )
1508 + num_seqs = attn_metadata.seq_lens.shape[0]
1509 + max_num_partitions = (
1510 + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
1511 + ) // _PARTITION_SIZE_ROCM
1512 + tmp_out = torch.empty(
1513 + (num_seqs, num_heads, max_num_partitions, head_size),
1514 + dtype=query.dtype,
1515 + device=query.device,
1516 + )
1517 + exp_sums = torch.empty(
1518 + (num_seqs, num_heads, max_num_partitions),
1519 + dtype=torch.float32,
1520 + device=query.device,
1521 + )
1522 + max_logits = torch.empty_like(exp_sums)
1523 + k_qscale = (
1524 + layer._k_scale
1525 + if attn_metadata.k_scale is None
1526 + else attn_metadata.k_scale
1527 + )
1528 + v_qscale = (
1529 + layer._v_scale
1530 + if attn_metadata.v_scale is None
1531 + else attn_metadata.v_scale
1532 + )
1533 + rocm_aiter_ops.paged_attention_common(
1534 + Q=query[:num_decode_tokens],
1535 + K=new_key_cache,
1536 + V=new_value_cache,
1537 + tmp_out=tmp_out,
1538 + max_logits=max_logits,
1539 + exp_sums=exp_sums,
1540 + max_seq_len=attn_metadata.max_seq_len,
1541 + block_tables=attn_metadata.block_table[:num_decodes],
1542 + context_lens=attn_metadata.seq_lens[:num_decodes],
1543 + block_tables_stride0=attn_metadata.block_table[
1544 + :num_decodes
1545 + ].stride(0),
1546 + scale=self.scale,
1547 + K_QScale_hip=k_qscale,
1548 + V_QScale_hip=v_qscale,
1549 + K_QScale_asm=k_qscale,
1550 + V_QScale_asm=v_qscale,
1551 + out_=output[:num_decode_tokens],
1552 + kv_cache_dtype=self.kv_cache_dtype,
1553 + )
1554 else:
1555 _, num_heads, head_size = query.shape
1556 nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
@@ -1393,6 +1603,14 @@ class AiterFlashAttentionImpl(AttentionImpl):
1603 def _split_kv_cache(
1604 self, kv_cache: torch.Tensor
1605 ) -> tuple[torch.Tensor, torch.Tensor]:
1606 + if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1607 + # (B, 2, N, H*hs) -> one (B, N, H, hs) per side, which is what the
1608 + # shuffle read/write kernels reinterpret in place. The heads are
1609 + # folded into the content dim by customize_spec.
1610 + key_cache, value_cache = kv_cache.unflatten(
1611 + -1, (self.num_kv_heads, self.head_size)
1612 + ).unbind(1)
1613 + return key_cache, value_cache
1614 # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1615 return kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1616