@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,194 lines 43.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Authors:
5 # - Burkhard Ringlein <ngl@zurich.ibm.com>
6 # - Jan van Lunteren <jvl@zurich.ibm.com>
7 # - Chih-Chieh Yang <chih.chieh.yang@ibm.com>
8 # - Thomas Parnell <tpa@zurich.ibm.com>
9
10 from typing import Any
11
12 import torch
13
14 import vllm.envs as envs
15 from vllm.logger import init_logger
16 from vllm.platforms import current_platform
17 from vllm.triton_utils import tl, triton
18 from vllm.v1.attention.ops.triton_attention_helpers import (
19 apply_alibi_to_score,
20 apply_softcap,
21 cdiv_fn,
22 compute_kv_seq_mask,
23 compute_tile_loop_bounds,
24 find_seq_idx,
25 init_softmax_M,
26 load_qq_bias_tile,
27 resolve_seq_and_query_len,
28 softmax_step,
29 store_segm_reduce_scalars,
30 )
31 from vllm.v1.kv_cache_interface import KVQuantMode
32
33 logger = init_logger(__name__)
34 is_batch_invariant = envs.VLLM_BATCH_INVARIANT
35 float8_info = torch.finfo(current_platform.fp8_dtype())
36
37
38 @triton.jit
39 def _cast_kv_tile(data, Q, tensor_scale, KV_QUANT_MODE: tl.constexpr):
40 """Cast a loaded KV tile to Q's dtype, dequantizing if needed.
41
42 Modes handled inside the core kernel:
43
44 - ``KV_QUANT_MODE == 0`` (NONE) and ``2`` (INT8 per-token-head) and
45 ``3`` (FP8 per-token-head): plain cast. Per-token-head modes apply
46 their scales separately on S/P inside the loop.
47 - ``KV_QUANT_MODE == 1`` (FP8 per-tensor): dequantize using the
48 tensor-wide scale, unless Q is also FP8 and the caller folds the scales
49 into the attention score and output accumulator.
50 """
51 if KV_QUANT_MODE == 1:
52 if Q.dtype.is_fp8():
53 return data.to(Q.dtype)
54 return (data.to(tl.float32) * tl.load(tensor_scale)).to(Q.dtype)
55 return data.to(Q.dtype)
56
57
58 # ---------------------------------------------------------------------------
59 # Tensor-descriptor (TD) helpers
60 #
61 # Used when the caller enables ``USE_TD`` (Intel Xe2/Xe3 HW 2D block reads).
62 # When ``USE_TD`` is False the helpers are dead-code-eliminated at Triton
63 # compile time, leaving the pointer-arithmetic load/store path untouched.
64 # ---------------------------------------------------------------------------
65
66
67 @triton.jit
68 def _load_q_td(
69 query_ptr,
70 q_block_local_len,
71 query_stride_0: tl.int64,
72 query_stride_1: tl.int64,
73 cur_batch_in_all_start_index,
74 q_block_local_idx,
75 kv_head_idx,
76 num_queries_per_kv: tl.constexpr,
77 BLOCK_Q: tl.constexpr,
78 BLOCK_M: tl.constexpr,
79 HEAD_SIZE: tl.constexpr,
80 HEAD_SIZE_PADDED: tl.constexpr,
81 ):
82 """Load Q via a 2D tensor descriptor.
83
84 Caller guarantees (via the wrapper's ``use_td_qo`` gate):
85 * ``HEAD_SIZE == HEAD_SIZE_PADDED`` (head_size is a power of 2),
86 * ``num_queries_per_kv`` is a power of 2,
87 * the ``num_queries_per_kv`` heads of the current KV group are
88 contiguous in memory (``query_stride_1 == HEAD_SIZE``, which is
89 the default vLLM query layout).
90
91 Under those preconditions the inner two axes are flattened into one
92 row of size ``num_queries_per_kv * HEAD_SIZE`` with stride 1, which
93 avoids the non-power-of-2 ``block_shape`` error from the Triton
94 tensor-descriptor validator. Returns (BLOCK_M, HEAD_SIZE_PADDED).
95 """
96 q_base = (
97 query_ptr
98 + (cur_batch_in_all_start_index + q_block_local_idx * BLOCK_Q) * query_stride_0
99 + (kv_head_idx * num_queries_per_kv) * query_stride_1
100 )
101 q_desc = tl.make_tensor_descriptor(
102 base=q_base,
103 shape=(q_block_local_len, num_queries_per_kv * HEAD_SIZE),
104 strides=(query_stride_0, 1),
105 block_shape=(BLOCK_Q, num_queries_per_kv * HEAD_SIZE_PADDED),
106 )
107 return q_desc.load([0, 0]).reshape(BLOCK_M, HEAD_SIZE_PADDED)
108
109
110 @triton.jit
111 def _load_kv_tile_td(
112 cache_ptr,
113 physical_block_idx_scalar,
114 kv_head_idx,
115 offset_in_block,
116 stride_cache_0: tl.int64,
117 stride_cache_1: tl.int64,
118 stride_cache_2: tl.int64,
119 stride_cache_3: tl.constexpr,
120 BLOCK_SIZE: tl.constexpr,
121 TILE_SIZE: tl.constexpr,
122 HEAD_SIZE: tl.constexpr,
123 HEAD_SIZE_PADDED: tl.constexpr,
124 ):
125 """Load a KV cache tile via tensor descriptor.
126
127 Returns shape (TILE_SIZE, HEAD_SIZE_PADDED). Caller transposes for K.
128 Tensor descriptors zero-pad reads beyond the shape boundary, so
129 ``HEAD_SIZE_PADDED > HEAD_SIZE`` is handled correctly.
130 """
131 base = (
132 cache_ptr
133 + physical_block_idx_scalar * stride_cache_0
134 + kv_head_idx * stride_cache_2
135 )
136 desc = tl.make_tensor_descriptor(
137 base=base,
138 shape=(BLOCK_SIZE, HEAD_SIZE),
139 strides=(stride_cache_1, stride_cache_3),
140 block_shape=(TILE_SIZE, HEAD_SIZE_PADDED),
141 )
142 return desc.load([offset_in_block, 0])
143
144
145 @triton.jit
146 def _store_output_td(
147 base_ptr,
148 acc,
149 q_block_local_len,
150 stride_token: tl.int64,
151 stride_head: tl.int64,
152 num_queries_per_kv: tl.constexpr,
153 BLOCK_Q: tl.constexpr,
154 HEAD_SIZE: tl.constexpr,
155 HEAD_SIZE_PADDED: tl.constexpr,
156 ):
157 """Store an output tile via a tensor descriptor.
158
159 The 2D and 3D epilogues differ only in ``base_ptr`` and the
160 ``(stride_token, stride_head)`` pair: 2D writes directly to the
161 flat output buffer, 3D writes to a single per-segment slice of
162 ``segm_output_ptr``. Descriptor shape / block_shape / reshape
163 are the same in both modes, so share one helper.
164 """
165 acc = acc.to(base_ptr.dtype.element_ty)
166 output_desc = tl.make_tensor_descriptor(
167 base=base_ptr,
168 shape=(q_block_local_len, num_queries_per_kv, HEAD_SIZE),
169 strides=(stride_token, stride_head, 1),
170 block_shape=(BLOCK_Q, num_queries_per_kv, HEAD_SIZE_PADDED),
171 )
172 output_desc.store(
173 [0, 0, 0],
174 acc.reshape(BLOCK_Q, num_queries_per_kv, HEAD_SIZE_PADDED),
175 )
176
177
178 @triton.jit
179 def kernel_unified_attention(
180 # Output destination for the 2D path. In 3D mode per-segment partials
181 # go to the ``segm_*`` tensors (see bottom of signature) and
182 # ``output_ptr`` is unused (callers may pass any non-null pointer).
183 output_ptr,
184 # Inputs
185 query_ptr,
186 key_cache_ptr,
187 value_cache_ptr,
188 sink_ptr,
189 block_tables_ptr,
190 seq_lens_ptr,
191 alibi_slopes_ptr,
192 qq_bias_ptr,
193 # Scalars
194 scale,
195 q_scale,
196 k_scale,
197 v_scale,
198 out_scale,
199 softcap,
200 num_query_heads: tl.constexpr, # int
201 num_queries_per_kv: tl.constexpr, # int
202 block_table_stride: tl.int64, # int
203 query_stride_0: tl.int64, # int
204 query_stride_1: tl.int64, # int, should be equal to head_size
205 output_stride_0: tl.int64, # int
206 output_stride_1: tl.int64, # int, should be equal to head_size
207 qq_bias_stride_0: tl.int64, # int
208 BLOCK_SIZE: tl.constexpr, # int
209 TILE_SIZE: tl.constexpr, # int must be power of 2
210 HEAD_SIZE: tl.constexpr, # int
211 HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
212 USE_ALIBI_SLOPES: tl.constexpr, # bool
213 USE_ALIBI_SQRT: tl.constexpr, # bool
214 USE_QQ_BIAS: tl.constexpr, # bool
215 USE_SOFTCAP: tl.constexpr, # bool
216 USE_SINKS: tl.constexpr, # bool
217 SLIDING_WINDOW: tl.constexpr, # int
218 USE_CAUSAL: tl.constexpr, # bool
219 USE_PER_SEQ_CAUSAL: tl.constexpr, # bool
220 per_seq_causal_ptr, # [num_seqs] int32, or None
221 USE_MM_PREFIX: tl.constexpr, # bool
222 MAX_MM_RANGES: tl.constexpr, # int
223 mm_prefix_range_ptr,
224 rswa_prefix_lens_ptr,
225 R_SWA_WINDOW: tl.constexpr, # int
226 USE_R_SWA: tl.constexpr, # bool
227 stride_k_cache_0: tl.int64, # int
228 stride_k_cache_1: tl.int64, # int
229 stride_k_cache_2: tl.int64, # int
230 stride_k_cache_3: tl.constexpr, # int
231 stride_v_cache_0: tl.int64, # int
232 stride_v_cache_1: tl.int64, # int
233 stride_v_cache_2: tl.int64, # int
234 stride_v_cache_3: tl.constexpr, # int
235 query_start_len_ptr,
236 BLOCK_Q: tl.constexpr,
237 num_seqs: tl.int32,
238 BLOCK_M: tl.constexpr,
239 NUM_SEGMENTS_PER_SEQ: tl.constexpr,
240 USE_FP8: tl.constexpr,
241 # Toggles 2D vs 3D layout. The 2D path runs the full sequence in one
242 # tile loop and writes to ``output_ptr``. The 3D path scopes the loop
243 # to ``[segm_idx, segm_idx+1) × tiles_per_segment`` and writes
244 # per-segment partials, finalized by ``reduce_segments``.
245 IS_3D: tl.constexpr,
246 # Parameters below default to None so Triton can skip materialising them
247 # on call sites where the corresponding constexpr branch is dead.
248 # Credit: @quinnlp identified this as a perf regression source in
249 # intel/intel-xpu-backend-for-triton#6758 (review comment r3204641104).
250 # Per-segment outputs: used in 3D mode; unused in 2D (IS_3D=False).
251 segm_output_ptr=None,
252 segm_max_ptr=None,
253 segm_expsum_ptr=None,
254 # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}.
255 k_scale_cache_ptr=None,
256 v_scale_cache_ptr=None,
257 # ``tl.int64`` cannot be combined with a ``None`` default — Triton's JIT
258 # rejects ``Optional[tl.int64]`` / ``tl.int64 | None`` at trace time, and
259 # plain ``tl.int64 = None`` raises ``TypeError: 'NoneType' object cannot
260 # be interpreted as an integer`` when callers omit these arguments.
261 # ``int | None`` is the only annotation that lets the wrapper pass
262 # ``None`` here so Triton can skip materialising the strides when the
263 # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead.
264 stride_ks_blk: int | None = None,
265 stride_ks_slot: int | None = None,
266 stride_ks_head: int | None = None,
267 stride_vs_blk: int | None = None,
268 stride_vs_slot: int | None = None,
269 stride_vs_head: int | None = None,
270 # KV cache quantization mode handled inside this kernel via constexpr
271 # branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2),
272 # FP8_PER_TOKEN_HEAD (3). Sub-byte INT4 (4) uses its own
273 # int4_per_token_head kernel, not this one.
274 KV_QUANT_MODE: tl.constexpr = 0,
275 FP8_MIN: tl.constexpr = float8_info.min,
276 FP8_MAX: tl.constexpr = float8_info.max,
277 # Chunked / block-local attention. ``CHUNK_LOOKBACK >= 0`` enables
278 # chunked masking (used by Gemma3 block-local layers); takes precedence
279 # over ``SLIDING_WINDOW`` inside the helpers. ``-1`` disables.
280 CHUNK_LOOKBACK: tl.constexpr = -1,
281 CHUNK_SIZE: tl.constexpr = -1,
282 # Tensor-descriptor load/store for HW 2D block reads on Intel Xe2/Xe3.
283 # ``USE_TD`` gates KV tile loads; ``USE_TD_QO`` separately gates Q/output
284 # (see ``unified_attention`` wrapper for the gating rules).
285 USE_TD: tl.constexpr = False,
286 USE_TD_QO: tl.constexpr = False,
287 Q_IS_FP8: tl.constexpr = False,
288 # Gemma4: clamp mm_prefix bidirectional ranges by the sliding window
289 # instead of letting them override it. Default False preserves the
290 # original (causal AND SW) OR mm_prefix behavior for all other models.
291 MM_PREFIX_CLAMP_SW: tl.constexpr = False,
292 ):
293 # Per-(token, head) scale caches: used iff KV_QUANT_MODE in {2, 3}.
294 USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = (KV_QUANT_MODE >= 2) and (
295 KV_QUANT_MODE <= 3
296 )
297 USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8
298
299 if USE_TD:
300 tl.static_assert(
301 BLOCK_SIZE % TILE_SIZE == 0,
302 "USE_TD requires BLOCK_SIZE to be a multiple of TILE_SIZE",
303 )
304
305 q_block_global_idx = tl.program_id(0)
306 kv_head_idx = tl.program_id(1)
307 segm_idx = tl.program_id(2) if IS_3D else 0
308
309 (
310 seq_idx,
311 q_block_local_idx,
312 cur_batch_in_all_start_index,
313 cur_batch_query_len,
314 seq_len,
315 ) = resolve_seq_and_query_len(
316 query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q
317 )
318
319 if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
320 return
321
322 if IS_3D:
323 tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
324 if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
325 return
326 else:
327 tiles_per_segment = 0
328
329 # Number of valid query rows in this block (used by TD descriptor
330 # shapes, but always computed so the variable stays in scope).
331 q_block_local_len = tl.minimum(
332 BLOCK_Q, cur_batch_query_len - q_block_local_idx * BLOCK_Q
333 )
334
335 offs_m = tl.arange(0, BLOCK_M)
336 offs_d = tl.arange(0, HEAD_SIZE_PADDED)
337 offs_t = tl.arange(0, TILE_SIZE)
338 query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
339
340 query_offset_0 = cur_batch_in_all_start_index + query_pos
341 query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv
342 query_offset = (
343 query_offset_0[:, None] * query_stride_0
344 + query_offset_1[:, None] * query_stride_1
345 + offs_d[None, :]
346 )
347
348 dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)
349 query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
350 query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
351
352 # Q : (BLOCK_M, HEAD_SIZE_PADDED)
353 if USE_TD_QO:
354 Q = _load_q_td(
355 query_ptr,
356 q_block_local_len,
357 query_stride_0,
358 query_stride_1,
359 cur_batch_in_all_start_index,
360 q_block_local_idx,
361 kv_head_idx,
362 num_queries_per_kv,
363 BLOCK_Q,
364 BLOCK_M,
365 HEAD_SIZE,
366 HEAD_SIZE_PADDED,
367 )
368 else:
369 Q = tl.load(
370 query_ptr + query_offset,
371 mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
372 other=0.0,
373 )
374
375 block_table_offset = seq_idx * block_table_stride
376
377 M = init_softmax_M(
378 sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D
379 )
380 L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
381 # acc : (BLOCK_M, HEAD_SIZE_PADDED)
382 acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
383 score_scale = scale
384 value_scale = 1.0
385 if USE_FP8_Q_DESCALE:
386 score_scale = scale * tl.load(q_scale) * tl.load(k_scale)
387 value_scale = tl.load(v_scale)
388
389 context_len = seq_len - cur_batch_query_len
390
391 if USE_ALIBI_SLOPES:
392 alibi_slope = tl.load(
393 alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0
394 )
395
396 if USE_QQ_BIAS:
397 qq_bias_row_ptrs = qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0
398
399 loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(
400 context_len,
401 seq_len,
402 cur_batch_query_len,
403 q_block_local_idx,
404 segm_idx,
405 tiles_per_segment,
406 TILE_SIZE,
407 BLOCK_M,
408 BLOCK_Q,
409 num_queries_per_kv,
410 SLIDING_WINDOW,
411 USE_MM_PREFIX,
412 IS_3D,
413 USE_CAUSAL,
414 USE_PER_SEQ_CAUSAL,
415 CHUNK_LOOKBACK,
416 CHUNK_SIZE,
417 USE_R_SWA,
418 MM_PREFIX_CLAMP_SW,
419 MAX_MM_RANGES,
420 mm_prefix_range_ptr,
421 seq_idx,
422 )
423
424 # iterate through tiles (now limited to the sliding window range)
425 for j in range(loop_lo, loop_hi):
426 seq_offset = j * TILE_SIZE + offs_t
427 tile_mask = seq_offset < max_seq_prefix_len
428
429 physical_block_idx = tl.load(
430 block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
431 ).to(tl.int64)
432
433 if USE_TD:
434 # All TILE_SIZE slots within a single KV tile map to one
435 # physical block (guaranteed by ``BLOCK_SIZE % TILE_SIZE == 0``
436 # from the static_assert above), so load the block index as
437 # a scalar instead of a broadcast reduction.
438 offset_in_block = (j * TILE_SIZE) % BLOCK_SIZE
439 physical_block_scalar = tl.load(
440 block_tables_ptr + block_table_offset + (j * TILE_SIZE) // BLOCK_SIZE
441 ).to(tl.int64)
442 # K : (HEAD_SIZE, TILE_SIZE)
443 K_load = _load_kv_tile_td(
444 key_cache_ptr,
445 physical_block_scalar,
446 kv_head_idx,
447 offset_in_block,
448 stride_k_cache_0,
449 stride_k_cache_1,
450 stride_k_cache_2,
451 stride_k_cache_3,
452 BLOCK_SIZE,
453 TILE_SIZE,
454 HEAD_SIZE,
455 HEAD_SIZE_PADDED,
456 ).T
457 # V : (TILE_SIZE, HEAD_SIZE)
458 V_load = _load_kv_tile_td(
459 value_cache_ptr,
460 physical_block_scalar,
461 kv_head_idx,
462 offset_in_block,
463 stride_v_cache_0,
464 stride_v_cache_1,
465 stride_v_cache_2,
466 stride_v_cache_3,
467 BLOCK_SIZE,
468 TILE_SIZE,
469 HEAD_SIZE,
470 HEAD_SIZE_PADDED,
471 )
472 else:
473 v_offset = (
474 physical_block_idx[:, None] * stride_v_cache_0
475 + kv_head_idx * stride_v_cache_2
476 + offs_d[None, :] * stride_v_cache_3
477 + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1
478 )
479 k_offset = (
480 physical_block_idx[None, :] * stride_k_cache_0
481 + kv_head_idx * stride_k_cache_2
482 + offs_d[:, None] * stride_k_cache_3
483 + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
484 )
485 # K : (HEAD_SIZE, TILE_SIZE)
486 K_load = tl.load(
487 key_cache_ptr + k_offset,
488 mask=dim_mask[:, None] & tile_mask[None, :],
489 other=0.0,
490 )
491 # V : (TILE_SIZE, HEAD_SIZE)
492 V_load = tl.load(
493 value_cache_ptr + v_offset,
494 mask=dim_mask[None, :] & tile_mask[:, None],
495 other=0.0,
496 )
497 K = _cast_kv_tile(K_load, Q, k_scale, KV_QUANT_MODE)
498 V = _cast_kv_tile(V_load, Q, v_scale, KV_QUANT_MODE)
499
500 # Per-(token, head) scales for INT8 / FP8 per-token-head modes.
501 if USE_PER_TOKEN_HEAD_SCALES:
502 scale_idx = (
503 physical_block_idx * stride_ks_blk
504 + (seq_offset % BLOCK_SIZE) * stride_ks_slot
505 + kv_head_idx * stride_ks_head
506 )
507 k_token_head_scales = tl.load(
508 k_scale_cache_ptr + scale_idx, mask=tile_mask, other=1.0
509 )
510 v_scale_idx = (
511 physical_block_idx * stride_vs_blk
512 + (seq_offset % BLOCK_SIZE) * stride_vs_slot
513 + kv_head_idx * stride_vs_head
514 )
515 v_token_head_scales = tl.load(
516 v_scale_cache_ptr + v_scale_idx, mask=tile_mask, other=1.0
517 )
518
519 query_abs_pos = context_len + query_pos[:, None]
520 seq_mask = compute_kv_seq_mask(
521 query_abs_pos,
522 seq_offset,
523 seq_idx,
524 seq_len,
525 mm_prefix_range_ptr,
526 SLIDING_WINDOW,
527 USE_MM_PREFIX,
528 MAX_MM_RANGES,
529 USE_CAUSAL,
530 USE_PER_SEQ_CAUSAL,
531 per_seq_causal_ptr,
532 rswa_prefix_lens_ptr,
533 R_SWA_WINDOW,
534 USE_R_SWA,
535 CHUNK_LOOKBACK,
536 CHUNK_SIZE,
537 MM_PREFIX_CLAMP_SW,
538 )
539
540 # S : (BLOCK_M, TILE_SIZE)
541 S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
542 if USE_PER_TOKEN_HEAD_SCALES:
543 # Per-token-head quant: fuse softmax_scale with per-head k_scale
544 # to avoid a separate BLOCK_M × TILE_SIZE multiply on S.
545 S += tl.dot(Q, K) * (score_scale * k_token_head_scales[None, :])
546 else:
547 S += score_scale * tl.dot(Q, K)
548
549 if USE_SOFTCAP:
550 S = apply_softcap(S, softcap)
551
552 S = tl.where(
553 query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")
554 )
555
556 if USE_ALIBI_SLOPES:
557 S = apply_alibi_to_score(
558 S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT
559 )
560
561 if USE_QQ_BIAS:
562 S += load_qq_bias_tile(
563 qq_bias_row_ptrs, seq_offset, context_len, qq_bias_stride_0
564 )
565
566 M, L, P, alpha = softmax_step(S, M, L)
567 acc = acc * alpha[:, None]
568
569 if SLIDING_WINDOW:
570 qpos_lo = q_block_local_idx * BLOCK_Q
571 dist = context_len + qpos_lo - seq_offset[:, None]
572 if USE_PER_SEQ_CAUSAL:
573 is_causal_seq = tl.load(per_seq_causal_ptr + seq_idx)
574 sw_mask_v = tl.where(
575 is_causal_seq,
576 dist < SLIDING_WINDOW,
577 (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW),
578 )
579 elif USE_CAUSAL:
580 sw_mask_v = dist < SLIDING_WINDOW
581 else:
582 sw_mask_v = (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW)
583 V = tl.where(sw_mask_v, V, 0.0)
584 if USE_PER_TOKEN_HEAD_SCALES:
585 # Per-token-head quant: apply v_scale to P instead of V.
586 P_v = (P * v_token_head_scales[None, :]).to(V.dtype)
587 acc += tl.dot(P_v, V)
588 else:
589 acc += tl.dot(P.to(V.dtype), V)
590
591 # ---- Epilogue ---------------------------------------------------------
592 if IS_3D:
593 if USE_FP8_Q_DESCALE:
594 acc *= value_scale
595 # Store per-segment partials; finalized by ``reduce_segments``.
596 if USE_TD_QO:
597 # 3D target: segm_output[token, head, segm_idx, :]. Advance
598 # the base to the correct (token-start, head-start, segm)
599 # slice; strides step between tokens / heads of the flattened
600 # (T, H, SEGS, PAD) layout.
601 segm_base = (
602 segm_output_ptr
603 + (cur_batch_in_all_start_index + q_block_local_idx * BLOCK_Q).to(
604 tl.int64
605 )
606 * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
607 + (kv_head_idx * num_queries_per_kv)
608 * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
609 + segm_idx * HEAD_SIZE_PADDED
610 )
611 _store_output_td(
612 segm_base,
613 acc,
614 q_block_local_len,
615 num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED,
616 NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED,
617 num_queries_per_kv,
618 BLOCK_Q,
619 HEAD_SIZE,
620 HEAD_SIZE_PADDED,
621 )
622 else:
623 segm_output_offset = (
624 query_offset_0[:, None].to(tl.int64)
625 * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
626 + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
627 + segm_idx * HEAD_SIZE_PADDED
628 + tl.arange(0, HEAD_SIZE_PADDED)[None, :]
629 )
630 tl.store(
631 segm_output_ptr + segm_output_offset,
632 acc,
633 mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
634 )
635 store_segm_reduce_scalars(
636 segm_max_ptr,
637 segm_expsum_ptr,
638 query_offset_0,
639 query_offset_1,
640 segm_idx,
641 M,
642 L,
643 query_mask_0,
644 query_mask_1,
645 num_query_heads,
646 NUM_SEGMENTS_PER_SEQ,
647 )
648 else:
649 acc = acc / L[:, None]
650 if USE_FP8_Q_DESCALE:
651 acc *= value_scale
652 if USE_FP8:
653 acc = acc * tl.load(out_scale)
654 acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
655 if USE_TD_QO:
656 # 2D target: flat output[token, head, :]. Strides come
657 # straight from the caller (``output_stride_0`` per token,
658 # ``output_stride_1`` per head).
659 output_base = (
660 output_ptr
661 + (cur_batch_in_all_start_index + q_block_local_idx * BLOCK_Q)
662 * output_stride_0
663 + (kv_head_idx * num_queries_per_kv) * output_stride_1
664 )
665 _store_output_td(
666 output_base,
667 acc,
668 q_block_local_len,
669 output_stride_0,
670 output_stride_1,
671 num_queries_per_kv,
672 BLOCK_Q,
673 HEAD_SIZE,
674 HEAD_SIZE_PADDED,
675 )
676 else:
677 output_offset = (
678 query_offset_0[:, None] * output_stride_0
679 + query_offset_1[:, None] * output_stride_1
680 + offs_d[None, :]
681 )
682 tl.store(
683 output_ptr + output_offset,
684 acc,
685 mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
686 )
687
688
689 @triton.jit
690 def reduce_segments(
691 output_ptr, # [num_tokens, num_query_heads, head_size]
692 segm_output_ptr,
693 # [num_tokens, num_query_heads, max_num_segments, head_size]
694 segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments]
695 segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments]
696 seq_lens_ptr, # [num_seqs]
697 num_seqs, # int
698 num_query_heads: tl.constexpr, # int
699 out_scale_inv, # float32
700 output_stride_0: tl.int64, # int
701 output_stride_1: tl.int64, # int, should be equal to head_size
702 block_table_stride: tl.int64, # int
703 TILE_SIZE: tl.constexpr, # int
704 HEAD_SIZE: tl.constexpr, # int, must be power of 2
705 HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
706 query_start_len_ptr, # [num_seqs+1]
707 BLOCK_Q: tl.constexpr, # int
708 NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int
709 USE_FP8: tl.constexpr, # bool
710 FP8_MIN: tl.constexpr = float8_info.min,
711 FP8_MAX: tl.constexpr = float8_info.max,
712 ):
713 query_token_idx = tl.program_id(0)
714 query_head_idx = tl.program_id(1)
715
716 seq_idx = find_seq_idx(
717 query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False
718 )
719
720 # sequence len for this particular sequence
721 seq_len = tl.load(seq_lens_ptr + seq_idx)
722
723 # number of segments for this particular sequence
724 num_segments = NUM_SEGMENTS_PER_SEQ
725 tiles_per_segment = cdiv_fn(seq_len, num_segments * TILE_SIZE)
726
727 # create masks for subsequent loads
728 act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE)
729 segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full(
730 [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32
731 )
732 dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1, 0).to(tl.int1)
733
734 # load segment maxima
735 segm_offset = (
736 query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
737 + query_head_idx * NUM_SEGMENTS_PER_SEQ
738 + tl.arange(0, NUM_SEGMENTS_PER_SEQ)
739 )
740 segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf"))
741 overall_max = tl.max(segm_max)
742
743 # load and rescale segment exp sums
744 segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0)
745 segm_expsum = segm_expsum * tl.exp(segm_max - overall_max)
746 overall_expsum = tl.sum(segm_expsum)
747
748 # load, rescale, and add segment attention outputs
749 segm_output_offset = (
750 query_token_idx.to(tl.int64)
751 * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
752 + query_head_idx * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
753 + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * HEAD_SIZE_PADDED
754 + tl.arange(0, HEAD_SIZE_PADDED)[None, :]
755 )
756 segm_output = tl.load(
757 segm_output_ptr + segm_output_offset,
758 mask=segm_mask[:, None] & dim_mask[None, :],
759 other=0.0,
760 )
761 segm_output *= tl.exp(segm_max - overall_max)[:, None]
762 acc_sum = tl.sum(segm_output, axis=0)
763 # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0
764 acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum)
765
766 if USE_FP8:
767 acc = acc * tl.load(out_scale_inv)
768 acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
769
770 # write result
771 output_offset = (
772 query_token_idx * output_stride_0
773 + query_head_idx * output_stride_1
774 + tl.arange(0, HEAD_SIZE_PADDED)
775 )
776 tl.store(output_ptr + output_offset, acc, mask=dim_mask)
777
778
779 def _is_gemma3_attention(head_size: int, sliding_window: int) -> bool:
780 """Detect Gemma3 models via unique (head_size, sliding_window) signature.
781
782 Gemma3 models are the only ones using sliding_window=1024 with
783 head_size 128 (27B) or 256 (1B, 4B, 12B). Other SWA models use
784 different window sizes (Mistral=4096, Phi-3=2047).
785 """
786 return sliding_window == 1024 and head_size in (128, 256)
787
788
789 def _get_tile_size(
790 head_size: int,
791 sliding_window: int,
792 element_size: int,
793 is_prefill: bool,
794 ) -> int:
795 """Select tile size with Gemma3-specific optimization."""
796 if _is_gemma3_attention(head_size, sliding_window):
797 # Gemma3: use 32 for decode (default is 16)
798 return 32
799
800 # Default behavior
801 if is_prefill:
802 return 32
803 # Note: tile size must be at least 32 for fp8 (element_size == 1).
804 return 16 if element_size >= 2 else 32
805
806
807 def unified_attention(
808 q,
809 k,
810 v,
811 out,
812 cu_seqlens_q,
813 max_seqlen_q,
814 seqused_k,
815 max_seqlen_k,
816 softmax_scale,
817 causal,
818 window_size,
819 block_table,
820 softcap,
821 q_descale,
822 k_descale,
823 v_descale,
824 seq_threshold_3D=None,
825 num_par_softmax_segments=None,
826 softmax_segm_output=None,
827 softmax_segm_max=None,
828 softmax_segm_expsum=None,
829 alibi_slopes=None,
830 output_scale=None,
831 qq_bias=None,
832 # Optional tensor for sinks
833 sinks=None,
834 # Optional tensor for prefix lengths (PrefixLM support)
835 mm_prefix_range=None,
836 # R-SWA support: prefix tokens stay globally visible, generated tokens use
837 # a fixed sliding window.
838 rswa_prefix_lens=None,
839 rswa_window: int | None = None,
840 use_alibi_sqrt=False,
841 # KV cache quantization mode and per-token-head scale caches.
842 kv_quant_mode: KVQuantMode = KVQuantMode.NONE,
843 k_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32
844 v_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32
845 # Chunked attention: restrict attention to aligned blocks with lookback.
846 chunk_lookback=-1,
847 # Tensor-descriptor mode: use ``tl.make_tensor_descriptor`` for Q/K/V
848 # loads and output stores. Enables HW 2D block reads on Intel Xe2/Xe3.
849 # The non-TD branch is dead-code-eliminated at Triton compile time so
850 # disabling this flag costs nothing.
851 use_td: bool = False,
852 # Gemma4: clamp mm_prefix bidirectional ranges by the sliding window.
853 # Default False keeps the original behavior for every other model.
854 mm_prefix_clamp_sliding_window: bool = False,
855 ):
856 # Resolve causal: bool or per-seq tensor.
857 use_per_seq_causal = isinstance(causal, torch.Tensor)
858 use_causal = bool(causal) if not use_per_seq_causal else True
859 per_seq_causal_ptr = causal if use_per_seq_causal else None
860
861 # Sub-byte packed mode (INT4) needs a bespoke kernel (split-dot +
862 # sub-byte unpack); everything else goes through the core kernel below.
863 if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
864 assert use_causal and not use_per_seq_causal, (
865 "INT4_PER_TOKEN_HEAD only supports causal attention"
866 )
867 from vllm.v1.attention.ops.int4_per_token_head import (
868 unified_attention_int4,
869 )
870
871 if sinks is not None:
872 assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
873 unified_attention_int4(
874 q=q,
875 k_cache=k,
876 v_cache=v,
877 out=out,
878 cu_seqlens_q=cu_seqlens_q,
879 max_seqlen_q=max_seqlen_q,
880 seqused_k=seqused_k,
881 max_seqlen_k=max_seqlen_k,
882 softmax_scale=softmax_scale,
883 window_size=window_size,
884 block_table=block_table,
885 softcap=softcap,
886 sinks=sinks,
887 alibi_slopes=alibi_slopes,
888 use_alibi_sqrt=use_alibi_sqrt,
889 qq_bias=qq_bias,
890 output_scale=output_scale,
891 mm_prefix_range=mm_prefix_range,
892 k_scale_cache=k_scale_cache,
893 v_scale_cache=v_scale_cache,
894 seq_threshold_3D=seq_threshold_3D,
895 num_par_softmax_segments=num_par_softmax_segments,
896 softmax_segm_output=softmax_segm_output,
897 softmax_segm_max=softmax_segm_max,
898 softmax_segm_expsum=softmax_segm_expsum,
899 )
900 return
901
902 if sinks is not None:
903 assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
904
905 use_per_token_head_scales = kv_quant_mode in (
906 KVQuantMode.INT8_PER_TOKEN_HEAD,
907 KVQuantMode.FP8_PER_TOKEN_HEAD,
908 )
909 if use_per_token_head_scales:
910 assert k_scale_cache is not None and v_scale_cache is not None, (
911 f"{kv_quant_mode.name} requires k_scale_cache / v_scale_cache"
912 )
913
914 use_mm_prefix = False
915 max_mm_ranges = 0
916 if mm_prefix_range is not None:
917 if mm_prefix_range.ndim == 3:
918 use_mm_prefix = True
919 max_mm_ranges = mm_prefix_range.shape[1]
920 else:
921 raise ValueError(
922 f"Unsupported mm_prefix_range shape: {mm_prefix_range.shape}"
923 )
924
925 use_rswa = rswa_window is not None and rswa_prefix_lens is not None
926
927 use_alibi_slopes = alibi_slopes is not None
928 use_qq_bias = qq_bias is not None
929
930 block_size = v.shape[1]
931 num_seqs = len(seqused_k)
932 num_query_heads = q.shape[1]
933 num_kv_heads = k.shape[2]
934 num_queries_per_kv = num_query_heads // num_kv_heads
935 head_size = q.shape[2]
936
937 BLOCK_M = (
938 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv)
939 )
940 BLOCK_Q = BLOCK_M // num_queries_per_kv
941
942 # Tuned launch parameters; ``None`` lets Triton pick its defaults.
943 launch_num_warps: int | None = None
944 launch_num_stages: int | None = None
945
946 # head_size 256 with many query rows per sequence (e.g. diffusion-gemma
947 # bidirectional canvas passes) is prefill-shaped, but the decode-oriented
948 # defaults (BLOCK_Q=8, TILE=32, 4 warps) under-tile it. A wider KV tile +
949 # more query rows per block + 8 warps is ~2x faster on B200.
950 tuned_large_head = (
951 head_size == 256
952 and max_seqlen_q > 1
953 and num_queries_per_kv <= 16
954 and current_platform.is_device_capability_family(100)
955 )
956 if tuned_large_head:
957 BLOCK_M = 32
958 BLOCK_Q = BLOCK_M // num_queries_per_kv
959 launch_num_warps = 8
960 launch_num_stages = 2
961
962 # Ideally we would launch with kernel with:
963 # \sum_i[ceil(query_len[i] / BLOCK_Q)] blocks.
964 # However, it is slow to realize the query_lens on cpu.
965 # Instead we use upper-bound:
966 # \sum_i[ceil(query_len[i] / BLOCK_Q)]
967 # <= \sum_i[floor(query_len[i] / BLOCK_Q) + 1]
968 # = \sum_i[floor(query_len[i] / BLOCK_Q)] + num_seqs
969 # <= floor(\sum_i(query_len[i]) / BLOCK_Q) + num_seqs
970 # = floor(q.shape[0] / BLOCK_Q) + num_seqs
971 total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
972
973 sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
974
975 # Compute chunked block size from sliding window if needed.
976 chunk_size = -1
977 if sliding_window_val > 0 and chunk_lookback > -1:
978 chunk_size = sliding_window_val // (chunk_lookback + 1)
979 assert chunk_size > 0, "sliding_window must be > chunk_lookback+1"
980 elif sliding_window_val <= 0:
981 chunk_lookback = -1
982
983 TILE_SIZE_PREFILL = _get_tile_size(
984 head_size, sliding_window_val, q.element_size(), is_prefill=True
985 )
986 TILE_SIZE_DECODE = _get_tile_size(
987 head_size, sliding_window_val, q.element_size(), is_prefill=False
988 )
989
990 # Wider KV tile for the tuned large-head path (see above). Only the 2D
991 # path (used when max_seqlen_q > 1) reads TILE_SIZE_PREFILL.
992 if tuned_large_head:
993 TILE_SIZE_PREFILL = 128
994
995 # USE_TD requires BLOCK_SIZE % TILE_SIZE == 0 (enforced by a
996 # ``tl.static_assert`` in the kernel). The default prefill tile
997 # size (32) is larger than a common ``block_size=16``, so clamp it
998 # down when TD is enabled. Zero overhead when disabled.
999 if use_td:
1000 TILE_SIZE_PREFILL = min(TILE_SIZE_PREFILL, block_size)
1001 TILE_SIZE_DECODE = min(TILE_SIZE_DECODE, block_size)
1002
1003 # Tensor descriptors for Q load / output store require every element
1004 # of ``block_shape`` to be a power of 2. ``num_queries_per_kv`` is
1005 # not always pow2 (e.g. Qwen2-7B: 28 / 4 = 7), so gate the Q/O paths
1006 # separately from the KV tile loads (whose ``block_shape`` does not
1007 # include ``num_queries_per_kv``).
1008 #
1009 # The Q/O descriptors also encode ``HEAD_SIZE_PADDED`` on the inner
1010 # axis while the backing buffers (both flat output and per-segment
1011 # output) are laid out with ``HEAD_SIZE``. When they differ (e.g.
1012 # Phi-3's head_size=96 → HEAD_SIZE_PADDED=128) the store would spill
1013 # padded lanes into neighbouring heads because tensor-descriptor
1014 # stores don't mask the padded tail. Fall back to the pointer path
1015 # for Q/O in that case — KV tile loads are unaffected because their
1016 # ``shape`` already matches ``block_shape`` on the inner axis.
1017 head_size_padded = triton.next_power_of_2(head_size)
1018 _is_pow2_nq = (num_queries_per_kv & (num_queries_per_kv - 1)) == 0
1019 _is_pow2_hs = head_size == head_size_padded
1020 use_td_qo = use_td and _is_pow2_nq and _is_pow2_hs
1021
1022 # ``_load_q_td`` / ``_store_output_td`` flatten ``(num_queries_per_kv,
1023 # HEAD_SIZE)`` into a single contiguous inner axis. That's only
1024 # equivalent to the pointer path when the ``num_queries_per_kv`` heads
1025 # for this KV group start at ``kv_head_idx * num_queries_per_kv`` and
1026 # lie exactly HEAD_SIZE apart — i.e. ``query_stride_1 == HEAD_SIZE``
1027 # and ``output_stride_1 == head_size``. This is the default vLLM
1028 # query/output layout; assert it explicitly so we fail fast if a
1029 # future caller passes a non-contiguous query tensor.
1030 if use_td_qo:
1031 assert q.stride(1) == head_size, (
1032 f"USE_TD_QO requires contiguous query heads "
1033 f"(q.stride(1) = {q.stride(1)} != head_size = {head_size}); "
1034 f"set VLLM_TRITON_USE_TD=0 or pad the query layout."
1035 )
1036 assert out.stride(1) == head_size, (
1037 f"USE_TD_QO requires contiguous output heads "
1038 f"(out.stride(1) = {out.stride(1)} != head_size = {head_size})."
1039 )
1040
1041 # Launch the 2D kernel if
1042 # 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or
1043 # 2. The batch includes at least one prefill request, or
1044 # 3. The number of sequences exceeds the configured threshold, or
1045 # 4. Batch invariance is enabled
1046 use_3d = not (
1047 seq_threshold_3D is None
1048 or num_par_softmax_segments is None
1049 or softmax_segm_output is None
1050 or softmax_segm_max is None
1051 or softmax_segm_expsum is None
1052 or max_seqlen_q > 1
1053 or num_seqs > seq_threshold_3D
1054 or is_batch_invariant
1055 )
1056
1057 # The kernel signature is the same for 2D and 3D — only the launch
1058 # grid + a handful of constexpr toggles differ. Per-token-head scale
1059 # caches and their strides are passed as ``None`` when the
1060 # ``USE_PER_TOKEN_HEAD_SCALES`` branch is dead so Triton can skip
1061 # materialising those arguments and the associated registers.
1062 if use_per_token_head_scales:
1063 ks_strides = k_scale_cache.stride()
1064 vs_strides = v_scale_cache.stride()
1065 ks_blk, ks_slot, ks_head = ks_strides[0], ks_strides[1], ks_strides[2]
1066 vs_blk, vs_slot, vs_head = vs_strides[0], vs_strides[1], vs_strides[2]
1067 k_scale_ptr = k_scale_cache
1068 v_scale_ptr = v_scale_cache
1069 else:
1070 ks_blk = ks_slot = ks_head = None
1071 vs_blk = vs_slot = vs_head = None
1072 k_scale_ptr = None
1073 v_scale_ptr = None
1074 # 3D needs real segm tensors; 2D never touches them. Pass ``None`` in
1075 # 2D mode so Triton can skip materialising these pointer arguments.
1076 segm_output_ptr = softmax_segm_output if use_3d else None
1077 segm_max_ptr = softmax_segm_max if use_3d else None
1078 segm_expsum_ptr = softmax_segm_expsum if use_3d else None
1079 num_segments = num_par_softmax_segments if use_3d else 1
1080
1081 grid: tuple[Any, ...]
1082 if not use_3d:
1083 grid = (total_num_q_blocks, num_kv_heads)
1084 tile_size = TILE_SIZE_PREFILL
1085 else:
1086 grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
1087 tile_size = TILE_SIZE_DECODE
1088
1089 launch_kwargs: dict[str, int] = {}
1090 if launch_num_warps is not None:
1091 launch_kwargs["num_warps"] = launch_num_warps
1092 if launch_num_stages is not None:
1093 launch_kwargs["num_stages"] = launch_num_stages
1094
1095 kernel_unified_attention[grid](
1096 output_ptr=out,
1097 segm_output_ptr=segm_output_ptr,
1098 segm_max_ptr=segm_max_ptr,
1099 segm_expsum_ptr=segm_expsum_ptr,
1100 query_ptr=q,
1101 key_cache_ptr=k,
1102 value_cache_ptr=v,
1103 sink_ptr=sinks,
1104 block_tables_ptr=block_table,
1105 seq_lens_ptr=seqused_k,
1106 alibi_slopes_ptr=alibi_slopes,
1107 qq_bias_ptr=qq_bias,
1108 k_scale_cache_ptr=k_scale_ptr,
1109 v_scale_cache_ptr=v_scale_ptr,
1110 scale=softmax_scale,
1111 q_scale=q_descale,
1112 k_scale=k_descale,
1113 v_scale=v_descale,
1114 out_scale=1 / output_scale if output_scale is not None else 1.0,
1115 softcap=softcap,
1116 num_query_heads=num_query_heads,
1117 num_queries_per_kv=num_queries_per_kv,
1118 block_table_stride=block_table.stride(0),
1119 query_stride_0=q.stride(0),
1120 query_stride_1=q.stride(1),
1121 output_stride_0=out.stride(0),
1122 output_stride_1=out.stride(1),
1123 qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
1124 BLOCK_SIZE=block_size,
1125 TILE_SIZE=tile_size,
1126 HEAD_SIZE=head_size,
1127 HEAD_SIZE_PADDED=head_size_padded,
1128 USE_ALIBI_SLOPES=use_alibi_slopes,
1129 USE_ALIBI_SQRT=use_alibi_sqrt,
1130 USE_QQ_BIAS=use_qq_bias,
1131 USE_SOFTCAP=(softcap > 0),
1132 USE_SINKS=(sinks is not None),
1133 SLIDING_WINDOW=(1 + window_size[0]),
1134 USE_CAUSAL=use_causal,
1135 USE_PER_SEQ_CAUSAL=use_per_seq_causal,
1136 per_seq_causal_ptr=per_seq_causal_ptr,
1137 USE_MM_PREFIX=use_mm_prefix,
1138 MAX_MM_RANGES=max_mm_ranges,
1139 mm_prefix_range_ptr=mm_prefix_range,
1140 rswa_prefix_lens_ptr=rswa_prefix_lens if use_rswa else seqused_k,
1141 R_SWA_WINDOW=rswa_window or 0,
1142 USE_R_SWA=use_rswa,
1143 stride_k_cache_0=k.stride(0),
1144 stride_k_cache_1=k.stride(1),
1145 stride_k_cache_2=k.stride(2),
1146 stride_k_cache_3=k.stride(3),
1147 stride_v_cache_0=v.stride(0),
1148 stride_v_cache_1=v.stride(1),
1149 stride_v_cache_2=v.stride(2),
1150 stride_v_cache_3=v.stride(3),
1151 stride_ks_blk=ks_blk,
1152 stride_ks_slot=ks_slot,
1153 stride_ks_head=ks_head,
1154 stride_vs_blk=vs_blk,
1155 stride_vs_slot=vs_slot,
1156 stride_vs_head=vs_head,
1157 query_start_len_ptr=cu_seqlens_q,
1158 BLOCK_Q=BLOCK_Q,
1159 num_seqs=num_seqs,
1160 BLOCK_M=BLOCK_M,
1161 NUM_SEGMENTS_PER_SEQ=num_segments,
1162 USE_FP8=output_scale is not None,
1163 IS_3D=use_3d,
1164 KV_QUANT_MODE=kv_quant_mode,
1165 Q_IS_FP8=(q.dtype == current_platform.fp8_dtype()),
1166 CHUNK_LOOKBACK=chunk_lookback,
1167 CHUNK_SIZE=chunk_size,
1168 USE_TD=use_td,
1169 USE_TD_QO=use_td_qo,
1170 MM_PREFIX_CLAMP_SW=mm_prefix_clamp_sliding_window,
1171 **launch_kwargs,
1172 )
1173
1174 if use_3d:
1175 reduce_segments[(q.shape[0], num_query_heads)](
1176 output_ptr=out,
1177 segm_output_ptr=softmax_segm_output,
1178 segm_max_ptr=softmax_segm_max,
1179 segm_expsum_ptr=softmax_segm_expsum,
1180 seq_lens_ptr=seqused_k,
1181 num_seqs=num_seqs,
1182 num_query_heads=num_query_heads,
1183 out_scale_inv=1 / output_scale if output_scale is not None else 1.0,
1184 output_stride_0=out.stride(0),
1185 output_stride_1=out.stride(1),
1186 block_table_stride=block_table.stride(0),
1187 TILE_SIZE=TILE_SIZE_DECODE,
1188 HEAD_SIZE=head_size,
1189 HEAD_SIZE_PADDED=head_size_padded,
1190 query_start_len_ptr=cu_seqlens_q,
1191 BLOCK_Q=BLOCK_Q,
1192 NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
1193 USE_FP8=output_scale is not None,
1194 )