| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm paged attention with Inkling's query-dependent relative bias. |
| 4 | |
| 5 | The NVIDIA implementation uses the score-mod hook in tml-fa4. ROCm Flash |
| 6 | Attention and AITER do not expose an equivalent hook, so this module implements |
| 7 | the same operation directly in Triton. Query heads belonging to one KV head |
| 8 | are processed together and KV pages are gathered through vLLM's block table. |
| 9 | """ |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | import os |
| 14 | from typing import cast |
| 15 | |
| 16 | import torch |
| 17 | |
| 18 | from vllm.models.inkling.amd.ops.rel_attention_decode import ( |
| 19 | inkling_rel_attention_split_kv_decode, |
| 20 | use_split_kv_decode, |
| 21 | ) |
| 22 | from vllm.platforms.rocm import on_gfx950 |
| 23 | from vllm.triton_utils import tl, triton |
| 24 | |
| 25 | |
| 26 | def bucket_max_seqlen_q(max_seqlen_q: int) -> int: |
| 27 | """Round the scheduling bound up to a power of two.""" |
| 28 | return 1 << max(0, max_seqlen_q - 1).bit_length() |
| 29 | |
| 30 | |
| 31 | def use_gfx950_gluon_decode( |
| 32 | *, max_query_len: int, page_size: int, head_dim: int |
| 33 | ) -> bool: |
| 34 | """Use the vendored TokenSpeed CDNA4 decode kernel where it is supported.""" |
| 35 | return ( |
| 36 | os.getenv("INKLING_GFX950_GLUON", "1") == "1" |
| 37 | and on_gfx950() |
| 38 | and max_query_len == 1 |
| 39 | and page_size in (64, 128, 256) |
| 40 | and head_dim in (64, 128) |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | def use_gfx950_gluon_extend( |
| 45 | *, |
| 46 | max_query_len: int, |
| 47 | max_kv_len: int, |
| 48 | page_size: int, |
| 49 | head_dim: int, |
| 50 | window_left: int, |
| 51 | ) -> bool: |
| 52 | """Use Gluon only for the long full-attention extend regime it wins.""" |
| 53 | return ( |
| 54 | os.getenv("INKLING_GFX950_GLUON", "1") == "1" |
| 55 | and on_gfx950() |
| 56 | and max_query_len > 1 |
| 57 | and max_kv_len >= 8192 |
| 58 | and window_left < 0 |
| 59 | and page_size in (64, 128, 256) |
| 60 | and head_dim in (64, 128) |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def inkling_fa4_num_splits( |
| 65 | *, |
| 66 | is_local: bool, |
| 67 | batch_size: int, |
| 68 | max_query_len: int, |
| 69 | num_heads: int, |
| 70 | num_kv_heads: int, |
| 71 | max_kv_len: int, |
| 72 | ) -> int: |
| 73 | """Keep the NVIDIA-facing split heuristic as API-compatible metadata. |
| 74 | |
| 75 | The ROCm Triton implementation performs online softmax in one program and |
| 76 | does not consume the result. Keeping this function unchanged avoids |
| 77 | platform-specific scheduling branches in :mod:`inkling.amd.attention`. |
| 78 | """ |
| 79 | if is_local: |
| 80 | return 1 |
| 81 | |
| 82 | q_rows = max_query_len * (num_heads // num_kv_heads) |
| 83 | q_tiles = (q_rows + 255) // 256 |
| 84 | base_ctas = batch_size * num_kv_heads * q_tiles |
| 85 | target_ctas = ( |
| 86 | 256 if q_tiles == 1 and batch_size == 1 else (128 if q_tiles == 1 else 64) |
| 87 | ) |
| 88 | max_splits = 128 |
| 89 | if q_tiles == 1 and batch_size == 1: |
| 90 | if num_kv_heads == 8: |
| 91 | max_splits = 16 |
| 92 | elif num_kv_heads == 4 or max_kv_len <= 8192: |
| 93 | max_splits = 32 |
| 94 | elif max_kv_len <= 65536: |
| 95 | max_splits = 64 |
| 96 | return max( |
| 97 | 1, |
| 98 | min(target_ctas // base_ctas, max_splits, (max_kv_len + 127) // 128), |
| 99 | ) |
| 100 | |
| 101 | |
| 102 | @triton.heuristics( |
| 103 | { |
| 104 | "BLOCK_D": lambda args: triton.next_power_of_2(args["head_dim"]), |
| 105 | "BLOCK_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]), |
| 106 | "BLOCK_QH": lambda args: ( |
| 107 | args["BLOCK_Q"] * triton.next_power_of_2(args["gqa_group_size"]) |
| 108 | ), |
| 109 | } |
| 110 | ) |
| 111 | @triton.jit(do_not_specialize_on_alignment=["cache_seqlens", "cu_seqlens_q"]) |
| 112 | def _inkling_rel_attention_kernel( |
| 113 | q_ptr, # [total_q, Hq, D] |
| 114 | k_ptr, # [blocks, page, Hkv, D] |
| 115 | v_ptr, # [blocks, page, Hkv, D] |
| 116 | rel_ptr, # [total_q, Hq, rel_extent] |
| 117 | out_ptr, # [total_q, Hq, D] |
| 118 | block_table_ptr, # [batch, max_pages] |
| 119 | cache_seqlens, |
| 120 | cu_seqlens_q, |
| 121 | gqa_group_size, |
| 122 | head_dim, |
| 123 | softmax_scale, |
| 124 | stride_q_t, |
| 125 | stride_q_h, |
| 126 | stride_q_d, |
| 127 | stride_k_b, |
| 128 | stride_k_p, |
| 129 | stride_k_h, |
| 130 | stride_k_d, |
| 131 | stride_v_b, |
| 132 | stride_v_p, |
| 133 | stride_v_h, |
| 134 | stride_v_d, |
| 135 | stride_r_t, |
| 136 | stride_r_h, |
| 137 | stride_r_e, |
| 138 | stride_o_t, |
| 139 | stride_o_h, |
| 140 | stride_o_d, |
| 141 | stride_bt_b, |
| 142 | page_size: tl.constexpr, |
| 143 | rel_extent: tl.constexpr, |
| 144 | window_left: tl.constexpr, |
| 145 | BLOCK_Q: tl.constexpr, |
| 146 | BLOCK_K: tl.constexpr, |
| 147 | BLOCK_D: tl.constexpr, |
| 148 | BLOCK_H: tl.constexpr, |
| 149 | BLOCK_QH: tl.constexpr, |
| 150 | ): |
| 151 | # A program owns BLOCK_Q query positions and every query head sharing one |
| 152 | # KV head. This makes the QK/PV operations MFMA-friendly on CDNA. |
| 153 | pid_q = tl.program_id(0) |
| 154 | pid_kh = tl.program_id(1) |
| 155 | pid_b = tl.program_id(2) |
| 156 | |
| 157 | q_start = tl.load(cu_seqlens_q + pid_b) |
| 158 | q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start |
| 159 | q_block = pid_q * BLOCK_Q |
| 160 | if q_block >= q_len: |
| 161 | return |
| 162 | |
| 163 | kv_len = tl.load(cache_seqlens + pid_b) |
| 164 | prefix_len = kv_len - q_len |
| 165 | q_head_start = pid_kh * gqa_group_size |
| 166 | bt_row = block_table_ptr + pid_b * stride_bt_b |
| 167 | |
| 168 | q_block_ptr = tl.make_block_ptr( |
| 169 | base=q_ptr + q_start * stride_q_t + q_head_start * stride_q_h, |
| 170 | shape=(q_len, gqa_group_size, head_dim), |
| 171 | strides=(stride_q_t, stride_q_h, stride_q_d), |
| 172 | offsets=(q_block, 0, 0), |
| 173 | block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), |
| 174 | order=(2, 1, 0), |
| 175 | ) |
| 176 | q = tl.load(q_block_ptr, boundary_check=(0, 1, 2), padding_option="zero") |
| 177 | q = tl.reshape(q, (BLOCK_QH, BLOCK_D)) |
| 178 | |
| 179 | q_rows = q_block + tl.arange(0, BLOCK_Q) |
| 180 | q_abs = prefix_len + q_rows |
| 181 | q_valid = q_rows < q_len |
| 182 | off_d = tl.arange(0, BLOCK_D) |
| 183 | d_valid = off_d < head_dim |
| 184 | |
| 185 | m_i = tl.full((BLOCK_QH,), float("-inf"), dtype=tl.float32) |
| 186 | l_i = tl.zeros((BLOCK_QH,), dtype=tl.float32) |
| 187 | acc = tl.zeros((BLOCK_QH, BLOCK_D), dtype=tl.float32) |
| 188 | log2e: tl.constexpr = 1.4426950408889634 |
| 189 | |
| 190 | for k_start in range(0, kv_len, BLOCK_K): |
| 191 | k_pos = k_start + tl.arange(0, BLOCK_K) |
| 192 | k_valid = k_pos < kv_len |
| 193 | # Masked lanes still need an in-range page-table load. |
| 194 | safe_pos = tl.minimum(k_pos, tl.maximum(kv_len - 1, 0)) |
| 195 | page = tl.load(bt_row + safe_pos // page_size).to(tl.int64) |
| 196 | page_off = safe_pos % page_size |
| 197 | |
| 198 | k = tl.load( |
| 199 | k_ptr |
| 200 | + page[None, :] * stride_k_b |
| 201 | + page_off[None, :] * stride_k_p |
| 202 | + pid_kh * stride_k_h |
| 203 | + off_d[:, None] * stride_k_d, |
| 204 | mask=d_valid[:, None] & k_valid[None, :], |
| 205 | other=0.0, |
| 206 | ) |
| 207 | scores = tl.dot(q, k) * (softmax_scale * log2e) |
| 208 | |
| 209 | dist = q_abs[:, None] - k_pos[None, :] |
| 210 | score_valid = q_valid[:, None] & k_valid[None, :] & (dist >= 0) |
| 211 | if window_left >= 0: |
| 212 | score_valid &= dist <= window_left |
| 213 | |
| 214 | # rel_logits is query- and head-dependent. Expand [Q, K] distance |
| 215 | # indices across the GQA heads, then flatten to match QK's row order. |
| 216 | off_h = tl.arange(0, BLOCK_H) |
| 217 | rel_dist = dist[:, None, :] |
| 218 | rel_valid = ( |
| 219 | q_valid[:, None, None] |
| 220 | & (off_h[None, :, None] < gqa_group_size) |
| 221 | & (rel_dist >= 0) |
| 222 | & (rel_dist < rel_extent) |
| 223 | ) |
| 224 | safe_dist = tl.maximum(0, tl.minimum(rel_dist, rel_extent - 1)) |
| 225 | bias = tl.load( |
| 226 | rel_ptr |
| 227 | + (q_start + q_rows[:, None, None]) * stride_r_t |
| 228 | + (q_head_start + off_h[None, :, None]) * stride_r_h |
| 229 | + safe_dist * stride_r_e, |
| 230 | mask=rel_valid, |
| 231 | other=0.0, |
| 232 | ) |
| 233 | bias = tl.reshape(bias, (BLOCK_QH, BLOCK_K)) |
| 234 | scores += bias.to(tl.float32) * log2e |
| 235 | score_valid = tl.reshape( |
| 236 | score_valid[:, None, :] & (off_h[None, :, None] < gqa_group_size), |
| 237 | (BLOCK_QH, BLOCK_K), |
| 238 | ) |
| 239 | scores = tl.where(score_valid, scores, float("-inf")) |
| 240 | |
| 241 | m_ij = tl.maximum(m_i, tl.max(scores, axis=1)) |
| 242 | has_scores = tl.sum(score_valid.to(tl.int32), axis=1) > 0 |
| 243 | # A sliding-window query can have entire early KV tiles masked. Avoid |
| 244 | # the -inf - -inf NaN in online softmax until its first live tile. |
| 245 | alpha = tl.where(has_scores, tl.exp2(m_i - m_ij), 1.0) |
| 246 | p = tl.where( |
| 247 | score_valid, |
| 248 | tl.exp2(scores - m_ij[:, None]), |
| 249 | 0.0, |
| 250 | ) |
| 251 | l_i = l_i * alpha + tl.sum(p, axis=1) |
| 252 | acc *= alpha[:, None] |
| 253 | |
| 254 | v = tl.load( |
| 255 | v_ptr |
| 256 | + page[:, None] * stride_v_b |
| 257 | + page_off[:, None] * stride_v_p |
| 258 | + pid_kh * stride_v_h |
| 259 | + off_d[None, :] * stride_v_d, |
| 260 | mask=k_valid[:, None] & d_valid[None, :], |
| 261 | other=0.0, |
| 262 | ) |
| 263 | acc += tl.dot(p.to(v.dtype), v) |
| 264 | m_i = m_ij |
| 265 | |
| 266 | acc /= l_i[:, None] |
| 267 | acc = tl.reshape(acc, (BLOCK_Q, BLOCK_H, BLOCK_D)) |
| 268 | out_block_ptr = tl.make_block_ptr( |
| 269 | base=out_ptr + q_start * stride_o_t + q_head_start * stride_o_h, |
| 270 | shape=(q_len, gqa_group_size, head_dim), |
| 271 | strides=(stride_o_t, stride_o_h, stride_o_d), |
| 272 | offsets=(q_block, 0, 0), |
| 273 | block_shape=(BLOCK_Q, BLOCK_H, BLOCK_D), |
| 274 | order=(2, 1, 0), |
| 275 | ) |
| 276 | tl.store( |
| 277 | out_block_ptr, |
| 278 | acc.to(out_ptr.dtype.element_ty), |
| 279 | boundary_check=(0, 1, 2), |
| 280 | ) |
| 281 | |
| 282 | |
| 283 | @torch.no_grad() |
| 284 | def inkling_fa4_rel_attention( |
| 285 | q: torch.Tensor, |
| 286 | key_cache: torch.Tensor, |
| 287 | value_cache: torch.Tensor, |
| 288 | *, |
| 289 | block_table: torch.Tensor, |
| 290 | cache_seqlens: torch.Tensor, |
| 291 | cu_seqlens_q: torch.Tensor, |
| 292 | max_seqlen_q: int, |
| 293 | softmax_scale: float, |
| 294 | causal: bool, |
| 295 | window_size: tuple[int, int], |
| 296 | rel_extent: int, |
| 297 | rel_logits: torch.Tensor, |
| 298 | num_splits: int = 32, |
| 299 | max_kv_len: int | None = None, |
| 300 | out: torch.Tensor | None = None, |
| 301 | ) -> torch.Tensor: |
| 302 | """Paged varlen attention with Inkling's relative score modification.""" |
| 303 | del num_splits |
| 304 | if not causal or window_size[1] not in (0, -1): |
| 305 | raise NotImplementedError("Inkling ROCm attention requires causal masking") |
| 306 | if q.ndim != 3 or key_cache.ndim != 4 or value_cache.ndim != 4: |
| 307 | raise ValueError("expected q [T,H,D] and paged K/V [B,P,Hkv,D]") |
| 308 | if rel_logits.shape != (q.shape[0], q.shape[1], rel_extent): |
| 309 | raise ValueError( |
| 310 | f"relative logits have shape {tuple(rel_logits.shape)}, expected " |
| 311 | f"{(q.shape[0], q.shape[1], rel_extent)}" |
| 312 | ) |
| 313 | |
| 314 | num_kv_heads = key_cache.shape[2] |
| 315 | if q.shape[1] % num_kv_heads: |
| 316 | raise ValueError("query heads must be divisible by KV heads") |
| 317 | if out is None: |
| 318 | out = torch.empty_like(q) |
| 319 | gqa_group_size = q.shape[1] // num_kv_heads |
| 320 | batch = cu_seqlens_q.shape[0] - 1 |
| 321 | if max_kv_len is None: |
| 322 | max_kv_len = key_cache.shape[0] * key_cache.shape[1] |
| 323 | if use_gfx950_gluon_decode( |
| 324 | max_query_len=max_seqlen_q, |
| 325 | page_size=key_cache.shape[1], |
| 326 | head_dim=q.shape[2], |
| 327 | ): |
| 328 | # Import only after the architecture guard. The implementation uses |
| 329 | # gfx950-only CDNA4 Gluon layouts and async-copy operations. |
| 330 | from vllm.models.inkling.amd.ops.gluon.rel_mha_decode_gfx950 import ( |
| 331 | gluon_rel_mha_decode_gfx950, |
| 332 | ) |
| 333 | |
| 334 | return gluon_rel_mha_decode_gfx950( |
| 335 | q, |
| 336 | key_cache, |
| 337 | value_cache, |
| 338 | block_table, |
| 339 | cache_seqlens, |
| 340 | max_kv_len, |
| 341 | rel_logits, |
| 342 | cu_seqlens_q, |
| 343 | max_seqlen_q=max_seqlen_q, |
| 344 | window_left=window_size[0], |
| 345 | softmax_scale=softmax_scale, |
| 346 | out=out, |
| 347 | ) |
| 348 | if use_gfx950_gluon_extend( |
| 349 | max_query_len=max_seqlen_q, |
| 350 | max_kv_len=max_kv_len, |
| 351 | page_size=key_cache.shape[1], |
| 352 | head_dim=q.shape[2], |
| 353 | window_left=window_size[0], |
| 354 | ): |
| 355 | from vllm.models.inkling.amd.ops.gluon.rel_mha_extend_gfx950 import ( |
| 356 | gluon_rel_mha_extend_gfx950, |
| 357 | ) |
| 358 | |
| 359 | # The copied kernel retains TokenSpeed's cu_seqlens_kv parameter for |
| 360 | # API compatibility but does not consume it. |
| 361 | return cast( |
| 362 | torch.Tensor, |
| 363 | gluon_rel_mha_extend_gfx950( |
| 364 | q, |
| 365 | cu_seqlens_q, |
| 366 | cu_seqlens_q, |
| 367 | key_cache, |
| 368 | value_cache, |
| 369 | block_table, |
| 370 | cache_seqlens, |
| 371 | window_left=window_size[0], |
| 372 | max_seqlen_q=max_seqlen_q, |
| 373 | max_seqlen_k=max_kv_len, |
| 374 | rel_logits=rel_logits, |
| 375 | softmax_scale=softmax_scale, |
| 376 | out=out, |
| 377 | ), |
| 378 | ) |
| 379 | if use_split_kv_decode( |
| 380 | max_query_len=max_seqlen_q, |
| 381 | max_kv_len=max_kv_len, |
| 382 | page_size=key_cache.shape[1], |
| 383 | window_left=window_size[0], |
| 384 | ): |
| 385 | return inkling_rel_attention_split_kv_decode( |
| 386 | q, |
| 387 | key_cache, |
| 388 | value_cache, |
| 389 | block_table=block_table, |
| 390 | cache_seqlens=cache_seqlens, |
| 391 | softmax_scale=softmax_scale, |
| 392 | window_left=window_size[0], |
| 393 | rel_extent=rel_extent, |
| 394 | rel_logits=rel_logits, |
| 395 | max_kv_len=max_kv_len, |
| 396 | out=out, |
| 397 | ) |
| 398 | block_q = 1 if max_seqlen_q == 1 else 4 |
| 399 | grid = (triton.cdiv(max_seqlen_q, block_q), num_kv_heads, batch) |
| 400 | _inkling_rel_attention_kernel[grid]( |
| 401 | q, |
| 402 | key_cache, |
| 403 | value_cache, |
| 404 | rel_logits, |
| 405 | out, |
| 406 | block_table, |
| 407 | cache_seqlens, |
| 408 | cu_seqlens_q, |
| 409 | gqa_group_size, |
| 410 | q.shape[2], |
| 411 | softmax_scale, |
| 412 | q.stride(0), |
| 413 | q.stride(1), |
| 414 | q.stride(2), |
| 415 | key_cache.stride(0), |
| 416 | key_cache.stride(1), |
| 417 | key_cache.stride(2), |
| 418 | key_cache.stride(3), |
| 419 | value_cache.stride(0), |
| 420 | value_cache.stride(1), |
| 421 | value_cache.stride(2), |
| 422 | value_cache.stride(3), |
| 423 | rel_logits.stride(0), |
| 424 | rel_logits.stride(1), |
| 425 | rel_logits.stride(2), |
| 426 | out.stride(0), |
| 427 | out.stride(1), |
| 428 | out.stride(2), |
| 429 | block_table.stride(0), |
| 430 | page_size=key_cache.shape[1], |
| 431 | rel_extent=rel_extent, |
| 432 | window_left=window_size[0], |
| 433 | BLOCK_Q=block_q, |
| 434 | BLOCK_K=64, |
| 435 | num_warps=4, |
| 436 | num_stages=1, |
| 437 | ) |
| 438 | return out |