@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 722 lines 25.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Triton kernels for MiniMax M3 block-sparse GQA attention.
4
5 The main heads attend only to the blocks selected by the lightning indexer (see
6 ``index_topk``). Adapted to vLLM's paged KV cache: the KV page size is forced to
7 equal the sparse block size (128), so one selected block maps to exactly one
8 page.
9
10 Main K/V cache layout (vLLM):
11 ``(num_blocks, num_kv_heads, 128, 2 * head_dim)``
12 K=[..., :head_dim] V=[..., head_dim:]
13
14 Only the paths MiniMax M3 uses are implemented: no attention sink, base-2
15 (exp2/log2) softmax. The decode kernels use split-K (flash-decoding) over the
16 selected blocks with a separate merge step, since one query token per request
17 leaves the prefill kernels (which parallelize over the query dim) idle.
18 """
19
20 import torch
21
22 from vllm.platforms import current_platform
23 from vllm.triton_utils import tl, triton
24
25 # One sparse block == one KV page.
26 SPARSE_BLOCK_SIZE = 128
27
28 _FP8_DTYPES = (
29 torch.float8_e4m3fn,
30 torch.float8_e4m3fnuz,
31 torch.float8_e5m2,
32 torch.float8_e5m2fnuz,
33 )
34
35
36 # ---------------------------------------------------------------------------
37 # GQA block-sparse attention (paged). Main heads attend only to the selected
38 # blocks. BLOCK_SIZE_K == 128 so each selected block is one page.
39 # ---------------------------------------------------------------------------
40 # since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens
41 # might lose pointer alignment, which trigger Triton recompiles. we don't actually
42 # need pointer alignment for those tensors anyway because we do scalar load.
43 @triton.heuristics(
44 {
45 "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
46 "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),
47 "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]
48 * triton.next_power_of_2(args["gqa_group_size"]),
49 }
50 )
51 @triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"])
52 def _gqa_sparse_fwd_kernel(
53 q_ptr, # [total_q, num_heads, head_dim]
54 kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
55 k_scale_ptr,
56 v_scale_ptr,
57 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
58 o_ptr, # [total_q, num_heads, head_dim]
59 block_table_ptr, # [num_reqs, max_blocks]
60 cu_seqlens_q,
61 cu_seqblocks_q,
62 seq_lens,
63 prefix_lens,
64 num_kv_heads,
65 gqa_group_size,
66 head_dim,
67 max_topk,
68 num_q_loop,
69 sm_scale,
70 stride_qn,
71 stride_qh,
72 stride_qd,
73 stride_kv_blk,
74 stride_kv_h,
75 stride_kv_pos,
76 stride_kv_d,
77 stride_ks_h,
78 stride_ks_t,
79 stride_vs_h,
80 stride_vs_t,
81 stride_th,
82 stride_tn,
83 stride_tk,
84 stride_on,
85 stride_oh,
86 stride_od,
87 stride_bt_b,
88 BLOCK_SIZE_Q: tl.constexpr,
89 BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
90 BLOCK_SIZE_D: tl.constexpr,
91 BLOCK_SIZE_H: tl.constexpr,
92 BLOCK_SIZE_QH: tl.constexpr,
93 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
94 KV_SCALE_MODE: tl.constexpr, # 0: none, 1: scalar, 2: [kv_head, token]
95 ):
96 sm_scale_log2e = sm_scale * 1.4426950409
97 pid_q = tl.program_id(0)
98 pid_kh = tl.program_id(1)
99 pid_b = tl.program_id(2)
100 pid_h = pid_kh * gqa_group_size
101 q_start = tl.load(cu_seqlens_q + pid_b)
102 q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start
103 q_block_start = tl.load(cu_seqblocks_q + pid_b)
104 q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start
105 seq_len = tl.load(seq_lens + pid_b)
106 prefix_len = tl.load(prefix_lens + pid_b)
107 if pid_q * num_q_loop >= q_block_len:
108 return
109 real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop)
110 bt_row = block_table_ptr + pid_b * stride_bt_b
111 off_n = tl.arange(0, BLOCK_SIZE_K)
112 off_d = tl.arange(0, BLOCK_SIZE_D)
113 d_mask = off_d < head_dim
114 for j in range(real_q_loop):
115 pid_q_j = pid_q * num_q_loop + j
116 t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th
117 # Valid block count from seq position (no sentinel): block_size_q == 1.
118 q_abs = prefix_len + pid_q_j * BLOCK_SIZE_Q
119 valid_blocks = (q_abs + BLOCK_SIZE_K) // BLOCK_SIZE_K
120 real_topk = tl.minimum(max_topk, valid_blocks)
121 q_ptrs = tl.make_block_ptr(
122 base=q_ptr + q_start * stride_qn + pid_h * stride_qh,
123 shape=(q_len, gqa_group_size, head_dim),
124 strides=(stride_qn, stride_qh, stride_qd),
125 offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
126 block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
127 order=(2, 1, 0),
128 )
129 q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero")
130 off_q = (
131 tl.arange(0, BLOCK_SIZE_Q)[:, None]
132 + pid_q_j * BLOCK_SIZE_Q
133 + prefix_len
134 - tl.arange(0, BLOCK_SIZE_K)[None, :]
135 )
136 m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
137 lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32)
138 acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32)
139 q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D)
140 for _ in range(real_topk):
141 blk = tl.load(t_ptr_j).to(tl.int32)
142 t_ptr_j = t_ptr_j + stride_tk
143 c = blk * BLOCK_SIZE_K
144 page = tl.load(bt_row + blk).to(tl.int64)
145 pos = c + off_n
146 pos_mask = pos < seq_len
147 k = tl.load(
148 kv_cache_ptr
149 + page * stride_kv_blk
150 + pid_kh * stride_kv_h
151 + off_n[None, :] * stride_kv_pos
152 + off_d[:, None] * stride_kv_d,
153 mask=d_mask[:, None] & pos_mask[None, :],
154 other=0.0,
155 )
156 if USE_FP8:
157 k = k.to(q.dtype)
158 if KV_SCALE_MODE == 1:
159 k = (k * tl.load(k_scale_ptr)).to(q.dtype)
160 elif KV_SCALE_MODE == 2:
161 k_scale = tl.load(
162 k_scale_ptr
163 + pid_kh * stride_ks_h
164 + (page * BLOCK_SIZE_K + off_n) * stride_ks_t,
165 mask=pos_mask,
166 other=1.0,
167 )
168 k = (k * k_scale[None, :]).to(q.dtype)
169 qk = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
170 # causal: q_abs_pos - k_off >= block_start (c)
171 qk += tl.where(off_q[:, None, :] >= c, 0, float("-inf"))
172 qk = tl.reshape(qk, BLOCK_SIZE_QH, BLOCK_SIZE_K)
173 qk += tl.dot(q, k) * sm_scale_log2e
174 qk += tl.where(pos_mask[None, :], 0, float("-inf"))
175 m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
176 p = tl.exp2(qk - m_ij[:, None])
177 l_ij = tl.sum(p, axis=1)
178 acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
179 v = tl.load(
180 kv_cache_ptr
181 + page * stride_kv_blk
182 + pid_kh * stride_kv_h
183 + off_n[:, None] * stride_kv_pos
184 + (head_dim + off_d[None, :]) * stride_kv_d,
185 mask=pos_mask[:, None] & d_mask[None, :],
186 other=0.0,
187 )
188 if USE_FP8:
189 v = v.to(q.dtype)
190 if KV_SCALE_MODE == 1:
191 v = (v * tl.load(v_scale_ptr)).to(q.dtype)
192 elif KV_SCALE_MODE == 2:
193 v_scale = tl.load(
194 v_scale_ptr
195 + pid_kh * stride_vs_h
196 + (page * BLOCK_SIZE_K + off_n) * stride_vs_t,
197 mask=pos_mask,
198 other=1.0,
199 )
200 v = (v * v_scale[:, None]).to(q.dtype)
201 acc_o += tl.dot(p.to(v.dtype), v)
202 m_i = m_ij
203 lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
204 acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None]
205 acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D)
206 o_ptrs = tl.make_block_ptr(
207 base=o_ptr + q_start * stride_on + pid_h * stride_oh,
208 shape=(q_len, gqa_group_size, head_dim),
209 strides=(stride_on, stride_oh, stride_od),
210 offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0),
211 block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D),
212 order=(2, 1, 0),
213 )
214 tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2))
215
216
217 # ---------------------------------------------------------------------------
218 # Decode kernels (split-K). Decode batches are flattened request-major, with a
219 # runtime query length used to map each query token back to its request metadata.
220 # This parallelizes over the selected top-k blocks, producing partials that the
221 # merge kernel combines (flash-decoding). All chunk counts depend only on shape
222 # constants so the grid is fixed within a cuda graph. Base-2 (exp2/log2)
223 # softmax matches the prefill kernel.
224 # ---------------------------------------------------------------------------
225 @triton.heuristics(
226 {
227 "BLOCK_SIZE_H": lambda args: max(
228 16, triton.next_power_of_2(args["gqa_group_size"])
229 ),
230 "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),
231 }
232 )
233 @triton.jit(do_not_specialize=["decode_query_len"])
234 def _gqa_sparse_decode_kernel(
235 q_ptr, # [total_q, num_heads, head_dim]
236 kv_cache_ptr, # main cache: [num_blocks, num_kv_heads, 128, 2*head_dim]
237 k_scale_ptr,
238 v_scale_ptr,
239 t_ptr, # topk_idx: [num_kv_heads, total_q, topk]
240 o_ptr, # partial out: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
241 lse_ptr, # partial lse (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
242 block_table_ptr, # [num_reqs, max_blocks]
243 seq_lens, # [num_reqs]
244 total_q,
245 gqa_group_size,
246 head_dim,
247 max_topk,
248 sm_scale,
249 decode_query_len,
250 stride_qn,
251 stride_qh,
252 stride_qd,
253 stride_kv_blk,
254 stride_kv_h,
255 stride_kv_pos,
256 stride_kv_d,
257 stride_ks_h,
258 stride_ks_t,
259 stride_vs_h,
260 stride_vs_t,
261 stride_th,
262 stride_tn,
263 stride_tk,
264 stride_o_c,
265 stride_o_b,
266 stride_o_h,
267 stride_o_d,
268 stride_l_c,
269 stride_l_b,
270 stride_l_h,
271 stride_bt_b,
272 BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)
273 NUM_TOPK_CHUNKS: tl.constexpr,
274 BLOCK_SIZE_H: tl.constexpr,
275 BLOCK_SIZE_D: tl.constexpr,
276 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load
277 KV_SCALE_MODE: tl.constexpr, # 0: none, 1: scalar, 2: [kv_head, token]
278 USE_PDL: tl.constexpr,
279 ):
280 sm_scale_log2e = sm_scale * 1.4426950409
281 # split-K over the topk dimension: pid(0) folds (query-token, chunk).
282 pid_bc, pid_kh = tl.program_id(0), tl.program_id(1)
283 pid_b = pid_bc % total_q
284 pid_c = pid_bc // total_q
285 req_id = pid_b // decode_query_len
286 q_offset = pid_b - req_id * decode_query_len
287 pid_h = pid_kh * gqa_group_size
288 chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS
289 chunk_start_topk = pid_c * chunk_size_topk
290 chunk_end_compiletime = chunk_start_topk + chunk_size_topk
291
292 if USE_PDL:
293 tl.extra.cuda.gdc_wait()
294
295 seq_len = tl.load(seq_lens + req_id)
296 query_pos = seq_len - decode_query_len + q_offset
297 # Full-CG padding uses zero-length request rows. Clamp to an empty
298 # attention range instead of letting padded rows produce negative lengths.
299 kv_len = tl.maximum(query_pos + 1, 0)
300
301 # Valid block count from seq_len (no sentinel): min(topk, cdiv(kv_len, blk)).
302 idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn
303 num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
304 real_topk = tl.minimum(max_topk, num_blocks)
305 chunk_end_topk = tl.minimum(chunk_end_compiletime, real_topk)
306
307 off_n = tl.arange(0, BLOCK_SIZE_K)
308 off_d = tl.arange(0, BLOCK_SIZE_D)
309 d_mask = off_d < head_dim
310 bt_row = block_table_ptr + req_id * stride_bt_b
311
312 m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
313 lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32)
314 acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32)
315 q_ptrs = tl.make_block_ptr(
316 base=q_ptr + pid_b * stride_qn + pid_h * stride_qh,
317 shape=(gqa_group_size, head_dim),
318 strides=(stride_qh, stride_qd),
319 offsets=(0, 0),
320 block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
321 order=(1, 0),
322 )
323 q = tl.load(q_ptrs, boundary_check=(0, 1), padding_option="zero")
324
325 cur_idx_ptr = idx_base + chunk_start_topk * stride_tk
326 for _ in tl.range(chunk_start_topk, chunk_end_topk):
327 blk = tl.load(cur_idx_ptr).to(tl.int32)
328 cur_idx_ptr = cur_idx_ptr + stride_tk
329 c = blk * BLOCK_SIZE_K
330 page = tl.load(bt_row + blk).to(tl.int64)
331 pos = c + off_n
332 pos_mask = pos < kv_len
333 k = tl.load(
334 kv_cache_ptr
335 + page * stride_kv_blk
336 + pid_kh * stride_kv_h
337 + off_n[None, :] * stride_kv_pos
338 + off_d[:, None] * stride_kv_d,
339 mask=d_mask[:, None] & pos_mask[None, :],
340 other=0.0,
341 )
342 if USE_FP8:
343 k = k.to(q.dtype)
344 if KV_SCALE_MODE == 1:
345 k = (k * tl.load(k_scale_ptr)).to(q.dtype)
346 elif KV_SCALE_MODE == 2:
347 k_scale = tl.load(
348 k_scale_ptr
349 + pid_kh * stride_ks_h
350 + (page * BLOCK_SIZE_K + off_n) * stride_ks_t,
351 mask=pos_mask,
352 other=1.0,
353 )
354 k = (k * k_scale[None, :]).to(q.dtype)
355 qk = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_K), dtype=tl.float32)
356 qk += tl.where(pos_mask[None, :], 0, float("-inf"))
357 qk += tl.dot(q, k) * sm_scale_log2e
358 m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
359 p = tl.exp2(qk - m_ij[:, None])
360 l_ij = tl.sum(p, axis=1)
361 acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None]
362 v = tl.load(
363 kv_cache_ptr
364 + page * stride_kv_blk
365 + pid_kh * stride_kv_h
366 + off_n[:, None] * stride_kv_pos
367 + (head_dim + off_d[None, :]) * stride_kv_d,
368 mask=pos_mask[:, None] & d_mask[None, :],
369 other=0.0,
370 )
371 if USE_FP8:
372 v = v.to(q.dtype)
373 if KV_SCALE_MODE == 1:
374 v = (v * tl.load(v_scale_ptr)).to(q.dtype)
375 elif KV_SCALE_MODE == 2:
376 v_scale = tl.load(
377 v_scale_ptr
378 + pid_kh * stride_vs_h
379 + (page * BLOCK_SIZE_K + off_n) * stride_vs_t,
380 mask=pos_mask,
381 other=1.0,
382 )
383 v = (v * v_scale[:, None]).to(q.dtype)
384 acc_o += tl.dot(p.to(v.dtype), v)
385 m_i = m_ij
386 lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij)
387
388 if USE_PDL:
389 tl.extra.cuda.gdc_launch_dependents()
390
391 # Empty chunks for active rows must store zero output; otherwise the merge
392 # can hit 0 * NaN. All-empty padded rows may still produce NaNs in merge.
393 scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), tl.zeros_like(lse_i))
394 acc_o = acc_o * scale[:, None]
395 o_ptrs = tl.make_block_ptr(
396 base=o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h,
397 shape=(gqa_group_size, head_dim),
398 strides=(stride_o_h, stride_o_d),
399 offsets=(0, 0),
400 block_shape=(BLOCK_SIZE_H, BLOCK_SIZE_D),
401 order=(1, 0),
402 )
403 tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1))
404 lse_ptrs = tl.make_block_ptr(
405 base=lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h,
406 shape=(gqa_group_size,),
407 strides=(stride_l_h,),
408 offsets=(0,),
409 block_shape=(BLOCK_SIZE_H,),
410 order=(0,),
411 )
412 tl.store(lse_ptrs, lse_i.to(lse_ptr.dtype.element_ty), boundary_check=(0,))
413
414
415 @triton.heuristics(
416 {"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}
417 )
418 @triton.jit
419 def _merge_topk_attn_out_kernel(
420 o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim]
421 lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads]
422 out_ptr, # merged out: [total_q, num_heads, head_dim]
423 head_dim,
424 stride_o_c,
425 stride_o_b,
426 stride_o_h,
427 stride_o_d,
428 stride_l_c,
429 stride_l_b,
430 stride_l_h,
431 stride_out_n,
432 stride_out_h,
433 stride_out_d,
434 NUM_TOPK_CHUNKS: tl.constexpr,
435 BLOCK_SIZE_D: tl.constexpr,
436 USE_PDL: tl.constexpr,
437 ):
438 pid_b, pid_h = tl.program_id(0), tl.program_id(1)
439
440 # NOTE: assume seq_lens is safe to load before gdc_wait()
441 if USE_PDL:
442 tl.extra.cuda.gdc_wait()
443 tl.extra.cuda.gdc_launch_dependents()
444
445 off_c = tl.arange(0, NUM_TOPK_CHUNKS)
446 off_d = tl.arange(0, BLOCK_SIZE_D)
447 o_ptrs = tl.make_block_ptr(
448 base=o_ptr + pid_b * stride_o_b + pid_h * stride_o_h,
449 shape=(NUM_TOPK_CHUNKS, head_dim),
450 strides=(stride_o_c, stride_o_d),
451 offsets=(0, 0),
452 block_shape=(NUM_TOPK_CHUNKS, BLOCK_SIZE_D),
453 order=(1, 0),
454 )
455 lse_ptrs = lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c
456 o = tl.load(o_ptrs, boundary_check=(0, 1), padding_option="zero")
457 lse = tl.load(lse_ptrs) # empty chunks contribute -inf -> weight 0
458 lse_max = tl.max(lse, axis=0)
459 weights = tl.exp2(lse - lse_max)
460 weights = weights / tl.sum(weights, axis=0)
461 o_merged = tl.sum(o * weights[:, None], axis=0)
462 out_ptrs = (
463 out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d
464 )
465 tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim)
466
467
468 # ---------------------------------------------------------------------------
469 # Python wrappers
470 # ---------------------------------------------------------------------------
471 _KV_SCALE_NONE = 0
472 _KV_SCALE_SCALAR = 1
473 _KV_SCALE_PER_TOKEN_HEAD = 2
474
475
476 def _kv_scale_args(
477 output: torch.Tensor,
478 num_kv_heads: int,
479 k_scale: torch.Tensor | None,
480 v_scale: torch.Tensor | None,
481 ) -> tuple[torch.Tensor, torch.Tensor, int, int, int, int, int]:
482 if k_scale is None and v_scale is None:
483 return output, output, 0, 0, 0, 0, _KV_SCALE_NONE
484 if k_scale is None or v_scale is None:
485 raise ValueError("k_scale and v_scale must be both provided or both None")
486 if k_scale.device != output.device or v_scale.device != output.device:
487 raise ValueError("k_scale and v_scale must be on the same device as output")
488 if k_scale.numel() == 1 and v_scale.numel() == 1:
489 return k_scale, v_scale, 0, 0, 0, 0, _KV_SCALE_SCALAR
490 if k_scale.dim() == 2 and v_scale.dim() == 2:
491 if k_scale.shape[0] != num_kv_heads or v_scale.shape[0] != num_kv_heads:
492 raise ValueError(
493 "per-token/head KV scales must have shape "
494 f"[{num_kv_heads}, max_kv_tokens]"
495 )
496 if k_scale.shape != v_scale.shape:
497 raise ValueError("k_scale and v_scale must have matching shapes")
498 return (
499 k_scale,
500 v_scale,
501 k_scale.stride(0),
502 k_scale.stride(1),
503 v_scale.stride(0),
504 v_scale.stride(1),
505 _KV_SCALE_PER_TOKEN_HEAD,
506 )
507 raise ValueError(
508 "MiniMax-M3 sparse attention supports scalar KV scales or "
509 "[num_kv_heads, max_kv_tokens] per-token/head scales"
510 )
511
512
513 @torch.no_grad()
514 def minimax_m3_sparse_attn(
515 q: torch.Tensor, # [total_q, num_heads, head_dim]
516 kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim]
517 topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
518 block_table: torch.Tensor, # [batch, max_blocks]
519 cu_seqlens_q: torch.Tensor, # [batch+1] int32
520 seq_lens: torch.Tensor, # [batch] int32
521 prefix_lens: torch.Tensor, # [batch] int32
522 max_query_len: int,
523 num_kv_heads: int,
524 sm_scale: float,
525 output: torch.Tensor, # [total_q, num_heads, head_dim]
526 k_scale: torch.Tensor | None = None,
527 v_scale: torch.Tensor | None = None,
528 ) -> None:
529 """GQA block-sparse attention over the selected blocks. block_size_q == 1."""
530 total_q, num_heads, head_dim = q.shape
531 batch = cu_seqlens_q.shape[0] - 1
532 topk = topk_idx.shape[-1]
533 gqa_group_size = num_heads // num_kv_heads
534 use_fp8 = kv_cache.dtype in _FP8_DTYPES
535 (
536 k_scale_arg,
537 v_scale_arg,
538 stride_ks_h,
539 stride_ks_t,
540 stride_vs_h,
541 stride_vs_t,
542 kv_scale_mode,
543 ) = (
544 _kv_scale_args(output, num_kv_heads, k_scale, v_scale)
545 if use_fp8
546 else (
547 output,
548 output,
549 0,
550 0,
551 0,
552 0,
553 _KV_SCALE_NONE,
554 )
555 )
556 grid = (max_query_len, num_kv_heads, batch)
557 _gqa_sparse_fwd_kernel[grid](
558 q,
559 kv_cache,
560 k_scale_arg,
561 v_scale_arg,
562 topk_idx,
563 output,
564 block_table,
565 cu_seqlens_q,
566 cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1
567 seq_lens,
568 prefix_lens,
569 num_kv_heads,
570 gqa_group_size,
571 head_dim,
572 topk,
573 1, # num_q_loop
574 sm_scale,
575 q.stride(0),
576 q.stride(1),
577 q.stride(2),
578 kv_cache.stride(0),
579 kv_cache.stride(1),
580 kv_cache.stride(2),
581 kv_cache.stride(3),
582 stride_ks_h,
583 stride_ks_t,
584 stride_vs_h,
585 stride_vs_t,
586 topk_idx.stride(0),
587 topk_idx.stride(1),
588 topk_idx.stride(2),
589 output.stride(0),
590 output.stride(1),
591 output.stride(2),
592 block_table.stride(0),
593 BLOCK_SIZE_Q=1,
594 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
595 USE_FP8=use_fp8,
596 KV_SCALE_MODE=kv_scale_mode,
597 )
598
599
600 @torch.no_grad()
601 def minimax_m3_sparse_attn_decode(
602 q: torch.Tensor, # [total_q, num_heads, head_dim]
603 kv_cache: torch.Tensor, # [num_blocks, num_kv_heads, 128, 2*head_dim]
604 topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk]
605 block_table: torch.Tensor, # [num_reqs, max_blocks]
606 seq_lens: torch.Tensor, # [num_reqs] int32
607 num_kv_heads: int,
608 sm_scale: float,
609 output: torch.Tensor, # [total_q, num_heads, head_dim]
610 decode_query_len: int,
611 k_scale: torch.Tensor | None = None,
612 v_scale: torch.Tensor | None = None,
613 ) -> None:
614 """GQA block-sparse attention for decode (split-K over the top-k blocks)."""
615 total_q, num_heads, head_dim = q.shape
616 assert total_q == seq_lens.shape[0] * decode_query_len
617 max_topk = topk_idx.shape[-1]
618 gqa_group_size = num_heads // num_kv_heads
619 use_fp8 = kv_cache.dtype in _FP8_DTYPES
620 (
621 k_scale_arg,
622 v_scale_arg,
623 stride_ks_h,
624 stride_ks_t,
625 stride_vs_h,
626 stride_vs_t,
627 kv_scale_mode,
628 ) = (
629 _kv_scale_args(output, num_kv_heads, k_scale, v_scale)
630 if use_fp8
631 else (
632 output,
633 output,
634 0,
635 0,
636 0,
637 0,
638 _KV_SCALE_NONE,
639 )
640 )
641 use_pdl = current_platform.is_arch_support_pdl()
642 # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA
643 # SM9+); this ROCm Triton rejects it even when False ("Keyword argument
644 # launch_pdl was specified but unrecognised"). Only pass it when PDL is
645 # actually supported -- on ROCm use_pdl is always False, so it's omitted.
646 pdl_launch = {"launch_pdl": True} if use_pdl else {}
647 # split-K over the selected blocks; chunk count is shape-constant (cuda graph).
648 TARGET_GRID = 256
649 target = max(1, min(max_topk, TARGET_GRID // max(1, total_q * num_kv_heads)))
650 num_topk_chunks = 1 << (target.bit_length() - 1)
651 o_partial = torch.empty(
652 num_topk_chunks, total_q, num_heads, head_dim, dtype=q.dtype, device=q.device
653 )
654 lse_partial = torch.empty(
655 num_topk_chunks, total_q, num_heads, dtype=torch.float32, device=q.device
656 )
657 grid = (total_q * num_topk_chunks, num_kv_heads)
658 _gqa_sparse_decode_kernel[grid](
659 q,
660 kv_cache,
661 k_scale_arg,
662 v_scale_arg,
663 topk_idx,
664 o_partial,
665 lse_partial,
666 block_table,
667 seq_lens,
668 total_q,
669 gqa_group_size,
670 head_dim,
671 max_topk,
672 sm_scale,
673 decode_query_len,
674 q.stride(0),
675 q.stride(1),
676 q.stride(2),
677 kv_cache.stride(0),
678 kv_cache.stride(1),
679 kv_cache.stride(2),
680 kv_cache.stride(3),
681 stride_ks_h,
682 stride_ks_t,
683 stride_vs_h,
684 stride_vs_t,
685 topk_idx.stride(0),
686 topk_idx.stride(1),
687 topk_idx.stride(2),
688 o_partial.stride(0),
689 o_partial.stride(1),
690 o_partial.stride(2),
691 o_partial.stride(3),
692 lse_partial.stride(0),
693 lse_partial.stride(1),
694 lse_partial.stride(2),
695 block_table.stride(0),
696 BLOCK_SIZE_K=SPARSE_BLOCK_SIZE,
697 NUM_TOPK_CHUNKS=num_topk_chunks,
698 USE_FP8=use_fp8,
699 KV_SCALE_MODE=kv_scale_mode,
700 USE_PDL=use_pdl,
701 **pdl_launch,
702 )
703 merge_grid = (total_q, num_heads)
704 _merge_topk_attn_out_kernel[merge_grid](
705 o_partial,
706 lse_partial,
707 output,
708 head_dim,
709 o_partial.stride(0),
710 o_partial.stride(1),
711 o_partial.stride(2),
712 o_partial.stride(3),
713 lse_partial.stride(0),
714 lse_partial.stride(1),
715 lse_partial.stride(2),
716 output.stride(0),
717 output.stride(1),
718 output.stride(2),
719 NUM_TOPK_CHUNKS=num_topk_chunks,
720 USE_PDL=use_pdl,
721 **pdl_launch,
722 )