| 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 decode Gluon kernel for AMD GFX950.""" |
| 25 | |
| 26 | import math |
| 27 | from typing import NamedTuple |
| 28 | |
| 29 | import torch |
| 30 | |
| 31 | from vllm.models.inkling.amd.ops.gluon.utils import ( |
| 32 | _INV_LN2_VALUE, |
| 33 | InputStrides, |
| 34 | PagedKVStrides, |
| 35 | max, |
| 36 | maximum, |
| 37 | ) |
| 38 | from vllm.triton_utils import aggregate, gl, gluon |
| 39 | |
| 40 | cdna4 = gl.amd.cdna4 |
| 41 | async_copy = cdna4.async_copy |
| 42 | cdiv = gl.cdiv |
| 43 | _GFX950_SM_COUNT = 256 |
| 44 | |
| 45 | |
| 46 | # ===-----------------------------------------------------------------------===# |
| 47 | # Kernel Config |
| 48 | # ===-----------------------------------------------------------------------===# |
| 49 | |
| 50 | |
| 51 | @aggregate |
| 52 | class AttentionConfig: |
| 53 | SM_SCALE: gl.constexpr |
| 54 | PAGE_TABLE_STRIDE: gl.constexpr |
| 55 | PAGE_SIZE: gl.constexpr |
| 56 | NUM_KV_SPLITS: gl.constexpr |
| 57 | MAX_SEQLEN_Q: gl.constexpr |
| 58 | NUM_Q_HEADS: gl.constexpr |
| 59 | NUM_KV_HEADS: gl.constexpr |
| 60 | HEAD_DIM: gl.constexpr |
| 61 | BLOCK_M: gl.constexpr |
| 62 | BLOCK_N: gl.constexpr |
| 63 | IS_SLIDING: gl.constexpr |
| 64 | WINDOW_LEFT: gl.constexpr |
| 65 | REL_EXTENT: gl.constexpr |
| 66 | REL_BIAS_QK_SCALE: gl.constexpr |
| 67 | IS_FP8: gl.constexpr |
| 68 | GROUP_SIZE: gl.constexpr |
| 69 | NUM_GROUPS: gl.constexpr |
| 70 | q_strides: InputStrides |
| 71 | rel_strides: InputStrides |
| 72 | k_strides: PagedKVStrides |
| 73 | v_strides: PagedKVStrides |
| 74 | qk_layout: gl.constexpr |
| 75 | pv_layout: gl.constexpr |
| 76 | q_layout: gl.constexpr |
| 77 | k_layout: gl.constexpr |
| 78 | p_layout: gl.constexpr |
| 79 | v_layout: gl.constexpr |
| 80 | load_layout: gl.constexpr |
| 81 | store_layout: gl.constexpr |
| 82 | reduce_layout: gl.constexpr |
| 83 | k_smem_layout: gl.constexpr |
| 84 | v_smem_layout: gl.constexpr |
| 85 | |
| 86 | @gluon.constexpr_function |
| 87 | def __init__( |
| 88 | self, |
| 89 | SM_SCALE, |
| 90 | PAGE_TABLE_STRIDE, |
| 91 | PAGE_SIZE, |
| 92 | NUM_KV_SPLITS, |
| 93 | MAX_SEQLEN_Q, |
| 94 | NUM_Q_HEADS, |
| 95 | NUM_KV_HEADS, |
| 96 | HEAD_DIM, |
| 97 | BLOCK_M, |
| 98 | BLOCK_N, |
| 99 | IS_SLIDING, |
| 100 | WINDOW_LEFT, |
| 101 | REL_EXTENT, |
| 102 | REL_BIAS_QK_SCALE, |
| 103 | IS_FP8, |
| 104 | q_strides, |
| 105 | rel_strides, |
| 106 | k_strides, |
| 107 | v_strides, |
| 108 | ): |
| 109 | assert NUM_Q_HEADS % NUM_KV_HEADS == 0 |
| 110 | assert HEAD_DIM in (64, 128) |
| 111 | assert BLOCK_N == PAGE_SIZE |
| 112 | if IS_SLIDING: |
| 113 | assert WINDOW_LEFT >= 0 |
| 114 | else: |
| 115 | assert WINDOW_LEFT == -1 |
| 116 | |
| 117 | mfma_layout = gl.amd.AMDMFMALayout( |
| 118 | version=4, |
| 119 | instr_shape=[16, 16, 32], |
| 120 | transposed=True, |
| 121 | warps_per_cta=[1, 1], |
| 122 | ) |
| 123 | qk_layout = mfma_layout |
| 124 | pv_layout = mfma_layout |
| 125 | # qk_kw is derived from a 128-bit load / dtype bitwidth. |
| 126 | # pv_kw is empirically tuned. |
| 127 | qk_kw = 16 if IS_FP8 else 8 |
| 128 | pv_kw = 8 if IS_FP8 else 4 |
| 129 | q_layout = gl.DotOperandLayout(0, qk_layout, k_width=qk_kw) |
| 130 | k_layout = gl.DotOperandLayout(1, qk_layout, k_width=qk_kw) |
| 131 | p_layout = gl.DotOperandLayout(0, pv_layout, k_width=pv_kw) |
| 132 | v_layout = gl.DotOperandLayout(1, pv_layout, k_width=pv_kw) |
| 133 | # Elements loaded per lane depend on the input dtype, matching qk_kw. |
| 134 | # load_threads is how many lanes span HEAD_DIM. |
| 135 | load_vec = 16 if IS_FP8 else 8 |
| 136 | load_threads = HEAD_DIM // load_vec |
| 137 | load_layout = gl.BlockedLayout( |
| 138 | [1, load_vec], [64 // load_threads, load_threads], [1, 1], [1, 0] |
| 139 | ) |
| 140 | # Output is always 16-bit, so a 128-bit store has 8 elements. |
| 141 | # store_threads is how many lanes span HEAD_DIM. |
| 142 | store_vec = 8 |
| 143 | store_threads = HEAD_DIM // store_vec |
| 144 | store_layout = gl.BlockedLayout( |
| 145 | [1, store_vec], [64 // store_threads, store_threads], [1, 1], [1, 0] |
| 146 | ) |
| 147 | reduce_layout = gl.BlockedLayout([1, HEAD_DIM // 64], [1, 64], [1, 1], [1, 0]) |
| 148 | # Padding interval is 64 lanes * load_vec elems. |
| 149 | pad_interval = 64 * load_vec |
| 150 | # Empirically tuned. |
| 151 | pad_k = 16 if IS_FP8 else 8 |
| 152 | pad_v = 32 |
| 153 | k_smem_layout = gl.PaddedSharedLayout.with_identity_for( |
| 154 | [[pad_interval, pad_k]], [BLOCK_N, HEAD_DIM], [1, 0] |
| 155 | ) |
| 156 | v_smem_layout = gl.PaddedSharedLayout.with_identity_for( |
| 157 | [[pad_interval, pad_v]], [BLOCK_N, HEAD_DIM], [1, 0] |
| 158 | ) |
| 159 | |
| 160 | self.SM_SCALE = gl.constexpr(SM_SCALE) |
| 161 | self.PAGE_TABLE_STRIDE = gl.constexpr(PAGE_TABLE_STRIDE) |
| 162 | self.PAGE_SIZE = gl.constexpr(PAGE_SIZE) |
| 163 | self.NUM_KV_SPLITS = gl.constexpr(NUM_KV_SPLITS) |
| 164 | self.MAX_SEQLEN_Q = gl.constexpr(MAX_SEQLEN_Q) |
| 165 | self.NUM_Q_HEADS = gl.constexpr(NUM_Q_HEADS) |
| 166 | self.NUM_KV_HEADS = gl.constexpr(NUM_KV_HEADS) |
| 167 | self.HEAD_DIM = gl.constexpr(HEAD_DIM) |
| 168 | self.BLOCK_M = gl.constexpr(BLOCK_M) |
| 169 | self.BLOCK_N = gl.constexpr(BLOCK_N) |
| 170 | self.IS_SLIDING = gl.constexpr(IS_SLIDING) |
| 171 | self.WINDOW_LEFT = gl.constexpr(WINDOW_LEFT) |
| 172 | self.REL_EXTENT = gl.constexpr(REL_EXTENT) |
| 173 | self.REL_BIAS_QK_SCALE = gl.constexpr(REL_BIAS_QK_SCALE) |
| 174 | self.IS_FP8 = gl.constexpr(IS_FP8) |
| 175 | self.GROUP_SIZE = gl.constexpr(NUM_Q_HEADS // NUM_KV_HEADS) |
| 176 | self.NUM_GROUPS = gl.constexpr((self.GROUP_SIZE + BLOCK_M - 1) // BLOCK_M) |
| 177 | self.q_strides = q_strides |
| 178 | self.rel_strides = rel_strides |
| 179 | self.k_strides = k_strides |
| 180 | self.v_strides = v_strides |
| 181 | self.qk_layout = gl.constexpr(qk_layout) |
| 182 | self.pv_layout = gl.constexpr(pv_layout) |
| 183 | self.q_layout = gl.constexpr(q_layout) |
| 184 | self.k_layout = gl.constexpr(k_layout) |
| 185 | self.p_layout = gl.constexpr(p_layout) |
| 186 | self.v_layout = gl.constexpr(v_layout) |
| 187 | self.load_layout = gl.constexpr(load_layout) |
| 188 | self.store_layout = gl.constexpr(store_layout) |
| 189 | self.reduce_layout = gl.constexpr(reduce_layout) |
| 190 | self.k_smem_layout = gl.constexpr(k_smem_layout) |
| 191 | self.v_smem_layout = gl.constexpr(v_smem_layout) |
| 192 | |
| 193 | |
| 194 | # ===-----------------------------------------------------------------------===# |
| 195 | # Kernel Program |
| 196 | # ===-----------------------------------------------------------------------===# |
| 197 | |
| 198 | |
| 199 | @aggregate |
| 200 | class AttentionProgram: |
| 201 | cfg: gl.constexpr |
| 202 | q_ptr: gl.tensor |
| 203 | rel_logits_ptr: gl.tensor |
| 204 | k_cache_ptr: gl.tensor |
| 205 | v_cache_ptr: gl.tensor |
| 206 | page_table_ptr: gl.tensor |
| 207 | cache_seqlens_ptr: gl.tensor |
| 208 | mid_o_ptr: gl.tensor |
| 209 | mid_lse_ptr: gl.tensor |
| 210 | q_index: gl.tensor |
| 211 | batch: gl.tensor |
| 212 | kv_head: gl.tensor |
| 213 | group_start: gl.tensor |
| 214 | split_id: gl.tensor |
| 215 | cache_len: gl.tensor |
| 216 | kv_start: gl.tensor |
| 217 | split_start: gl.tensor |
| 218 | split_end: gl.tensor |
| 219 | |
| 220 | @gluon.constexpr_function |
| 221 | def __init__( |
| 222 | self, |
| 223 | cfg, |
| 224 | q_ptr, |
| 225 | rel_logits_ptr, |
| 226 | k_cache_ptr, |
| 227 | v_cache_ptr, |
| 228 | page_table_ptr, |
| 229 | cache_seqlens_ptr, |
| 230 | mid_o_ptr, |
| 231 | mid_lse_ptr, |
| 232 | q_index, |
| 233 | batch, |
| 234 | kv_head, |
| 235 | group_start, |
| 236 | split_id, |
| 237 | cache_len, |
| 238 | kv_start, |
| 239 | split_start, |
| 240 | split_end, |
| 241 | ): |
| 242 | self.cfg = gl.constexpr(cfg) |
| 243 | self.q_ptr = q_ptr |
| 244 | self.rel_logits_ptr = rel_logits_ptr |
| 245 | self.k_cache_ptr = k_cache_ptr |
| 246 | self.v_cache_ptr = v_cache_ptr |
| 247 | self.page_table_ptr = page_table_ptr |
| 248 | self.cache_seqlens_ptr = cache_seqlens_ptr |
| 249 | self.mid_o_ptr = mid_o_ptr |
| 250 | self.mid_lse_ptr = mid_lse_ptr |
| 251 | self.q_index = q_index |
| 252 | self.batch = batch |
| 253 | self.kv_head = kv_head |
| 254 | self.group_start = group_start |
| 255 | self.split_id = split_id |
| 256 | self.cache_len = cache_len |
| 257 | self.kv_start = kv_start |
| 258 | self.split_start = split_start |
| 259 | self.split_end = split_end |
| 260 | |
| 261 | @gluon.jit |
| 262 | def create( |
| 263 | cfg, |
| 264 | q_ptr, |
| 265 | rel_logits_ptr, |
| 266 | k_cache_ptr, |
| 267 | v_cache_ptr, |
| 268 | page_table_ptr, |
| 269 | cache_seqlens_ptr, |
| 270 | mid_o_ptr, |
| 271 | mid_lse_ptr, |
| 272 | ): |
| 273 | # Gluon treats ``cfg`` as the constexpr factory argument, while mypy |
| 274 | # models the first parameter of this aggregate method as ``self``. |
| 275 | q_index = gl.program_id(0) |
| 276 | batch = q_index // cfg.MAX_SEQLEN_Q # type: ignore[attr-defined] |
| 277 | q_pos = q_index - batch * cfg.MAX_SEQLEN_Q # type: ignore[attr-defined] |
| 278 | head_block = gl.program_id(1) |
| 279 | kv_head = head_block // cfg.NUM_GROUPS # type: ignore[attr-defined] |
| 280 | group_block = head_block - kv_head * cfg.NUM_GROUPS # type: ignore[attr-defined] |
| 281 | group_start = group_block * cfg.BLOCK_M # type: ignore[attr-defined] |
| 282 | split_id = gl.program_id(2) |
| 283 | cache_len = gl.load(cache_seqlens_ptr + batch) |
| 284 | cache_len = cache_len - ( |
| 285 | cfg.MAX_SEQLEN_Q - 1 - q_pos # type: ignore[attr-defined] |
| 286 | ) |
| 287 | cache_len = maximum(cache_len, 0) |
| 288 | if cfg.IS_SLIDING: # type: ignore[attr-defined] |
| 289 | # WINDOW_LEFT is defined as exclusive (keys strictly to the left, not |
| 290 | # counting the current token), so the window is WINDOW_LEFT + 1 keys |
| 291 | # once the current token is included. E.g. with |
| 292 | # WINDOW_LEFT = 127 and cache_len = 500 (current token at index 499), |
| 293 | # kv_start = 500 - (127 + 1) = 372 keeps keys [372, 499] = 128 keys. |
| 294 | # Without the + 1, kv_start = 373 would drop the leftmost key. |
| 295 | kv_start = cache_len - min( |
| 296 | cache_len, |
| 297 | cfg.WINDOW_LEFT + 1, # type: ignore[attr-defined] |
| 298 | ) |
| 299 | else: |
| 300 | kv_start = cache_len - cache_len |
| 301 | first_page = kv_start // cfg.PAGE_SIZE # type: ignore[attr-defined] |
| 302 | end_page = cdiv(cache_len, cfg.PAGE_SIZE) # type: ignore[attr-defined] |
| 303 | num_pages = end_page - first_page |
| 304 | pages_per_split = cdiv( |
| 305 | num_pages, |
| 306 | cfg.NUM_KV_SPLITS, # type: ignore[attr-defined] |
| 307 | ) |
| 308 | split_start_page = first_page + split_id * pages_per_split |
| 309 | split_end_page = min(split_start_page + pages_per_split, end_page) |
| 310 | split_start = ( |
| 311 | split_start_page * cfg.PAGE_SIZE # type: ignore[attr-defined] |
| 312 | ) |
| 313 | split_end = min( |
| 314 | split_end_page * cfg.PAGE_SIZE, # type: ignore[attr-defined] |
| 315 | cache_len, |
| 316 | ) |
| 317 | return AttentionProgram( |
| 318 | gl.constexpr(cfg), |
| 319 | q_ptr, |
| 320 | rel_logits_ptr, |
| 321 | k_cache_ptr, |
| 322 | v_cache_ptr, |
| 323 | page_table_ptr, |
| 324 | cache_seqlens_ptr, |
| 325 | mid_o_ptr, |
| 326 | mid_lse_ptr, |
| 327 | q_index, |
| 328 | batch, |
| 329 | kv_head, |
| 330 | group_start, |
| 331 | split_id, |
| 332 | cache_len, |
| 333 | kv_start, |
| 334 | split_start, |
| 335 | split_end, |
| 336 | ) |
| 337 | |
| 338 | @gluon.jit |
| 339 | def load_q(self): |
| 340 | cfg = self.cfg |
| 341 | offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.q_layout)) |
| 342 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.q_layout)) |
| 343 | q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m |
| 344 | valid = (self.group_start + offs_m) < cfg.GROUP_SIZE |
| 345 | offsets = cfg.q_strides.offsets(self.q_index, q_heads[:, None], offs_d[None, :]) |
| 346 | return cdna4.buffer_load(self.q_ptr, offsets, mask=valid[:, None], other=0.0) |
| 347 | |
| 348 | @gluon.jit |
| 349 | def init_state(self): |
| 350 | cfg = self.cfg |
| 351 | m_i = gl.full( |
| 352 | [cfg.BLOCK_M], |
| 353 | value=-float("inf"), |
| 354 | dtype=gl.float32, |
| 355 | layout=gl.SliceLayout(1, cfg.pv_layout), |
| 356 | ) |
| 357 | l_i = gl.full( |
| 358 | [cfg.BLOCK_M], |
| 359 | value=0.0, |
| 360 | dtype=gl.float32, |
| 361 | layout=gl.SliceLayout(1, cfg.pv_layout), |
| 362 | ) |
| 363 | acc = gl.zeros( |
| 364 | [cfg.BLOCK_M, cfg.HEAD_DIM], dtype=gl.float32, layout=cfg.pv_layout |
| 365 | ) |
| 366 | return m_i, l_i, acc |
| 367 | |
| 368 | @gluon.jit |
| 369 | def load_page(self, start_n): |
| 370 | cfg = self.cfg |
| 371 | page_index = start_n // cfg.PAGE_SIZE |
| 372 | valid = start_n < self.split_end |
| 373 | return gl.load( |
| 374 | self.page_table_ptr + self.batch * cfg.PAGE_TABLE_STRIDE + page_index, |
| 375 | mask=valid, |
| 376 | other=0, |
| 377 | ) |
| 378 | |
| 379 | @gluon.jit |
| 380 | def issue_load_k(self, physical_page, k_smem): |
| 381 | cfg = self.cfg |
| 382 | offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) |
| 383 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) |
| 384 | offsets = cfg.k_strides.offsets( |
| 385 | physical_page, |
| 386 | offs_n[:, None], |
| 387 | self.kv_head, |
| 388 | offs_d[None, :], |
| 389 | ) |
| 390 | # can't use buffer_load: paged KV offsets may exceed its 32-bit range. |
| 391 | async_copy.global_load_to_shared(k_smem, self.k_cache_ptr + offsets) |
| 392 | async_copy.commit_group() |
| 393 | |
| 394 | @gluon.jit |
| 395 | def issue_load_v(self, physical_page, v_smem): |
| 396 | cfg = self.cfg |
| 397 | offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(1, cfg.load_layout)) |
| 398 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.load_layout)) |
| 399 | offsets = cfg.v_strides.offsets( |
| 400 | physical_page, |
| 401 | offs_n[:, None], |
| 402 | self.kv_head, |
| 403 | offs_d[None, :], |
| 404 | ) |
| 405 | # can't use buffer_load: paged KV offsets may exceed its 32-bit range. |
| 406 | async_copy.global_load_to_shared(v_smem, self.v_cache_ptr + offsets) |
| 407 | async_copy.commit_group() |
| 408 | |
| 409 | @gluon.jit |
| 410 | def shared_load_k(self, k_smem): |
| 411 | return k_smem.permute([1, 0]).load(self.cfg.k_layout) |
| 412 | |
| 413 | @gluon.jit |
| 414 | def shared_load_v(self, v_smem): |
| 415 | return v_smem.load(self.cfg.v_layout) |
| 416 | |
| 417 | @gluon.jit |
| 418 | def compute_qk(self, q, k): |
| 419 | cfg = self.cfg |
| 420 | qk = gl.zeros( |
| 421 | [cfg.BLOCK_M, cfg.BLOCK_N], dtype=gl.float32, layout=cfg.qk_layout |
| 422 | ) |
| 423 | return cdna4.mfma(q, k, qk) |
| 424 | |
| 425 | @gluon.jit |
| 426 | def apply_rel_bias(self, qk, start_n): |
| 427 | cfg = self.cfg |
| 428 | offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.qk_layout)) |
| 429 | offs_n = start_n + gl.arange( |
| 430 | 0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout) |
| 431 | ) |
| 432 | q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m |
| 433 | rel_dist = (self.cache_len - 1) - offs_n |
| 434 | rel_valid = (rel_dist >= 0) & (rel_dist < cfg.REL_EXTENT) |
| 435 | rel_idx = gl.where(rel_dist >= 0, rel_dist, 0) |
| 436 | rel_idx = gl.where(rel_idx < cfg.REL_EXTENT, rel_idx, cfg.REL_EXTENT - 1) |
| 437 | offsets = cfg.rel_strides.offsets( |
| 438 | self.q_index, q_heads[:, None], rel_idx[None, :] |
| 439 | ) |
| 440 | head_valid = (self.group_start + offs_m) < cfg.GROUP_SIZE |
| 441 | rel_bias = cdna4.buffer_load( |
| 442 | self.rel_logits_ptr, |
| 443 | offsets, |
| 444 | mask=head_valid[:, None] & rel_valid[None, :], |
| 445 | other=0.0, |
| 446 | ).to(gl.float32) |
| 447 | return qk + rel_bias * cfg.REL_BIAS_QK_SCALE |
| 448 | |
| 449 | @gluon.jit |
| 450 | def apply_kv_mask(self, qk, start_n): |
| 451 | cfg = self.cfg |
| 452 | offs_n = gl.arange(0, cfg.BLOCK_N, layout=gl.SliceLayout(0, cfg.qk_layout)) |
| 453 | tokens = start_n + offs_n[None, :] |
| 454 | mask = (tokens >= self.kv_start) & (tokens < self.split_end) |
| 455 | return gl.where(mask, qk, -float("inf")) |
| 456 | |
| 457 | @gluon.jit |
| 458 | def softmax(self, qk, m_i, l_i, acc): |
| 459 | cfg = self.cfg |
| 460 | row_max = max(qk, axis=1) |
| 461 | row_max = gl.convert_layout(row_max, gl.SliceLayout(1, cfg.pv_layout)) |
| 462 | m_new = maximum(m_i, row_max) |
| 463 | m_new_scaled = m_new * cfg.SM_SCALE |
| 464 | qk_shifted = qk * cfg.SM_SCALE - m_new_scaled[:, None] |
| 465 | p = gl.exp2(qk_shifted) |
| 466 | m_diff = m_i * cfg.SM_SCALE - m_new_scaled |
| 467 | alpha = gl.exp2(m_diff) |
| 468 | l_ij = gl.sum(p, axis=1) |
| 469 | l_i = l_i * alpha + l_ij |
| 470 | acc = acc * alpha[:, None] |
| 471 | p = p.to(self.q_ptr.dtype.element_ty) |
| 472 | p = gl.convert_layout(p, cfg.p_layout) |
| 473 | return p, m_new, l_i, acc |
| 474 | |
| 475 | @gluon.jit |
| 476 | def compute_pv(self, p, v, acc): |
| 477 | return cdna4.mfma(p, v, acc) |
| 478 | |
| 479 | @gluon.jit |
| 480 | def store_split(self, acc, l_i, m_i): |
| 481 | cfg = self.cfg |
| 482 | offs_m = gl.arange(0, cfg.BLOCK_M, layout=gl.SliceLayout(1, cfg.store_layout)) |
| 483 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.store_layout)) |
| 484 | q_heads = self.kv_head * cfg.GROUP_SIZE + self.group_start + offs_m |
| 485 | valid = ((self.group_start + offs_m) < cfg.GROUP_SIZE) & ( |
| 486 | self.split_start < self.split_end |
| 487 | ) |
| 488 | acc = gl.convert_layout(acc, cfg.store_layout) |
| 489 | l_i = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.store_layout)) |
| 490 | m_i = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.store_layout)) |
| 491 | recip_l_i = 1.0 / l_i |
| 492 | part_o = acc * recip_l_i[:, None] |
| 493 | part_lse = m_i * cfg.SM_SCALE + gl.log2(l_i) |
| 494 | mid_o_offsets = ( |
| 495 | (self.q_index * cfg.NUM_Q_HEADS + q_heads[:, None]) * cfg.NUM_KV_SPLITS |
| 496 | + self.split_id |
| 497 | ) * cfg.HEAD_DIM + offs_d[None, :] |
| 498 | mid_lse_offsets = ( |
| 499 | self.q_index * cfg.NUM_Q_HEADS + q_heads |
| 500 | ) * cfg.NUM_KV_SPLITS + self.split_id |
| 501 | cdna4.buffer_store(part_o, self.mid_o_ptr, mid_o_offsets, mask=valid[:, None]) |
| 502 | cdna4.buffer_store(part_lse, self.mid_lse_ptr, mid_lse_offsets, mask=valid) |
| 503 | |
| 504 | |
| 505 | # ===-----------------------------------------------------------------------===# |
| 506 | # Entry Point |
| 507 | # ===-----------------------------------------------------------------------===# |
| 508 | |
| 509 | |
| 510 | @gluon.jit |
| 511 | def _rel_mha_decode_fp16( |
| 512 | q_ptr, |
| 513 | rel_logits_ptr, |
| 514 | k_cache_ptr, |
| 515 | v_cache_ptr, |
| 516 | page_table_ptr, |
| 517 | cache_seqlens_ptr, |
| 518 | mid_o_ptr, |
| 519 | mid_lse_ptr, |
| 520 | Q_STRIDE_B: gl.constexpr, |
| 521 | Q_STRIDE_H: gl.constexpr, |
| 522 | Q_STRIDE_D: gl.constexpr, |
| 523 | K_STRIDE_B: gl.constexpr, |
| 524 | K_STRIDE_P: gl.constexpr, |
| 525 | K_STRIDE_H: gl.constexpr, |
| 526 | K_STRIDE_D: gl.constexpr, |
| 527 | V_STRIDE_B: gl.constexpr, |
| 528 | V_STRIDE_P: gl.constexpr, |
| 529 | V_STRIDE_H: gl.constexpr, |
| 530 | V_STRIDE_D: gl.constexpr, |
| 531 | SM_SCALE: gl.constexpr, |
| 532 | PAGE_TABLE_STRIDE: gl.constexpr, |
| 533 | PAGE_SIZE: gl.constexpr, |
| 534 | NUM_KV_SPLITS: gl.constexpr, |
| 535 | MAX_SEQLEN_Q: gl.constexpr, |
| 536 | NUM_Q_HEADS: gl.constexpr, |
| 537 | NUM_KV_HEADS: gl.constexpr, |
| 538 | HEAD_DIM: gl.constexpr, |
| 539 | BLOCK_M: gl.constexpr, |
| 540 | BLOCK_N: gl.constexpr, |
| 541 | IS_SLIDING: gl.constexpr, |
| 542 | WINDOW_LEFT: gl.constexpr, |
| 543 | REL_STRIDE_T: gl.constexpr, |
| 544 | REL_STRIDE_H: gl.constexpr, |
| 545 | REL_STRIDE_E: gl.constexpr, |
| 546 | REL_EXTENT: gl.constexpr, |
| 547 | REL_BIAS_QK_SCALE: gl.constexpr, |
| 548 | IS_FP8: gl.constexpr, |
| 549 | ): |
| 550 | cfg = AttentionConfig( |
| 551 | SM_SCALE, |
| 552 | PAGE_TABLE_STRIDE, |
| 553 | PAGE_SIZE, |
| 554 | NUM_KV_SPLITS, |
| 555 | MAX_SEQLEN_Q, |
| 556 | NUM_Q_HEADS, |
| 557 | NUM_KV_HEADS, |
| 558 | HEAD_DIM, |
| 559 | BLOCK_M, |
| 560 | BLOCK_N, |
| 561 | IS_SLIDING, |
| 562 | WINDOW_LEFT, |
| 563 | REL_EXTENT, |
| 564 | REL_BIAS_QK_SCALE, |
| 565 | IS_FP8, |
| 566 | InputStrides(Q_STRIDE_B, Q_STRIDE_H, Q_STRIDE_D), |
| 567 | InputStrides(REL_STRIDE_T, REL_STRIDE_H, REL_STRIDE_E), |
| 568 | PagedKVStrides(K_STRIDE_B, K_STRIDE_P, K_STRIDE_H, K_STRIDE_D), |
| 569 | PagedKVStrides(V_STRIDE_B, V_STRIDE_P, V_STRIDE_H, V_STRIDE_D), |
| 570 | ) |
| 571 | program = AttentionProgram.create( |
| 572 | cfg, |
| 573 | q_ptr, |
| 574 | rel_logits_ptr, |
| 575 | k_cache_ptr, |
| 576 | v_cache_ptr, |
| 577 | page_table_ptr, |
| 578 | cache_seqlens_ptr, |
| 579 | mid_o_ptr, |
| 580 | mid_lse_ptr, |
| 581 | ) |
| 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_state() |
| 591 | |
| 592 | physical_page = program.load_page(program.split_start) |
| 593 | |
| 594 | for start_n in range(program.split_start, program.split_end, cfg.BLOCK_N): |
| 595 | program.issue_load_k(physical_page, k_smem) |
| 596 | program.issue_load_v(physical_page, v_smem) |
| 597 | physical_page = program.load_page(start_n + cfg.BLOCK_N) |
| 598 | |
| 599 | async_copy.wait_group(1) |
| 600 | k = program.shared_load_k(k_smem) |
| 601 | qk = program.compute_qk(q, k) |
| 602 | qk = program.apply_rel_bias(qk, start_n) |
| 603 | qk = program.apply_kv_mask(qk, start_n) |
| 604 | p, m_i, l_i, acc = program.softmax(qk, m_i, l_i, acc) |
| 605 | |
| 606 | async_copy.wait_group(0) |
| 607 | v = program.shared_load_v(v_smem) |
| 608 | acc = program.compute_pv(p, v, acc) |
| 609 | |
| 610 | program.store_split(acc, l_i, m_i) |
| 611 | |
| 612 | |
| 613 | @gluon.jit |
| 614 | def _rel_mha_decode_reduce_fp16( |
| 615 | mid_o_ptr, |
| 616 | mid_lse_ptr, |
| 617 | out_ptr, |
| 618 | cache_seqlens_ptr, |
| 619 | SM_SCALE: gl.constexpr, |
| 620 | PAGE_TABLE_STRIDE: gl.constexpr, |
| 621 | NUM_KV_SPLITS: gl.constexpr, |
| 622 | MAX_SEQLEN_Q: gl.constexpr, |
| 623 | PAGE_SIZE: gl.constexpr, |
| 624 | NUM_Q_HEADS: gl.constexpr, |
| 625 | NUM_KV_HEADS: gl.constexpr, |
| 626 | HEAD_DIM: gl.constexpr, |
| 627 | BLOCK_M: gl.constexpr, |
| 628 | BLOCK_N: gl.constexpr, |
| 629 | IS_SLIDING: gl.constexpr, |
| 630 | WINDOW_LEFT: gl.constexpr, |
| 631 | IS_FP8: gl.constexpr, |
| 632 | ): |
| 633 | cfg = AttentionConfig( |
| 634 | SM_SCALE, |
| 635 | PAGE_TABLE_STRIDE, |
| 636 | PAGE_SIZE, |
| 637 | NUM_KV_SPLITS, |
| 638 | MAX_SEQLEN_Q, |
| 639 | NUM_Q_HEADS, |
| 640 | NUM_KV_HEADS, |
| 641 | HEAD_DIM, |
| 642 | BLOCK_M, |
| 643 | BLOCK_N, |
| 644 | IS_SLIDING, |
| 645 | WINDOW_LEFT, |
| 646 | 1, # REL_EXTENT |
| 647 | 1.0, # REL_BIAS_QK_SCALE |
| 648 | IS_FP8, |
| 649 | InputStrides(1, 1, 1), |
| 650 | InputStrides(1, 1, 1), |
| 651 | PagedKVStrides(1, 1, 1, 1), |
| 652 | PagedKVStrides(1, 1, 1, 1), |
| 653 | ) |
| 654 | q_index = gl.program_id(0) |
| 655 | batch = q_index // MAX_SEQLEN_Q |
| 656 | q_pos = q_index - batch * MAX_SEQLEN_Q |
| 657 | q_head = gl.program_id(1) |
| 658 | cache_len = gl.load(cache_seqlens_ptr + batch) |
| 659 | cache_len = cache_len - (MAX_SEQLEN_Q - 1 - q_pos) |
| 660 | cache_len = maximum(cache_len, 0) |
| 661 | if cfg.IS_SLIDING: |
| 662 | kv_start = cache_len - min(cache_len, cfg.WINDOW_LEFT + 1) |
| 663 | else: |
| 664 | kv_start = cache_len - cache_len |
| 665 | first_page = kv_start // cfg.PAGE_SIZE |
| 666 | end_page = cdiv(cache_len, cfg.PAGE_SIZE) |
| 667 | num_pages = end_page - first_page |
| 668 | pages_per_split = cdiv(num_pages, cfg.NUM_KV_SPLITS) |
| 669 | |
| 670 | # SPLIT_TILE pads NUM_KV_SPLITS up to a power of 2. |
| 671 | SPLIT_TILE: gl.constexpr = 1 << (NUM_KV_SPLITS - 1).bit_length() |
| 672 | offs_s = gl.arange(0, SPLIT_TILE, layout=gl.SliceLayout(1, cfg.reduce_layout)) |
| 673 | offs_d = gl.arange(0, cfg.HEAD_DIM, layout=gl.SliceLayout(0, cfg.reduce_layout)) |
| 674 | # split_valid masks out empty splits and the power-of-2 padding tail. |
| 675 | split_start_page = first_page + offs_s * pages_per_split |
| 676 | split_end_page_raw = split_start_page + pages_per_split |
| 677 | split_end_page = gl.where( |
| 678 | split_end_page_raw < end_page, split_end_page_raw, end_page |
| 679 | ) |
| 680 | split_start_tok = split_start_page * cfg.PAGE_SIZE |
| 681 | split_end_raw = split_end_page * cfg.PAGE_SIZE |
| 682 | split_end_tok = gl.where(split_end_raw < cache_len, split_end_raw, cache_len) |
| 683 | split_valid = (split_start_tok < split_end_tok) & (offs_s < cfg.NUM_KV_SPLITS) |
| 684 | # Load every split's partial output and lse. |
| 685 | base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.NUM_KV_SPLITS + offs_s |
| 686 | part_lse = gl.load(mid_lse_ptr + base, mask=split_valid, other=-float("inf")) |
| 687 | o_off = base[:, None] * cfg.HEAD_DIM + offs_d[None, :] |
| 688 | part_o = cdna4.buffer_load(mid_o_ptr, o_off, mask=split_valid[:, None], other=0.0) |
| 689 | |
| 690 | # Global softmax max over all splits. |
| 691 | m_i = max(part_lse, axis=0) |
| 692 | # Weighted sum of the split partials, normalized by the total softmax mass. |
| 693 | beta = gl.exp2(part_lse - m_i) |
| 694 | l_i = gl.sum(beta, axis=0) |
| 695 | acc = gl.sum(part_o * beta[:, None], axis=0) |
| 696 | |
| 697 | out_base = (q_index * cfg.NUM_Q_HEADS + q_head) * cfg.HEAD_DIM |
| 698 | output = acc * (1.0 / l_i) |
| 699 | output = output.to(out_ptr.dtype.element_ty) |
| 700 | cdna4.buffer_store(output, out_ptr, out_base + offs_d) |
| 701 | |
| 702 | |
| 703 | def _select_num_kv_splits( |
| 704 | *, |
| 705 | batch: int, |
| 706 | num_kv_heads: int, |
| 707 | num_groups: int, |
| 708 | num_pages: int, |
| 709 | sm_count: int, |
| 710 | ) -> int: |
| 711 | """Pick num_kv_splits to balance occupancy against reduce overhead. |
| 712 | |
| 713 | The launch grid is (batch * num_kv_heads * num_groups) * num_kv_splits |
| 714 | work-groups. Too few splits under-fill the machine at low batch; too many |
| 715 | leave each split with a handful of pages, so the reduce kernel dominates. |
| 716 | |
| 717 | Return the smaller of two candidate counts: splits_for_occupancy (enough to |
| 718 | fill ~wave_target waves of CUs) and splits_for_pages (~min_pages_per_split |
| 719 | pages per split), with the pages candidate clamped to [min_page_splits, |
| 720 | max_page_splits] so a short context still splits without launching empty work |
| 721 | and a long one does not over-split where reduce cost outgrows the decode win. |
| 722 | """ |
| 723 | wave_target = 2 |
| 724 | min_pages_per_split = 2 |
| 725 | min_page_splits = 8 |
| 726 | max_page_splits = 32 |
| 727 | |
| 728 | base_ctas = batch * num_kv_heads * num_groups |
| 729 | target_ctas = sm_count * wave_target |
| 730 | splits_for_occupancy = (target_ctas + base_ctas - 1) // base_ctas |
| 731 | |
| 732 | splits_for_pages = num_pages // min_pages_per_split |
| 733 | min_page_splits = min(min_page_splits, num_pages) |
| 734 | if splits_for_pages < min_page_splits: |
| 735 | splits_for_pages = min_page_splits |
| 736 | if splits_for_pages > max_page_splits: |
| 737 | splits_for_pages = max_page_splits |
| 738 | return min(splits_for_occupancy, splits_for_pages) |
| 739 | |
| 740 | |
| 741 | class LaunchConfig(NamedTuple): |
| 742 | num_q_heads: int |
| 743 | num_kv_heads: int |
| 744 | num_groups: int |
| 745 | head_dim: int |
| 746 | page_size: int |
| 747 | num_kv_splits: int |
| 748 | block_m: int |
| 749 | block_n: int |
| 750 | sm_scale: float |
| 751 | rel_bias_qk_scale: float |
| 752 | is_sliding: bool |
| 753 | window_left: int |
| 754 | |
| 755 | |
| 756 | def get_config( |
| 757 | *, |
| 758 | q: torch.Tensor, |
| 759 | k_cache: torch.Tensor, |
| 760 | max_seqlen_k: int, |
| 761 | window_left: int, |
| 762 | softmax_scale: float | None, |
| 763 | ) -> LaunchConfig: |
| 764 | head_dim = q.shape[2] |
| 765 | page_size = k_cache.shape[1] |
| 766 | block_m = 16 |
| 767 | block_n = page_size |
| 768 | group_size = q.shape[1] // k_cache.shape[2] |
| 769 | num_groups = math.ceil(group_size / block_m) |
| 770 | is_sliding = window_left >= 0 |
| 771 | window_left = window_left if is_sliding else -1 |
| 772 | if softmax_scale is None: |
| 773 | softmax_scale = 1.0 / math.sqrt(head_dim) |
| 774 | sm_scale = softmax_scale |
| 775 | effective_seqlen_k = ( |
| 776 | min(max_seqlen_k, window_left + 1) if is_sliding else max_seqlen_k |
| 777 | ) |
| 778 | num_pages = (effective_seqlen_k + page_size - 1) // page_size |
| 779 | if is_sliding: |
| 780 | num_kv_splits = 4 |
| 781 | else: |
| 782 | num_kv_splits = _select_num_kv_splits( |
| 783 | batch=q.shape[0], |
| 784 | num_kv_heads=k_cache.shape[2], |
| 785 | num_groups=num_groups, |
| 786 | num_pages=num_pages, |
| 787 | sm_count=_GFX950_SM_COUNT, |
| 788 | ) |
| 789 | return LaunchConfig( |
| 790 | num_q_heads=q.shape[1], |
| 791 | num_kv_heads=k_cache.shape[2], |
| 792 | num_groups=num_groups, |
| 793 | head_dim=head_dim, |
| 794 | page_size=page_size, |
| 795 | num_kv_splits=num_kv_splits, |
| 796 | block_m=block_m, |
| 797 | block_n=block_n, |
| 798 | sm_scale=sm_scale * _INV_LN2_VALUE, |
| 799 | rel_bias_qk_scale=1.0 / softmax_scale, |
| 800 | is_sliding=is_sliding, |
| 801 | window_left=window_left, |
| 802 | ) |
| 803 | |
| 804 | |
| 805 | def gluon_rel_mha_decode_gfx950( |
| 806 | q: torch.Tensor, |
| 807 | k_cache: torch.Tensor, |
| 808 | v_cache: torch.Tensor, |
| 809 | page_table: torch.Tensor, |
| 810 | cache_seqlens: torch.Tensor, |
| 811 | max_seqlen_k: int, |
| 812 | rel_logits: torch.Tensor, |
| 813 | cu_seqlens_q: torch.Tensor, |
| 814 | max_seqlen_q: int = 1, |
| 815 | window_left: int = -1, |
| 816 | softmax_scale: float | None = None, |
| 817 | q_scale: torch.Tensor | None = None, |
| 818 | k_scale: torch.Tensor | None = None, |
| 819 | v_scale: torch.Tensor | None = None, |
| 820 | out: torch.Tensor | None = None, |
| 821 | ) -> torch.Tensor: |
| 822 | total_q = q.shape[0] |
| 823 | |
| 824 | config = get_config( |
| 825 | q=q, |
| 826 | k_cache=k_cache, |
| 827 | max_seqlen_k=max_seqlen_k, |
| 828 | window_left=window_left, |
| 829 | softmax_scale=softmax_scale, |
| 830 | ) |
| 831 | |
| 832 | is_fp8 = q.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) |
| 833 | out_dtype = torch.bfloat16 if is_fp8 else q.dtype |
| 834 | if k_cache.shape[1] not in (64, 128, 256): |
| 835 | raise ValueError( |
| 836 | "gfx950 Gluon relative-attention decode requires page size " |
| 837 | f"64, 128, or 256; got {k_cache.shape[1]}" |
| 838 | ) |
| 839 | output = ( |
| 840 | torch.empty(q.shape, device=q.device, dtype=out_dtype) if out is None else out |
| 841 | ) |
| 842 | |
| 843 | # Always use split-k for full attention and for the small-batch sliding |
| 844 | # decode path. Sliding uses a fixed 8 splits, one page per split for the |
| 845 | # TP-4 local-attention shape with 512-token window and 64-token pages. |
| 846 | mid_o = torch.empty( |
| 847 | (total_q, config.num_q_heads, config.num_kv_splits, config.head_dim), |
| 848 | device=q.device, |
| 849 | dtype=torch.float32, |
| 850 | ) |
| 851 | mid_lse = torch.empty( |
| 852 | (total_q, config.num_q_heads, config.num_kv_splits), |
| 853 | device=q.device, |
| 854 | dtype=torch.float32, |
| 855 | ) |
| 856 | |
| 857 | grid = ( |
| 858 | total_q, |
| 859 | config.num_kv_heads * config.num_groups, |
| 860 | config.num_kv_splits, |
| 861 | ) |
| 862 | _rel_mha_decode_fp16[grid]( |
| 863 | q, |
| 864 | rel_logits, |
| 865 | k_cache, |
| 866 | v_cache, |
| 867 | page_table, |
| 868 | cache_seqlens, |
| 869 | mid_o, |
| 870 | mid_lse, |
| 871 | q.stride(0), |
| 872 | q.stride(1), |
| 873 | q.stride(2), |
| 874 | k_cache.stride(0), |
| 875 | k_cache.stride(1), |
| 876 | k_cache.stride(2), |
| 877 | k_cache.stride(3), |
| 878 | v_cache.stride(0), |
| 879 | v_cache.stride(1), |
| 880 | v_cache.stride(2), |
| 881 | v_cache.stride(3), |
| 882 | config.sm_scale, |
| 883 | page_table.stride(0), |
| 884 | config.page_size, |
| 885 | config.num_kv_splits, |
| 886 | max_seqlen_q, |
| 887 | config.num_q_heads, |
| 888 | config.num_kv_heads, |
| 889 | config.head_dim, |
| 890 | config.block_m, |
| 891 | config.block_n, |
| 892 | config.is_sliding, |
| 893 | config.window_left, |
| 894 | rel_logits.stride(0), |
| 895 | rel_logits.stride(1), |
| 896 | rel_logits.stride(2), |
| 897 | rel_logits.shape[2], |
| 898 | config.rel_bias_qk_scale, |
| 899 | is_fp8, |
| 900 | num_warps=1, |
| 901 | ) |
| 902 | |
| 903 | reduce_grid = (total_q, config.num_q_heads) |
| 904 | _rel_mha_decode_reduce_fp16[reduce_grid]( |
| 905 | mid_o, |
| 906 | mid_lse, |
| 907 | output, |
| 908 | cache_seqlens, |
| 909 | config.sm_scale, |
| 910 | page_table.stride(0), |
| 911 | config.num_kv_splits, |
| 912 | max_seqlen_q, |
| 913 | config.page_size, |
| 914 | config.num_q_heads, |
| 915 | config.num_kv_heads, |
| 916 | config.head_dim, |
| 917 | config.block_m, |
| 918 | config.block_n, |
| 919 | config.is_sliding, |
| 920 | config.window_left, |
| 921 | is_fp8, |
| 922 | num_warps=1, |
| 923 | ) |
| 924 | return output |