A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 16:15 UTC
[Refactor] Unify 2D/3D kernels in triton_unified_attention (#40631)
Signed-off-by: JartX <sagformas@epdcenter.es>
JartX committed
Apr 24, 2026 at 17:18 UTC
914d0464c1b1ec77d1560b485624f32002532b83
2 files changed
+704
-888
vllm/v1/attention/ops/triton_attention_helpers.py
new
+383
@@ -0,0 +1,383 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared ``@triton.jit`` helpers used by the unified attention kernel
4
+and ``reduce_segments``.
5
+
6
+These are plain attention-loop helpers — mask building, ALiBi / QQ-bias
7
+score post-processing, online-softmax bookkeeping, tile-loop bounds,
8
+sequence lookup — extracted so the 2D and 3D paths of the unified
9
+kernel (and any future consumer) share a single implementation.
10
+"""
11
+
12
+from __future__ import annotations
13
+
14
+from vllm.triton_utils import tl, triton
15
+
16
+# ===========================================================================
17
+# Scalar helpers (reused by every kernel + reduce_segments)
18
+# ===========================================================================
19
+
20
+
21
+@triton.jit
22
+def cdiv_fn(x, y):
23
+ """Ceiling division. Kept as a helper to keep kernel bodies terse."""
24
+ return (x + y - 1) // y
25
+
26
+
27
+@triton.jit
28
+def apply_softcap(S, x):
29
+ """Softcap (aka tanh-style clamp) used to bound attention scores.
30
+
31
+ ``x * tanh(S / x)`` rewritten to avoid a direct ``tanh`` call.
32
+ """
33
+ Sdiv = S / x
34
+ p1 = tl.exp(Sdiv)
35
+ p2 = tl.exp(-Sdiv)
36
+ return x * (p1 - p2) / (p1 + p2)
37
+
38
+
39
+# ===========================================================================
40
+# Attention loop
41
+# ===========================================================================
42
+
43
+
44
+@triton.jit
45
+def resolve_seq_and_query_len(
46
+ query_start_len_ptr,
47
+ seq_lens_ptr,
48
+ q_block_global_idx,
49
+ num_seqs,
50
+ BLOCK_Q: tl.constexpr,
51
+):
52
+ """Resolve the (sequence, q-block-within-sequence) pair and load the
53
+ per-sequence lengths.
54
+
55
+ Shared across every attention kernel — the ``q_block_global_idx``
56
+ program id indexes into the flattened ``(seq, q_block_in_seq)``
57
+ space, and a binary search over ``query_start_len_ptr`` recovers
58
+ the (seq, local-q-block) pair.
59
+
60
+ Returns ``(seq_idx, q_block_local_idx, cur_batch_in_all_start_index,
61
+ cur_batch_query_len, seq_len)``. Callers must still early-return
62
+ when ``q_block_local_idx * BLOCK_Q >= cur_batch_query_len`` (Triton
63
+ helpers cannot return from the caller).
64
+ """
65
+ # find_seq_idx is defined below; forward use is fine inside @triton.jit.
66
+ seq_idx = find_seq_idx(
67
+ query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True
68
+ )
69
+ q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx
70
+ q_block_local_idx = q_block_global_idx - q_block_start_idx
71
+ cur_start = tl.load(query_start_len_ptr + seq_idx)
72
+ cur_stop = tl.load(query_start_len_ptr + seq_idx + 1)
73
+ cur_batch_query_len = cur_stop - cur_start
74
+ seq_len = tl.load(seq_lens_ptr + seq_idx)
75
+ return seq_idx, q_block_local_idx, cur_start, cur_batch_query_len, seq_len
76
+
77
+
78
+@triton.jit
79
+def find_seq_idx(
80
+ query_start_len_ptr,
81
+ target_idx,
82
+ num_seqs,
83
+ BLOCK_Q: tl.constexpr,
84
+ use_q_block_mode: tl.constexpr,
85
+):
86
+ """Binary search over the cumulative query-length prefix.
87
+
88
+ When ``use_q_block_mode`` is True, the prefix values are reshaped
89
+ into units of ``BLOCK_Q`` plus one entry per boundary — matching
90
+ the q-block grid laid out by the attention kernels. When False
91
+ we search the plain cumulative-length prefix (used by
92
+ ``reduce_segments`` which iterates over raw query tokens).
93
+ """
94
+ left: tl.int32 = 0
95
+ right = num_seqs
96
+ while left < right:
97
+ mid = (left + right) // 2
98
+ val = tl.load(query_start_len_ptr + mid)
99
+ mid_val = val // BLOCK_Q + mid if use_q_block_mode else val
100
+
101
+ if mid_val <= target_idx:
102
+ left = mid + 1
103
+ else:
104
+ right = mid
105
+
106
+ return left - 1
107
+
108
+
109
+@triton.jit
110
+def init_softmax_M(
111
+ sink_ptr,
112
+ query_offset_1,
113
+ query_mask_1,
114
+ segm_idx_or_0,
115
+ BLOCK_M: tl.constexpr,
116
+ USE_SINKS: tl.constexpr,
117
+ IS_3D: tl.constexpr,
118
+):
119
+ """Initial row-max ``M`` for the online softmax.
120
+
121
+ Without sinks: ``-inf``. With sinks: load the per-head sink bias
122
+ once. In 3D mode only segment 0 loads — ``reduce_segments`` adds
123
+ the sink contribution exactly once across segments, so other
124
+ segments must start from ``-inf``.
125
+
126
+ ``segm_idx_or_0`` is the 3D segment index or 0 for 2D (caller
127
+ passes ``0`` when ``IS_3D`` is False).
128
+ """
129
+ M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
130
+ if USE_SINKS:
131
+ load_sinks = (not IS_3D) or (segm_idx_or_0 == 0)
132
+ if load_sinks:
133
+ M = tl.load(
134
+ sink_ptr + query_offset_1,
135
+ mask=query_mask_1,
136
+ other=float("-inf"),
137
+ ).to(tl.float32)
138
+ return M
139
+
140
+
141
+@triton.jit
142
+def compute_tile_loop_bounds(
143
+ context_len,
144
+ seq_len,
145
+ cur_batch_query_len,
146
+ q_block_local_idx,
147
+ segm_idx_or_0,
148
+ tiles_per_segment_or_0,
149
+ TILE_SIZE: tl.constexpr,
150
+ BLOCK_M: tl.constexpr,
151
+ BLOCK_Q: tl.constexpr,
152
+ num_queries_per_kv: tl.constexpr,
153
+ SLIDING_WINDOW: tl.constexpr,
154
+ USE_MM_PREFIX: tl.constexpr,
155
+ IS_3D: tl.constexpr,
156
+ CHUNK_LOOKBACK: tl.constexpr = -1,
157
+ CHUNK_SIZE: tl.constexpr = -1,
158
+):
159
+ """Compute the tile-loop bounds ``(loop_lo, loop_hi)`` and the
160
+ derived ``max_seq_prefix_len`` used for per-tile masking.
161
+
162
+ Combines three concerns into one helper:
163
+
164
+ 1. Longest prefix spanned by any query token in this q-block.
165
+ Clamped to ``seq_len`` (causal) or extended to it when
166
+ mm_prefix is active (bidirectional ranges can reach past the
167
+ causal prefix).
168
+ 2. Sliding-window pruning: narrows ``[tile_start, tile_end)`` to
169
+ only tiles that can contain an allowed key under SWA.
170
+ 3. 3D scoping: when ``IS_3D`` is True, further narrows to the
171
+ segment's slice via ``(segm_idx * tiles_per_segment,
172
+ (segm_idx + 1) * tiles_per_segment)``.
173
+ """
174
+ # compute the length of the longest sequence prefix spanned by any
175
+ # query token in the current q_block (q_block_local_idx)
176
+ max_seq_prefix_len = (
177
+ context_len
178
+ + q_block_local_idx * BLOCK_Q
179
+ + (BLOCK_M - 1) // num_queries_per_kv
180
+ + 1
181
+ )
182
+ if USE_MM_PREFIX:
183
+ # image bidirectional attention ranges require a full range
184
+ # including q_block padding to make sure doc mask is correct
185
+ max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len)
186
+ else:
187
+ max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
188
+
189
+ num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)
190
+
191
+ # ---- Sliding-window tile pruning --------------------
192
+ # Default: keep previous global behavior
193
+ tile_start = 0
194
+ tile_end = num_tiles
195
+ # TODO(Isotr0py): sliding window pruning with image bidirectional mask
196
+ if SLIDING_WINDOW > 0 and not USE_MM_PREFIX:
197
+ # Query rows covered by this Q-block
198
+ qpos_lo = q_block_local_idx * BLOCK_Q
199
+ qpos_hi = tl.minimum(
200
+ qpos_lo + (BLOCK_M - 1) // num_queries_per_kv,
201
+ cur_batch_query_len - 1,
202
+ )
203
+ # For sliding window, each query position q can only attend to
204
+ # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs]
205
+ # where q_abs = context_len + q
206
+ # The union of allowed key positions for this Q-block is:
207
+ # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
208
+ q_abs = context_len + qpos_lo
209
+ if CHUNK_LOOKBACK > -1:
210
+ # Chunked attention: align lower bound to the start of the
211
+ # lookback'th previous chunk.
212
+ first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE
213
+ else:
214
+ first_allowed_key = q_abs - SLIDING_WINDOW + 1
215
+ last_allowed_key = context_len + qpos_hi
216
+ # Convert to tile indices and clamp
217
+ tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
218
+ tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles)
219
+
220
+ if IS_3D:
221
+ loop_lo = max(segm_idx_or_0 * tiles_per_segment_or_0, tile_start)
222
+ loop_hi = min((segm_idx_or_0 + 1) * tiles_per_segment_or_0, tile_end)
223
+ else:
224
+ loop_lo = tile_start
225
+ loop_hi = tile_end
226
+
227
+ return loop_lo, loop_hi, max_seq_prefix_len
228
+
229
+
230
+@triton.jit
231
+def store_segm_reduce_scalars(
232
+ segm_max_ptr,
233
+ segm_expsum_ptr,
234
+ query_offset_0,
235
+ query_offset_1,
236
+ segm_idx,
237
+ M,
238
+ L,
239
+ query_mask_0,
240
+ query_mask_1,
241
+ num_query_heads: tl.constexpr,
242
+ NUM_SEGMENTS_PER_SEQ: tl.constexpr,
243
+):
244
+ """Store per-segment ``M`` and ``L`` for ``reduce_segments`` to
245
+ combine into the final softmax.
246
+
247
+ Shared across every 3D attention epilogue; the per-token output
248
+ stripes are mode-specific (flat / 2-stream split / 4-stream split)
249
+ and stay inlined.
250
+ """
251
+ segm_offset = (
252
+ query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
253
+ + query_offset_1 * NUM_SEGMENTS_PER_SEQ
254
+ + segm_idx
255
+ )
256
+ tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1)
257
+ tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1)
258
+
259
+
260
+@triton.jit
261
+def compute_kv_seq_mask(
262
+ query_abs_pos,
263
+ seq_offset,
264
+ seq_idx,
265
+ mm_prefix_range_ptr,
266
+ SLIDING_WINDOW: tl.constexpr,
267
+ USE_MM_PREFIX: tl.constexpr,
268
+ MAX_MM_RANGES: tl.constexpr,
269
+ CHUNK_LOOKBACK: tl.constexpr = -1,
270
+ CHUNK_SIZE: tl.constexpr = -1,
271
+):
272
+ """Build the KV mask for one tile.
273
+
274
+ Causal (key <= query) by default; AND-ed with either chunked
275
+ attention (``CHUNK_LOOKBACK >= 0``) or sliding window
276
+ (``SLIDING_WINDOW > 0``); OR-ed with the bidirectional ranges from
277
+ ``mm_prefix_range`` when PrefixLM / multimodal attention is active.
278
+ Order matches FlexAttention: ``(causal AND window) OR mm_prefix``.
279
+ Chunked attention takes precedence over sliding window when both
280
+ are non-default — the launcher zeros ``CHUNK_LOOKBACK`` whenever
281
+ sliding window is disabled.
282
+ """
283
+ # Compute attention mask: causal by default (key <= query)
284
+ seq_mask = seq_offset[None, :] <= query_abs_pos
285
+
286
+ # Apply sliding window / chunked attention to base mask
287
+ # BEFORE mm_prefix OR.
288
+ # Order must match FlexAttention:
289
+ # (causal AND sliding_window) OR mm_prefix
290
+ if CHUNK_LOOKBACK > -1:
291
+ seq_mask = seq_mask & (
292
+ (query_abs_pos // CHUNK_SIZE - seq_offset[None, :] // CHUNK_SIZE)
293
+ <= CHUNK_LOOKBACK
294
+ )
295
+ elif SLIDING_WINDOW > 0:
296
+ seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
297
+
298
+ # PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
299
+ # Applied AFTER sliding window so mm_prefix ranges override SW restriction.
300
+ if USE_MM_PREFIX:
301
+ for i in range(MAX_MM_RANGES):
302
+ range_start = tl.load(
303
+ mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2
304
+ )
305
+ range_end = tl.load(
306
+ mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1
307
+ )
308
+ is_valid = range_start < range_end
309
+ q_in_range = (
310
+ (query_abs_pos >= range_start) & (query_abs_pos <= range_end) & is_valid
311
+ )
312
+ k_in_range = (
313
+ (seq_offset[None, :] >= range_start)
314
+ & (seq_offset[None, :] <= range_end)
315
+ & is_valid
316
+ )
317
+ seq_mask |= q_in_range & k_in_range
318
+ return seq_mask
319
+
320
+
321
+@triton.jit
322
+def apply_alibi_to_score(
323
+ S,
324
+ alibi_slope,
325
+ seq_offset,
326
+ context_len,
327
+ query_pos,
328
+ USE_ALIBI_SQRT: tl.constexpr,
329
+):
330
+ """Add the ALiBi positional bias (linear or sqrt variant) to S in-place."""
331
+ if USE_ALIBI_SQRT:
332
+ relative_pos = seq_offset - (context_len + query_pos[:, None])
333
+ alibi_offset = tl.where(
334
+ relative_pos <= 0,
335
+ -tl.sqrt((-relative_pos).to(tl.float32)),
336
+ 0.0,
337
+ )
338
+ else:
339
+ alibi_offset = seq_offset - context_len
340
+ return S + alibi_slope[:, None] * alibi_offset
341
+
342
+
343
+@triton.jit
344
+def load_qq_bias_tile(
345
+ qq_bias_row_ptrs,
346
+ seq_offset,
347
+ context_len,
348
+ qq_bias_stride_0,
349
+):
350
+ """Load the qq-bias slice for keys that correspond to query rows."""
351
+ key_rel_pos = seq_offset - context_len
352
+ is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0
353
+ return tl.load(
354
+ qq_bias_row_ptrs + key_rel_pos[None, :],
355
+ mask=is_query_key[None, :],
356
+ other=0.0,
357
+ )
358
+
359
+
360
+@triton.jit
361
+def softmax_step(S, M, L):
362
+ """Online softmax update for one tile.
363
+
364
+ Returns ``(M_new, L_new, P, alpha)``. Caller is responsible for
365
+ rescaling its accumulator(s) by ``alpha[:, None]`` — done outside so
366
+ kernels with a different number / shape of accumulators can reuse
367
+ the same step.
368
+ """
369
+ # compute running maximum
370
+ # m_j : (BLOCK_M,)
371
+ m_j = tl.maximum(M, tl.max(S, axis=1))
372
+ # For sliding window there's a chance the max is -inf due to masking of
373
+ # the entire row. In this case we need to set m_j 0 to avoid NaN
374
+ m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
375
+ # P : (BLOCK_M, TILE_SIZE)
376
+ P = tl.exp(S - m_j[:, None])
377
+ # l_j : (BLOCK_M,)
378
+ l_j = tl.sum(P, axis=1)
379
+ # alpha : (BLOCK_M, )
380
+ alpha = tl.exp(M - m_j)
381
+ # update constants
382
+ L_new = L * alpha + l_j
383
+ return m_j, L_new, P, alpha
vllm/v1/attention/ops/triton_unified_attention.py
+321
-888
@@ -7,12 +7,27 @@
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__)
@@ -21,114 +36,53 @@ float8_info = torch.finfo(current_platform.fp8_dtype())
36
37
38
@triton.jit
24
-def cdiv_fn(x, y):
25
- return (x + y - 1) // y
26
-
27
-
28
-@triton.jit
29
-def apply_softcap(S, x):
30
- Sdiv = S / x
31
- p1 = tl.exp(Sdiv)
32
- p2 = tl.exp(-Sdiv)
33
- return x * (p1 - p2) / (p1 + p2)
34
-
35
-
36
-@triton.jit
37
-def _prepare_kv_tile(
38
- data,
39
- Q,
40
- tensor_scale,
41
- scale_cache_ptr,
42
- physical_block_idx,
43
- seq_offset,
44
- kv_head_idx,
45
- stride_s_blk,
46
- stride_s_slot,
47
- stride_s_head,
48
- tile_mask,
49
- BLOCK_SIZE: tl.constexpr,
50
- KV_QUANT_MODE: tl.constexpr,
51
-):
52
- """Prepare a loaded KV tile for attention computation.
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
54
- Casts the raw KV data to Q's dtype and loads per-token-head scales
55
- when applicable:
42
+ Modes handled inside the core kernel:
43
57
- - ``KV_QUANT_MODE == 0``: cast only (no-op for bf16/fp16).
58
- - ``KV_QUANT_MODE == 1`` (FP8 per-tensor): dequantize inline
59
- using the tensor-wide scale.
60
- - ``KV_QUANT_MODE >= 2`` (per-token-head int8/fp8): cast to Q's
61
- dtype and return per-head scales separately — the caller applies
62
- them after the dot product for better numerical efficiency.
63
-
64
- Returns ``(data, token_head_scales)``. *token_head_scales* is only
65
- meaningful when ``KV_QUANT_MODE >= 2``; callers gate its use on
66
- the same constexpr so the compiler eliminates dead code.
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.
49
"""
68
- # KV_QUANT_MODE values: 0=none, 1=fp8 per-tensor,
69
- # 2=int8 per-token-head, 3=fp8 per-token-head
70
-
71
- # Placeholder scales (float32) — never read when KV_QUANT_MODE < 2.
72
- unused_scales = tile_mask.to(tl.float32)
73
-
74
- if KV_QUANT_MODE == 1: # FP8 per-tensor
50
+ if KV_QUANT_MODE == 1:
51
if Q.dtype.is_fp8():
76
- return data.to(Q.dtype), unused_scales
77
- return (data.to(tl.float32) * tl.load(tensor_scale)).to(Q.dtype), unused_scales
78
- if KV_QUANT_MODE >= 2: # per-token-head (int8 or fp8)
79
- scale_idx = (
80
- physical_block_idx * stride_s_blk
81
- + (seq_offset % BLOCK_SIZE) * stride_s_slot
82
- + kv_head_idx * stride_s_head
83
- )
84
- token_head_scales = tl.load(
85
- scale_cache_ptr + scale_idx, mask=tile_mask, other=1.0
86
- )
87
- return data.to(Q.dtype), token_head_scales
88
- # .to(Q.dtype) is a no-op when data is already Q's type (bf16/fp16),
89
- # but required so Triton sees consistent return types across branches.
90
- return data.to(Q.dtype), unused_scales
52
+ return data.to(Q.dtype)
53
+ return (data.to(tl.float32) * tl.load(tensor_scale)).to(Q.dtype)
54
+ return data.to(Q.dtype)
55
56
57
@triton.jit
94
-def find_seq_idx(
95
- query_start_len_ptr,
96
- target_idx,
97
- num_seqs,
98
- BLOCK_Q: tl.constexpr,
99
- use_q_block_mode: tl.constexpr,
100
-):
101
- left: tl.int32 = 0
102
- right = num_seqs
103
- while left < right:
104
- mid = (left + right) // 2
105
- val = tl.load(query_start_len_ptr + mid)
106
- mid_val = val // BLOCK_Q + mid if use_q_block_mode else val
107
-
108
- if mid_val <= target_idx:
109
- left = mid + 1
110
- else:
111
- right = mid
112
-
113
- return left - 1
114
-
115
-
116
-@triton.jit
117
-def kernel_unified_attention_2d(
118
- output_ptr, # [num_tokens, num_query_heads, head_size]
119
- query_ptr, # [num_tokens, num_query_heads, head_size]
120
- key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
121
- value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
122
- sink_ptr, # [num_query_heads]
123
- block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
124
- seq_lens_ptr, # [num_seqs]
125
- alibi_slopes_ptr, # [num_query_heads]
126
- qq_bias_ptr, # [num_query_tokens, num_query_tokens]
127
- scale, # float32
128
- k_scale, # float32
129
- v_scale, # float32
130
- out_scale, # float32
131
- softcap, # float32
58
+def kernel_unified_attention(
59
+ # Output destinations. In 2D mode we write the final result into
60
+ # ``output_ptr``; in 3D mode we write per-segment partials into the
61
+ # three ``segm_*`` tensors and ``output_ptr`` is unused (callers may
62
+ # pass any non-null pointer).
63
+ output_ptr,
64
+ segm_output_ptr,
65
+ segm_max_ptr,
66
+ segm_expsum_ptr,
67
+ # Inputs
68
+ query_ptr,
69
+ key_cache_ptr,
70
+ value_cache_ptr,
71
+ sink_ptr,
72
+ block_tables_ptr,
73
+ seq_lens_ptr,
74
+ alibi_slopes_ptr,
75
+ qq_bias_ptr,
76
+ # Per-(token, head) scale caches (used iff KV_QUANT_MODE in {2, 3}).
77
+ # For other modes callers may pass any non-null pointer.
78
+ k_scale_cache_ptr,
79
+ v_scale_cache_ptr,
80
+ # Scalars
81
+ scale,
82
+ k_scale,
83
+ v_scale,
84
+ out_scale,
85
+ softcap,
86
num_query_heads: tl.constexpr, # int
87
num_queries_per_kv: tl.constexpr, # int
88
block_table_stride: tl.int64, # int
@@ -149,7 +103,7 @@ def kernel_unified_attention_2d(
103
SLIDING_WINDOW: tl.constexpr, # int
104
USE_MM_PREFIX: tl.constexpr, # bool
105
MAX_MM_RANGES: tl.constexpr, # int
152
- mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence
106
+ mm_prefix_range_ptr,
107
stride_k_cache_0: tl.int64, # int
108
stride_k_cache_1: tl.int64, # int
109
stride_k_cache_2: tl.int64, # int
@@ -158,47 +112,61 @@ def kernel_unified_attention_2d(
112
stride_v_cache_1: tl.int64, # int
113
stride_v_cache_2: tl.int64, # int
114
stride_v_cache_3: tl.constexpr, # int
161
- query_start_len_ptr, # [num_seqs+1]
162
- BLOCK_Q: tl.constexpr, # int
115
+ stride_ks_blk: tl.int64,
116
+ stride_ks_slot: tl.int64,
117
+ stride_ks_head: tl.int64,
118
+ stride_vs_blk: tl.int64,
119
+ stride_vs_slot: tl.int64,
120
+ stride_vs_head: tl.int64,
121
+ query_start_len_ptr,
122
+ BLOCK_Q: tl.constexpr,
123
num_seqs: tl.int32,
164
- BLOCK_M: tl.constexpr, # int
165
- USE_FP8: tl.constexpr, # bool
166
- # KV cache quantization: 0=none, 1=fp8, 2=per-token-head
124
+ BLOCK_M: tl.constexpr,
125
+ NUM_SEGMENTS_PER_SEQ: tl.constexpr,
126
+ USE_FP8: tl.constexpr,
127
+ # Toggles 2D vs 3D layout. The 2D path runs the full sequence in one
128
+ # tile loop and writes to ``output_ptr``. The 3D path scopes the loop
129
+ # to ``[segm_idx, segm_idx+1) × tiles_per_segment`` and writes
130
+ # per-segment partials, finalized by ``reduce_segments``.
131
+ IS_3D: tl.constexpr,
132
+ # KV cache quantization mode handled inside this kernel via constexpr
133
+ # branches: NONE (0), FP8_PER_TENSOR (1), INT8_PER_TOKEN_HEAD (2),
134
+ # FP8_PER_TOKEN_HEAD (3).
135
KV_QUANT_MODE: tl.constexpr = 0,
136
FP8_MIN: tl.constexpr = float8_info.min,
137
FP8_MAX: tl.constexpr = float8_info.max,
170
- # Per-token-head scale caches (KV_QUANT_MODE >= 2)
171
- # Shape: [num_blocks, block_size, num_kv_heads]
172
- k_scale_cache_ptr=None,
173
- v_scale_cache_ptr=None,
174
- stride_ks_blk=0,
175
- stride_ks_slot=0,
176
- stride_ks_head=0,
177
- stride_vs_blk=0,
178
- stride_vs_slot=0,
179
- stride_vs_head=0,
138
+ # Chunked / block-local attention. ``CHUNK_LOOKBACK >= 0`` enables
139
+ # chunked masking (used by Gemma3 block-local layers); takes precedence
140
+ # over ``SLIDING_WINDOW`` inside the helpers. ``-1`` disables.
141
CHUNK_LOOKBACK: tl.constexpr = -1,
142
CHUNK_SIZE: tl.constexpr = -1,
143
):
144
+ USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2
145
+
146
q_block_global_idx = tl.program_id(0)
147
kv_head_idx = tl.program_id(1)
185
-
186
- seq_idx = find_seq_idx(
187
- query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True
148
+ segm_idx = tl.program_id(2) if IS_3D else 0
149
+
150
+ (
151
+ seq_idx,
152
+ q_block_local_idx,
153
+ cur_batch_in_all_start_index,
154
+ cur_batch_query_len,
155
+ seq_len,
156
+ ) = resolve_seq_and_query_len(
157
+ query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q
158
)
159
190
- q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx
191
-
192
- q_block_local_idx = q_block_global_idx - q_block_start_idx
193
-
194
- cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
195
- cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)
196
-
197
- cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index
198
-
160
if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
161
return
162
163
+ if IS_3D:
164
+ tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
165
+ if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
166
+ return
167
+ else:
168
+ tiles_per_segment = 0
169
+
170
offs_m = tl.arange(0, BLOCK_M)
171
offs_d = tl.arange(0, HEAD_SIZE_PADDED)
172
offs_t = tl.arange(0, TILE_SIZE)
@@ -225,88 +193,43 @@ def kernel_unified_attention_2d(
193
194
block_table_offset = seq_idx * block_table_stride
195
228
- if not USE_SINKS:
229
- M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
230
- else:
231
- M = tl.load(
232
- sink_ptr + query_offset_1,
233
- mask=query_mask_1,
234
- other=float("-inf"),
235
- ).to(dtype=tl.float32)
236
-
196
+ M = init_softmax_M(
197
+ sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D
198
+ )
199
L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
200
+ # acc : (BLOCK_M, HEAD_SIZE_PADDED)
201
acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
202
240
- # sequence len for this particular sequence
241
- seq_len = tl.load(seq_lens_ptr + seq_idx)
242
-
243
- # context length for this particular sequences
203
context_len = seq_len - cur_batch_query_len
204
246
- # alibi slope for this head
205
if USE_ALIBI_SLOPES:
206
alibi_slope = tl.load(
207
alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0
208
)
209
252
- # query-query attention bias
210
if USE_QQ_BIAS:
254
- qq_bias_row_ptrs = (
255
- qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0
256
- ) # shape: [BLOCK_M]
257
-
258
- # compute the length of the longest sequence prefix spanned by any
259
- # query token in the current q_block (q_block_local_idx)
260
- max_seq_prefix_len = (
261
- context_len
262
- + q_block_local_idx * BLOCK_Q
263
- + (BLOCK_M - 1) // num_queries_per_kv
264
- + 1
211
+ qq_bias_row_ptrs = qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0
212
+
213
+ loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds(
214
+ context_len,
215
+ seq_len,
216
+ cur_batch_query_len,
217
+ q_block_local_idx,
218
+ segm_idx,
219
+ tiles_per_segment,
220
+ TILE_SIZE,
221
+ BLOCK_M,
222
+ BLOCK_Q,
223
+ num_queries_per_kv,
224
+ SLIDING_WINDOW,
225
+ USE_MM_PREFIX,
226
+ IS_3D,
227
+ CHUNK_LOOKBACK,
228
+ CHUNK_SIZE,
229
)
230
267
- if USE_MM_PREFIX:
268
- # image bidirectional attention ranges require a full range
269
- # including q_block padding to make sure doc mask is correct
270
- max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len)
271
- else:
272
- # adjust for potential padding in the last q_block by considering the
273
- # actual sequence length
274
- max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
275
-
276
- # calculate the number of tiles that need to be processed to
277
- # cover the longest sequence prefix (due to causal masking, tiles beyond
278
- # this prefix can be skipped)
279
- num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)
280
-
281
- # ---- Sliding-window tile pruning --------------------
282
- # Default: keep previous global behavior
283
- tile_start = 0
284
- tile_end = num_tiles
285
- # TODO(Isotr0py): sliding window pruning with image bidirectional mask
286
- if SLIDING_WINDOW > 0 and not USE_MM_PREFIX:
287
- # Query rows covered by this Q-block
288
- qpos_lo = q_block_local_idx * BLOCK_Q
289
- qpos_hi = tl.minimum(
290
- qpos_lo + (BLOCK_M - 1) // num_queries_per_kv,
291
- cur_batch_query_len - 1,
292
- )
293
- # For sliding window, each query position q can only attend to
294
- # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs]
295
- # where q_abs = context_len + q
296
- # The union of allowed key positions for this Q-block is:
297
- # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
298
- q_abs = context_len + qpos_lo
299
- if CHUNK_LOOKBACK > -1:
300
- first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE
301
- else:
302
- first_allowed_key = q_abs - SLIDING_WINDOW + 1
303
- last_allowed_key = context_len + qpos_hi
304
- # Convert to tile indices and clamp
305
- tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
306
- tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles)
307
-
231
# iterate through tiles (now limited to the sliding window range)
309
- for j in range(tile_start, tile_end):
232
+ for j in range(loop_lo, loop_hi):
233
seq_offset = j * TILE_SIZE + offs_t
234
tile_mask = seq_offset < max_seq_prefix_len
235
@@ -320,107 +243,64 @@ def kernel_unified_attention_2d(
243
+ offs_d[None, :] * stride_v_cache_3
244
+ (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1
245
)
323
-
246
k_offset = (
247
physical_block_idx[None, :] * stride_k_cache_0
248
+ kv_head_idx * stride_k_cache_2
249
+ offs_d[:, None] * stride_k_cache_3
250
+ (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
251
)
330
-
252
# K : (HEAD_SIZE, TILE_SIZE)
253
K_load = tl.load(
254
key_cache_ptr + k_offset,
255
mask=dim_mask[:, None] & tile_mask[None, :],
256
other=0.0,
257
)
337
- K, k_token_head_scales = _prepare_kv_tile(
338
- K_load,
339
- Q,
340
- k_scale,
341
- k_scale_cache_ptr,
342
- physical_block_idx,
343
- seq_offset,
344
- kv_head_idx,
345
- stride_ks_blk,
346
- stride_ks_slot,
347
- stride_ks_head,
348
- tile_mask,
349
- BLOCK_SIZE,
350
- KV_QUANT_MODE,
351
- )
352
-
258
+ K = _cast_kv_tile(K_load, Q, k_scale, KV_QUANT_MODE)
259
# V : (TILE_SIZE, HEAD_SIZE)
260
V_load = tl.load(
261
value_cache_ptr + v_offset,
262
mask=dim_mask[None, :] & tile_mask[:, None],
263
other=0.0,
264
)
359
- V, v_token_head_scales = _prepare_kv_tile(
360
- V_load,
361
- Q,
362
- v_scale,
363
- v_scale_cache_ptr,
364
- physical_block_idx,
365
- seq_offset,
366
- kv_head_idx,
367
- stride_vs_blk,
368
- stride_vs_slot,
369
- stride_vs_head,
370
- tile_mask,
371
- BLOCK_SIZE,
372
- KV_QUANT_MODE,
373
- )
265
+ V = _cast_kv_tile(V_load, Q, v_scale, KV_QUANT_MODE)
266
+
267
+ # Per-(token, head) scales for INT8 / FP8 per-token-head modes.
268
+ if USE_PER_TOKEN_HEAD_SCALES:
269
+ scale_idx = (
270
+ physical_block_idx * stride_ks_blk
271
+ + (seq_offset % BLOCK_SIZE) * stride_ks_slot
272
+ + kv_head_idx * stride_ks_head
273
+ )
274
+ k_token_head_scales = tl.load(
275
+ k_scale_cache_ptr + scale_idx, mask=tile_mask, other=1.0
276
+ )
277
+ v_scale_idx = (
278
+ physical_block_idx * stride_vs_blk
279
+ + (seq_offset % BLOCK_SIZE) * stride_vs_slot
280
+ + kv_head_idx * stride_vs_head
281
+ )
282
+ v_token_head_scales = tl.load(
283
+ v_scale_cache_ptr + v_scale_idx, mask=tile_mask, other=1.0
284
+ )
285
375
- # Compute attention mask: causal by default (key <= query)
286
query_abs_pos = context_len + query_pos[:, None]
377
- seq_mask = seq_offset[None, :] <= query_abs_pos
378
-
379
- # Apply sliding window / chunked attention to base mask
380
- # BEFORE mm_prefix OR.
381
- # Order must match FlexAttention:
382
- # (causal AND sliding_window) OR mm_prefix
383
- if CHUNK_LOOKBACK > -1:
384
- seq_mask = seq_mask & (
385
- (
386
- (context_len + query_pos[:, None]) // CHUNK_SIZE
387
- - (seq_offset[None, :] // CHUNK_SIZE)
388
- )
389
- <= CHUNK_LOOKBACK
390
- )
391
- elif SLIDING_WINDOW > 0:
392
- seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
393
-
394
- # PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
395
- # Applied AFTER sliding window so mm_prefix ranges override SW restriction.
396
- if USE_MM_PREFIX:
397
- for i in range(MAX_MM_RANGES):
398
- range_start = tl.load(
399
- mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2
400
- )
401
- range_end = tl.load(
402
- mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1
403
- )
404
-
405
- is_valid = range_start < range_end
406
- q_in_range = (
407
- (query_abs_pos >= range_start)
408
- & (query_abs_pos <= range_end)
409
- & is_valid
410
- )
411
- k_in_range = (
412
- (seq_offset[None, :] >= range_start)
413
- & (seq_offset[None, :] <= range_end)
414
- & is_valid
415
- )
416
- seq_mask |= q_in_range & k_in_range
287
+ seq_mask = compute_kv_seq_mask(
288
+ query_abs_pos,
289
+ seq_offset,
290
+ seq_idx,
291
+ mm_prefix_range_ptr,
292
+ SLIDING_WINDOW,
293
+ USE_MM_PREFIX,
294
+ MAX_MM_RANGES,
295
+ CHUNK_LOOKBACK,
296
+ CHUNK_SIZE,
297
+ )
298
299
# S : (BLOCK_M, TILE_SIZE)
300
S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
420
-
421
- # Per-token-head quant: fuse softmax_scale with per-head k_scale
422
- # to avoid a separate BLOCK_M × TILE_SIZE multiply on S.
423
- if KV_QUANT_MODE >= 2:
301
+ if USE_PER_TOKEN_HEAD_SCALES:
302
+ # Per-token-head quant: fuse softmax_scale with per-head k_scale
303
+ # to avoid a separate BLOCK_M × TILE_SIZE multiply on S.
304
S += tl.dot(Q, K) * (scale * k_token_head_scales[None, :])
305
else:
306
S += scale * tl.dot(Q, K)
@@ -433,494 +313,76 @@ def kernel_unified_attention_2d(
313
)
314
315
if USE_ALIBI_SLOPES:
436
- if USE_ALIBI_SQRT:
437
- relative_pos = seq_offset - (context_len + query_pos[:, None])
438
- alibi_offset = tl.where(
439
- relative_pos <= 0,
440
- -tl.sqrt((-relative_pos).to(tl.float32)),
441
- 0.0,
442
- )
443
- else:
444
- alibi_offset = seq_offset - context_len
445
- S += alibi_slope[:, None] * alibi_offset
316
+ S = apply_alibi_to_score(
317
+ S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT
318
+ )
319
320
if USE_QQ_BIAS:
448
- # compute key positions relative to query section
449
- key_rel_pos = seq_offset - context_len # shape: [BLOCK_SIZE]
450
- # load bias only for keys that correspond to queries
451
- is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0
452
- qq_bias = tl.load(
453
- qq_bias_row_ptrs + key_rel_pos[None, :],
454
- mask=is_query_key[None, :], # avoid OOB for context keys
455
- other=0.0,
321
+ S += load_qq_bias_tile(
322
+ qq_bias_row_ptrs, seq_offset, context_len, qq_bias_stride_0
323
)
457
- S += qq_bias
458
-
459
- # compute running maximum
460
- # m_j : (BLOCK_M,)
461
- m_j = tl.maximum(M, tl.max(S, axis=1))
462
-
463
- # For sliding window there's a chance the max is -inf due to masking of
464
- # the entire row. In this case we need to set m_j 0 to avoid NaN
465
- m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
466
-
467
- # P : (BLOCK_M, TILE_SIZE)
468
- P = tl.exp(S - m_j[:, None])
469
-
470
- # l_j : (BLOCK_M,)
471
- l_j = tl.sum(P, axis=1)
324
473
- # alpha : (BLOCK_M, )
474
- alpha = tl.exp(M - m_j)
475
-
476
- # acc : (BLOCK_M, HEAD_SIZE_PADDED)
325
+ M, L, P, alpha = softmax_step(S, M, L)
326
acc = acc * alpha[:, None]
327
479
- # update constants
480
- L = L * alpha + l_j
481
- M = m_j
482
-
328
if SLIDING_WINDOW:
329
qpos_lo = q_block_local_idx * BLOCK_Q
330
V = tl.where(
486
- (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0
331
+ (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW,
332
+ V,
333
+ 0.0,
334
)
488
-
489
- # acc : (BLOCK_M, HEAD_SIZE_PADDED)
490
- # Per-token-head quant: apply v_scale to P instead of V.
491
- if KV_QUANT_MODE >= 2:
335
+ if USE_PER_TOKEN_HEAD_SCALES:
336
+ # Per-token-head quant: apply v_scale to P instead of V.
337
P_v = (P * v_token_head_scales[None, :]).to(V.dtype)
338
acc += tl.dot(P_v, V)
339
else:
340
acc += tl.dot(P.to(V.dtype), V)
341
497
- # epilogue
498
- acc = acc / L[:, None]
499
- if USE_FP8:
500
- acc = acc * tl.load(out_scale)
501
- acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
502
-
503
- output_offset = (
504
- query_offset_0[:, None] * output_stride_0
505
- + query_offset_1[:, None] * output_stride_1
506
- + offs_d[None, :]
507
- )
508
-
509
- tl.store(
510
- output_ptr + output_offset,
511
- acc,
512
- mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
513
- )
514
-
515
-
516
-@triton.jit
517
-def kernel_unified_attention_3d(
518
- segm_output_ptr,
519
- # [num_tokens, num_query_heads, num_segments, head_size_padded]
520
- segm_max_ptr, # [num_tokens, num_query_heads, num_segments]
521
- segm_expsum_ptr, # [num_tokens, num_query_heads, num_segments]
522
- query_ptr, # [num_tokens, num_query_heads, head_size]
523
- key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]
524
- value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]
525
- sink_ptr, # [num_query_heads]
526
- block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
527
- seq_lens_ptr, # [num_seqs]
528
- alibi_slopes_ptr, # [num_query_heads]
529
- qq_bias_ptr, # [num_query_tokens, num_query_tokens]
530
- scale, # float32
531
- k_scale, # float32
532
- v_scale, # float32
533
- softcap, # float32
534
- num_query_heads: tl.constexpr, # int
535
- num_queries_per_kv: tl.constexpr, # int
536
- block_table_stride: tl.int64, # int
537
- query_stride_0: tl.int64, # int
538
- query_stride_1: tl.int64, # int, should be equal to head_size
539
- qq_bias_stride_0: tl.int64, # int
540
- BLOCK_SIZE: tl.constexpr, # int
541
- TILE_SIZE: tl.constexpr, # int, must be power of 2
542
- HEAD_SIZE: tl.constexpr, # int
543
- HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
544
- USE_ALIBI_SLOPES: tl.constexpr, # bool
545
- USE_ALIBI_SQRT: tl.constexpr, # bool
546
- USE_QQ_BIAS: tl.constexpr, # bool
547
- USE_SOFTCAP: tl.constexpr, # bool
548
- USE_SINKS: tl.constexpr, # bool
549
- SLIDING_WINDOW: tl.constexpr, # int
550
- stride_k_cache_0: tl.int64, # int
551
- stride_k_cache_1: tl.int64, # int
552
- stride_k_cache_2: tl.int64, # int
553
- stride_k_cache_3: tl.constexpr, # int
554
- stride_v_cache_0: tl.int64, # int
555
- stride_v_cache_1: tl.int64, # int
556
- stride_v_cache_2: tl.int64, # int
557
- stride_v_cache_3: tl.constexpr, # int
558
- query_start_len_ptr, # [num_seqs+1]
559
- BLOCK_Q: tl.constexpr, # int
560
- num_seqs: tl.int32,
561
- BLOCK_M: tl.constexpr, # int
562
- NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int
563
- USE_MM_PREFIX: tl.constexpr, # bool
564
- MAX_MM_RANGES: tl.constexpr, # int
565
- mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence
566
- # KV cache quantization: 0=none, 1=fp8, 2=per-token-head
567
- KV_QUANT_MODE: tl.constexpr = 0,
568
- # Per-token-head scale caches (KV_QUANT_MODE >= 2)
569
- # Shape: [num_blocks, block_size, num_kv_heads]
570
- k_scale_cache_ptr=None,
571
- v_scale_cache_ptr=None,
572
- stride_ks_blk=0,
573
- stride_ks_slot=0,
574
- stride_ks_head=0,
575
- stride_vs_blk=0,
576
- stride_vs_slot=0,
577
- stride_vs_head=0,
578
- CHUNK_LOOKBACK: tl.constexpr = -1,
579
- CHUNK_SIZE: tl.constexpr = -1,
580
-):
581
- q_block_global_idx = tl.program_id(0)
582
- kv_head_idx = tl.program_id(1)
583
- segm_idx = tl.program_id(2)
584
-
585
- seq_idx = find_seq_idx(
586
- query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True
587
- )
588
-
589
- q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx
590
-
591
- q_block_local_idx = q_block_global_idx - q_block_start_idx
592
-
593
- cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
594
- cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)
595
-
596
- cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index
597
-
598
- if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
599
- return
600
-
601
- # sequence len for this particular sequence
602
- seq_len = tl.load(seq_lens_ptr + seq_idx)
603
-
604
- # number of segments for this particular sequence
605
- num_segments = NUM_SEGMENTS_PER_SEQ
606
- tiles_per_segment = cdiv_fn(seq_len, num_segments * TILE_SIZE)
607
-
608
- if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
609
- return
610
-
611
- offs_m = tl.arange(0, BLOCK_M)
612
- offs_d = tl.arange(0, HEAD_SIZE_PADDED)
613
- offs_t = tl.arange(0, TILE_SIZE)
614
- query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
615
-
616
- query_offset_0 = cur_batch_in_all_start_index + query_pos
617
- query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv
618
- query_offset = (
619
- query_offset_0[:, None] * query_stride_0
620
- + query_offset_1[:, None] * query_stride_1
621
- + offs_d[None, :]
622
- )
623
-
624
- dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)
625
- query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
626
- query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
627
-
628
- # Q : (BLOCK_M, HEAD_SIZE_PADDED)
629
- Q = tl.load(
630
- query_ptr + query_offset,
631
- mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
632
- other=0.0,
633
- )
634
-
635
- block_table_offset = seq_idx * block_table_stride
636
-
637
- if USE_SINKS:
638
- if segm_idx == 0:
639
- M = tl.load(
640
- sink_ptr + query_offset_1,
641
- mask=query_mask_1,
642
- other=float("-inf"),
643
- ).to(dtype=tl.float32)
644
- else:
645
- M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
646
- else:
647
- M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
648
-
649
- L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
650
- acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
651
-
652
- # context length for this particular sequences
653
- context_len = seq_len - cur_batch_query_len
654
-
655
- # alibi slope for this head
656
- if USE_ALIBI_SLOPES:
657
- alibi_slope = tl.load(
658
- alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0
659
- )
660
-
661
- # query-query attention bias
662
- if USE_QQ_BIAS:
663
- qq_bias_row_ptrs = (
664
- qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0
665
- ) # shape: [BLOCK_M]
666
-
667
- # compute the length of the longest sequence prefix spanned by any
668
- # query token in the current q_block (q_block_local_idx)
669
- max_seq_prefix_len = (
670
- context_len
671
- + q_block_local_idx * BLOCK_Q
672
- + (BLOCK_M - 1) // num_queries_per_kv
673
- + 1
674
- )
675
-
676
- # adjust for potential padding in the last q_block by considering the
677
- # actual sequence length
678
- max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)
679
-
680
- # calculate the number of tiles that need to be processed to
681
- # cover the longest sequence prefix (due to causal masking, tiles beyond
682
- # this prefix can be skipped)
683
- num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)
684
-
685
- # ---- Sliding-window tile pruning --------------------
686
- # Default: keep previous global behavior
687
- tile_start = 0
688
- tile_end = num_tiles
689
- # TODO(Isotr0py): sliding window pruning with image bidirectional mask
690
- if SLIDING_WINDOW > 0 and not USE_MM_PREFIX:
691
- # Query rows covered by this Q-block
692
- qpos_lo = q_block_local_idx * BLOCK_Q
693
- qpos_hi = tl.minimum(
694
- qpos_lo + (BLOCK_M - 1) // num_queries_per_kv,
695
- cur_batch_query_len - 1,
696
- )
697
- # For sliding window, each query position q can only attend to
698
- # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs]
699
- # where q_abs = context_len + q
700
- # The union of allowed key positions for this Q-block is:
701
- # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
702
- q_abs = context_len + qpos_lo
703
- if CHUNK_LOOKBACK > -1:
704
- first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE
705
- else:
706
- first_allowed_key = q_abs - SLIDING_WINDOW + 1
707
- last_allowed_key = context_len + qpos_hi
708
- # Convert to tile indices and clamp
709
- tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
710
- tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles)
711
-
712
- # iterate through tiles (now limited to the sliding window range)
713
- for j in range(
714
- max(segm_idx * tiles_per_segment, tile_start),
715
- min((segm_idx + 1) * tiles_per_segment, tile_end),
716
- ):
717
- seq_offset = j * TILE_SIZE + offs_t
718
- tile_mask = seq_offset < max_seq_prefix_len
719
-
720
- physical_block_idx = tl.load(
721
- block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
722
- ).to(tl.int64)
723
-
724
- v_offset = (
725
- physical_block_idx[:, None] * stride_v_cache_0
726
- + kv_head_idx * stride_v_cache_2
727
- + offs_d[None, :] * stride_v_cache_3
728
- + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1
729
- )
730
-
731
- k_offset = (
732
- physical_block_idx[None, :] * stride_k_cache_0
733
- + kv_head_idx * stride_k_cache_2
734
- + offs_d[:, None] * stride_k_cache_3
735
- + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
342
+ # ---- Epilogue ---------------------------------------------------------
343
+ if IS_3D:
344
+ # Store per-segment partials; finalized by ``reduce_segments``.
345
+ segm_output_offset = (
346
+ query_offset_0[:, None].to(tl.int64)
347
+ * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
348
+ + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
349
+ + segm_idx * HEAD_SIZE_PADDED
350
+ + tl.arange(0, HEAD_SIZE_PADDED)[None, :]
351
)
737
-
738
- # K : (HEAD_SIZE, TILE_SIZE)
739
- K_load = tl.load(
740
- key_cache_ptr + k_offset,
741
- mask=dim_mask[:, None] & tile_mask[None, :],
742
- other=0.0,
352
+ tl.store(
353
+ segm_output_ptr + segm_output_offset,
354
+ acc,
355
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
356
)
744
- K, k_token_head_scales = _prepare_kv_tile(
745
- K_load,
746
- Q,
747
- k_scale,
748
- k_scale_cache_ptr,
749
- physical_block_idx,
750
- seq_offset,
751
- kv_head_idx,
752
- stride_ks_blk,
753
- stride_ks_slot,
754
- stride_ks_head,
755
- tile_mask,
756
- BLOCK_SIZE,
757
- KV_QUANT_MODE,
357
+ store_segm_reduce_scalars(
358
+ segm_max_ptr,
359
+ segm_expsum_ptr,
360
+ query_offset_0,
361
+ query_offset_1,
362
+ segm_idx,
363
+ M,
364
+ L,
365
+ query_mask_0,
366
+ query_mask_1,
367
+ num_query_heads,
368
+ NUM_SEGMENTS_PER_SEQ,
369
)
759
-
760
- # V : (TILE_SIZE, HEAD_SIZE)
761
- V_load = tl.load(
762
- value_cache_ptr + v_offset,
763
- mask=dim_mask[None, :] & tile_mask[:, None],
764
- other=0.0,
370
+ else:
371
+ acc = acc / L[:, None]
372
+ if USE_FP8:
373
+ acc = acc * tl.load(out_scale)
374
+ acc = tl.clamp(acc, FP8_MIN, FP8_MAX)
375
+ output_offset = (
376
+ query_offset_0[:, None] * output_stride_0
377
+ + query_offset_1[:, None] * output_stride_1
378
+ + offs_d[None, :]
379
)
766
- V, v_token_head_scales = _prepare_kv_tile(
767
- V_load,
768
- Q,
769
- v_scale,
770
- v_scale_cache_ptr,
771
- physical_block_idx,
772
- seq_offset,
773
- kv_head_idx,
774
- stride_vs_blk,
775
- stride_vs_slot,
776
- stride_vs_head,
777
- tile_mask,
778
- BLOCK_SIZE,
779
- KV_QUANT_MODE,
380
+ tl.store(
381
+ output_ptr + output_offset,
382
+ acc,
383
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
384
)
385
782
- # Compute attention mask: causal by default (key <= query)
783
- query_abs_pos = context_len + query_pos[:, None]
784
- seq_mask = seq_offset[None, :] <= query_abs_pos
785
-
786
- # Apply sliding window / chunked attention to base mask
787
- # BEFORE mm_prefix OR.
788
- # Order must match FlexAttention:
789
- # (causal AND sliding_window) OR mm_prefix
790
- if CHUNK_LOOKBACK > -1:
791
- seq_mask = seq_mask & (
792
- (
793
- (context_len + query_pos[:, None]) // CHUNK_SIZE
794
- - (seq_offset[None, :] // CHUNK_SIZE)
795
- )
796
- <= CHUNK_LOOKBACK
797
- )
798
- elif SLIDING_WINDOW > 0:
799
- seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)
800
-
801
- # PrefixLM: extend mask with bidirectional ranges for multimodal tokens.
802
- # Applied AFTER sliding window so mm_prefix ranges override SW restriction.
803
- if USE_MM_PREFIX:
804
- for i in range(MAX_MM_RANGES):
805
- range_start = tl.load(
806
- mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2
807
- )
808
- range_end = tl.load(
809
- mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1
810
- )
811
-
812
- is_valid = range_start < range_end
813
- q_in_range = (
814
- (query_abs_pos >= range_start)
815
- & (query_abs_pos <= range_end)
816
- & is_valid
817
- )
818
- k_in_range = (
819
- (seq_offset[None, :] >= range_start)
820
- & (seq_offset[None, :] <= range_end)
821
- & is_valid
822
- )
823
- seq_mask |= q_in_range & k_in_range
824
-
825
- # S : (BLOCK_M, TILE_SIZE)
826
- S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)
827
-
828
- # Per-token-head quant: fuse softmax_scale with per-head k_scale
829
- # to avoid a separate BLOCK_M × TILE_SIZE multiply on S.
830
- if KV_QUANT_MODE >= 2:
831
- S += tl.dot(Q, K) * (scale * k_token_head_scales[None, :])
832
- else:
833
- S += scale * tl.dot(Q, K)
834
-
835
- if USE_SOFTCAP:
836
- S = apply_softcap(S, softcap)
837
-
838
- S = tl.where(
839
- query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf")
840
- )
841
-
842
- if USE_ALIBI_SLOPES:
843
- if USE_ALIBI_SQRT:
844
- relative_pos = seq_offset - (context_len + query_pos[:, None])
845
- alibi_offset = tl.where(
846
- relative_pos <= 0,
847
- -tl.sqrt((-relative_pos).to(tl.float32)),
848
- 0.0,
849
- )
850
- else:
851
- alibi_offset = seq_offset - context_len
852
- S += alibi_slope[:, None] * alibi_offset
853
-
854
- if USE_QQ_BIAS:
855
- # compute key positions relative to query section
856
- key_rel_pos = seq_offset - context_len # shape: [BLOCK_SIZE]
857
- # load bias only for keys that correspond to queries
858
- is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0
859
- qq_bias = tl.load(
860
- qq_bias_row_ptrs + key_rel_pos[None, :],
861
- mask=is_query_key[None, :], # avoid OOB for context keys
862
- other=0.0,
863
- )
864
- S += qq_bias
865
-
866
- # compute running maximum
867
- # m_j : (BLOCK_M,)
868
- m_j = tl.maximum(M, tl.max(S, axis=1))
869
-
870
- # For sliding window there's a chance the max is -inf due to masking of
871
- # the entire row. In this case we need to set m_j 0 to avoid NaN
872
- m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
873
-
874
- # P : (BLOCK_M, TILE_SIZE,)
875
- P = tl.exp(S - m_j[:, None])
876
-
877
- # l_j : (BLOCK_M,)
878
- l_j = tl.sum(P, axis=1)
879
-
880
- # alpha : (BLOCK_M, )
881
- alpha = tl.exp(M - m_j)
882
-
883
- # acc : (BLOCK_M, HEAD_SIZE_PADDED)
884
- acc = acc * alpha[:, None]
885
-
886
- # update constants
887
- L = L * alpha + l_j
888
- M = m_j
889
-
890
- if SLIDING_WINDOW:
891
- qpos_lo = q_block_local_idx * BLOCK_Q
892
- V = tl.where(
893
- (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, V, 0.0
894
- )
895
-
896
- # acc : (BLOCK_M, HEAD_SIZE_PADDED)
897
- # Per-token-head quant: apply v_scale to P instead of V.
898
- if KV_QUANT_MODE >= 2:
899
- P_v = (P * v_token_head_scales[None, :]).to(V.dtype)
900
- acc += tl.dot(P_v, V)
901
- else:
902
- acc += tl.dot(P.to(V.dtype), V)
903
-
904
- segm_output_offset = (
905
- query_offset_0[:, None].to(tl.int64)
906
- * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
907
- + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)
908
- + segm_idx * HEAD_SIZE_PADDED
909
- + tl.arange(0, HEAD_SIZE_PADDED)[None, :]
910
- )
911
- tl.store(
912
- segm_output_ptr + segm_output_offset,
913
- acc,
914
- mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
915
- )
916
- segm_offset = (
917
- query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
918
- + query_offset_1 * NUM_SEGMENTS_PER_SEQ
919
- + segm_idx
920
- )
921
- tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1)
922
- tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1)
923
-
386
387
@triton.jit
388
def reduce_segments(
@@ -1028,12 +490,7 @@ def _get_tile_size(
490
element_size: int,
491
is_prefill: bool,
492
) -> int:
1031
- """Select tile size with Gemma3-specific optimization.
1032
-
1033
- For Gemma3, use 32 for both prefill and decode to better utilize
1034
- the larger head dimension (128/256). For other models, use
1035
- the default vLLM behavior.
1036
- """
493
+ """Select tile size with Gemma3-specific optimization."""
494
if _is_gemma3_attention(head_size, sliding_window):
495
# Gemma3: use 32 for decode (default is 16)
496
return 32
@@ -1041,6 +498,7 @@ def _get_tile_size(
498
# Default behavior
499
if is_prefill:
500
return 32
501
+ # Note: tile size must be at least 32 for fp8 (element_size == 1).
502
return 16 if element_size >= 2 else 32
503
504
@@ -1087,6 +545,15 @@ def unified_attention(
545
if sinks is not None:
546
assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"
547
548
+ use_per_token_head_scales = kv_quant_mode in (
549
+ KVQuantMode.INT8_PER_TOKEN_HEAD,
550
+ KVQuantMode.FP8_PER_TOKEN_HEAD,
551
+ )
552
+ if use_per_token_head_scales:
553
+ assert k_scale_cache is not None and v_scale_cache is not None, (
554
+ f"{kv_quant_mode.name} requires k_scale_cache / v_scale_cache"
555
+ )
556
+
557
use_mm_prefix = False
558
max_mm_ranges = 0
559
if mm_prefix_range is not None:
@@ -1124,8 +591,6 @@ def unified_attention(
591
# = floor(q.shape[0] / BLOCK_Q) + num_seqs
592
total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
593
1127
- # Tile sizes for prefill and decode. Gemma3 models use optimized values.
1128
- # Note: tile size must be at least 32 for fp8 (element_size == 1).
594
sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
595
596
# Compute chunked block size from sliding window if needed.
@@ -1137,16 +602,10 @@ def unified_attention(
602
chunk_lookback = -1
603
604
TILE_SIZE_PREFILL = _get_tile_size(
1140
- head_size,
1141
- sliding_window_val,
1142
- q.element_size(),
1143
- is_prefill=True,
605
+ head_size, sliding_window_val, q.element_size(), is_prefill=True
606
)
607
TILE_SIZE_DECODE = _get_tile_size(
1146
- head_size,
1147
- sliding_window_val,
1148
- q.element_size(),
1149
- is_prefill=False,
608
+ head_size, sliding_window_val, q.element_size(), is_prefill=False
609
)
610
611
# Launch the 2D kernel if
@@ -1154,7 +613,7 @@ def unified_attention(
613
# 2. The batch includes at least one prefill request, or
614
# 3. The number of sequences exceeds the configured threshold, or
615
# 4. Batch invariance is enabled
1157
- if (
616
+ use_3d = not (
617
seq_threshold_3D is None
618
or num_par_softmax_segments is None
619
or softmax_segm_output is None
@@ -1163,136 +622,110 @@ def unified_attention(
622
or max_seqlen_q > 1
623
or num_seqs > seq_threshold_3D
624
or is_batch_invariant
1166
- ):
1167
- kernel_unified_attention_2d[
1168
- (
1169
- total_num_q_blocks,
1170
- num_kv_heads,
1171
- )
1172
- ](
1173
- output_ptr=out,
1174
- query_ptr=q,
1175
- key_cache_ptr=k,
1176
- value_cache_ptr=v,
1177
- sink_ptr=sinks,
1178
- block_tables_ptr=block_table,
1179
- seq_lens_ptr=seqused_k,
1180
- alibi_slopes_ptr=alibi_slopes,
1181
- qq_bias_ptr=qq_bias,
1182
- scale=softmax_scale,
1183
- k_scale=k_descale,
1184
- v_scale=v_descale,
1185
- out_scale=1 / output_scale if output_scale is not None else 1.0,
1186
- softcap=softcap,
1187
- num_query_heads=num_query_heads,
1188
- num_queries_per_kv=num_queries_per_kv,
1189
- block_table_stride=block_table.stride(0),
1190
- query_stride_0=q.stride(0),
1191
- query_stride_1=q.stride(1),
1192
- output_stride_0=out.stride(0),
1193
- output_stride_1=out.stride(1),
1194
- qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
1195
- BLOCK_SIZE=block_size,
1196
- TILE_SIZE=TILE_SIZE_PREFILL,
1197
- HEAD_SIZE=head_size,
1198
- HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
1199
- USE_ALIBI_SLOPES=use_alibi_slopes,
1200
- USE_ALIBI_SQRT=use_alibi_sqrt,
1201
- USE_QQ_BIAS=use_qq_bias,
1202
- USE_SOFTCAP=(softcap > 0),
1203
- USE_SINKS=(sinks is not None),
1204
- USE_MM_PREFIX=use_mm_prefix,
1205
- MAX_MM_RANGES=max_mm_ranges,
1206
- mm_prefix_range_ptr=mm_prefix_range,
1207
- SLIDING_WINDOW=(1 + window_size[0]),
1208
- stride_k_cache_0=k.stride(0),
1209
- stride_k_cache_1=k.stride(1),
1210
- stride_k_cache_2=k.stride(2),
1211
- stride_k_cache_3=k.stride(3),
1212
- stride_v_cache_0=v.stride(0),
1213
- stride_v_cache_1=v.stride(1),
1214
- stride_v_cache_2=v.stride(2),
1215
- stride_v_cache_3=v.stride(3),
1216
- query_start_len_ptr=cu_seqlens_q,
1217
- BLOCK_Q=BLOCK_Q,
1218
- num_seqs=num_seqs,
1219
- BLOCK_M=BLOCK_M,
1220
- USE_FP8=output_scale is not None,
1221
- KV_QUANT_MODE=kv_quant_mode,
1222
- k_scale_cache_ptr=k_scale_cache,
1223
- v_scale_cache_ptr=v_scale_cache,
1224
- stride_ks_blk=k_scale_cache.stride(0) if k_scale_cache is not None else 0,
1225
- stride_ks_slot=k_scale_cache.stride(1) if k_scale_cache is not None else 0,
1226
- stride_ks_head=k_scale_cache.stride(2) if k_scale_cache is not None else 0,
1227
- stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0,
1228
- stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0,
1229
- stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0,
1230
- CHUNK_LOOKBACK=chunk_lookback,
1231
- CHUNK_SIZE=chunk_size,
1232
- )
625
+ )
626
+
627
+ # The kernel signature is the same for 2D and 3D — only the launch
628
+ # grid + a handful of constexpr toggles differ. Per-token-head scale
629
+ # caches and their strides are required arguments; non-per-token-head
630
+ # modes pass dummy zeros (the code path is dead-code eliminated by
631
+ # the ``USE_PER_TOKEN_HEAD_SCALES`` constexpr branch in the kernel).
632
+ if use_per_token_head_scales:
633
+ ks_strides = k_scale_cache.stride()
634
+ vs_strides = v_scale_cache.stride()
635
+ ks_blk, ks_slot, ks_head = ks_strides[0], ks_strides[1], ks_strides[2]
636
+ vs_blk, vs_slot, vs_head = vs_strides[0], vs_strides[1], vs_strides[2]
637
+ k_scale_ptr = k_scale_cache
638
+ v_scale_ptr = v_scale_cache
639
else:
1234
- kernel_unified_attention_3d[
1235
- (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
1236
- ](
1237
- segm_output_ptr=softmax_segm_output,
1238
- segm_max_ptr=softmax_segm_max,
1239
- segm_expsum_ptr=softmax_segm_expsum,
1240
- query_ptr=q,
1241
- key_cache_ptr=k,
1242
- value_cache_ptr=v,
1243
- sink_ptr=sinks,
1244
- block_tables_ptr=block_table,
1245
- seq_lens_ptr=seqused_k,
1246
- alibi_slopes_ptr=alibi_slopes,
1247
- qq_bias_ptr=qq_bias,
1248
- scale=softmax_scale,
1249
- k_scale=k_descale,
1250
- v_scale=v_descale,
1251
- softcap=softcap,
1252
- num_query_heads=num_query_heads,
1253
- num_queries_per_kv=num_queries_per_kv,
1254
- block_table_stride=block_table.stride(0),
1255
- query_stride_0=q.stride(0),
1256
- query_stride_1=q.stride(1),
1257
- qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
1258
- BLOCK_SIZE=block_size,
1259
- TILE_SIZE=TILE_SIZE_DECODE,
1260
- HEAD_SIZE=head_size,
1261
- HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
1262
- USE_ALIBI_SLOPES=use_alibi_slopes,
1263
- USE_ALIBI_SQRT=use_alibi_sqrt,
1264
- USE_QQ_BIAS=use_qq_bias,
1265
- USE_SOFTCAP=(softcap > 0),
1266
- USE_SINKS=(sinks is not None),
1267
- USE_MM_PREFIX=use_mm_prefix,
1268
- MAX_MM_RANGES=max_mm_ranges,
1269
- mm_prefix_range_ptr=mm_prefix_range,
1270
- SLIDING_WINDOW=(1 + window_size[0]),
1271
- stride_k_cache_0=k.stride(0),
1272
- stride_k_cache_1=k.stride(1),
1273
- stride_k_cache_2=k.stride(2),
1274
- stride_k_cache_3=k.stride(3),
1275
- stride_v_cache_0=v.stride(0),
1276
- stride_v_cache_1=v.stride(1),
1277
- stride_v_cache_2=v.stride(2),
1278
- stride_v_cache_3=v.stride(3),
1279
- query_start_len_ptr=cu_seqlens_q,
1280
- BLOCK_Q=BLOCK_Q,
1281
- num_seqs=num_seqs,
1282
- BLOCK_M=BLOCK_M,
1283
- NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
1284
- KV_QUANT_MODE=kv_quant_mode,
1285
- k_scale_cache_ptr=k_scale_cache,
1286
- v_scale_cache_ptr=v_scale_cache,
1287
- stride_ks_blk=k_scale_cache.stride(0) if k_scale_cache is not None else 0,
1288
- stride_ks_slot=k_scale_cache.stride(1) if k_scale_cache is not None else 0,
1289
- stride_ks_head=k_scale_cache.stride(2) if k_scale_cache is not None else 0,
1290
- stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0,
1291
- stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0,
1292
- stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0,
1293
- CHUNK_LOOKBACK=chunk_lookback,
1294
- CHUNK_SIZE=chunk_size,
1295
- )
640
+ ks_blk = ks_slot = ks_head = 0
641
+ vs_blk = vs_slot = vs_head = 0
642
+ # Pass the K cache as a stand-in pointer; never dereferenced.
643
+ k_scale_ptr = k
644
+ v_scale_ptr = v
645
+
646
+ # 3D needs real segm tensors; 2D never touches them but Triton wants
647
+ # a non-null pointer. Reuse ``out`` as the placeholder.
648
+ segm_output_ptr = softmax_segm_output if use_3d else out
649
+ segm_max_ptr = softmax_segm_max if use_3d else out
650
+ segm_expsum_ptr = softmax_segm_expsum if use_3d else out
651
+ num_segments = num_par_softmax_segments if use_3d else 1
652
+
653
+ grid: tuple[Any, ...]
654
+ if not use_3d:
655
+ grid = (total_num_q_blocks, num_kv_heads)
656
+ tile_size = TILE_SIZE_PREFILL
657
+ else:
658
+ grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
659
+ tile_size = TILE_SIZE_DECODE
660
+
661
+ kernel_unified_attention[grid](
662
+ output_ptr=out,
663
+ segm_output_ptr=segm_output_ptr,
664
+ segm_max_ptr=segm_max_ptr,
665
+ segm_expsum_ptr=segm_expsum_ptr,
666
+ query_ptr=q,
667
+ key_cache_ptr=k,
668
+ value_cache_ptr=v,
669
+ sink_ptr=sinks,
670
+ block_tables_ptr=block_table,
671
+ seq_lens_ptr=seqused_k,
672
+ alibi_slopes_ptr=alibi_slopes,
673
+ qq_bias_ptr=qq_bias,
674
+ k_scale_cache_ptr=k_scale_ptr,
675
+ v_scale_cache_ptr=v_scale_ptr,
676
+ scale=softmax_scale,
677
+ k_scale=k_descale,
678
+ v_scale=v_descale,
679
+ out_scale=1 / output_scale if output_scale is not None else 1.0,
680
+ softcap=softcap,
681
+ num_query_heads=num_query_heads,
682
+ num_queries_per_kv=num_queries_per_kv,
683
+ block_table_stride=block_table.stride(0),
684
+ query_stride_0=q.stride(0),
685
+ query_stride_1=q.stride(1),
686
+ output_stride_0=out.stride(0),
687
+ output_stride_1=out.stride(1),
688
+ qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,
689
+ BLOCK_SIZE=block_size,
690
+ TILE_SIZE=tile_size,
691
+ HEAD_SIZE=head_size,
692
+ HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
693
+ USE_ALIBI_SLOPES=use_alibi_slopes,
694
+ USE_ALIBI_SQRT=use_alibi_sqrt,
695
+ USE_QQ_BIAS=use_qq_bias,
696
+ USE_SOFTCAP=(softcap > 0),
697
+ USE_SINKS=(sinks is not None),
698
+ USE_MM_PREFIX=use_mm_prefix,
699
+ MAX_MM_RANGES=max_mm_ranges,
700
+ mm_prefix_range_ptr=mm_prefix_range,
701
+ SLIDING_WINDOW=(1 + window_size[0]),
702
+ stride_k_cache_0=k.stride(0),
703
+ stride_k_cache_1=k.stride(1),
704
+ stride_k_cache_2=k.stride(2),
705
+ stride_k_cache_3=k.stride(3),
706
+ stride_v_cache_0=v.stride(0),
707
+ stride_v_cache_1=v.stride(1),
708
+ stride_v_cache_2=v.stride(2),
709
+ stride_v_cache_3=v.stride(3),
710
+ stride_ks_blk=ks_blk,
711
+ stride_ks_slot=ks_slot,
712
+ stride_ks_head=ks_head,
713
+ stride_vs_blk=vs_blk,
714
+ stride_vs_slot=vs_slot,
715
+ stride_vs_head=vs_head,
716
+ query_start_len_ptr=cu_seqlens_q,
717
+ BLOCK_Q=BLOCK_Q,
718
+ num_seqs=num_seqs,
719
+ BLOCK_M=BLOCK_M,
720
+ NUM_SEGMENTS_PER_SEQ=num_segments,
721
+ USE_FP8=output_scale is not None,
722
+ IS_3D=use_3d,
723
+ KV_QUANT_MODE=kv_quant_mode,
724
+ CHUNK_LOOKBACK=chunk_lookback,
725
+ CHUNK_SIZE=chunk_size,
726
+ )
727
+
728
+ if use_3d:
729
reduce_segments[(q.shape[0], num_query_heads)](
730
output_ptr=out,
731
segm_output_ptr=softmax_segm_output,