| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Correctness test for the Inkling FA4 relative-attention score-mod kernel. |
| 4 | |
| 5 | Checks ``_INKLING_FA4_REL_ATTENTION_KERNEL`` against a pure-PyTorch reference |
| 6 | that implements the relative bias exactly as documented in the Inkling architecture |
| 7 | guide:: |
| 8 | |
| 9 | logit(i, j, h) = (1 / head_dim) * dot(q[i, h], k[j, h]) + rel_bias(i, j, h) |
| 10 | rel_bias(i, j, h) = rel_logits[i, h, i - j] if 0 <= i - j < rel_extent |
| 11 | = 0 otherwise |
| 12 | |
| 13 | with causal (and optionally sliding-window) masking handled by the backend. |
| 14 | """ |
| 15 | |
| 16 | import importlib |
| 17 | |
| 18 | import pytest |
| 19 | import torch |
| 20 | |
| 21 | from vllm.models.inkling.nvidia.attention import ( |
| 22 | InklingAttention, |
| 23 | compute_log_scaling_tau, |
| 24 | ) |
| 25 | from vllm.models.inkling.nvidia.ops.fa4_rel_attention import ( |
| 26 | _INKLING_FA4_REL_ATTENTION_KERNEL, |
| 27 | _use_sheared_bias, |
| 28 | inkling_fa4_num_splits, |
| 29 | ) |
| 30 | from vllm.platforms import current_platform |
| 31 | from vllm.platforms.interface import DeviceCapability |
| 32 | |
| 33 | _cap = current_platform.get_device_capability() if current_platform.is_cuda() else None |
| 34 | |
| 35 | NUM_HEADS = [(4, 4), (8, 2)] # (num_heads, num_kv_heads) |
| 36 | GLOBAL_REL_EXTENTS = [128, 1024] |
| 37 | LOCAL_REL_EXTENTS = [128, 256] |
| 38 | HEAD_DIM = 128 |
| 39 | BLOCK_SIZE = 16 |
| 40 | DTYPE = torch.bfloat16 |
| 41 | |
| 42 | |
| 43 | def test_log_scaling_tau_matches_reference(): |
| 44 | positions = torch.tensor([0, 127999, 128000, 999999], dtype=torch.int64) |
| 45 | actual = compute_log_scaling_tau(positions, 128000, 0.1) |
| 46 | expected = 1.0 + 0.1 * torch.log( |
| 47 | torch.clamp((positions + 1).float() / 128000.0, min=1.0) |
| 48 | ) |
| 49 | torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
| 50 | |
| 51 | |
| 52 | def test_split_packed_kv_cache(): |
| 53 | attention = InklingAttention.__new__(InklingAttention) |
| 54 | torch.nn.Module.__init__(attention) |
| 55 | attention.head_dim = 8 |
| 56 | attention.kv_cache = torch.arange(2 * 3 * 4 * 16).reshape(2, 3, 4, 16) |
| 57 | |
| 58 | key_cache, value_cache = attention._split_kv_cache() |
| 59 | |
| 60 | assert key_cache.shape == value_cache.shape == (2, 4, 3, 8) |
| 61 | torch.testing.assert_close(key_cache, attention.kv_cache[..., :8].transpose(1, 2)) |
| 62 | torch.testing.assert_close(value_cache, attention.kv_cache[..., 8:].transpose(1, 2)) |
| 63 | |
| 64 | |
| 65 | def test_num_splits_hopper_is_unsplit(monkeypatch): |
| 66 | monkeypatch.setattr( |
| 67 | current_platform, |
| 68 | "get_device_capability", |
| 69 | lambda: DeviceCapability(major=9, minor=0), |
| 70 | ) |
| 71 | assert ( |
| 72 | inkling_fa4_num_splits( |
| 73 | is_local=False, |
| 74 | batch_size=1, |
| 75 | max_query_len=1, |
| 76 | num_heads=16, |
| 77 | num_kv_heads=2, |
| 78 | max_kv_len=1_048_576, |
| 79 | ) |
| 80 | == 1 |
| 81 | ) |
| 82 | |
| 83 | |
| 84 | @pytest.mark.parametrize( |
| 85 | ("major", "expected"), |
| 86 | [(9, False), (10, True), (11, True), (12, False)], |
| 87 | ) |
| 88 | def test_sheared_bias_architecture_selection(monkeypatch, major, expected): |
| 89 | monkeypatch.setattr( |
| 90 | current_platform, |
| 91 | "get_device_capability", |
| 92 | lambda: DeviceCapability(major=major, minor=0), |
| 93 | ) |
| 94 | _use_sheared_bias.cache_clear() |
| 95 | try: |
| 96 | assert _use_sheared_bias() is expected |
| 97 | finally: |
| 98 | _use_sheared_bias.cache_clear() |
| 99 | |
| 100 | |
| 101 | @pytest.fixture |
| 102 | def blackwell_platform(monkeypatch): |
| 103 | monkeypatch.setattr( |
| 104 | current_platform, |
| 105 | "get_device_capability", |
| 106 | lambda: DeviceCapability(major=10, minor=0), |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | @pytest.mark.parametrize( |
| 111 | ("batch_size", "max_query_len", "expected"), |
| 112 | [ |
| 113 | (1, 1, (16, 32, 128, 128)), |
| 114 | (8, 1, (2, 4, 8, 16)), |
| 115 | (32, 1, (1, 1, 2, 4)), |
| 116 | (1, 128, (2, 4, 8, 16)), |
| 117 | (1, 2048, (1, 1, 1, 1)), |
| 118 | ], |
| 119 | ) |
| 120 | def test_num_splits_all_tp(blackwell_platform, batch_size, max_query_len, expected): |
| 121 | actual = tuple( |
| 122 | inkling_fa4_num_splits( |
| 123 | is_local=False, |
| 124 | batch_size=batch_size, |
| 125 | max_query_len=max_query_len, |
| 126 | num_heads=64 // tp, |
| 127 | num_kv_heads=8 // tp, |
| 128 | max_kv_len=131072, |
| 129 | ) |
| 130 | for tp in (1, 2, 4, 8) |
| 131 | ) |
| 132 | assert actual == expected |
| 133 | |
| 134 | |
| 135 | @pytest.mark.parametrize("tp", [1, 2, 4, 8]) |
| 136 | def test_num_splits_local_is_unsplit(tp): |
| 137 | assert ( |
| 138 | inkling_fa4_num_splits( |
| 139 | is_local=True, |
| 140 | batch_size=1, |
| 141 | max_query_len=1, |
| 142 | num_heads=64 // tp, |
| 143 | num_kv_heads=16 // tp, |
| 144 | max_kv_len=512, |
| 145 | ) |
| 146 | == 1 |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | @pytest.mark.parametrize( |
| 151 | ("max_kv_len", "expected"), |
| 152 | [(8192, 32), (65536, 64), (1048576, 128)], |
| 153 | ) |
| 154 | @pytest.mark.parametrize("tp", [4, 8]) |
| 155 | def test_num_splits_long_context_bound(blackwell_platform, tp, max_kv_len, expected): |
| 156 | assert ( |
| 157 | inkling_fa4_num_splits( |
| 158 | is_local=False, |
| 159 | batch_size=1, |
| 160 | max_query_len=1, |
| 161 | num_heads=64 // tp, |
| 162 | num_kv_heads=8 // tp, |
| 163 | max_kv_len=max_kv_len, |
| 164 | ) |
| 165 | == expected |
| 166 | ) |
| 167 | |
| 168 | |
| 169 | def _ref_rel_attn( |
| 170 | q: torch.Tensor, # [total_q, H, D] |
| 171 | key_cache: torch.Tensor, # [num_blocks, block, Hkv, D] |
| 172 | value_cache: torch.Tensor, |
| 173 | rel_logits: torch.Tensor, # [total_q, H, rel_extent] |
| 174 | *, |
| 175 | q_lens: list[int], |
| 176 | kv_lens: list[int], |
| 177 | block_table: torch.Tensor, |
| 178 | scale: float, |
| 179 | rel_extent: int, |
| 180 | window_left: int | None, |
| 181 | ) -> torch.Tensor: |
| 182 | num_kv_heads = key_cache.shape[2] |
| 183 | num_heads = q.shape[1] |
| 184 | g = num_heads // num_kv_heads |
| 185 | bt = block_table.cpu().numpy() |
| 186 | out = torch.empty_like(q) |
| 187 | |
| 188 | start = 0 |
| 189 | for i, (ql, kl) in enumerate(zip(q_lens, kv_lens)): |
| 190 | qi = q[start : start + ql].float() # [ql, H, D] |
| 191 | rl = rel_logits[start : start + ql].float() # [ql, H, rel_extent] |
| 192 | |
| 193 | nblk = (kl + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 194 | blk = bt[i, :nblk] |
| 195 | k = key_cache[blk].reshape(-1, num_kv_heads, HEAD_DIM)[:kl].float() |
| 196 | v = value_cache[blk].reshape(-1, num_kv_heads, HEAD_DIM)[:kl].float() |
| 197 | k = k.repeat_interleave(g, dim=1) # [kl, H, D] |
| 198 | v = v.repeat_interleave(g, dim=1) |
| 199 | |
| 200 | # [H, ql, kl] |
| 201 | scores = torch.einsum("qhd,khd->hqk", qi, k) * scale |
| 202 | |
| 203 | dev = q.device |
| 204 | qpos = torch.arange(ql, device=dev).view(ql, 1) + (kl - ql) # query pos |
| 205 | kpos = torch.arange(kl, device=dev).view(1, kl) |
| 206 | dist = qpos - kpos # [ql, kl] = i - j |
| 207 | |
| 208 | # Relative bias: rel_logits[i, h, dist] when 0 <= dist < rel_extent. |
| 209 | in_rng = (dist >= 0) & (dist < rel_extent) # [ql, kl] |
| 210 | idx = dist.clamp(0, rel_extent - 1) |
| 211 | # gather per head: bias[h, i, j] = rl[i, h, idx[i, j]] |
| 212 | bias = rl.permute(1, 0, 2).gather( # [H, ql, rel_extent] |
| 213 | 2, idx.unsqueeze(0).expand(num_heads, -1, -1) |
| 214 | ) # [H, ql, kl] |
| 215 | bias = torch.where(in_rng.unsqueeze(0), bias, torch.zeros_like(bias)) |
| 216 | scores = scores + bias |
| 217 | |
| 218 | mask = dist < 0 # causal |
| 219 | if window_left is not None: |
| 220 | mask = mask | (dist > window_left) |
| 221 | scores.masked_fill_(mask.unsqueeze(0), float("-inf")) |
| 222 | |
| 223 | probs = torch.softmax(scores, dim=-1) |
| 224 | out[start : start + ql] = torch.einsum("hqk,khd->qhd", probs, v).to(q.dtype) |
| 225 | start += ql |
| 226 | return out |
| 227 | |
| 228 | |
| 229 | def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0): |
| 230 | torch.manual_seed(seed) |
| 231 | device = "cuda" |
| 232 | q_lens = [s[0] for s in seq_lens] |
| 233 | kv_lens = [s[1] for s in seq_lens] |
| 234 | total_q = sum(q_lens) |
| 235 | num_seqs = len(seq_lens) |
| 236 | |
| 237 | scale = 1.0 / HEAD_DIM |
| 238 | |
| 239 | # q/k are RMS-normed in the model (unit-ish norm); normalize here so the |
| 240 | # logit magnitudes are realistic and the bias is not numerically dwarfed. |
| 241 | q = torch.randn(total_q, num_heads, HEAD_DIM, device=device, dtype=DTYPE) |
| 242 | q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) |
| 243 | |
| 244 | # Paged KV cache. |
| 245 | max_blocks = (max(kv_lens) + BLOCK_SIZE - 1) // BLOCK_SIZE |
| 246 | num_blocks = num_seqs * max_blocks + 1 |
| 247 | key_cache = torch.randn( |
| 248 | num_blocks, BLOCK_SIZE, num_kv_heads, HEAD_DIM, device=device, dtype=DTYPE |
| 249 | ) |
| 250 | key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) |
| 251 | value_cache = torch.randn( |
| 252 | num_blocks, BLOCK_SIZE, num_kv_heads, HEAD_DIM, device=device, dtype=DTYPE |
| 253 | ) |
| 254 | |
| 255 | # Distinct blocks per sequence (block 0 left as a never-referenced pad). |
| 256 | block_table = torch.zeros(num_seqs, max_blocks, dtype=torch.int32, device=device) |
| 257 | for i in range(num_seqs): |
| 258 | block_table[i] = torch.arange( |
| 259 | 1 + i * max_blocks, 1 + (i + 1) * max_blocks, dtype=torch.int32 |
| 260 | ) |
| 261 | |
| 262 | cu_seqlens_q = torch.tensor( |
| 263 | [0, *torch.cumsum(torch.tensor(q_lens), 0).tolist()], |
| 264 | dtype=torch.int32, |
| 265 | device=device, |
| 266 | ) |
| 267 | cache_seqlens = torch.tensor(kv_lens, dtype=torch.int32, device=device) |
| 268 | |
| 269 | rel_logits = torch.randn(total_q, num_heads, rel_extent, device=device, dtype=DTYPE) |
| 270 | |
| 271 | window_size = (-1, -1) if window_left is None else (window_left, 0) |
| 272 | |
| 273 | preallocated_out = torch.empty_like(q) |
| 274 | num_splits = inkling_fa4_num_splits( |
| 275 | is_local=window_left is not None, |
| 276 | batch_size=num_seqs, |
| 277 | max_query_len=max(q_lens), |
| 278 | num_heads=num_heads, |
| 279 | num_kv_heads=num_kv_heads, |
| 280 | max_kv_len=max(kv_lens), |
| 281 | ) |
| 282 | out = _INKLING_FA4_REL_ATTENTION_KERNEL( |
| 283 | q, |
| 284 | key_cache, |
| 285 | value_cache, |
| 286 | block_table=block_table, |
| 287 | cache_seqlens=cache_seqlens, |
| 288 | cu_seqlens_q=cu_seqlens_q, |
| 289 | max_seqlen_q=max(q_lens), |
| 290 | softmax_scale=scale, |
| 291 | causal=True, |
| 292 | window_size=window_size, |
| 293 | rel_extent=rel_extent, |
| 294 | rel_logits=rel_logits, |
| 295 | num_splits=num_splits, |
| 296 | out=preallocated_out, |
| 297 | ) |
| 298 | assert out.data_ptr() == preallocated_out.data_ptr() |
| 299 | out = out.view(total_q, num_heads, HEAD_DIM) |
| 300 | |
| 301 | ref = _ref_rel_attn( |
| 302 | q, |
| 303 | key_cache, |
| 304 | value_cache, |
| 305 | rel_logits, |
| 306 | q_lens=q_lens, |
| 307 | kv_lens=kv_lens, |
| 308 | block_table=block_table, |
| 309 | scale=scale, |
| 310 | rel_extent=rel_extent, |
| 311 | window_left=window_left, |
| 312 | ) |
| 313 | |
| 314 | torch.testing.assert_close(out.float(), ref.float(), atol=2e-2, rtol=2e-2) |
| 315 | |
| 316 | |
| 317 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") |
| 318 | @pytest.mark.skipif( |
| 319 | _cap is None or _cap.major < 9, |
| 320 | reason="FA4 score-mod requires Hopper+ (SM90+)", |
| 321 | ) |
| 322 | @torch.inference_mode() |
| 323 | def test_score_mod_relative_attention(monkeypatch): |
| 324 | module = importlib.import_module("vllm.models.inkling.nvidia.ops.fa4_rel_attention") |
| 325 | monkeypatch.setattr(module, "_use_sheared_bias", lambda: False) |
| 326 | _run_case( |
| 327 | [(64, 64), (1, 80)], |
| 328 | num_heads=4, |
| 329 | num_kv_heads=4, |
| 330 | rel_extent=128, |
| 331 | window_left=None, |
| 332 | ) |
| 333 | |
| 334 | |
| 335 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") |
| 336 | @pytest.mark.skipif( |
| 337 | _cap is None or _cap.major < 9, |
| 338 | reason="FA4 score-mod requires Hopper+ (SM90+)", |
| 339 | ) |
| 340 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 341 | @pytest.mark.parametrize( |
| 342 | "seq_lens", |
| 343 | [ |
| 344 | [(64, 64)], # single full prefill |
| 345 | [(64, 64), (33, 33), (17, 17)], # ragged prefill batch |
| 346 | [(512, 512)], # seq_len >> rel_extent (most keys get zero bias) |
| 347 | [(300, 300), (512, 512), (129, 129)], # large ragged batch |
| 348 | ], |
| 349 | ) |
| 350 | @pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS) |
| 351 | @torch.inference_mode() |
| 352 | def test_full_attention(seq_lens, num_heads, rel_extent): |
| 353 | # rel_extent=128 exercises the out-of-range (zero bias) path; 1024 covers all. |
| 354 | # With the 512-token cases and rel_extent=128, query/seq lengths are far |
| 355 | # larger than rel_extent so the vast majority of (i, j) pairs are out of |
| 356 | # range and must contribute zero bias. |
| 357 | _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None) |
| 358 | |
| 359 | |
| 360 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") |
| 361 | @pytest.mark.skipif( |
| 362 | _cap is None or _cap.major < 9, |
| 363 | reason="FA4 score-mod requires Hopper+ (SM90+)", |
| 364 | ) |
| 365 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 366 | @pytest.mark.parametrize( |
| 367 | "seq_lens", |
| 368 | [ |
| 369 | [(200, 512)], # chunked prefill: q_len=200 (> rel_extent), 312 cached |
| 370 | [(200, 512), (50, 300), (1, 400)], # mixed chunked + decode |
| 371 | ], |
| 372 | ) |
| 373 | @pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS) |
| 374 | @torch.inference_mode() |
| 375 | def test_chunked_prefill(seq_lens, num_heads, rel_extent): |
| 376 | # q_len < kv_len with q_len itself larger than rel_extent (for the 128 case): |
| 377 | # exercises the seqlen_k - seqlen_q offset together with the out-of-range path. |
| 378 | _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None) |
| 379 | |
| 380 | |
| 381 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") |
| 382 | @pytest.mark.skipif( |
| 383 | _cap is None or _cap.major < 9, |
| 384 | reason="FA4 score-mod requires Hopper+ (SM90+)", |
| 385 | ) |
| 386 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 387 | @pytest.mark.parametrize( |
| 388 | "seq_lens", |
| 389 | [ |
| 390 | [(64, 64), (40, 40)], # seq_len > window |
| 391 | [(512, 512), (300, 300)], # seq_len/query_len >> window |
| 392 | [(1, 512)], # decode with kv_len >> window |
| 393 | ], |
| 394 | ) |
| 395 | @pytest.mark.parametrize("local_extent", LOCAL_REL_EXTENTS) |
| 396 | @torch.inference_mode() |
| 397 | def test_sliding_window(seq_lens, num_heads, local_extent): |
| 398 | # Local layers use window_size=(local_extent-1, 0) and rel_extent==local_extent. |
| 399 | # With the 512-token cases, query/seq lengths far exceed the window so most |
| 400 | # keys are masked out by the sliding window. |
| 401 | _run_case( |
| 402 | seq_lens, |
| 403 | num_heads[0], |
| 404 | num_heads[1], |
| 405 | rel_extent=local_extent, |
| 406 | window_left=local_extent - 1, |
| 407 | ) |
| 408 | |
| 409 | |
| 410 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") |
| 411 | @pytest.mark.skipif( |
| 412 | _cap is None or _cap.major < 9, |
| 413 | reason="FA4 score-mod requires Hopper+ (SM90+)", |
| 414 | ) |
| 415 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 416 | @pytest.mark.parametrize( |
| 417 | "seq_lens", |
| 418 | [ |
| 419 | [(1, 50)], |
| 420 | [(1, 50), (1, 7), (1, 200)], |
| 421 | [(1, 512), (1, 333)], # kv_len >> rel_extent |
| 422 | ], |
| 423 | ) |
| 424 | @pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS) |
| 425 | @torch.inference_mode() |
| 426 | def test_decode(seq_lens, num_heads, rel_extent): |
| 427 | # q_len=1 with kv_len>q_len: the score-mod's seqlen_k - seqlen_q offset path. |
| 428 | _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None) |