@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 328 lines 12.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """ROCm gfx942/gfx950 block-sparse GQA prefill kernel for MiniMax-M3.
4
5 Only the prefill path is specialized on CDNA: each 128-token KV block is split
6 into SUB_K-token sub-tiles to right-size the per-block QK/PV MFMAs. Everything
7 else -- the decode split-K kernels, the FP8 dtype set, the sparse block size --
8 is reused unchanged from ``common.ops.sparse_attn``.
9 """
10
11 import torch
12
13 from vllm.models.minimax_m3.common.ops.sparse_attn import (
14 _FP8_DTYPES,
15 _KV_SCALE_NONE,
16 SPARSE_BLOCK_SIZE,
17 _kv_scale_args,
18 minimax_m3_sparse_attn_decode,
19 )
20 from vllm.platforms.rocm import on_gfx950, on_mi3xx
21 from vllm.triton_utils import tl, triton
22
23 __all__ = ["minimax_m3_sparse_attn", "minimax_m3_sparse_attn_decode"]
24
25
26 # Sub-tile width for the prefill kernel's per-block QK/PV GEMMs. gfx950 -> 64,
27 # gfx942 -> 32 (re-tune with tune_sparse_attn.py). Must divide SPARSE_BLOCK_SIZE.
28 _SPARSE_ATTN_SUB_K = SPARSE_BLOCK_SIZE // 2 if on_gfx950() else SPARSE_BLOCK_SIZE // 4
29
30 _SPARSE_ATTN_PREFILL_KWARG: dict | None = None
31
32
33 def _sparse_attn_prefill_kwargs() -> dict:
34 """MFMA + pipeline launch params for the sub-tiled prefill kernel.
35
36 gfx942 and gfx950 share the same params: ``num_warps=1`` keeps one wave
37 resident on the small per-sub-tile GEMM, ``matrix_instr_nonkdim=16`` /
38 ``kpack=2`` select the MFMA_16x16 path, and ``num_stages=1`` fits LDS and is
39 fastest in the sweep. Only the sub-tile width (``_SPARSE_ATTN_SUB_K``)
40 differs by arch. Empty on other AMD archs. Cached: arch is fixed per process.
41 """
42 global _SPARSE_ATTN_PREFILL_KWARG
43 if _SPARSE_ATTN_PREFILL_KWARG is None:
44 kwarg: dict = {}
45 if on_mi3xx():
46 kwarg = {
47 "num_warps": 1,
48 "matrix_instr_nonkdim": 16,
49 "kpack": 2,
50 "num_stages": 1,
51 }
52 _SPARSE_ATTN_PREFILL_KWARG = kwarg
53 return _SPARSE_ATTN_PREFILL_KWARG
54
55
56 # ---------------------------------------------------------------------------
57 # GQA block-sparse attention (paged). Main heads attend only to the selected
58 # blocks. BLOCK_SIZE_K == 128 so each selected block is one page.
59 # ---------------------------------------------------------------------------
60 # since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
61 # might lose pointer alignment, which trigger Triton recompiles. we don't actually
62 # need pointer alignment for those tensors anyway because we do scalar load.
63 @triton.heuristics(
64 {
65 "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
66 "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
67 "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
68 "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]
69 * triton.next_power_of_2(args["gqa_group_size"]),
70 }
71 )
72 @triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
73 def _gqa_sparse_fwd_kernel(
74 q_ptr, # [total_q, num_heads, head_dim]
75 kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
76 k_scale_ptr,
77 v_scale_ptr,
78 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
79 o_ptr, # [total_q, num_heads, head_dim]
80 block_table_ptr, # [num_reqs, max_blocks]
81 cu_seqlens_q,
82 cu_seqblocks_q,
83 seq_lens,
84 prefix_lens,
85 num_kv_heads,
86 gqa_group_size,
87 head_dim,
88 max_topk,
89 num_q_loop,
90 sm_scale,
91 stride_qn,
92 stride_qh,
93 stride_qd,
94 stride_kv_blk,
95 stride_kv_h,
96 stride_kv_pos,
97 stride_kv_d,
98 stride_ks_h,
99 stride_ks_t,
100 stride_vs_h,
101 stride_vs_t,
102 stride_th,
103 stride_tn,
104 stride_tk,
105 stride_on,
106 stride_oh,
107 stride_od,
108 stride_bt_b,
109 BLOCK_SIZE_Q: tl.constexpr,
110 BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
111 BLOCK_SIZE_D: tl.constexpr,
112 BLOCK_SIZE_H: tl.constexpr,
113 BLOCK_SIZE_T: tl.constexpr,
114 BLOCK_SIZE_QH: tl.constexpr,
115 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
116 KV_SCALE_MODE: tl.constexpr, # 0: none, 1: scalar, 2: [kv_head, token]
117 SUB_K: tl.constexpr, # CDNA only: KV sub-tile width (see _IS_MI3XX)
118 ):
119 sm_scale_log2e = sm_scale * 1.4426950409
120 pid_q = tl.program_id(0)
121 pid_kh = tl.program_id(1)
122 pid_b = tl.program_id(2)
123 pid_h = pid_kh * gqa_group_size
124 q_start = tl.load(cu_seqlens_q + pid_b)
125 q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start
126 q_block_start = tl.load(cu_seqblocks_q + pid_b)
127 q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start
128 seq_len = tl.load(seq_lens + pid_b)
129 prefix_len = tl.load(prefix_lens + pid_b)
130 if pid_q * num_q_loop >= q_block_len:
131 return
132 real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop)
133 bt_row = block_table_ptr + pid_b * stride_bt_b
134 off_d = tl.arange(0, BLOCK_SIZE_D)
135 d_mask = off_d < head_dim
136 for j in range(real_q_loop):
137 pid_q_j = pid_q * num_q_loop + j
138 t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
139 off_t = tl.arange(0, BLOCK_SIZE_T)
140 topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)
141 real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)
142 q_ptrs = tl.make_block_ptr(
143 base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
144 shape=(q_len, gqa_group_size, head_dim),
145 strides=(stride_qn, stride_qh, stride_qd),
146 offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
147 block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
148 order=(2, 1, 0),
149 )
150 q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
151 m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
152 lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
153 acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32)
154 q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D)
155
156 # CDNA: process each 128-token KV block in SUB_K-token sub-tiles so
157 # each QK/PV MFMA is right-sized. Numerically equivalent to the dense
158 # path below (flash-softmax reassociation).
159 NUM_SUB: tl.constexpr = BLOCK_SIZE_K // SUB_K
160 for _ in tl.range(real_topk):
161 blk = tl.load(t_ptr_j).to(tl.int32)
162 t_ptr_j = t_ptr_j + stride_tk
163 c = blk * BLOCK_SIZE_K
164 page = tl.load(bt_row + blk).to(tl.int64)
165 kv_base = kv_cache_ptr + page * stride_kv_blk + pid_kh * stride_kv_h
166 for sub_i in range(NUM_SUB):
167 off_sub = tl.arange(0, SUB_K) + sub_i * SUB_K
168 pos_sub = c + off_sub
169 pos_mask_sub = pos_sub < seq_len
170 k_sub = tl.load(
171 kv_base
172 + off_sub[None, :] * stride_kv_pos
173 + off_d[:, None] * stride_kv_d,
174 mask=d_mask[:, None] & pos_mask_sub[None, :],
175 other=0.0,
176 )
177 if USE_FP8:
178 k_sub = k_sub.to(q.dtype)
179 if KV_SCALE_MODE == 1:
180 k_sub = (k_sub * tl.load(k_scale_ptr)).to(q.dtype)
181 elif KV_SCALE_MODE == 2:
182 k_scale = tl.load(
183 k_scale_ptr
184 + pid_kh * stride_ks_h
185 + (page * BLOCK_SIZE_K + off_sub) * stride_ks_t,
186 mask=pos_mask_sub,
187 other=1.0,
188 )
189 k_sub = (k_sub * k_scale[None, :]).to(q.dtype)
190 off_q_sub = (
191 tl.arange(0, BLOCK_SIZE_Q)[:, None]
192 + pid_q_j * BLOCK_SIZE_Q
193 + prefix_len
194 - off_sub[None, :]
195 )
196 qk_sub = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, SUB_K), dtype=tl.float32)
197 # causal: q_abs_pos - k_off >= block_start (c)
198 qk_sub += tl.where(off_q_sub[:, None, :] >= c, 0, float("-inf"))
199 qk_sub = tl.reshape(qk_sub, BLOCK_SIZE_QH, SUB_K)
200 qk_sub += tl.dot(q, k_sub) * sm_scale_log2e
201 qk_sub += tl.where(pos_mask_sub[None, :], 0, float("-inf"))
202 m_ij = tl.maximum(m_i, tl.max(qk_sub, axis=1))
203 p_sub = tl.exp2(qk_sub - m_ij[:, None])
204 l_ij = tl.sum(p_sub, axis=1)
205 acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
206 v_sub = tl.load(
207 kv_base
208 + off_sub[:, None] * stride_kv_pos
209 + (head_dim + off_d[None, :]) * stride_kv_d,
210 mask=pos_mask_sub[:, None] & d_mask[None, :],
211 other=0.0,
212 )
213 if USE_FP8:
214 v_sub = v_sub.to(q.dtype)
215 if KV_SCALE_MODE == 1:
216 v_sub = (v_sub * tl.load(v_scale_ptr)).to(q.dtype)
217 elif KV_SCALE_MODE == 2:
218 v_scale = tl.load(
219 v_scale_ptr
220 + pid_kh * stride_vs_h
221 + (page * BLOCK_SIZE_K + off_sub) * stride_vs_t,
222 mask=pos_mask_sub,
223 other=1.0,
224 )
225 v_sub = (v_sub * v_scale[:, None]).to(q.dtype)
226 acc_o += tl.dot(p_sub.to(v_sub.dtype), v_sub)
227 m_i = m_ij
228 lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
229 acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
230 acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
231 o_ptrs = tl.make_block_ptr(
232 base=o_ptr + q_start * stride_on + pid_h * stride_oh,
233 shape=(q_len, gqa_group_size, head_dim),
234 strides=(stride_on, stride_oh, stride_od),
235 offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
236 block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
237 order=(2, 1, 0),
238 )
239 tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
240
241
242 @torch.no_grad()
243 def minimax_m3_sparse_attn(
244 q: torch.Tensor, # [total_q, num_heads, head_dim]
245 kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim]
246 topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
247 block_table: torch.Tensor, # [batch, max_blocks]
248 cu_seqlens_q: torch.Tensor, # [batch+1] int32
249 seq_lens: torch.Tensor, # [batch] int32
250 prefix_lens: torch.Tensor, # [batch] int32
251 max_query_len: int,
252 num_kv_heads: int,
253 sm_scale: float,
254 output: torch.Tensor, # [total_q, num_heads, head_dim]
255 k_scale: torch.Tensor | None = None,
256 v_scale: torch.Tensor | None = None,
257 ) -> None:
258 """GQA block-sparse attention over the selected blocks. block_size_q == 1."""
259 total_q, num_heads, head_dim = q.shape
260 batch = cu_seqlens_q.shape[0] - 1
261 topk = topk_idx.shape[-1]
262 gqa_group_size = num_heads // num_kv_heads
263 use_fp8 = kv_cache.dtype in _FP8_DTYPES
264 (
265 k_scale_arg,
266 v_scale_arg,
267 stride_ks_h,
268 stride_ks_t,
269 stride_vs_h,
270 stride_vs_t,
271 kv_scale_mode,
272 ) = (
273 _kv_scale_args(output, num_kv_heads, k_scale, v_scale)
274 if use_fp8
275 else (
276 output,
277 output,
278 0,
279 0,
280 0,
281 0,
282 _KV_SCALE_NONE,
283 )
284 )
285 grid = (max_query_len, num_kv_heads, batch)
286 _gqa_sparse_fwd_kernel[grid](
287 q,
288 kv_cache,
289 k_scale_arg,
290 v_scale_arg,
291 topk_idx,
292 output,
293 block_table,
294 cu_seqlens_q,
295 cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
296 seq_lens,
297 prefix_lens,
298 num_kv_heads,
299 gqa_group_size,
300 head_dim,
301 topk,
302 1, # num_q_loop
303 sm_scale,
304 q.stride(0),
305 q.stride(1),
306 q.stride(2),
307 kv_cache.stride(0),
308 kv_cache.stride(1),
309 kv_cache.stride(2),
310 kv_cache.stride(3),
311 stride_ks_h,
312 stride_ks_t,
313 stride_vs_h,
314 stride_vs_t,
315 topk_idx.stride(0),
316 topk_idx.stride(1),
317 topk_idx.stride(2),
318 output.stride(0),
319 output.stride(1),
320 output.stride(2),
321 block_table.stride(0),
322 BLOCK_SIZE_Q=1,
323 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
324 USE_FP8=use_fp8,
325 KV_SCALE_MODE=kv_scale_mode,
326 SUB_K=_SPARSE_ATTN_SUB_K,
327 **_sparse_attn_prefill_kwargs(),
328 )