| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Correctness tests for Inkling's ROCm relative-bias paged attention.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.models.inkling.amd.ops.fa4_rel_attention import ( |
| 9 | bucket_max_seqlen_q, |
| 10 | inkling_fa4_rel_attention, |
| 11 | use_gfx950_gluon_decode, |
| 12 | use_gfx950_gluon_extend, |
| 13 | ) |
| 14 | from vllm.models.inkling.amd.ops.rel_attention_decode import ( |
| 15 | decode_split_count, |
| 16 | use_split_kv_decode, |
| 17 | ) |
| 18 | |
| 19 | HEAD_DIM = 128 |
| 20 | DTYPE = torch.bfloat16 |
| 21 | |
| 22 | |
| 23 | def _reference( |
| 24 | q: torch.Tensor, |
| 25 | key_cache: torch.Tensor, |
| 26 | value_cache: torch.Tensor, |
| 27 | rel_logits: torch.Tensor, |
| 28 | block_table: torch.Tensor, |
| 29 | q_lens: list[int], |
| 30 | kv_lens: list[int], |
| 31 | rel_extent: int, |
| 32 | window_left: int | None, |
| 33 | ) -> torch.Tensor: |
| 34 | num_heads = q.shape[1] |
| 35 | num_kv_heads = key_cache.shape[2] |
| 36 | gqa_group = num_heads // num_kv_heads |
| 37 | out: list[torch.Tensor] = [] |
| 38 | q_start = 0 |
| 39 | |
| 40 | for req, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)): |
| 41 | page_size = key_cache.shape[1] |
| 42 | num_pages = (kv_len + page_size - 1) // page_size |
| 43 | page_ids = block_table[req, :num_pages].long() |
| 44 | k = key_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len] |
| 45 | v = value_cache[page_ids].reshape(-1, num_kv_heads, HEAD_DIM)[:kv_len] |
| 46 | k = k.float().repeat_interleave(gqa_group, dim=1) |
| 47 | v = v.float().repeat_interleave(gqa_group, dim=1) |
| 48 | |
| 49 | q_req = q[q_start : q_start + q_len].float() |
| 50 | rel_req = rel_logits[q_start : q_start + q_len].float() |
| 51 | scores = torch.einsum("qhd,khd->hqk", q_req, k) / HEAD_DIM |
| 52 | |
| 53 | q_pos = torch.arange(q_len, device=q.device)[:, None] + kv_len - q_len |
| 54 | k_pos = torch.arange(kv_len, device=q.device)[None, :] |
| 55 | distance = q_pos - k_pos |
| 56 | rel_idx = distance.clamp(0, rel_extent - 1) |
| 57 | bias = rel_req.permute(1, 0, 2).gather( |
| 58 | 2, rel_idx[None].expand(num_heads, -1, -1) |
| 59 | ) |
| 60 | in_rel_extent = (distance >= 0) & (distance < rel_extent) |
| 61 | scores += torch.where(in_rel_extent[None], bias, 0.0) |
| 62 | |
| 63 | masked = distance < 0 |
| 64 | if window_left is not None: |
| 65 | masked |= distance > window_left |
| 66 | scores.masked_fill_(masked[None], float("-inf")) |
| 67 | out.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), v)) |
| 68 | q_start += q_len |
| 69 | |
| 70 | return torch.cat(out).to(DTYPE) |
| 71 | |
| 72 | |
| 73 | def test_query_length_bucket(): |
| 74 | assert bucket_max_seqlen_q(1) == 1 |
| 75 | assert bucket_max_seqlen_q(17) == 32 |
| 76 | assert bucket_max_seqlen_q(33) == 64 |
| 77 | |
| 78 | |
| 79 | def test_split_kv_dispatch_thresholds(): |
| 80 | assert not use_split_kv_decode( |
| 81 | max_query_len=4, |
| 82 | max_kv_len=65536, |
| 83 | page_size=128, |
| 84 | window_left=-1, |
| 85 | ) |
| 86 | assert not use_split_kv_decode( |
| 87 | max_query_len=1, |
| 88 | max_kv_len=4096, |
| 89 | page_size=16, |
| 90 | window_left=-1, |
| 91 | ) |
| 92 | assert use_split_kv_decode( |
| 93 | max_query_len=1, |
| 94 | max_kv_len=8192, |
| 95 | page_size=16, |
| 96 | window_left=-1, |
| 97 | ) |
| 98 | assert use_split_kv_decode( |
| 99 | max_query_len=1, |
| 100 | max_kv_len=512, |
| 101 | page_size=128, |
| 102 | window_left=511, |
| 103 | ) |
| 104 | assert decode_split_count(512, 511) == 4 |
| 105 | assert decode_split_count(65536, -1) == 32 |
| 106 | |
| 107 | |
| 108 | def test_gfx950_gluon_dispatch_is_strict(monkeypatch: pytest.MonkeyPatch): |
| 109 | import vllm.models.inkling.amd.ops.fa4_rel_attention as rel_attention |
| 110 | |
| 111 | monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True) |
| 112 | assert use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128) |
| 113 | assert not use_gfx950_gluon_decode(max_query_len=1, page_size=16, head_dim=128) |
| 114 | assert not use_gfx950_gluon_decode(max_query_len=4, page_size=128, head_dim=128) |
| 115 | assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=256) |
| 116 | |
| 117 | monkeypatch.setattr(rel_attention, "on_gfx950", lambda: False) |
| 118 | assert not use_gfx950_gluon_decode(max_query_len=1, page_size=128, head_dim=128) |
| 119 | |
| 120 | monkeypatch.setattr(rel_attention, "on_gfx950", lambda: True) |
| 121 | assert use_gfx950_gluon_extend( |
| 122 | max_query_len=4, |
| 123 | max_kv_len=8192, |
| 124 | page_size=128, |
| 125 | head_dim=128, |
| 126 | window_left=-1, |
| 127 | ) |
| 128 | assert not use_gfx950_gluon_extend( |
| 129 | max_query_len=4, |
| 130 | max_kv_len=8192, |
| 131 | page_size=128, |
| 132 | head_dim=128, |
| 133 | window_left=511, |
| 134 | ) |
| 135 | assert not use_gfx950_gluon_extend( |
| 136 | max_query_len=4, |
| 137 | max_kv_len=8192, |
| 138 | page_size=16, |
| 139 | head_dim=128, |
| 140 | window_left=-1, |
| 141 | ) |
| 142 | |
| 143 | |
| 144 | @pytest.mark.parametrize( |
| 145 | ("page_size", "window_left", "max_kv_len"), |
| 146 | [ |
| 147 | (16, None, None), |
| 148 | (16, 15, None), |
| 149 | (128, None, 8192), |
| 150 | ], |
| 151 | ) |
| 152 | @torch.inference_mode() |
| 153 | def test_ragged_multi_page_relative_attention( |
| 154 | page_size: int, |
| 155 | window_left: int | None, |
| 156 | max_kv_len: int | None, |
| 157 | ): |
| 158 | """Covers full/chunked prefill, decode, GQA, and the local window.""" |
| 159 | torch.manual_seed(19 + int(window_left is not None)) |
| 160 | device = "cuda" |
| 161 | q_lens = [17, 1] |
| 162 | kv_lens = [35, 80] |
| 163 | num_heads = 8 |
| 164 | num_kv_heads = 2 |
| 165 | rel_extent = 16 |
| 166 | max_pages = max((length + page_size - 1) // page_size for length in kv_lens) |
| 167 | |
| 168 | q = torch.randn(sum(q_lens), num_heads, HEAD_DIM, device=device) |
| 169 | q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) |
| 170 | key_cache = torch.randn( |
| 171 | 1 + len(q_lens) * max_pages, |
| 172 | page_size, |
| 173 | num_kv_heads, |
| 174 | HEAD_DIM, |
| 175 | device=device, |
| 176 | ) |
| 177 | key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) |
| 178 | value_cache = torch.randn_like(key_cache) |
| 179 | rel_logits = torch.randn( |
| 180 | sum(q_lens), num_heads, rel_extent, device=device, dtype=DTYPE |
| 181 | ) |
| 182 | |
| 183 | block_table = torch.stack( |
| 184 | [ |
| 185 | torch.arange(1 + req * max_pages, 1 + (req + 1) * max_pages) |
| 186 | for req in range(len(q_lens)) |
| 187 | ] |
| 188 | ).to(device=device, dtype=torch.int32) |
| 189 | cu_seqlens_q = torch.tensor( |
| 190 | [0, *torch.tensor(q_lens).cumsum(0).tolist()], |
| 191 | device=device, |
| 192 | dtype=torch.int32, |
| 193 | ) |
| 194 | cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) |
| 195 | window_size = (-1, -1) if window_left is None else (window_left, 0) |
| 196 | |
| 197 | preallocated = torch.empty_like(q) |
| 198 | actual = inkling_fa4_rel_attention( |
| 199 | q, |
| 200 | key_cache, |
| 201 | value_cache, |
| 202 | block_table=block_table, |
| 203 | cache_seqlens=cache_seqlens, |
| 204 | cu_seqlens_q=cu_seqlens_q, |
| 205 | max_seqlen_q=bucket_max_seqlen_q(max(q_lens)), |
| 206 | softmax_scale=1.0 / HEAD_DIM, |
| 207 | causal=True, |
| 208 | window_size=window_size, |
| 209 | rel_extent=rel_extent, |
| 210 | rel_logits=rel_logits, |
| 211 | max_kv_len=max_kv_len, |
| 212 | out=preallocated, |
| 213 | ) |
| 214 | assert actual.data_ptr() == preallocated.data_ptr() |
| 215 | |
| 216 | expected = _reference( |
| 217 | q, |
| 218 | key_cache, |
| 219 | value_cache, |
| 220 | rel_logits, |
| 221 | block_table, |
| 222 | q_lens, |
| 223 | kv_lens, |
| 224 | rel_extent, |
| 225 | window_left, |
| 226 | ) |
| 227 | torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) |
| 228 | |
| 229 | |
| 230 | @pytest.mark.parametrize( |
| 231 | ("page_size", "kv_lens", "num_kv_heads", "rel_extent", "window_left"), |
| 232 | [ |
| 233 | (16, [257, 513], 1, 64, None), |
| 234 | (128, [257, 512], 2, 512, 511), |
| 235 | ], |
| 236 | ) |
| 237 | @torch.inference_mode() |
| 238 | def test_split_kv_decode_matches_reference( |
| 239 | page_size: int, |
| 240 | kv_lens: list[int], |
| 241 | num_kv_heads: int, |
| 242 | rel_extent: int, |
| 243 | window_left: int | None, |
| 244 | ): |
| 245 | """Covers page-16 split-KV and page-128 gfx950 Gluon decode.""" |
| 246 | torch.manual_seed(41 + page_size) |
| 247 | device = "cuda" |
| 248 | batch_size = len(kv_lens) |
| 249 | q_lens = [1] * batch_size |
| 250 | num_heads = 8 |
| 251 | max_pages = max((length + page_size - 1) // page_size for length in kv_lens) |
| 252 | |
| 253 | q = torch.randn(batch_size, num_heads, HEAD_DIM, device=device) |
| 254 | q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) |
| 255 | key_cache = torch.randn( |
| 256 | batch_size * max_pages, |
| 257 | page_size, |
| 258 | num_kv_heads, |
| 259 | HEAD_DIM, |
| 260 | device=device, |
| 261 | ) |
| 262 | key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) |
| 263 | value_cache = torch.randn_like(key_cache) |
| 264 | rel_logits = torch.randn( |
| 265 | batch_size, |
| 266 | num_heads, |
| 267 | rel_extent, |
| 268 | device=device, |
| 269 | dtype=DTYPE, |
| 270 | ) |
| 271 | block_table = torch.arange( |
| 272 | batch_size * max_pages, |
| 273 | device=device, |
| 274 | dtype=torch.int32, |
| 275 | ).view(batch_size, max_pages) |
| 276 | cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) |
| 277 | cu_seqlens_q = torch.arange( |
| 278 | batch_size + 1, |
| 279 | device=device, |
| 280 | dtype=torch.int32, |
| 281 | ) |
| 282 | window_size = (-1, -1) if window_left is None else (window_left, 0) |
| 283 | max_kv_len = 8192 if page_size == 16 else max(kv_lens) |
| 284 | |
| 285 | actual = inkling_fa4_rel_attention( |
| 286 | q, |
| 287 | key_cache, |
| 288 | value_cache, |
| 289 | block_table=block_table, |
| 290 | cache_seqlens=cache_seqlens, |
| 291 | cu_seqlens_q=cu_seqlens_q, |
| 292 | max_seqlen_q=1, |
| 293 | softmax_scale=1.0 / HEAD_DIM, |
| 294 | causal=True, |
| 295 | window_size=window_size, |
| 296 | rel_extent=rel_extent, |
| 297 | rel_logits=rel_logits, |
| 298 | max_kv_len=max_kv_len, |
| 299 | ) |
| 300 | expected = _reference( |
| 301 | q, |
| 302 | key_cache, |
| 303 | value_cache, |
| 304 | rel_logits, |
| 305 | block_table, |
| 306 | q_lens, |
| 307 | kv_lens, |
| 308 | rel_extent, |
| 309 | window_left, |
| 310 | ) |
| 311 | torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) |
| 312 | |
| 313 | |
| 314 | @pytest.mark.parametrize( |
| 315 | ( |
| 316 | "q_len", |
| 317 | "kv_len", |
| 318 | "num_kv_heads", |
| 319 | "rel_extent", |
| 320 | "window_left", |
| 321 | "gluon_enabled", |
| 322 | ), |
| 323 | [ |
| 324 | (1, 513, 2, 512, 511, True), |
| 325 | (1, 513, 2, 512, 511, False), |
| 326 | (4, 8192, 1, 1024, None, True), |
| 327 | ], |
| 328 | ) |
| 329 | @torch.inference_mode() |
| 330 | def test_gfx950_page128_packed_kv_views_match_reference( |
| 331 | monkeypatch: pytest.MonkeyPatch, |
| 332 | q_len: int, |
| 333 | kv_len: int, |
| 334 | num_kv_heads: int, |
| 335 | rel_extent: int, |
| 336 | window_left: int | None, |
| 337 | gluon_enabled: bool, |
| 338 | ): |
| 339 | """Cover both Gluon and split-KV against vLLM's packed KV allocation.""" |
| 340 | monkeypatch.setenv( |
| 341 | "INKLING_GFX950_GLUON", |
| 342 | "1" if gluon_enabled else "0", |
| 343 | ) |
| 344 | torch.manual_seed(71 + q_len) |
| 345 | device = "cuda" |
| 346 | page_size = 128 |
| 347 | batch_size = 2 |
| 348 | num_heads = 8 |
| 349 | q_lens = [q_len] * batch_size |
| 350 | kv_lens = [kv_len] * batch_size |
| 351 | pages_per_req = (kv_len + page_size - 1) // page_size |
| 352 | |
| 353 | q = torch.randn( |
| 354 | batch_size * q_len, |
| 355 | num_heads, |
| 356 | HEAD_DIM, |
| 357 | device=device, |
| 358 | dtype=DTYPE, |
| 359 | ) |
| 360 | q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) |
| 361 | |
| 362 | # FlashAttentionBackend allocates logical [block, head, page, 2 * dim]. |
| 363 | # Inkling transposes to [block, page, head, 2 * dim] and splits K/V, |
| 364 | # producing non-contiguous views whose page/head strides include both. |
| 365 | packed_kv = torch.randn( |
| 366 | batch_size * pages_per_req, |
| 367 | num_kv_heads, |
| 368 | page_size, |
| 369 | 2 * HEAD_DIM, |
| 370 | device=device, |
| 371 | dtype=DTYPE, |
| 372 | ) |
| 373 | key_cache, value_cache = packed_kv.transpose(1, 2).split(HEAD_DIM, dim=-1) |
| 374 | assert not key_cache.is_contiguous() |
| 375 | assert key_cache.stride() == value_cache.stride() |
| 376 | |
| 377 | rel_logits = torch.randn( |
| 378 | batch_size * q_len, |
| 379 | num_heads, |
| 380 | rel_extent, |
| 381 | device=device, |
| 382 | dtype=DTYPE, |
| 383 | ) |
| 384 | block_table = torch.arange( |
| 385 | batch_size * pages_per_req, |
| 386 | device=device, |
| 387 | dtype=torch.int32, |
| 388 | ).view(batch_size, pages_per_req) |
| 389 | cache_seqlens = torch.tensor(kv_lens, device=device, dtype=torch.int32) |
| 390 | cu_seqlens_q = torch.arange( |
| 391 | 0, |
| 392 | batch_size * q_len + 1, |
| 393 | q_len, |
| 394 | device=device, |
| 395 | dtype=torch.int32, |
| 396 | ) |
| 397 | window_size = (-1, -1) if window_left is None else (window_left, 0) |
| 398 | |
| 399 | actual = inkling_fa4_rel_attention( |
| 400 | q, |
| 401 | key_cache, |
| 402 | value_cache, |
| 403 | block_table=block_table, |
| 404 | cache_seqlens=cache_seqlens, |
| 405 | cu_seqlens_q=cu_seqlens_q, |
| 406 | max_seqlen_q=q_len, |
| 407 | softmax_scale=1.0 / HEAD_DIM, |
| 408 | causal=True, |
| 409 | window_size=window_size, |
| 410 | rel_extent=rel_extent, |
| 411 | rel_logits=rel_logits, |
| 412 | max_kv_len=kv_len, |
| 413 | ) |
| 414 | expected = _reference( |
| 415 | q, |
| 416 | key_cache, |
| 417 | value_cache, |
| 418 | rel_logits, |
| 419 | block_table, |
| 420 | q_lens, |
| 421 | kv_lens, |
| 422 | rel_extent, |
| 423 | window_left, |
| 424 | ) |
| 425 | torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) |