| 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)`` computed without a direct ``tanh`` call. |
| 32 | |
| 33 | Use ``tanh(y) = 1 - 2 / (exp(2y) + 1)`` rather than |
| 34 | ``(exp(y) - exp(-y)) / (exp(y) + exp(-y))``. The symmetric-exponential |
| 35 | form overflows to ``inf`` for ``|y| > ~88`` and then evaluates to |
| 36 | ``inf / inf = NaN``, poisoning the attention row. The single-exponential |
| 37 | form saturates cleanly at +/-1 for either sign. |
| 38 | """ |
| 39 | y = S / x |
| 40 | return x * (1.0 - 2.0 / (tl.exp(2.0 * y) + 1.0)) |
| 41 | |
| 42 | |
| 43 | # =========================================================================== |
| 44 | # Attention loop |
| 45 | # =========================================================================== |
| 46 | |
| 47 | |
| 48 | @triton.jit |
| 49 | def resolve_seq_and_query_len( |
| 50 | query_start_len_ptr, |
| 51 | seq_lens_ptr, |
| 52 | q_block_global_idx, |
| 53 | num_seqs, |
| 54 | BLOCK_Q: tl.constexpr, |
| 55 | ): |
| 56 | """Resolve the (sequence, q-block-within-sequence) pair and load the |
| 57 | per-sequence lengths. |
| 58 | |
| 59 | Shared across every attention kernel — the ``q_block_global_idx`` |
| 60 | program id indexes into the flattened ``(seq, q_block_in_seq)`` |
| 61 | space, and a binary search over ``query_start_len_ptr`` recovers |
| 62 | the (seq, local-q-block) pair. |
| 63 | |
| 64 | Returns ``(seq_idx, q_block_local_idx, cur_batch_in_all_start_index, |
| 65 | cur_batch_query_len, seq_len)``. Callers must still early-return |
| 66 | when ``q_block_local_idx * BLOCK_Q >= cur_batch_query_len`` (Triton |
| 67 | helpers cannot return from the caller). |
| 68 | """ |
| 69 | # find_seq_idx is defined below; forward use is fine inside @triton.jit. |
| 70 | seq_idx = find_seq_idx( |
| 71 | query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True |
| 72 | ) |
| 73 | q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx |
| 74 | q_block_local_idx = q_block_global_idx - q_block_start_idx |
| 75 | cur_start = tl.load(query_start_len_ptr + seq_idx) |
| 76 | cur_stop = tl.load(query_start_len_ptr + seq_idx + 1) |
| 77 | cur_batch_query_len = cur_stop - cur_start |
| 78 | seq_len = tl.load(seq_lens_ptr + seq_idx) |
| 79 | return seq_idx, q_block_local_idx, cur_start, cur_batch_query_len, seq_len |
| 80 | |
| 81 | |
| 82 | @triton.jit |
| 83 | def find_seq_idx( |
| 84 | query_start_len_ptr, |
| 85 | target_idx, |
| 86 | num_seqs, |
| 87 | BLOCK_Q: tl.constexpr, |
| 88 | use_q_block_mode: tl.constexpr, |
| 89 | ): |
| 90 | """Binary search over the cumulative query-length prefix. |
| 91 | |
| 92 | When ``use_q_block_mode`` is True, the prefix values are reshaped |
| 93 | into units of ``BLOCK_Q`` plus one entry per boundary — matching |
| 94 | the q-block grid laid out by the attention kernels. When False |
| 95 | we search the plain cumulative-length prefix (used by |
| 96 | ``reduce_segments`` which iterates over raw query tokens). |
| 97 | """ |
| 98 | left: tl.int32 = 0 |
| 99 | right = num_seqs |
| 100 | while left < right: |
| 101 | mid = (left + right) // 2 |
| 102 | val = tl.load(query_start_len_ptr + mid) |
| 103 | mid_val = val // BLOCK_Q + mid if use_q_block_mode else val |
| 104 | |
| 105 | if mid_val <= target_idx: |
| 106 | left = mid + 1 |
| 107 | else: |
| 108 | right = mid |
| 109 | |
| 110 | return left - 1 |
| 111 | |
| 112 | |
| 113 | @triton.jit |
| 114 | def init_softmax_M( |
| 115 | sink_ptr, |
| 116 | query_offset_1, |
| 117 | query_mask_1, |
| 118 | segm_idx_or_0, |
| 119 | BLOCK_M: tl.constexpr, |
| 120 | USE_SINKS: tl.constexpr, |
| 121 | IS_3D: tl.constexpr, |
| 122 | ): |
| 123 | """Initial row-max ``M`` for the online softmax. |
| 124 | |
| 125 | Without sinks: ``-inf``. With sinks: load the per-head sink bias |
| 126 | once. In 3D mode only segment 0 loads — ``reduce_segments`` adds |
| 127 | the sink contribution exactly once across segments, so other |
| 128 | segments must start from ``-inf``. |
| 129 | |
| 130 | ``segm_idx_or_0`` is the 3D segment index or 0 for 2D (caller |
| 131 | passes ``0`` when ``IS_3D`` is False). |
| 132 | """ |
| 133 | M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) |
| 134 | if USE_SINKS: |
| 135 | load_sinks = (not IS_3D) or (segm_idx_or_0 == 0) |
| 136 | if load_sinks: |
| 137 | M = tl.load( |
| 138 | sink_ptr + query_offset_1, |
| 139 | mask=query_mask_1, |
| 140 | other=float("-inf"), |
| 141 | ).to(tl.float32) |
| 142 | return M |
| 143 | |
| 144 | |
| 145 | @triton.jit |
| 146 | def compute_tile_loop_bounds( |
| 147 | context_len, |
| 148 | seq_len, |
| 149 | cur_batch_query_len, |
| 150 | q_block_local_idx, |
| 151 | segm_idx_or_0, |
| 152 | tiles_per_segment_or_0, |
| 153 | TILE_SIZE: tl.constexpr, |
| 154 | BLOCK_M: tl.constexpr, |
| 155 | BLOCK_Q: tl.constexpr, |
| 156 | num_queries_per_kv: tl.constexpr, |
| 157 | SLIDING_WINDOW: tl.constexpr, |
| 158 | USE_MM_PREFIX: tl.constexpr, |
| 159 | IS_3D: tl.constexpr, |
| 160 | USE_CAUSAL: tl.constexpr = True, |
| 161 | USE_PER_SEQ_CAUSAL: tl.constexpr = False, |
| 162 | CHUNK_LOOKBACK: tl.constexpr = -1, |
| 163 | CHUNK_SIZE: tl.constexpr = -1, |
| 164 | USE_R_SWA: tl.constexpr = False, |
| 165 | MM_PREFIX_CLAMP_SW: tl.constexpr = False, |
| 166 | MAX_MM_RANGES: tl.constexpr = 0, |
| 167 | mm_prefix_range_ptr=None, |
| 168 | seq_idx=0, |
| 169 | ): |
| 170 | """Compute the tile-loop bounds ``(loop_lo, loop_hi)`` and the |
| 171 | derived ``max_seq_prefix_len`` used for per-tile masking. |
| 172 | |
| 173 | Combines three concerns into one helper: |
| 174 | |
| 175 | 1. Longest prefix spanned by any query token in this q-block. |
| 176 | Clamped to ``seq_len`` (causal) or extended to it when |
| 177 | mm_prefix is active or non-causal sequences need the full |
| 178 | sequence. |
| 179 | 2. Sliding-window pruning: narrows ``[tile_start, tile_end)`` to |
| 180 | only tiles that can contain an allowed key under SWA. For non-causal |
| 181 | sequences, the window extends in both directions. For Gemma4's |
| 182 | window-clamped multimodal prefix mask, the bounds include the union of |
| 183 | the base mask and each image range intersecting the query block. |
| 184 | 3. 3D scoping: when ``IS_3D`` is True, further narrows to the |
| 185 | segment's slice via ``(segm_idx * tiles_per_segment, |
| 186 | (segm_idx + 1) * tiles_per_segment)``. |
| 187 | """ |
| 188 | # compute the length of the longest sequence prefix spanned by any |
| 189 | # query token in the current q_block (q_block_local_idx) |
| 190 | max_seq_prefix_len = ( |
| 191 | context_len |
| 192 | + q_block_local_idx * BLOCK_Q |
| 193 | + (BLOCK_M - 1) // num_queries_per_kv |
| 194 | + 1 |
| 195 | ) |
| 196 | if USE_MM_PREFIX or USE_R_SWA or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): |
| 197 | # Read the full sequence but never past seq_len: the causal-style |
| 198 | # formula above can overshoot for non-causal sequences, and slots |
| 199 | # >= seq_len are unwritten KV (last-block tail) that may hold NaN |
| 200 | # (0 * NaN poisons the output). Per-element masking in |
| 201 | # compute_kv_seq_mask handles the causal/non-causal boundary. |
| 202 | max_seq_prefix_len = seq_len |
| 203 | else: |
| 204 | max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) |
| 205 | |
| 206 | num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) |
| 207 | |
| 208 | # ---- Sliding-window tile pruning -------------------- |
| 209 | # Default: keep previous global behavior |
| 210 | tile_start = 0 |
| 211 | tile_end = num_tiles |
| 212 | # Prefix ranges normally override the sliding window, so they require the |
| 213 | # complete sequence. Gemma4 instead clamps prefix attention to the left |
| 214 | # edge of the sliding window. In that case, the union of possible keys is |
| 215 | # bounded by the normal left window and the end of any multimodal range |
| 216 | # intersecting this query block. |
| 217 | can_prune_sliding = (not USE_R_SWA) and ((not USE_MM_PREFIX) or MM_PREFIX_CLAMP_SW) |
| 218 | if SLIDING_WINDOW > 0 and can_prune_sliding: |
| 219 | # Query rows covered by this Q-block |
| 220 | qpos_lo = q_block_local_idx * BLOCK_Q |
| 221 | qpos_hi = tl.minimum( |
| 222 | qpos_lo + (BLOCK_M - 1) // num_queries_per_kv, |
| 223 | cur_batch_query_len - 1, |
| 224 | ) |
| 225 | # For sliding window, each query position q can only attend to |
| 226 | # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs] |
| 227 | # where q_abs = context_len + q |
| 228 | # The union of allowed key positions for this Q-block is: |
| 229 | # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi] |
| 230 | q_abs = context_len + qpos_lo |
| 231 | if CHUNK_LOOKBACK > -1: |
| 232 | first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE |
| 233 | else: |
| 234 | first_allowed_key = q_abs - SLIDING_WINDOW + 1 |
| 235 | if USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): |
| 236 | # Non-causal: keys can be AHEAD of query within the window |
| 237 | last_allowed_key = context_len + qpos_hi + SLIDING_WINDOW - 1 |
| 238 | else: |
| 239 | last_allowed_key = context_len + qpos_hi |
| 240 | if USE_MM_PREFIX and MM_PREFIX_CLAMP_SW: |
| 241 | query_abs_lo = context_len + qpos_lo |
| 242 | query_abs_hi = context_len + qpos_hi |
| 243 | for i in range(MAX_MM_RANGES): |
| 244 | range_start = tl.load( |
| 245 | mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 |
| 246 | ) |
| 247 | range_end = tl.load( |
| 248 | mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1 |
| 249 | ) |
| 250 | intersects_query_block = ( |
| 251 | (range_start < range_end) |
| 252 | & (range_start <= query_abs_hi) |
| 253 | & (range_end >= query_abs_lo) |
| 254 | ) |
| 255 | mm_first_allowed_key = tl.maximum( |
| 256 | range_start, query_abs_lo - SLIDING_WINDOW + 1 |
| 257 | ) |
| 258 | first_allowed_key = tl.minimum( |
| 259 | first_allowed_key, |
| 260 | tl.where( |
| 261 | intersects_query_block, |
| 262 | mm_first_allowed_key, |
| 263 | first_allowed_key, |
| 264 | ), |
| 265 | ) |
| 266 | last_allowed_key = tl.maximum( |
| 267 | last_allowed_key, |
| 268 | tl.where(intersects_query_block, range_end, last_allowed_key), |
| 269 | ) |
| 270 | last_allowed_key = tl.minimum(last_allowed_key, seq_len - 1) |
| 271 | # Convert to tile indices and clamp |
| 272 | tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) |
| 273 | tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) |
| 274 | |
| 275 | if IS_3D: |
| 276 | loop_lo = max(segm_idx_or_0 * tiles_per_segment_or_0, tile_start) |
| 277 | loop_hi = min((segm_idx_or_0 + 1) * tiles_per_segment_or_0, tile_end) |
| 278 | else: |
| 279 | loop_lo = tile_start |
| 280 | loop_hi = tile_end |
| 281 | |
| 282 | return loop_lo, loop_hi, max_seq_prefix_len |
| 283 | |
| 284 | |
| 285 | @triton.jit |
| 286 | def store_segm_reduce_scalars( |
| 287 | segm_max_ptr, |
| 288 | segm_expsum_ptr, |
| 289 | query_offset_0, |
| 290 | query_offset_1, |
| 291 | segm_idx, |
| 292 | M, |
| 293 | L, |
| 294 | query_mask_0, |
| 295 | query_mask_1, |
| 296 | num_query_heads: tl.constexpr, |
| 297 | NUM_SEGMENTS_PER_SEQ: tl.constexpr, |
| 298 | ): |
| 299 | """Store per-segment ``M`` and ``L`` for ``reduce_segments`` to |
| 300 | combine into the final softmax. |
| 301 | |
| 302 | Shared across every 3D attention epilogue; the per-token output |
| 303 | stripes are mode-specific (flat / 2-stream split / 4-stream split) |
| 304 | and stay inlined. |
| 305 | """ |
| 306 | segm_offset = ( |
| 307 | query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) |
| 308 | + query_offset_1 * NUM_SEGMENTS_PER_SEQ |
| 309 | + segm_idx |
| 310 | ) |
| 311 | tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1) |
| 312 | tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1) |
| 313 | |
| 314 | |
| 315 | @triton.jit |
| 316 | def compute_kv_seq_mask( |
| 317 | query_abs_pos, |
| 318 | seq_offset, |
| 319 | seq_idx, |
| 320 | seq_len, |
| 321 | mm_prefix_range_ptr, |
| 322 | SLIDING_WINDOW: tl.constexpr, |
| 323 | USE_MM_PREFIX: tl.constexpr, |
| 324 | MAX_MM_RANGES: tl.constexpr, |
| 325 | USE_CAUSAL: tl.constexpr = True, |
| 326 | USE_PER_SEQ_CAUSAL: tl.constexpr = False, |
| 327 | per_seq_causal_ptr=None, |
| 328 | rswa_prefix_lens_ptr=None, |
| 329 | R_SWA_WINDOW: tl.constexpr = 0, |
| 330 | USE_R_SWA: tl.constexpr = False, |
| 331 | CHUNK_LOOKBACK: tl.constexpr = -1, |
| 332 | CHUNK_SIZE: tl.constexpr = -1, |
| 333 | MM_PREFIX_CLAMP_SW: tl.constexpr = False, |
| 334 | ): |
| 335 | """Build the KV mask for one tile. |
| 336 | |
| 337 | Causal (key <= query) by default; AND-ed with either chunked |
| 338 | attention (``CHUNK_LOOKBACK >= 0``) or sliding window |
| 339 | (``SLIDING_WINDOW > 0``); OR-ed with the bidirectional ranges from |
| 340 | ``mm_prefix_range`` when PrefixLM / multimodal attention is active. |
| 341 | Order matches FlexAttention: ``(causal AND window) OR mm_prefix``. |
| 342 | Chunked attention takes precedence over sliding window when both |
| 343 | are non-default — the launcher zeros ``CHUNK_LOOKBACK`` whenever |
| 344 | sliding window is disabled. |
| 345 | |
| 346 | When ``USE_PER_SEQ_CAUSAL`` is set, each sequence carries its own |
| 347 | causal flag via ``per_seq_causal_ptr``; non-causal sequences use a |
| 348 | simple ``key < seq_len`` bound instead. ``USE_CAUSAL=False`` |
| 349 | disables causal masking entirely. |
| 350 | """ |
| 351 | if USE_PER_SEQ_CAUSAL: |
| 352 | is_causal = tl.load(per_seq_causal_ptr + seq_idx) |
| 353 | seq_mask = tl.where( |
| 354 | is_causal, |
| 355 | seq_offset[None, :] <= query_abs_pos, |
| 356 | seq_offset[None, :] < seq_len, |
| 357 | ) |
| 358 | elif USE_CAUSAL: |
| 359 | seq_mask = seq_offset[None, :] <= query_abs_pos |
| 360 | else: |
| 361 | seq_mask = seq_offset[None, :] < seq_len |
| 362 | |
| 363 | # Apply sliding window / chunked attention to base mask |
| 364 | # BEFORE mm_prefix OR. |
| 365 | # Order must match FlexAttention: |
| 366 | # (causal AND sliding_window) OR mm_prefix |
| 367 | if CHUNK_LOOKBACK > -1: |
| 368 | seq_mask = seq_mask & ( |
| 369 | (query_abs_pos // CHUNK_SIZE - seq_offset[None, :] // CHUNK_SIZE) |
| 370 | <= CHUNK_LOOKBACK |
| 371 | ) |
| 372 | elif SLIDING_WINDOW > 0 and not USE_R_SWA: |
| 373 | sw_left = (query_abs_pos - seq_offset) < SLIDING_WINDOW |
| 374 | if USE_PER_SEQ_CAUSAL: |
| 375 | sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW |
| 376 | seq_mask = seq_mask & tl.where(is_causal, sw_left, sw_left & sw_right) |
| 377 | elif not USE_CAUSAL: |
| 378 | sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW |
| 379 | seq_mask = seq_mask & sw_left & sw_right |
| 380 | else: |
| 381 | seq_mask = seq_mask & sw_left |
| 382 | |
| 383 | if USE_R_SWA: |
| 384 | prefix_len = tl.load(rswa_prefix_lens_ptr + seq_idx) |
| 385 | in_prefix = seq_offset[None, :] < prefix_len |
| 386 | in_window = (query_abs_pos - seq_offset) < R_SWA_WINDOW |
| 387 | seq_mask = seq_mask & (in_prefix | in_window) |
| 388 | |
| 389 | # PrefixLM: extend mask with bidirectional ranges for multimodal tokens. |
| 390 | # Default (MM_PREFIX_CLAMP_SW=False): applied AFTER sliding window so |
| 391 | # mm_prefix ranges override the SW restriction -> (causal AND SW) OR mm. |
| 392 | # Gemma4 (MM_PREFIX_CLAMP_SW=True): the bidirectional image block must stay |
| 393 | # within the sliding window, matching HF's (causal OR blockwise) AND |
| 394 | # sliding_window. We AND each range with the SW past-bound |
| 395 | # (query_abs_pos - seq_offset) < SLIDING_WINDOW (== HF sliding_window_overlay |
| 396 | # kv > q - sw; future kv passes trivially). Inert for full-attention layers |
| 397 | # (SLIDING_WINDOW <= 0). |
| 398 | if USE_MM_PREFIX: |
| 399 | for i in range(MAX_MM_RANGES): |
| 400 | range_start = tl.load( |
| 401 | mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 |
| 402 | ) |
| 403 | range_end = tl.load( |
| 404 | mm_prefix_range_ptr + seq_idx * MAX_MM_RANGES * 2 + i * 2 + 1 |
| 405 | ) |
| 406 | is_valid = range_start < range_end |
| 407 | q_in_range = ( |
| 408 | (query_abs_pos >= range_start) & (query_abs_pos <= range_end) & is_valid |
| 409 | ) |
| 410 | k_in_range = ( |
| 411 | (seq_offset[None, :] >= range_start) |
| 412 | & (seq_offset[None, :] <= range_end) |
| 413 | & is_valid |
| 414 | ) |
| 415 | mm_mask = q_in_range & k_in_range |
| 416 | if MM_PREFIX_CLAMP_SW and SLIDING_WINDOW > 0: |
| 417 | mm_mask = mm_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW) |
| 418 | seq_mask |= mm_mask |
| 419 | return seq_mask |
| 420 | |
| 421 | |
| 422 | @triton.jit |
| 423 | def apply_alibi_to_score( |
| 424 | S, |
| 425 | alibi_slope, |
| 426 | seq_offset, |
| 427 | context_len, |
| 428 | query_pos, |
| 429 | USE_ALIBI_SQRT: tl.constexpr, |
| 430 | ): |
| 431 | """Add the ALiBi positional bias (linear or sqrt variant) to S in-place.""" |
| 432 | if USE_ALIBI_SQRT: |
| 433 | relative_pos = seq_offset - (context_len + query_pos[:, None]) |
| 434 | alibi_offset = tl.where( |
| 435 | relative_pos <= 0, |
| 436 | -tl.sqrt((-relative_pos).to(tl.float32)), |
| 437 | 0.0, |
| 438 | ) |
| 439 | else: |
| 440 | alibi_offset = seq_offset - context_len |
| 441 | return S + alibi_slope[:, None] * alibi_offset |
| 442 | |
| 443 | |
| 444 | @triton.jit |
| 445 | def load_qq_bias_tile( |
| 446 | qq_bias_row_ptrs, |
| 447 | seq_offset, |
| 448 | context_len, |
| 449 | qq_bias_stride_0, |
| 450 | ): |
| 451 | """Load the qq-bias slice for keys that correspond to query rows.""" |
| 452 | key_rel_pos = seq_offset - context_len |
| 453 | is_query_key = key_rel_pos >= 0 and key_rel_pos < qq_bias_stride_0 |
| 454 | return tl.load( |
| 455 | qq_bias_row_ptrs + key_rel_pos[None, :], |
| 456 | mask=is_query_key[None, :], |
| 457 | other=0.0, |
| 458 | ) |
| 459 | |
| 460 | |
| 461 | @triton.jit |
| 462 | def softmax_step(S, M, L): |
| 463 | """Online softmax update for one tile. |
| 464 | |
| 465 | Returns ``(M_new, L_new, P, alpha)``. Caller is responsible for |
| 466 | rescaling its accumulator(s) by ``alpha[:, None]`` — done outside so |
| 467 | kernels with a different number / shape of accumulators can reuse |
| 468 | the same step. |
| 469 | """ |
| 470 | # compute running maximum |
| 471 | # m_j : (BLOCK_M,) |
| 472 | m_j = tl.maximum(M, tl.max(S, axis=1)) |
| 473 | # For sliding window there's a chance the max is -inf due to masking of |
| 474 | # the entire row. In this case we need to set m_j 0 to avoid NaN |
| 475 | m_j = tl.where(m_j > float("-inf"), m_j, 0.0) |
| 476 | # P : (BLOCK_M, TILE_SIZE) |
| 477 | P = tl.exp(S - m_j[:, None]) |
| 478 | # l_j : (BLOCK_M,) |
| 479 | l_j = tl.sum(P, axis=1) |
| 480 | # alpha : (BLOCK_M, ) |
| 481 | alpha = tl.exp(M - m_j) |
| 482 | # update constants |
| 483 | L_new = L * alpha + l_j |
| 484 | return m_j, L_new, P, alpha |