| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Copyright (c) 2026 LightSeek Foundation |
| 5 | # |
| 6 | # Permission is hereby granted, free of charge, to any person obtaining a copy |
| 7 | # of this software and associated documentation files (the "Software"), to deal |
| 8 | # in the Software without restriction, including without limitation the rights |
| 9 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 10 | # copies of the Software, and to permit persons to whom the Software is |
| 11 | # furnished to do so, subject to the following conditions: |
| 12 | # |
| 13 | # The above copyright notice and this permission notice shall be included in |
| 14 | # all copies or substantial portions of the Software. |
| 15 | # |
| 16 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 17 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 18 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 19 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 20 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 21 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 22 | # SOFTWARE. |
| 23 | |
| 24 | """rel_mha extend Gluon kernel for AMD GFX950. |
| 25 | |
| 26 | This handles ragged, multi-token queries against a paged KV cache. The query |
| 27 | axis is tiled into the MFMA ``M`` dimension (prefill-style): ``BLOCK_M`` query |
| 28 | rows of a single q-head share each paged KV tile, so every KV tile is loaded |
| 29 | once and reused across all rows in the tile. The grid is |
| 30 | ``(blocks_per_req, batch, n_heads)`` -- all host-known sizes -- and each program |
| 31 | self-locates its request from ``cu_seqlens_q`` / ``cache_seqlens`` in-kernel, so |
| 32 | the launch stays CUDA-graph static with no device->host sync. |
| 33 | |
| 34 | Visibility per query row depends on ``is_causal`` and the optional sliding |
| 35 | window: |
| 36 | |
| 37 | * ``is_causal=False``: every query token attends the full visible cache, i.e. |
| 38 | ``visible_kv = cache_seqlens[batch]``. |
| 39 | * ``is_causal=True``: query tokens are a causal suffix, so the ``i``-th query |
| 40 | token of a request (0-indexed) attends ``prefix + i + 1`` tokens, where |
| 41 | ``prefix = cache_seqlens[batch] - query_len[batch]``. |
| 42 | |
| 43 | Causal masking only touches the few KV tiles that reach the diagonal; the long |
| 44 | prefix is a mask-free fast path. Sinks and sliding windows (causal or not) are |
| 45 | applied per tile. This is the sole Gluon extend implementation. |
| 46 | """ |
| 47 | |
| 48 | import math |
| 49 | |
| 50 | import torch |
| 51 | |
| 52 | from vllm.models.inkling.amd.ops.gluon.utils import ( |
| 53 | _INV_LN2_VALUE, |
| 54 | _LN2, |
| 55 | InputStrides, |
| 56 | PagedKVStrides, |
| 57 | max, |
| 58 | maximum, |
| 59 | ) |
| 60 | from vllm.triton_utils import aggregate, gl, gluon |
| 61 | |
| 62 | cdna4 = gl.amd.cdna4 |
| 63 | async_copy = cdna4.async_copy |
| 64 | |
| 65 | |
| 66 | # ===-----------------------------------------------------------------------===# |
| 67 | # Query-batched path |
| 68 | # |
| 69 | # This path tiles BLOCK_M query rows (of a single q-head) into the MFMA M |
| 70 | # dimension -- like the prefill kernel -- so each paged KV tile is loaded once |
| 71 | # and reused across all BLOCK_M rows. KV still comes from the paged cache |
| 72 | # (decode-style async page loads). Causal masking only touches the few KV tiles |
| 73 | # that reach the diagonal; the long prefix is a mask-free fast path. |
| 74 | # ===-----------------------------------------------------------------------===# |
| 75 | |
| 76 | _EXTEND_SHORT_Q_BLOCK_M = 64 |
| 77 | _EXTEND_SHORT_Q_NUM_WARPS = 2 |
| 78 | _EXTEND_LONG_Q_BLOCK_M = 128 |
| 79 | _EXTEND_LONG_Q_NUM_WARPS = 4 |
| 80 | |
| 81 | |
| 82 | def _select_extend_tile(max_seqlen_q: int, page_size: int) -> tuple[int, int, int]: |
| 83 | """Return (BLOCK_M, BLOCK_N, NUM_WARPS) for the given max query length. |
| 84 | |
| 85 | Queries that fit in a single short tile use it (least padding, most |
| 86 | occupancy); longer ones use the tall tile that covers more rows per |
| 87 | shared-KV pass. |
| 88 | """ |
| 89 | if max_seqlen_q <= _EXTEND_SHORT_Q_BLOCK_M: |
| 90 | return _EXTEND_SHORT_Q_BLOCK_M, page_size, _EXTEND_SHORT_Q_NUM_WARPS |
| 91 | return _EXTEND_LONG_Q_BLOCK_M, page_size, _EXTEND_LONG_Q_NUM_WARPS |
| 92 | |
| 93 | |
| 94 | @aggregate |
| 95 | class ExtendConfig: |
| 96 | N_HEADS: gl.constexpr |
| 97 | N_KV_HEADS: gl.constexpr |
| 98 | GROUP_SIZE: gl.constexpr |
| 99 | HEAD_DIM: gl.constexpr |
| 100 | SM_SCALE: gl.constexpr |
| 101 | BLOCK_M: gl.constexpr |
| 102 | BLOCK_N: gl.constexpr |
| 103 | NUM_WARPS: gl.constexpr |
| 104 | PAGE_SIZE: gl.constexpr |
| 105 | PAGE_TABLE_STRIDE: gl.constexpr |
| 106 | IS_CAUSAL: gl.constexpr |
| 107 | HAS_LSE: gl.constexpr |
| 108 | WINDOW_LEFT: gl.constexpr |
| 109 | REL_EXTENT: gl.constexpr |
| 110 | REL_BIAS_QK_SCALE: gl.constexpr |
| 111 | IS_FP8: gl.constexpr |
| 112 | q_strides: InputStrides |
| 113 | rel_strides: InputStrides |
| 114 | k_strides: PagedKVStrides |
| 115 | v_strides: PagedKVStrides |
| 116 | qk_layout: gl.constexpr |
| 117 | pv_layout: gl.constexpr |
| 118 | q_layout: gl.constexpr |
| 119 | k_layout: gl.constexpr |
| 120 | p_layout: gl.constexpr |
| 121 | v_layout: gl.constexpr |
| 122 | load_layout: gl.constexpr |
| 123 | store_layout: gl.constexpr |
| 124 | k_smem_layout: gl.constexpr |
| 125 | v_smem_layout: gl.constexpr |
| 126 | |
| 127 | @gluon.constexpr_function |
| 128 | def __init__( |
| 129 | self, |
| 130 | N_HEADS, |
| 131 | N_KV_HEADS, |
| 132 | HEAD_DIM, |
| 133 | SM_SCALE, |
| 134 | BLOCK_M, |
| 135 | BLOCK_N, |
| 136 | NUM_WARPS, |
| 137 | PAGE_SIZE, |
| 138 | PAGE_TABLE_STRIDE, |
| 139 | IS_CAUSAL, |
| 140 | HAS_LSE, |
| 141 | WINDOW_LEFT, |
| 142 | REL_EXTENT, |
| 143 | REL_BIAS_QK_SCALE, |
| 144 | IS_FP8, |
| 145 | q_strides, |
| 146 | rel_strides, |
| 147 | k_strides, |
| 148 | v_strides, |
| 149 | ): |
| 150 | assert HEAD_DIM in (64, 128) |
| 151 | assert BLOCK_N == PAGE_SIZE |
| 152 | |
| 153 | instr_shape = [32, 32, 16] |
| 154 | mfma_layout = gl.amd.AMDMFMALayout( |
| 155 | version=4, |
| 156 | instr_shape=instr_shape, |
| 157 | transposed=True, |
| 158 | warps_per_cta=[NUM_WARPS, 1], |
| 159 | ) |
| 160 | qk_layout = mfma_layout |
| 161 | pv_layout = mfma_layout |
| 162 | # Elements loaded per lane depend on the input dtype, matching qk_kw. |
| 163 | # load_threads is how many lanes span HEAD_DIM. |
| 164 | load_vec = 16 if IS_FP8 else 8 |
| 165 | load_threads = HEAD_DIM // load_vec |
| 166 | load_layout = gl.BlockedLayout( |
| 167 | [1, load_vec], [64 // load_threads, load_threads], [NUM_WARPS, 1], [1, 0] |
| 168 | ) |
| 169 | # Output is always 16-bit, so a 128-bit store has 8 elements. |
| 170 | # store_threads is how many lanes span HEAD_DIM. |
| 171 | store_vec = 8 |
| 172 | store_threads = HEAD_DIM // store_vec |
| 173 | store_layout = gl.BlockedLayout( |
| 174 | [1, store_vec], [64 // store_threads, store_threads], [NUM_WARPS, 1], [1, 0] |
| 175 | ) |
| 176 | # Padding interval is 64 lanes * load_vec elems. |
| 177 | pad_interval = 64 * load_vec |
| 178 | # Empirically tuned. |
| 179 | pad_k = 16 if IS_FP8 else 8 |
| 180 | pad_v = 16 if IS_FP8 else 32 |
| 181 | k_smem_layout = gl.PaddedSharedLayout.with_identity_for( |
| 182 | [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0] |
| 183 | ) |
| 184 | v_smem_layout = gl.PaddedSharedLayout.with_identity_for( |
| 185 | [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0] |
| 186 | ) |
| 187 | |
| 188 | self.N_HEADS = gl.constexpr(N_HEADS) |
| 189 | self.N_KV_HEADS = gl.constexpr(N_KV_HEADS) |
| 190 | self.GROUP_SIZE = gl.constexpr(N_HEADS // N_KV_HEADS) |
| 191 | self.HEAD_DIM = gl.constexpr(HEAD_DIM) |
| 192 | self.SM_SCALE = gl.constexpr(SM_SCALE) |
| 193 | self.BLOCK_M = gl.constexpr(BLOCK_M) |
| 194 | self.BLOCK_N = gl.constexpr(BLOCK_N) |
| 195 | self.NUM_WARPS = gl.constexpr(NUM_WARPS) |
| 196 | self.PAGE_SIZE = gl.constexpr(PAGE_SIZE) |
| 197 | self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE) |
| 198 | self.IS_CAUSAL = gl.constexpr(IS_CAUSAL) |
| 199 | self.HAS_LSE = gl.constexpr(HAS_LSE) |
| 200 | self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT) |
| 201 | self.REL_EXTENT = gl.constexpr(REL_EXTENT) |
| 202 | self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE) |
| 203 | self.IS_FP8 = gl.constexpr(IS_FP8) |
| 204 | self.q_strides = q_strides |
| 205 | self.rel_strides = rel_strides |
| 206 | self.k_strides = k_strides |
| 207 | self.v_strides = v_strides |
| 208 | self.qk_layout = gl.constexpr(qk_layout) |
| 209 | self.pv_layout = gl.constexpr(pv_layout) |
| 210 | # qk_kw is derived from a 128-bit load / dtype bitwidth. |
| 211 | # pv_kw is empirically tuned. |
| 212 | qk_kw = 16 if IS_FP8 else 8 |
| 213 | pv_kw = 16 if IS_FP8 else 4 |
| 214 | self.q_layout = gl.constexpr(gl.DotOperandLayout(0, qk_layout, k_width=qk_kw)) |
| 215 | self.k_layout = gl.constexpr(gl.DotOperandLayout(1, qk_layout, k_width=qk_kw)) |
| 216 | self.p_layout = gl.constexpr(gl.DotOperandLayout(0, pv_layout, k_width=pv_kw)) |
| 217 | self.v_layout = gl.constexpr(gl.DotOperandLayout(1, pv_layout, k_width=pv_kw)) |
| 218 | self.load_layout = gl.constexpr(load_layout) |
| 219 | self.store_layout = gl.constexpr(store_layout) |
| 220 | self.k_smem_layout = gl.constexpr(k_smem_layout) |
| 221 | self.v_smem_layout = gl.constexpr(v_smem_layout) |
| 222 | |
| 223 | |
| 224 | @aggregate |
| 225 | class ExtendProgram: |
| 226 | cfg: gl.constexpr |
| 227 | q_ptr: gl.tensor |
| 228 | rel_logits_ptr: gl.tensor |
| 229 | k_cache_ptr: gl.tensor |
| 230 | v_cache_ptr: gl.tensor |
| 231 | page_table_ptr: gl.tensor |
| 232 | output_ptr: gl.tensor |
| 233 | lse_ptr: gl.tensor |
| 234 | batch: gl.tensor |
| 235 | q_head: gl.tensor |
| 236 | kv_head: gl.tensor |
| 237 | q_start: gl.tensor |
| 238 | seq_base: gl.tensor |
| 239 | seq_len: gl.tensor |
| 240 | prefix: gl.tensor |
| 241 | cache_len: gl.tensor |
| 242 | |
| 243 | @gluon.constexpr_function |
| 244 | def __init__( |
| 245 | self, |
| 246 | cfg, |
| 247 | q_ptr, |
| 248 | rel_logits_ptr, |
| 249 | k_cache_ptr, |
| 250 | v_cache_ptr, |
| 251 | page_table_ptr, |
| 252 | output_ptr, |
| 253 | lse_ptr, |
| 254 | batch, |
| 255 | q_head, |
| 256 | kv_head, |
| 257 | q_start, |
| 258 | seq_base, |
| 259 | seq_len, |
| 260 | prefix, |
| 261 | cache_len, |
| 262 | ): |
| 263 | self.cfg = gl.constexpr(cfg) |
| 264 | self.q_ptr = q_ptr |
| 265 | self.rel_logits_ptr = rel_logits_ptr |
| 266 | self.k_cache_ptr = k_cache_ptr |
| 267 | self.v_cache_ptr = v_cache_ptr |
| 268 | self.page_table_ptr = page_table_ptr |
| 269 | self.output_ptr = output_ptr |
| 270 | self.lse_ptr = lse_ptr |
| 271 | self.batch = batch |
| 272 | self.q_head = q_head |
| 273 | self.kv_head = kv_head |
| 274 | self.q_start = q_start |
| 275 | self.seq_base = seq_base |
| 276 | self.seq_len = seq_len |
| 277 | self.prefix = prefix |
| 278 | self.cache_len = cache_len |
| 279 | |
| 280 | @gluon.jit |
| 281 | def create( |
| 282 | cfg, |
| 283 | q_ptr, |
| 284 | rel_logits_ptr, |
| 285 | k_cache_ptr, |
| 286 | v_cache_ptr, |
| 287 | page_table_ptr, |
| 288 | output_ptr, |
| 289 | lse_ptr, |
| 290 | cu_seqlens_q_ptr, |
| 291 | cache_seqlens_ptr, |
| 292 | ): |
| 293 | # Gluon treats ``cfg`` as the constexpr factory argument, while mypy |
| 294 | # models the first parameter of this aggregate method as ``self``. |
| 295 | block_in_req = gl.program_id(0) |
| 296 | batch = gl.program_id(1) |
| 297 | q_head = gl.program_id(2) |
| 298 | q_start = block_in_req * cfg.BLOCK_M # type: ignore[attr-defined] |
| 299 | seq_base = gl.load(cu_seqlens_q_ptr + batch) |
| 300 | seq_end = gl.load(cu_seqlens_q_ptr + batch + 1) |
| 301 | seq_len = seq_end - seq_base |
| 302 | cache_len = gl.load(cache_seqlens_ptr + batch) |
| 303 | prefix = cache_len - seq_len |
| 304 | kv_head = q_head // cfg.GROUP_SIZE # type: ignore[attr-defined] |
| 305 | return ExtendProgram( |
| 306 | gl.constexpr(cfg), |
| 307 | q_ptr, |
| 308 | rel_logits_ptr, |
| 309 | k_cache_ptr, |
| 310 | v_cache_ptr, |
| 311 | page_table_ptr, |
| 312 | output_ptr, |
| 313 | lse_ptr, |
| 314 | batch, |
| 315 | q_head, |
| 316 | kv_head, |
| 317 | q_start, |
| 318 | seq_base, |
| 319 | seq_len, |
| 320 | prefix, |
| 321 | cache_len, |
| 322 | ) |
| 323 | |
| 324 | @gluon.jit |
| 325 | def load_q(self): |
| 326 | cfg = self.cfg |
| 327 | offs_m = self.q_start + gl.arange( |
| 328 | 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout) |
| 329 | ) |
| 330 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout)) |
| 331 | row = self.seq_base + offs_m |
| 332 | offsets = cfg.q_strides.offsets(row[:, None], self.q_head, offs_d[None, :]) |
| 333 | mask = offs_m[:, None] < self.seq_len |
| 334 | return cdna4.buffer_load(self.q_ptr, offsets, mask=mask, other=0.0) |
| 335 | |
| 336 | @gluon.jit |
| 337 | def load_page(self, start_n): |
| 338 | cfg = self.cfg |
| 339 | page_index = start_n // cfg.PAGE_SIZE |
| 340 | valid = start_n < self.cache_len |
| 341 | return gl.load( |
| 342 | self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index, |
| 343 | mask=valid, |
| 344 | other=0, |
| 345 | ) |
| 346 | |
| 347 | @gluon.jit |
| 348 | def issue_load_k(self, physical_page, k_smem): |
| 349 | cfg = self.cfg |
| 350 | offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) |
| 351 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) |
| 352 | offsets = cfg.k_strides.offsets( |
| 353 | physical_page, |
| 354 | offs_n[:, None], |
| 355 | self.kv_head, |
| 356 | offs_d[None, :], |
| 357 | ) |
| 358 | async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets) |
| 359 | async_copy.commit_group() |
| 360 | |
| 361 | @gluon.jit |
| 362 | def issue_load_v(self, physical_page, v_smem): |
| 363 | cfg = self.cfg |
| 364 | offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) |
| 365 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) |
| 366 | offsets = cfg.v_strides.offsets( |
| 367 | physical_page, |
| 368 | offs_n[:, None], |
| 369 | self.kv_head, |
| 370 | offs_d[None, :], |
| 371 | ) |
| 372 | async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets) |
| 373 | async_copy.commit_group() |
| 374 | |
| 375 | @gluon.jit |
| 376 | def shared_load_k(self, k_smem): |
| 377 | cfg = self.cfg |
| 378 | k_buffer = k_smem.permute([1, 0]) |
| 379 | return k_buffer.load(cfg.k_layout) |
| 380 | |
| 381 | @gluon.jit |
| 382 | def shared_load_v(self, v_smem): |
| 383 | cfg = self.cfg |
| 384 | return v_smem.load(cfg.v_layout) |
| 385 | |
| 386 | @gluon.jit |
| 387 | def compute_qk(self, q, k): |
| 388 | cfg = self.cfg |
| 389 | qk = gl.zeros( |
| 390 | [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout |
| 391 | ) |
| 392 | return cdna4.mfma(q, k, qk) |
| 393 | |
| 394 | @gluon.jit |
| 395 | def apply_rel_bias(self, qk, start_n): |
| 396 | cfg = self.cfg |
| 397 | offs_m = self.q_start + gl.arange( |
| 398 | 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout) |
| 399 | ) |
| 400 | offs_n_abs = start_n + gl.arange( |
| 401 | 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout) |
| 402 | ) |
| 403 | q_pos = self.prefix + offs_m |
| 404 | rel_dist = q_pos[:, None] - offs_n_abs[None, :] |
| 405 | rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT) |
| 406 | rel_idx = gl.where(rel_dist >= 0, rel_dist, 0) |
| 407 | rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1) |
| 408 | offsets = cfg.rel_strides.offsets( |
| 409 | self.seq_base + offs_m[:, None], self.q_head, rel_idx |
| 410 | ) |
| 411 | mask = (offs_m[:, None] < self.seq_len) & rel_valid |
| 412 | rel_bias = cdna4.buffer_load( |
| 413 | self.rel_logits_ptr, offsets, mask=mask, other=0.0 |
| 414 | ).to(gl.float32) |
| 415 | return qk + rel_bias * cfg.REL_BIAS_QK_SCALE |
| 416 | |
| 417 | @gluon.jit |
| 418 | def compute_pv(self, p, v, acc): |
| 419 | return cdna4.mfma(p, v, acc) |
| 420 | |
| 421 | @gluon.jit |
| 422 | def init_attention_state(self): |
| 423 | cfg = self.cfg |
| 424 | m_i = gl.full( |
| 425 | [cfg.BLOCK_M], |
| 426 | value=-float("inf"), |
| 427 | dtype=gl.float32, |
| 428 | layout=gl.SliceLayout(1, cfg.pv_layout), |
| 429 | ) |
| 430 | l_i = gl.full( |
| 431 | [cfg.BLOCK_M], |
| 432 | value=0, |
| 433 | dtype=gl.float32, |
| 434 | layout=gl.SliceLayout(1, cfg.pv_layout), |
| 435 | ) |
| 436 | acc = gl.zeros( |
| 437 | [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout |
| 438 | ) |
| 439 | return m_i, l_i, acc |
| 440 | |
| 441 | @gluon.jit |
| 442 | def softmax(self, qk, m_i, l_i, acc): |
| 443 | cfg = self.cfg |
| 444 | # In sliding window case, some rows can see fully masked tiles before |
| 445 | # any valid KV. Guard the online softmax state so `-inf - -inf` does not |
| 446 | # produce NaNs. |
| 447 | HAS_INVALID: gl.constexpr = cfg.WINDOW_LEFT >= 0 |
| 448 | |
| 449 | row_max = max(qk, 1) |
| 450 | m_new = maximum(m_i, row_max) |
| 451 | m_new_scaled = m_new * cfg.SM_SCALE |
| 452 | if HAS_INVALID: |
| 453 | invalid = m_new == -float("inf") |
| 454 | m_new_scaled = gl.where(invalid, 0.0, m_new_scaled) |
| 455 | |
| 456 | qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None] |
| 457 | p = gl.exp2(qk_shifted) |
| 458 | m_diff = m_i * cfg.SM_SCALE - m_new_scaled |
| 459 | if HAS_INVALID: |
| 460 | m_diff = gl.where(invalid, 0.0, m_diff) |
| 461 | |
| 462 | alpha = gl.exp2(m_diff) |
| 463 | l_ij = gl.sum(p, axis=1) |
| 464 | l_i = l_i * alpha + l_ij |
| 465 | acc = acc * alpha[:, None] |
| 466 | p = p.to(self.q_ptr.dtype.element_ty) |
| 467 | p = gl.convert_layout(p, cfg.p_layout) |
| 468 | return p, m_new, l_i, acc |
| 469 | |
| 470 | @gluon.jit |
| 471 | def store_output(self, output): |
| 472 | cfg = self.cfg |
| 473 | offs_m = self.q_start + gl.arange( |
| 474 | 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout) |
| 475 | ) |
| 476 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) |
| 477 | offsets = ( |
| 478 | ((self.seq_base + offs_m[:, None]) * cfg.N_HEADS + self.q_head) |
| 479 | * cfg.HEAD_DIM |
| 480 | + offs_d[None, :] |
| 481 | ).to(gl.int32) |
| 482 | mask = offs_m[:, None] < self.seq_len |
| 483 | output = output.to(self.output_ptr.dtype.element_ty) |
| 484 | cdna4.buffer_store(output, self.output_ptr, offsets, mask=mask) |
| 485 | |
| 486 | @gluon.jit |
| 487 | def store_lse(self, l_i, m_i): |
| 488 | cfg = self.cfg |
| 489 | if cfg.HAS_LSE: |
| 490 | offs_m = self.q_start + gl.arange( |
| 491 | 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.pv_layout) |
| 492 | ) |
| 493 | offsets = ((self.seq_base + offs_m) * cfg.N_HEADS + self.q_head).to( |
| 494 | gl.int32 |
| 495 | ) |
| 496 | mask = offs_m < self.seq_len |
| 497 | lse_l_i = gl.where(l_i > 0.0, l_i, 1.0) |
| 498 | # Softmax runs in base-2 (exp2 hardware fast path), so m_i*SM_SCALE + |
| 499 | # log2(l_i) is the LSE in base-2 units. Convert to natural log (the |
| 500 | # public op contract / torch.logsumexp convention) by scaling by ln2. |
| 501 | lse = (m_i * cfg.SM_SCALE + gl.log2(lse_l_i)) * _LN2 |
| 502 | cdna4.buffer_store(lse, self.lse_ptr, offsets, mask=mask) |
| 503 | |
| 504 | |
| 505 | @gluon.jit |
| 506 | def _rel_mha_extend_fp16( |
| 507 | q_ptr, |
| 508 | rel_logits_ptr, |
| 509 | k_cache_ptr, |
| 510 | v_cache_ptr, |
| 511 | page_table_ptr, |
| 512 | output_ptr, |
| 513 | lse_ptr, |
| 514 | cu_seqlens_q_ptr, |
| 515 | cache_seqlens_ptr, |
| 516 | Q_STRIDE_T: gl.constexpr, |
| 517 | Q_STRIDE_H: gl.constexpr, |
| 518 | Q_STRIDE_D: gl.constexpr, |
| 519 | K_STRIDE_B: gl.constexpr, |
| 520 | K_STRIDE_P: gl.constexpr, |
| 521 | K_STRIDE_H: gl.constexpr, |
| 522 | K_STRIDE_D: gl.constexpr, |
| 523 | V_STRIDE_B: gl.constexpr, |
| 524 | V_STRIDE_P: gl.constexpr, |
| 525 | V_STRIDE_H: gl.constexpr, |
| 526 | V_STRIDE_D: gl.constexpr, |
| 527 | SM_SCALE: gl.constexpr, |
| 528 | PAGE_TABLE_STRIDE: gl.constexpr, |
| 529 | PAGE_SIZE: gl.constexpr, |
| 530 | N_HEADS: gl.constexpr, |
| 531 | N_KV_HEADS: gl.constexpr, |
| 532 | HEAD_DIM: gl.constexpr, |
| 533 | BLOCK_M: gl.constexpr, |
| 534 | BLOCK_N: gl.constexpr, |
| 535 | NUM_WARPS: gl.constexpr, |
| 536 | IS_CAUSAL: gl.constexpr, |
| 537 | HAS_LSE: gl.constexpr, |
| 538 | WINDOW_LEFT: gl.constexpr, |
| 539 | REL_STRIDE_T: gl.constexpr, |
| 540 | REL_STRIDE_H: gl.constexpr, |
| 541 | REL_STRIDE_E: gl.constexpr, |
| 542 | REL_EXTENT: gl.constexpr, |
| 543 | REL_BIAS_QK_SCALE: gl.constexpr, |
| 544 | IS_FP8: gl.constexpr, |
| 545 | ): |
| 546 | cfg = ExtendConfig( |
| 547 | N_HEADS, |
| 548 | N_KV_HEADS, |
| 549 | HEAD_DIM, |
| 550 | SM_SCALE, |
| 551 | BLOCK_M, |
| 552 | BLOCK_N, |
| 553 | NUM_WARPS, |
| 554 | PAGE_SIZE, |
| 555 | PAGE_TABLE_STRIDE, |
| 556 | IS_CAUSAL, |
| 557 | HAS_LSE, |
| 558 | WINDOW_LEFT, |
| 559 | REL_EXTENT, |
| 560 | REL_BIAS_QK_SCALE, |
| 561 | IS_FP8, |
| 562 | InputStrides(Q_STRIDE_T, Q_STRIDE_H, Q_STRIDE_D), |
| 563 | InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), |
| 564 | PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), |
| 565 | PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), |
| 566 | ) |
| 567 | program = ExtendProgram.create( |
| 568 | cfg, |
| 569 | q_ptr, |
| 570 | rel_logits_ptr, |
| 571 | k_cache_ptr, |
| 572 | v_cache_ptr, |
| 573 | page_table_ptr, |
| 574 | output_ptr, |
| 575 | lse_ptr, |
| 576 | cu_seqlens_q_ptr, |
| 577 | cache_seqlens_ptr, |
| 578 | ) |
| 579 | # Over-provisioned tile past this request's real query rows: nothing to do. |
| 580 | if program.q_start >= program.seq_len: |
| 581 | return |
| 582 | k_smem = gl.allocate_shared_memory( |
| 583 | k_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.k_smem_layout |
| 584 | ) |
| 585 | v_smem = gl.allocate_shared_memory( |
| 586 | v_cache_ptr.dtype.element_ty, [cfg.BLOCK_N, cfg.HEAD_DIM], cfg.v_smem_layout |
| 587 | ) |
| 588 | |
| 589 | q = program.load_q() |
| 590 | m_i, l_i, acc = program.init_attention_state() |
| 591 | |
| 592 | offs_m_q = program.q_start + gl.arange( |
| 593 | 0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout) |
| 594 | ) |
| 595 | offs_n_q = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)) |
| 596 | diag_row = program.prefix + offs_m_q |
| 597 | |
| 598 | if IS_CAUSAL: |
| 599 | kv_end = min(program.cache_len, program.prefix + program.q_start + cfg.BLOCK_M) |
| 600 | else: |
| 601 | kv_end = program.cache_len |
| 602 | |
| 603 | # Sliding window (inclusive-left, matches flash-attn window_size=(W, 0)): |
| 604 | # skip KV tiles entirely below the window's lower edge. The tile's top query |
| 605 | # row sits at absolute position pos_top; its window opens at |
| 606 | # pos_top - WINDOW_LEFT (that key is visible -> W + 1 keys total). The min() |
| 607 | # form clamps to 0 without the shadowed builtin max(). |
| 608 | if cfg.WINDOW_LEFT >= 0: |
| 609 | pos_top = program.prefix + program.q_start |
| 610 | kv_start = pos_top - min(pos_top, cfg.WINDOW_LEFT) |
| 611 | kv_start = (kv_start // cfg.BLOCK_N) * cfg.BLOCK_N |
| 612 | else: |
| 613 | kv_start = 0 |
| 614 | |
| 615 | for start_n in range(kv_start, kv_end, cfg.BLOCK_N): |
| 616 | physical_page = program.load_page(start_n) |
| 617 | program.issue_load_k(physical_page, k_smem) |
| 618 | program.issue_load_v(physical_page, v_smem) |
| 619 | |
| 620 | async_copy.wait_group(1) |
| 621 | k = program.shared_load_k(k_smem) |
| 622 | qk = program.compute_qk(q, k) |
| 623 | qk = program.apply_rel_bias(qk, start_n) |
| 624 | |
| 625 | if cfg.WINDOW_LEFT >= 0: |
| 626 | # Window lower edge + cache bound always apply; causal upper edge only |
| 627 | # when IS_CAUSAL (independent layering keeps non-causal + window correct). |
| 628 | offs_n_abs = start_n + offs_n_q |
| 629 | mask = (offs_n_abs[None, :] >= diag_row[:, None] - cfg.WINDOW_LEFT) & ( |
| 630 | offs_n_abs[None, :] < program.cache_len |
| 631 | ) |
| 632 | if IS_CAUSAL: |
| 633 | mask &= offs_n_abs[None, :] <= diag_row[:, None] |
| 634 | qk = gl.where(mask, qk, -float("inf")) |
| 635 | elif IS_CAUSAL: |
| 636 | if start_n + cfg.BLOCK_N > program.prefix + program.q_start: |
| 637 | offs_n_abs = start_n + offs_n_q |
| 638 | mask = (offs_n_abs[None, :] <= diag_row[:, None]) & ( |
| 639 | offs_n_abs[None, :] < program.cache_len |
| 640 | ) |
| 641 | qk = gl.where(mask, qk, -float("inf")) |
| 642 | else: |
| 643 | if start_n + cfg.BLOCK_N > program.cache_len: |
| 644 | offs_n_abs = start_n + offs_n_q |
| 645 | qk = gl.where( |
| 646 | offs_n_abs[None, :] < program.cache_len, qk, -float("inf") |
| 647 | ) |
| 648 | |
| 649 | p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) |
| 650 | |
| 651 | async_copy.wait_group(0) |
| 652 | v = program.shared_load_v(v_smem) |
| 653 | acc = program.compute_pv(p, v, acc) |
| 654 | |
| 655 | denom = gl.where(l_i > 0.0, l_i, 1.0) |
| 656 | output = acc * (1.0 / denom)[:, None] |
| 657 | output = gl.convert_layout(output, cfg.store_layout) |
| 658 | program.store_output(output) |
| 659 | program.store_lse(l_i, m_i) |
| 660 | |
| 661 | |
| 662 | def gluon_rel_mha_extend_gfx950( |
| 663 | q: torch.Tensor, |
| 664 | cu_seqlens_q: torch.Tensor, |
| 665 | cu_seqlens_kv: torch.Tensor, |
| 666 | k_cache: torch.Tensor, |
| 667 | v_cache: torch.Tensor, |
| 668 | page_table: torch.Tensor, |
| 669 | cache_seqlens: torch.Tensor, |
| 670 | window_left: int = -1, |
| 671 | return_lse: bool = False, |
| 672 | max_seqlen_q: int = 1, |
| 673 | max_seqlen_k: int = 1, |
| 674 | rel_logits: torch.Tensor | None = None, |
| 675 | softmax_scale: float | None = None, |
| 676 | q_scale: torch.Tensor | None = None, |
| 677 | k_scale: torch.Tensor | None = None, |
| 678 | v_scale: torch.Tensor | None = None, |
| 679 | out: torch.Tensor | None = None, |
| 680 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 681 | assert rel_logits is not None |
| 682 | head_dim = q.shape[2] |
| 683 | n_heads = q.shape[1] |
| 684 | n_kv_heads = k_cache.shape[2] |
| 685 | page_size = k_cache.shape[1] |
| 686 | block_m, block_n, num_warps = _select_extend_tile(max_seqlen_q, page_size) |
| 687 | if softmax_scale is None: |
| 688 | softmax_scale = 1.0 / math.sqrt(head_dim) |
| 689 | sm_scale = softmax_scale * _INV_LN2_VALUE |
| 690 | rel_bias_qk_scale = 1.0 / softmax_scale |
| 691 | |
| 692 | # max_seqlen_q must be >= the true max query length; extra tiles early-exit. |
| 693 | batch = cu_seqlens_q.shape[0] - 1 |
| 694 | safe_max_q = max_seqlen_q if max_seqlen_q > 0 else 1 |
| 695 | blocks_per_req = (safe_max_q + block_m - 1) // block_m |
| 696 | cu_q_i32 = cu_seqlens_q.to(torch.int32).contiguous() |
| 697 | cache_i32 = cache_seqlens.to(torch.int32).contiguous() |
| 698 | |
| 699 | is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) |
| 700 | out_dtype = torch.bfloat16 if is_fp8 else q.dtype |
| 701 | if page_size not in (64, 128, 256): |
| 702 | raise ValueError( |
| 703 | "gfx950 Gluon relative-attention extend requires page size " |
| 704 | f"64, 128, or 256; got {page_size}" |
| 705 | ) |
| 706 | output = ( |
| 707 | torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out |
| 708 | ) |
| 709 | if return_lse: |
| 710 | lse = torch.empty((q.shape[0], n_heads), device=q.device, dtype=torch.float32) |
| 711 | lse_arg = lse |
| 712 | else: |
| 713 | lse = None |
| 714 | lse_arg = q |
| 715 | grid = (blocks_per_req, batch, n_heads) |
| 716 | _rel_mha_extend_fp16[grid]( |
| 717 | q, |
| 718 | rel_logits, |
| 719 | k_cache, |
| 720 | v_cache, |
| 721 | page_table, |
| 722 | output, |
| 723 | lse_arg, |
| 724 | cu_q_i32, |
| 725 | cache_i32, |
| 726 | q.stride(0), |
| 727 | q.stride(1), |
| 728 | q.stride(2), |
| 729 | k_cache.stride(0), |
| 730 | k_cache.stride(1), |
| 731 | k_cache.stride(2), |
| 732 | k_cache.stride(3), |
| 733 | v_cache.stride(0), |
| 734 | v_cache.stride(1), |
| 735 | v_cache.stride(2), |
| 736 | v_cache.stride(3), |
| 737 | sm_scale, |
| 738 | page_table.stride(0), |
| 739 | page_size, |
| 740 | n_heads, |
| 741 | n_kv_heads, |
| 742 | head_dim, |
| 743 | block_m, |
| 744 | block_n, |
| 745 | num_warps, |
| 746 | True, |
| 747 | return_lse, |
| 748 | window_left, |
| 749 | rel_logits.stride(0), |
| 750 | rel_logits.stride(1), |
| 751 | rel_logits.stride(2), |
| 752 | rel_logits.shape[2], |
| 753 | rel_bias_qk_scale, |
| 754 | is_fp8, |
| 755 | num_warps=num_warps, |
| 756 | ) |
| 757 | if return_lse: |
| 758 | assert lse is not None |
| 759 | return output, lse |
| 760 | return output |