| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # SPDX-License-Identifier: MIT |
| 5 | # Copyright (c) 2026 LightSeek Foundation |
| 6 | # |
| 7 | # Permission is hereby granted, free of charge, to any person obtaining a copy |
| 8 | # of this software and associated documentation files (the "Software"), to deal |
| 9 | # in the Software without restriction, including without limitation the rights |
| 10 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 11 | # copies of the Software, and to permit persons to whom the Software is |
| 12 | # furnished to do so, subject to the following conditions: |
| 13 | # |
| 14 | # The above copyright notice and this permission notice shall be included in |
| 15 | # all copies or substantial portions of the Software. |
| 16 | # |
| 17 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 18 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 19 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 20 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 21 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 22 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 23 | # SOFTWARE. |
| 24 | """Split-KV decode for Inkling relative attention on ROCm. |
| 25 | |
| 26 | Adapted from LightSeek TokenSpeed's portable Triton relative-MHA decode. |
| 27 | """ |
| 28 | |
| 29 | from __future__ import annotations |
| 30 | |
| 31 | import os |
| 32 | |
| 33 | import torch |
| 34 | |
| 35 | from vllm.triton_utils import tl, triton |
| 36 | |
| 37 | _MIN_BLOCK_KV = tl.constexpr(32) |
| 38 | |
| 39 | |
| 40 | def decode_split_count(max_kv_len: int, window_left: int) -> int: |
| 41 | """Return the number of parallel KV partitions for decode.""" |
| 42 | if window_left >= 0: |
| 43 | effective_len = max(1, min(max_kv_len, window_left + 1)) |
| 44 | return min(4, max(1, triton.cdiv(effective_len, 128))) |
| 45 | return min(32, max(1, triton.cdiv(max(1, max_kv_len), 2048))) |
| 46 | |
| 47 | |
| 48 | def use_split_kv_decode( |
| 49 | *, |
| 50 | max_query_len: int, |
| 51 | max_kv_len: int, |
| 52 | page_size: int, |
| 53 | window_left: int, |
| 54 | ) -> bool: |
| 55 | """Select split-KV only where it outperforms the single-pass kernel.""" |
| 56 | if os.getenv("INKLING_SPLIT_KV", "1") != "1": |
| 57 | return False |
| 58 | if max_query_len != 1: |
| 59 | return False |
| 60 | if page_size >= 64: |
| 61 | return True |
| 62 | if window_left >= 0: |
| 63 | return page_size >= 32 |
| 64 | return max_kv_len >= 8192 |
| 65 | |
| 66 | |
| 67 | @triton.jit |
| 68 | def _split_kv_stage1( |
| 69 | q_ptr, |
| 70 | rel_ptr, |
| 71 | k_ptr, |
| 72 | v_ptr, |
| 73 | block_table_ptr, |
| 74 | cache_seqlens, |
| 75 | mid_out_ptr, |
| 76 | mid_lse_ptr, |
| 77 | softmax_scale, |
| 78 | stride_q_t, |
| 79 | stride_q_h, |
| 80 | stride_k_b, |
| 81 | stride_k_p, |
| 82 | stride_k_h, |
| 83 | stride_k_d, |
| 84 | stride_v_b, |
| 85 | stride_v_p, |
| 86 | stride_v_h, |
| 87 | stride_v_d, |
| 88 | stride_r_t, |
| 89 | stride_r_h, |
| 90 | stride_r_e, |
| 91 | stride_mo_t, |
| 92 | stride_mo_h, |
| 93 | stride_mo_s, |
| 94 | stride_ml_t, |
| 95 | stride_ml_h, |
| 96 | stride_ml_s, |
| 97 | stride_bt_b: tl.constexpr, |
| 98 | page_size: tl.constexpr, |
| 99 | window_left: tl.constexpr, |
| 100 | gqa_group_size: tl.constexpr, |
| 101 | num_q_heads: tl.constexpr, |
| 102 | head_dim: tl.constexpr, |
| 103 | rel_extent: tl.constexpr, |
| 104 | max_kv_splits: tl.constexpr, |
| 105 | BLOCK_D: tl.constexpr, |
| 106 | BLOCK_H: tl.constexpr, |
| 107 | BLOCK_N: tl.constexpr, |
| 108 | ): |
| 109 | pid_t = tl.program_id(0) |
| 110 | pid_h = tl.program_id(1) |
| 111 | pid_s = tl.program_id(2) |
| 112 | |
| 113 | head_blocks_per_kv = tl.cdiv(gqa_group_size, BLOCK_H) |
| 114 | kv_head = pid_h // head_blocks_per_kv |
| 115 | valid_block_h: tl.constexpr = min(BLOCK_H, gqa_group_size) |
| 116 | off_h = pid_h * valid_block_h + tl.arange(0, BLOCK_H) |
| 117 | h_valid = (off_h < (pid_h + 1) * valid_block_h) & (off_h < num_q_heads) |
| 118 | off_d = tl.arange(0, BLOCK_D) |
| 119 | d_valid = off_d < head_dim |
| 120 | |
| 121 | cache_len = tl.load(cache_seqlens + pid_t) |
| 122 | effective_len = ( |
| 123 | tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len |
| 124 | ) |
| 125 | kv_offset = cache_len - effective_len |
| 126 | split_len = ( |
| 127 | tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV |
| 128 | ) |
| 129 | split_start = split_len * pid_s |
| 130 | split_end = tl.minimum(split_start + split_len, effective_len) |
| 131 | |
| 132 | q_offsets = pid_t * stride_q_t + off_h[:, None] * stride_q_h + off_d[None, :] |
| 133 | row_max = tl.full((BLOCK_H,), float("-inf"), dtype=tl.float32) |
| 134 | row_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) |
| 135 | acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32) |
| 136 | |
| 137 | if split_end > split_start: |
| 138 | q = tl.load( |
| 139 | q_ptr + q_offsets, |
| 140 | mask=h_valid[:, None] & d_valid[None, :], |
| 141 | other=0.0, |
| 142 | ) |
| 143 | for start_n in range(split_start, split_end, BLOCK_N): |
| 144 | off_n = start_n + tl.arange(0, BLOCK_N) |
| 145 | n_valid = off_n < split_end |
| 146 | token_idx = kv_offset + off_n |
| 147 | logical_page = token_idx // page_size |
| 148 | page_offset = token_idx % page_size |
| 149 | physical_page = tl.load( |
| 150 | block_table_ptr + pid_t * stride_bt_b + logical_page, |
| 151 | mask=n_valid, |
| 152 | other=0, |
| 153 | ) |
| 154 | k_offsets = ( |
| 155 | physical_page[None, :].to(tl.int64) * stride_k_b |
| 156 | + page_offset[None, :] * stride_k_p |
| 157 | + kv_head * stride_k_h |
| 158 | + off_d[:, None] * stride_k_d |
| 159 | ) |
| 160 | k = tl.load( |
| 161 | k_ptr + k_offsets, |
| 162 | mask=d_valid[:, None] & n_valid[None, :], |
| 163 | other=0.0, |
| 164 | ) |
| 165 | scores = tl.dot(q, k.to(q.dtype)) * softmax_scale |
| 166 | |
| 167 | rel_dist = cache_len - 1 - token_idx |
| 168 | rel_valid = (rel_dist >= 0) & (rel_dist < rel_extent) |
| 169 | rel_idx = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1)) |
| 170 | rel_offsets = ( |
| 171 | pid_t * stride_r_t |
| 172 | + off_h[:, None] * stride_r_h |
| 173 | + rel_idx[None, :] * stride_r_e |
| 174 | ) |
| 175 | rel_bias = tl.load( |
| 176 | rel_ptr + rel_offsets, |
| 177 | mask=h_valid[:, None] & rel_valid[None, :] & n_valid[None, :], |
| 178 | other=0.0, |
| 179 | ) |
| 180 | scores += rel_bias.to(tl.float32) |
| 181 | scores = tl.where( |
| 182 | h_valid[:, None] & n_valid[None, :], |
| 183 | scores, |
| 184 | float("-inf"), |
| 185 | ) |
| 186 | |
| 187 | v_offsets = ( |
| 188 | physical_page[:, None].to(tl.int64) * stride_v_b |
| 189 | + page_offset[:, None] * stride_v_p |
| 190 | + kv_head * stride_v_h |
| 191 | + off_d[None, :] * stride_v_d |
| 192 | ) |
| 193 | v = tl.load( |
| 194 | v_ptr + v_offsets, |
| 195 | mask=n_valid[:, None] & d_valid[None, :], |
| 196 | other=0.0, |
| 197 | ) |
| 198 | |
| 199 | next_max = tl.maximum(tl.max(scores, axis=1), row_max) |
| 200 | old_scale = tl.exp(row_max - next_max) |
| 201 | probs = tl.exp(scores - next_max[:, None]) |
| 202 | acc *= old_scale[:, None] |
| 203 | acc += tl.dot(probs.to(v.dtype), v) |
| 204 | row_sum = row_sum * old_scale + tl.sum(probs, axis=1) |
| 205 | row_max = next_max |
| 206 | |
| 207 | mid_out_offsets = ( |
| 208 | pid_t * stride_mo_t |
| 209 | + off_h[:, None] * stride_mo_h |
| 210 | + pid_s * stride_mo_s |
| 211 | + off_d[None, :] |
| 212 | ) |
| 213 | tl.store( |
| 214 | mid_out_ptr + mid_out_offsets, |
| 215 | acc / row_sum[:, None], |
| 216 | mask=h_valid[:, None] & d_valid[None, :], |
| 217 | ) |
| 218 | mid_lse_offsets = ( |
| 219 | pid_t * stride_ml_t + off_h * stride_ml_h + pid_s * stride_ml_s |
| 220 | ) |
| 221 | tl.store( |
| 222 | mid_lse_ptr + mid_lse_offsets, |
| 223 | row_max + tl.log(row_sum), |
| 224 | mask=h_valid, |
| 225 | ) |
| 226 | |
| 227 | |
| 228 | @triton.jit |
| 229 | def _split_kv_stage2( |
| 230 | mid_out_ptr, |
| 231 | mid_lse_ptr, |
| 232 | out_ptr, |
| 233 | cache_seqlens, |
| 234 | stride_mo_t, |
| 235 | stride_mo_h, |
| 236 | stride_mo_s, |
| 237 | stride_ml_t, |
| 238 | stride_ml_h, |
| 239 | stride_ml_s, |
| 240 | stride_o_t, |
| 241 | stride_o_h, |
| 242 | window_left: tl.constexpr, |
| 243 | head_dim: tl.constexpr, |
| 244 | max_kv_splits: tl.constexpr, |
| 245 | BLOCK_D: tl.constexpr, |
| 246 | ): |
| 247 | pid_t = tl.program_id(0) |
| 248 | pid_h = tl.program_id(1) |
| 249 | |
| 250 | cache_len = tl.load(cache_seqlens + pid_t) |
| 251 | effective_len = ( |
| 252 | tl.minimum(cache_len, window_left + 1) if window_left >= 0 else cache_len |
| 253 | ) |
| 254 | split_len = ( |
| 255 | tl.cdiv(tl.cdiv(effective_len, max_kv_splits), _MIN_BLOCK_KV) * _MIN_BLOCK_KV |
| 256 | ) |
| 257 | |
| 258 | off_d = tl.arange(0, BLOCK_D) |
| 259 | d_valid = off_d < head_dim |
| 260 | row_max = -float("inf") |
| 261 | row_sum = 0.0 |
| 262 | acc = tl.zeros((BLOCK_D,), dtype=tl.float32) |
| 263 | |
| 264 | for split_id in range(max_kv_splits): |
| 265 | split_start = split_len * split_id |
| 266 | split_end = tl.minimum(split_start + split_len, effective_len) |
| 267 | if split_end > split_start: |
| 268 | value = tl.load( |
| 269 | mid_out_ptr |
| 270 | + pid_t * stride_mo_t |
| 271 | + pid_h * stride_mo_h |
| 272 | + split_id * stride_mo_s |
| 273 | + off_d, |
| 274 | mask=d_valid, |
| 275 | other=0.0, |
| 276 | ) |
| 277 | split_lse = tl.load( |
| 278 | mid_lse_ptr |
| 279 | + pid_t * stride_ml_t |
| 280 | + pid_h * stride_ml_h |
| 281 | + split_id * stride_ml_s |
| 282 | ) |
| 283 | next_max = tl.maximum(split_lse, row_max) |
| 284 | old_scale = tl.exp(row_max - next_max) |
| 285 | split_scale = tl.exp(split_lse - next_max) |
| 286 | acc = acc * old_scale + value * split_scale |
| 287 | row_sum = row_sum * old_scale + split_scale |
| 288 | row_max = next_max |
| 289 | |
| 290 | tl.store( |
| 291 | out_ptr + pid_t * stride_o_t + pid_h * stride_o_h + off_d, |
| 292 | acc / row_sum, |
| 293 | mask=d_valid, |
| 294 | ) |
| 295 | |
| 296 | |
| 297 | @torch.no_grad() |
| 298 | def inkling_rel_attention_split_kv_decode( |
| 299 | q: torch.Tensor, |
| 300 | key_cache: torch.Tensor, |
| 301 | value_cache: torch.Tensor, |
| 302 | *, |
| 303 | block_table: torch.Tensor, |
| 304 | cache_seqlens: torch.Tensor, |
| 305 | softmax_scale: float, |
| 306 | window_left: int, |
| 307 | rel_extent: int, |
| 308 | rel_logits: torch.Tensor, |
| 309 | max_kv_len: int, |
| 310 | out: torch.Tensor, |
| 311 | ) -> torch.Tensor: |
| 312 | """Run split-KV relative attention for single-token decode.""" |
| 313 | num_kv_heads = key_cache.shape[2] |
| 314 | gqa_group_size = q.shape[1] // num_kv_heads |
| 315 | max_kv_splits = decode_split_count(max_kv_len, window_left) |
| 316 | block_d = triton.next_power_of_2(q.shape[2]) |
| 317 | block_h = min(8, gqa_group_size) |
| 318 | |
| 319 | mid_out = torch.empty( |
| 320 | q.shape[0], |
| 321 | q.shape[1], |
| 322 | max_kv_splits, |
| 323 | q.shape[2], |
| 324 | dtype=torch.float32, |
| 325 | device=q.device, |
| 326 | ) |
| 327 | mid_lse = torch.empty( |
| 328 | q.shape[0], |
| 329 | q.shape[1], |
| 330 | max_kv_splits, |
| 331 | dtype=torch.float32, |
| 332 | device=q.device, |
| 333 | ) |
| 334 | stage1_grid = ( |
| 335 | q.shape[0], |
| 336 | triton.cdiv(q.shape[1], block_h), |
| 337 | max_kv_splits, |
| 338 | ) |
| 339 | _split_kv_stage1[stage1_grid]( |
| 340 | q, |
| 341 | rel_logits, |
| 342 | key_cache, |
| 343 | value_cache, |
| 344 | block_table, |
| 345 | cache_seqlens, |
| 346 | mid_out, |
| 347 | mid_lse, |
| 348 | softmax_scale, |
| 349 | q.stride(0), |
| 350 | q.stride(1), |
| 351 | key_cache.stride(0), |
| 352 | key_cache.stride(1), |
| 353 | key_cache.stride(2), |
| 354 | key_cache.stride(3), |
| 355 | value_cache.stride(0), |
| 356 | value_cache.stride(1), |
| 357 | value_cache.stride(2), |
| 358 | value_cache.stride(3), |
| 359 | rel_logits.stride(0), |
| 360 | rel_logits.stride(1), |
| 361 | rel_logits.stride(2), |
| 362 | mid_out.stride(0), |
| 363 | mid_out.stride(1), |
| 364 | mid_out.stride(2), |
| 365 | mid_lse.stride(0), |
| 366 | mid_lse.stride(1), |
| 367 | mid_lse.stride(2), |
| 368 | block_table.stride(0), |
| 369 | page_size=key_cache.shape[1], |
| 370 | window_left=window_left, |
| 371 | gqa_group_size=gqa_group_size, |
| 372 | num_q_heads=q.shape[1], |
| 373 | head_dim=q.shape[2], |
| 374 | rel_extent=rel_extent, |
| 375 | max_kv_splits=max_kv_splits, |
| 376 | BLOCK_D=block_d, |
| 377 | BLOCK_H=block_h, |
| 378 | BLOCK_N=key_cache.shape[1], |
| 379 | num_warps=4, |
| 380 | num_stages=1, |
| 381 | ) |
| 382 | stage2_grid = (q.shape[0], q.shape[1]) |
| 383 | _split_kv_stage2[stage2_grid]( |
| 384 | mid_out, |
| 385 | mid_lse, |
| 386 | out, |
| 387 | cache_seqlens, |
| 388 | mid_out.stride(0), |
| 389 | mid_out.stride(1), |
| 390 | mid_out.stride(2), |
| 391 | mid_lse.stride(0), |
| 392 | mid_lse.stride(1), |
| 393 | mid_lse.stride(2), |
| 394 | out.stride(0), |
| 395 | out.stride(1), |
| 396 | window_left=window_left, |
| 397 | head_dim=q.shape[2], |
| 398 | max_kv_splits=max_kv_splits, |
| 399 | BLOCK_D=block_d, |
| 400 | num_warps=4, |
| 401 | num_stages=2, |
| 402 | ) |
| 403 | return out |