@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 530 lines 18 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Triton unified attention with different K/V head dimensions (DiffKV).
4
5 This is a slimmed fork of ``triton_unified_attention.py`` for models like
6 MiMo-V2.5 where the V tensor's head dimension differs from K's. The KV cache
7 is the same packed layout used by ``FlashAttentionDiffKVBackend``:
8
9 kv_cache: [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v]
10
11 We slice ``key_cache = kv_cache[..., :head_size_qk]`` and
12 ``value_cache = kv_cache[..., head_size_qk:]`` on the host, so the kernel
13 takes two cache pointers but with two distinct head sizes.
14
15 Both 2D and 3D launches are supported:
16 - 2D: one program per (q-block, kv-head); tile-loop walks the full KV
17 sequence; final output written directly. Used for prefill and large
18 decode batches.
19 - 3D: one program per (q-block, kv-head, segm); each program covers a
20 KV slice and writes per-segment partials (max/expsum/output). A
21 follow-up ``kernel_reduce_segments_diffkv`` combines them. Selected
22 for decode-only batches whose 2D grid would under-fill the GPU.
23 """
24
25 from typing import Any
26
27 import torch
28
29 import vllm.envs as envs
30 from vllm.logger import init_logger
31 from vllm.triton_utils import tl, triton
32 from vllm.v1.attention.ops.triton_attention_helpers import (
33 apply_alibi_to_score,
34 apply_softcap,
35 cdiv_fn,
36 compute_kv_seq_mask,
37 compute_tile_loop_bounds,
38 find_seq_idx,
39 init_softmax_M,
40 resolve_seq_and_query_len,
41 softmax_step,
42 store_segm_reduce_scalars,
43 )
44
45 logger = init_logger(__name__)
46
47 is_batch_invariant = envs.VLLM_BATCH_INVARIANT
48
49
50 @triton.jit
51 def kernel_unified_attention_diffkv(
52 # Output destinations. In 2D mode we write the final result into
53 # ``output_ptr``; in 3D mode we write per-segment partials into
54 # ``segm_*`` and ``output_ptr`` is unused (callers may pass any
55 # non-null pointer).
56 output_ptr,
57 segm_output_ptr,
58 segm_max_ptr,
59 segm_expsum_ptr,
60 query_ptr,
61 key_cache_ptr, # view of packed cache: [..., :head_size_qk]
62 value_cache_ptr, # view of packed cache: [..., head_size_qk:hqk+hv]
63 sink_ptr,
64 block_tables_ptr,
65 seq_lens_ptr,
66 alibi_slopes_ptr,
67 scale,
68 softcap,
69 num_query_heads: tl.constexpr,
70 num_queries_per_kv: tl.constexpr,
71 block_table_stride: tl.int64,
72 query_stride_0: tl.int64,
73 query_stride_1: tl.int64, # == HEAD_SIZE_QK
74 output_stride_0: tl.int64,
75 output_stride_1: tl.int64, # == HEAD_SIZE_V
76 BLOCK_SIZE: tl.constexpr,
77 TILE_SIZE: tl.constexpr,
78 HEAD_SIZE_QK: tl.constexpr,
79 HEAD_SIZE_QK_PADDED: tl.constexpr,
80 HEAD_SIZE_V: tl.constexpr,
81 HEAD_SIZE_V_PADDED: tl.constexpr,
82 USE_ALIBI_SLOPES: tl.constexpr,
83 USE_ALIBI_SQRT: tl.constexpr,
84 USE_SOFTCAP: tl.constexpr,
85 USE_SINKS: tl.constexpr,
86 SLIDING_WINDOW: tl.constexpr,
87 # Strides for both cache views (they share the same packed buffer, so
88 # dims 0/1/2 strides match; only the per-head extent differs).
89 stride_k_cache_0: tl.int64,
90 stride_k_cache_1: tl.int64,
91 stride_k_cache_2: tl.int64,
92 stride_k_cache_3: tl.constexpr,
93 stride_v_cache_0: tl.int64,
94 stride_v_cache_1: tl.int64,
95 stride_v_cache_2: tl.int64,
96 stride_v_cache_3: tl.constexpr,
97 query_start_len_ptr,
98 BLOCK_Q: tl.constexpr,
99 num_seqs: tl.int32,
100 BLOCK_M: tl.constexpr,
101 NUM_SEGMENTS_PER_SEQ: tl.constexpr,
102 # ``IS_3D`` toggles between 2D layout (one program walks the full KV
103 # sequence) and 3D layout (split-KV / FlashDecoding-style: per-segm
104 # programs write partials, finalized by ``kernel_reduce_segments_diffkv``).
105 IS_3D: tl.constexpr,
106 ):
107 q_block_global_idx = tl.program_id(0)
108 kv_head_idx = tl.program_id(1)
109 segm_idx = tl.program_id(2) if IS_3D else 0
110
111 (
112 seq_idx,
113 q_block_local_idx,
114 cur_batch_in_all_start_index,
115 cur_batch_query_len,
116 seq_len,
117 ) = resolve_seq_and_query_len(
118 query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q
119 )
120
121 if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
122 return
123
124 if IS_3D:
125 tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
126 if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
127 return
128 else:
129 tiles_per_segment = 0
130
131 offs_m = tl.arange(0, BLOCK_M)
132 offs_d_qk = tl.arange(0, HEAD_SIZE_QK_PADDED)
133 offs_d_v = tl.arange(0, HEAD_SIZE_V_PADDED)
134 offs_t = tl.arange(0, TILE_SIZE)
135 query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
136
137 query_offset_0 = cur_batch_in_all_start_index + query_pos
138 query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv
139 query_offset = (
140 query_offset_0[:, None] * query_stride_0
141 + query_offset_1[:, None] * query_stride_1
142 + offs_d_qk[None, :]
143 )
144
145 dim_mask_qk = tl.where(offs_d_qk < HEAD_SIZE_QK, 1, 0).to(tl.int1)
146 dim_mask_v = tl.where(offs_d_v < HEAD_SIZE_V, 1, 0).to(tl.int1)
147 query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
148 query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
149
150 # Q : (BLOCK_M, HEAD_SIZE_QK_PADDED)
151 Q = tl.load(
152 query_ptr + query_offset,
153 mask=dim_mask_qk[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
154 other=0.0,
155 )
156
157 block_table_offset = seq_idx * block_table_stride
158
159 M = init_softmax_M(
160 sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D
161 )
162 L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
163 # acc : (BLOCK_M, HEAD_SIZE_V_PADDED)
164 acc = tl.zeros([BLOCK_M, HEAD_SIZE_V_PADDED], dtype=tl.float32)
165
166 context_len = seq_len - cur_batch_query_len
167
168 if USE_ALIBI_SLOPES:
169 alibi_slope = tl.load(
170 alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0
171 )
172
173 loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(
174 context_len,
175 seq_len,
176 cur_batch_query_len,
177 q_block_local_idx,
178 segm_idx,
179 tiles_per_segment,
180 TILE_SIZE,
181 BLOCK_M,
182 BLOCK_Q,
183 num_queries_per_kv,
184 SLIDING_WINDOW,
185 False, # USE_MM_PREFIX
186 IS_3D,
187 )
188
189 for j in range(loop_lo, loop_hi):
190 seq_offset = j * TILE_SIZE + offs_t
191 tile_mask = seq_offset < max_seq_prefix_len
192
193 physical_block_idx = tl.load(
194 block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
195 ).to(tl.int64)
196
197 v_offset = (
198 physical_block_idx[:, None] * stride_v_cache_0
199 + kv_head_idx * stride_v_cache_2
200 + offs_d_v[None, :] * stride_v_cache_3
201 + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1
202 )
203 k_offset = (
204 physical_block_idx[None, :] * stride_k_cache_0
205 + kv_head_idx * stride_k_cache_2
206 + offs_d_qk[:, None] * stride_k_cache_3
207 + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
208 )
209 # K : (HEAD_SIZE_QK_PADDED, TILE_SIZE)
210 K_load = tl.load(
211 key_cache_ptr + k_offset,
212 mask=dim_mask_qk[:, None] & tile_mask[None, :],
213 other=0.0,
214 )
215 K = K_load.to(Q.dtype)
216 # V : (TILE_SIZE, HEAD_SIZE_V_PADDED)
217 V_load = tl.load(
218 value_cache_ptr + v_offset,
219 mask=dim_mask_v[None, :] & tile_mask[:, None],
220 other=0.0,
221 )
222 V = V_load.to(Q.dtype)
223
224 query_abs_pos = context_len + query_pos[:, None]
225 seq_mask = compute_kv_seq_mask(
226 query_abs_pos,
227 seq_offset,
228 seq_idx,
229 seq_len,
230 None, # mm_prefix_range_ptr
231 SLIDING_WINDOW,
232 False, # USE_MM_PREFIX
233 0, # MAX_MM_RANGES
234 )
235
236 # S : (BLOCK_M, TILE_SIZE)
237 S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
238 S += scale * tl.dot(Q, K)
239
240 if USE_SOFTCAP:
241 S = apply_softcap(S, softcap)
242
243 S = tl.where(
244 query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")
245 )
246
247 if USE_ALIBI_SLOPES:
248 S = apply_alibi_to_score(
249 S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT
250 )
251
252 M, L, P, alpha = softmax_step(S, M, L)
253 acc = acc * alpha[:, None]
254
255 if SLIDING_WINDOW:
256 qpos_lo = q_block_local_idx * BLOCK_Q
257 V = tl.where(
258 (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW,
259 V,
260 0.0,
261 )
262 acc += tl.dot(P.to(V.dtype), V)
263
264 # ---- Epilogue --------------------------------------------------------
265 if IS_3D:
266 # Store per-segment partials; finalized by reduce_segments_diffkv.
267 segm_output_offset = (
268 query_offset_0[:, None].to(tl.int64)
269 * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
270 + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
271 + segm_idx * HEAD_SIZE_V_PADDED
272 + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :]
273 )
274 tl.store(
275 segm_output_ptr + segm_output_offset,
276 acc,
277 mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
278 )
279 store_segm_reduce_scalars(
280 segm_max_ptr,
281 segm_expsum_ptr,
282 query_offset_0,
283 query_offset_1,
284 segm_idx,
285 M,
286 L,
287 query_mask_0,
288 query_mask_1,
289 num_query_heads,
290 NUM_SEGMENTS_PER_SEQ,
291 )
292 else:
293 acc = acc / L[:, None]
294 output_offset = (
295 query_offset_0[:, None] * output_stride_0
296 + query_offset_1[:, None] * output_stride_1
297 + offs_d_v[None, :]
298 )
299 tl.store(
300 output_ptr + output_offset,
301 acc,
302 mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
303 )
304
305
306 @triton.jit
307 def kernel_reduce_segments_diffkv(
308 output_ptr, # [num_tokens, num_query_heads, head_size_v]
309 segm_output_ptr,
310 # [num_tokens, num_query_heads, max_num_segments, head_size_v]
311 segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments]
312 segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments]
313 seq_lens_ptr, # [num_seqs]
314 num_seqs,
315 num_query_heads: tl.constexpr,
316 output_stride_0: tl.int64,
317 output_stride_1: tl.int64, # == HEAD_SIZE_V
318 TILE_SIZE: tl.constexpr,
319 HEAD_SIZE_V: tl.constexpr,
320 HEAD_SIZE_V_PADDED: tl.constexpr,
321 query_start_len_ptr, # [num_seqs+1]
322 BLOCK_Q: tl.constexpr,
323 NUM_SEGMENTS_PER_SEQ: tl.constexpr,
324 ):
325 """Combine per-segment partials into the final softmax output.
326
327 Mirrors ``reduce_segments`` from triton_unified_attention.py but
328 indexes V's head size (``HEAD_SIZE_V``) instead of the shared one.
329 """
330 query_token_idx = tl.program_id(0)
331 query_head_idx = tl.program_id(1)
332
333 seq_idx = find_seq_idx(
334 query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False
335 )
336 seq_len = tl.load(seq_lens_ptr + seq_idx)
337
338 tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
339 act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE)
340 segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full(
341 [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32
342 )
343 dim_mask = tl.where(tl.arange(0, HEAD_SIZE_V_PADDED) < HEAD_SIZE_V, 1, 0).to(
344 tl.int1
345 )
346
347 segm_offset = (
348 query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
349 + query_head_idx * NUM_SEGMENTS_PER_SEQ
350 + tl.arange(0, NUM_SEGMENTS_PER_SEQ)
351 )
352 segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf"))
353 overall_max = tl.max(segm_max)
354
355 segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0)
356 segm_expsum = segm_expsum * tl.exp(segm_max - overall_max)
357 overall_expsum = tl.sum(segm_expsum)
358
359 segm_output_offset = (
360 query_token_idx.to(tl.int64)
361 * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
362 + query_head_idx * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED)
363 + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * HEAD_SIZE_V_PADDED
364 + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :]
365 )
366 segm_output = tl.load(
367 segm_output_ptr + segm_output_offset,
368 mask=segm_mask[:, None] & dim_mask[None, :],
369 other=0.0,
370 )
371 segm_output *= tl.exp(segm_max - overall_max)[:, None]
372 acc_sum = tl.sum(segm_output, axis=0)
373 acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum)
374
375 output_offset = (
376 query_token_idx * output_stride_0
377 + query_head_idx * output_stride_1
378 + tl.arange(0, HEAD_SIZE_V_PADDED)
379 )
380 tl.store(output_ptr + output_offset, acc, mask=dim_mask)
381
382
383 def unified_attention_diffkv(
384 q, # [num_tokens, num_query_heads, head_size_qk]
385 k, # view: [num_blocks, block_size, num_kv_heads, head_size_qk]
386 v, # view: [num_blocks, block_size, num_kv_heads, head_size_v]
387 out, # [num_tokens, num_query_heads, head_size_v]
388 cu_seqlens_q,
389 seqused_k,
390 softmax_scale,
391 causal,
392 window_size,
393 block_table,
394 softcap,
395 max_seqlen_q: int = 1,
396 alibi_slopes=None,
397 sinks=None,
398 use_alibi_sqrt=False,
399 # 3D / split-KV softmax buffers. When all four are provided and the
400 # batch is decode-only with few sequences, the 3D path is taken.
401 seq_threshold_3D: int | None = None,
402 num_par_softmax_segments: int | None = None,
403 softmax_segm_output: torch.Tensor | None = None,
404 softmax_segm_max: torch.Tensor | None = None,
405 softmax_segm_expsum: torch.Tensor | None = None,
406 ):
407 assert causal, "Only causal attention is supported"
408
409 if sinks is not None:
410 assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
411
412 use_alibi_slopes = alibi_slopes is not None
413
414 block_size = v.shape[1]
415 num_seqs = len(seqused_k)
416 num_query_heads = q.shape[1]
417 num_kv_heads = k.shape[2]
418 num_queries_per_kv = num_query_heads // num_kv_heads
419 head_size_qk = q.shape[2]
420 head_size_v = v.shape[3]
421
422 BLOCK_M = (
423 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv)
424 )
425 BLOCK_Q = BLOCK_M // num_queries_per_kv
426
427 total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
428
429 sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
430
431 # Decide between 2D and 3D launch. Mirrors the standard launcher:
432 # 3D requires preallocated softmax buffers, decode-only batches, and
433 # a small number of sequences (otherwise 2D already saturates the SM).
434 use_3d = not (
435 seq_threshold_3D is None
436 or num_par_softmax_segments is None
437 or softmax_segm_output is None
438 or softmax_segm_max is None
439 or softmax_segm_expsum is None
440 or max_seqlen_q > 1
441 or num_seqs > seq_threshold_3D
442 or is_batch_invariant
443 )
444
445 # Tile size: 32 for prefill-class kernels. Decode (small Q) prefers
446 # smaller tiles to expose more parallelism along the KV dim.
447 tile_size = 32 if not use_3d else (16 if q.element_size() >= 2 else 32)
448
449 grid: tuple[Any, ...]
450 if use_3d:
451 grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
452 segm_output_ptr = softmax_segm_output
453 segm_max_ptr = softmax_segm_max
454 segm_expsum_ptr = softmax_segm_expsum
455 num_segments = num_par_softmax_segments
456 else:
457 grid = (total_num_q_blocks, num_kv_heads)
458 # 2D never touches the segm tensors but Triton wants a non-null
459 # pointer; reuse ``out``.
460 segm_output_ptr = out
461 segm_max_ptr = out
462 segm_expsum_ptr = out
463 num_segments = 1
464
465 kernel_unified_attention_diffkv[grid](
466 output_ptr=out,
467 segm_output_ptr=segm_output_ptr,
468 segm_max_ptr=segm_max_ptr,
469 segm_expsum_ptr=segm_expsum_ptr,
470 query_ptr=q,
471 key_cache_ptr=k,
472 value_cache_ptr=v,
473 sink_ptr=sinks,
474 block_tables_ptr=block_table,
475 seq_lens_ptr=seqused_k,
476 alibi_slopes_ptr=alibi_slopes,
477 scale=softmax_scale,
478 softcap=softcap,
479 num_query_heads=num_query_heads,
480 num_queries_per_kv=num_queries_per_kv,
481 block_table_stride=block_table.stride(0),
482 query_stride_0=q.stride(0),
483 query_stride_1=q.stride(1),
484 output_stride_0=out.stride(0),
485 output_stride_1=out.stride(1),
486 BLOCK_SIZE=block_size,
487 TILE_SIZE=tile_size,
488 HEAD_SIZE_QK=head_size_qk,
489 HEAD_SIZE_QK_PADDED=triton.next_power_of_2(head_size_qk),
490 HEAD_SIZE_V=head_size_v,
491 HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v),
492 USE_ALIBI_SLOPES=use_alibi_slopes,
493 USE_ALIBI_SQRT=use_alibi_sqrt,
494 USE_SOFTCAP=(softcap > 0),
495 USE_SINKS=(sinks is not None),
496 SLIDING_WINDOW=sliding_window_val,
497 stride_k_cache_0=k.stride(0),
498 stride_k_cache_1=k.stride(1),
499 stride_k_cache_2=k.stride(2),
500 stride_k_cache_3=k.stride(3),
501 stride_v_cache_0=v.stride(0),
502 stride_v_cache_1=v.stride(1),
503 stride_v_cache_2=v.stride(2),
504 stride_v_cache_3=v.stride(3),
505 query_start_len_ptr=cu_seqlens_q,
506 BLOCK_Q=BLOCK_Q,
507 num_seqs=num_seqs,
508 BLOCK_M=BLOCK_M,
509 NUM_SEGMENTS_PER_SEQ=num_segments,
510 IS_3D=use_3d,
511 )
512
513 if use_3d:
514 kernel_reduce_segments_diffkv[(q.shape[0], num_query_heads)](
515 output_ptr=out,
516 segm_output_ptr=softmax_segm_output,
517 segm_max_ptr=softmax_segm_max,
518 segm_expsum_ptr=softmax_segm_expsum,
519 seq_lens_ptr=seqused_k,
520 num_seqs=num_seqs,
521 num_query_heads=num_query_heads,
522 output_stride_0=out.stride(0),
523 output_stride_1=out.stride(1),
524 TILE_SIZE=tile_size,
525 HEAD_SIZE_V=head_size_v,
526 HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v),
527 query_start_len_ptr=cu_seqlens_q,
528 BLOCK_Q=BLOCK_Q,
529 NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
530 )