| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | import torch.nn.functional as F |
| 7 | |
| 8 | import vllm.v1.attention.ops.triton_prefill_attention as prefill_ops |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd |
| 11 | |
| 12 | DEVICE_TYPE = current_platform.device_type |
| 13 | |
| 14 | |
| 15 | def ref_masked_attention( |
| 16 | q: torch.Tensor, |
| 17 | k: torch.Tensor, |
| 18 | v: torch.Tensor, |
| 19 | is_causal: bool = True, |
| 20 | sliding_window_q: int | None = None, |
| 21 | sliding_window_k: int | None = None, |
| 22 | ) -> torch.Tensor: |
| 23 | """Reference implementation using PyTorch SDPA.""" |
| 24 | # q, k, v: [total_tokens, num_heads, head_dim] |
| 25 | # SDPA expects [batch, num_heads, seq_len, head_dim] |
| 26 | |
| 27 | total_tokens = q.shape[0] |
| 28 | |
| 29 | # Add batch dimension and transpose |
| 30 | q = q.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim] |
| 31 | k = k.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim] |
| 32 | v = v.unsqueeze(0).transpose(1, 2) # [1, num_heads, total_tokens, head_dim] |
| 33 | |
| 34 | # Create attention mask if needed |
| 35 | attn_mask = None |
| 36 | use_causal = is_causal |
| 37 | |
| 38 | # If we have sliding window or need custom masking, create explicit mask |
| 39 | sliding_window_q = sliding_window_q if sliding_window_q is not None else 0 |
| 40 | sliding_window_k = sliding_window_k if sliding_window_k is not None else 0 |
| 41 | if (sliding_window_q > 0) or (sliding_window_k > 0): |
| 42 | # Position indices |
| 43 | pos_q = torch.arange(total_tokens, device=q.device).unsqueeze(1) |
| 44 | pos_k = torch.arange(total_tokens, device=q.device).unsqueeze(0) |
| 45 | |
| 46 | # Start with valid mask (False = no masking) |
| 47 | mask = torch.ones( |
| 48 | (total_tokens, total_tokens), dtype=torch.bool, device=q.device |
| 49 | ) |
| 50 | |
| 51 | # Apply causal mask |
| 52 | if is_causal: |
| 53 | mask = mask & (pos_q >= pos_k) |
| 54 | |
| 55 | # Apply sliding window masks |
| 56 | sliding_window_mask = torch.ones_like(mask) |
| 57 | if sliding_window_q > 0: |
| 58 | sliding_window_mask &= pos_q - pos_k <= sliding_window_q |
| 59 | |
| 60 | if sliding_window_k > 0: |
| 61 | sliding_window_mask &= pos_k - pos_q <= sliding_window_k |
| 62 | |
| 63 | mask = mask & sliding_window_mask |
| 64 | |
| 65 | attn_mask = torch.where(mask, 0.0, float("-inf")).to(q.dtype) |
| 66 | use_causal = False # Don't use is_causal when providing explicit mask |
| 67 | |
| 68 | # Use SDPA |
| 69 | output = F.scaled_dot_product_attention( |
| 70 | q, k, v, attn_mask=attn_mask, is_causal=use_causal, dropout_p=0.0 |
| 71 | ) |
| 72 | |
| 73 | # Convert back to original shape: [total_tokens, num_heads, head_dim] |
| 74 | output = output.transpose(1, 2).squeeze(0) |
| 75 | |
| 76 | return output |
| 77 | |
| 78 | |
| 79 | @pytest.mark.parametrize("B", [5]) |
| 80 | @pytest.mark.parametrize("max_seq_len", [1024]) |
| 81 | @pytest.mark.parametrize("H_Q", [32]) |
| 82 | @pytest.mark.parametrize("H_KV", [32, 8]) |
| 83 | @pytest.mark.parametrize("D", [128]) |
| 84 | @pytest.mark.parametrize("is_causal", [True, False]) |
| 85 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 86 | def test_context_attention( |
| 87 | B: int, |
| 88 | max_seq_len: int, |
| 89 | H_Q: int, |
| 90 | H_KV: int, |
| 91 | D: int, |
| 92 | is_causal: bool, |
| 93 | dtype: torch.dtype, |
| 94 | ): |
| 95 | """Test basic context attention without sliding window.""" |
| 96 | torch.manual_seed(42) |
| 97 | |
| 98 | # Generate random sequence lengths for each batch |
| 99 | seq_lens = torch.randint( |
| 100 | max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE |
| 101 | ) |
| 102 | total_tokens = seq_lens.sum().item() |
| 103 | |
| 104 | # Create batch start locations |
| 105 | b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE) |
| 106 | b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0) |
| 107 | |
| 108 | # Create input tensors |
| 109 | q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE) |
| 110 | k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) |
| 111 | v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) |
| 112 | o = torch.zeros_like(q) |
| 113 | |
| 114 | # Call Triton kernel |
| 115 | context_attention_fwd( |
| 116 | q, |
| 117 | k, |
| 118 | v, |
| 119 | o, |
| 120 | b_start_loc, |
| 121 | seq_lens, |
| 122 | max_seq_len, |
| 123 | is_causal=is_causal, |
| 124 | sliding_window_q=None, |
| 125 | sliding_window_k=None, |
| 126 | ) |
| 127 | |
| 128 | # Compute reference output for each sequence in batch |
| 129 | o_ref = torch.zeros_like(q) |
| 130 | for i in range(B): |
| 131 | start = b_start_loc[i].item() |
| 132 | end = start + seq_lens[i].item() |
| 133 | |
| 134 | q_seq = q[start:end] |
| 135 | k_seq = k[start:end] |
| 136 | v_seq = v[start:end] |
| 137 | |
| 138 | # Expand KV heads if using GQA |
| 139 | if H_Q != H_KV: |
| 140 | kv_group_num = H_Q // H_KV |
| 141 | k_seq = k_seq.repeat_interleave(kv_group_num, dim=1) |
| 142 | v_seq = v_seq.repeat_interleave(kv_group_num, dim=1) |
| 143 | |
| 144 | o_ref[start:end] = ref_masked_attention( |
| 145 | q_seq, |
| 146 | k_seq, |
| 147 | v_seq, |
| 148 | is_causal=is_causal, |
| 149 | sliding_window_q=None, |
| 150 | sliding_window_k=None, |
| 151 | ) |
| 152 | |
| 153 | # Compare outputs |
| 154 | torch.testing.assert_close(o, o_ref, rtol=1e-2, atol=1e-2) |
| 155 | |
| 156 | |
| 157 | @pytest.mark.parametrize("B", [4]) |
| 158 | @pytest.mark.parametrize("max_seq_len", [1024]) |
| 159 | @pytest.mark.parametrize("H_Q", [32]) |
| 160 | @pytest.mark.parametrize("H_KV", [32, 8]) |
| 161 | @pytest.mark.parametrize("D", [128]) |
| 162 | @pytest.mark.parametrize("sliding_window", [(32, 32), (32, 0), (0, 32)]) |
| 163 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 164 | def test_context_attention_sliding_window( |
| 165 | B: int, |
| 166 | max_seq_len: int, |
| 167 | H_Q: int, |
| 168 | H_KV: int, |
| 169 | D: int, |
| 170 | sliding_window: tuple[int, int], |
| 171 | dtype: torch.dtype, |
| 172 | ): |
| 173 | sliding_window_q, sliding_window_k = sliding_window |
| 174 | """Test context attention with sliding window.""" |
| 175 | torch.manual_seed(42) |
| 176 | |
| 177 | # Generate random sequence lengths for each batch |
| 178 | seq_lens = torch.randint( |
| 179 | max_seq_len // 2, max_seq_len + 1, (B,), device=DEVICE_TYPE |
| 180 | ) |
| 181 | total_tokens = seq_lens.sum().item() |
| 182 | |
| 183 | # Create batch start locations |
| 184 | b_start_loc = torch.zeros(B, dtype=torch.int32, device=DEVICE_TYPE) |
| 185 | b_start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0) |
| 186 | |
| 187 | # Create input tensors |
| 188 | q = torch.randn(total_tokens, H_Q, D, dtype=dtype, device=DEVICE_TYPE) |
| 189 | k = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) |
| 190 | v = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=DEVICE_TYPE) |
| 191 | o = torch.zeros_like(q) |
| 192 | |
| 193 | # Call Triton kernel |
| 194 | context_attention_fwd( |
| 195 | q, |
| 196 | k, |
| 197 | v, |
| 198 | o, |
| 199 | b_start_loc, |
| 200 | seq_lens, |
| 201 | max_seq_len, |
| 202 | is_causal=False, |
| 203 | sliding_window_q=sliding_window_q, |
| 204 | sliding_window_k=sliding_window_k, |
| 205 | ) |
| 206 | |
| 207 | # Compute reference output for each sequence in batch |
| 208 | o_ref = torch.zeros_like(q) |
| 209 | for i in range(B): |
| 210 | start = b_start_loc[i].item() |
| 211 | end = start + seq_lens[i].item() |
| 212 | |
| 213 | q_seq = q[start:end] |
| 214 | k_seq = k[start:end] |
| 215 | v_seq = v[start:end] |
| 216 | |
| 217 | # Expand KV heads if using GQA |
| 218 | if H_Q != H_KV: |
| 219 | kv_group_num = H_Q // H_KV |
| 220 | k_seq = k_seq.repeat_interleave(kv_group_num, dim=1) |
| 221 | v_seq = v_seq.repeat_interleave(kv_group_num, dim=1) |
| 222 | |
| 223 | o_ref[start:end] = ref_masked_attention( |
| 224 | q_seq, |
| 225 | k_seq, |
| 226 | v_seq, |
| 227 | is_causal=False, |
| 228 | sliding_window_q=sliding_window_q if sliding_window_q > 0 else None, |
| 229 | sliding_window_k=sliding_window_k if sliding_window_k > 0 else None, |
| 230 | ) |
| 231 | |
| 232 | # Compare outputs |
| 233 | torch.testing.assert_close(o, o_ref, rtol=2e-2, atol=2e-2) |
| 234 | |
| 235 | |
| 236 | class _LaunchCapture: |
| 237 | """Records the launch configuration in place of ``_fwd_kernel``.""" |
| 238 | |
| 239 | grid: tuple |
| 240 | kwargs: dict |
| 241 | |
| 242 | def __getitem__(self, grid): |
| 243 | self.grid = grid |
| 244 | return self._record |
| 245 | |
| 246 | def _record(self, *args, **kwargs) -> None: |
| 247 | self.kwargs = kwargs |
| 248 | |
| 249 | |
| 250 | def _capture_tile_config( |
| 251 | monkeypatch, *, is_rocm: bool, on_gfx1x: bool, dtype=torch.bfloat16 |
| 252 | ) -> _LaunchCapture: |
| 253 | """Capture the tile configuration with the platform predicates mocked. |
| 254 | |
| 255 | Both tile widths are numerically correct, so the tests above pass whichever |
| 256 | one is selected. This needs no RDNA part and allocates no device memory, so |
| 257 | it covers the RDNA branch on the CDNA and NVIDIA agents CI actually runs. |
| 258 | """ |
| 259 | platform = prefill_ops.current_platform |
| 260 | monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm) |
| 261 | # Every device this kernel targets is cuda-alike at capability 80 or better, |
| 262 | # so the stock tile is 128 for 16-bit dtypes. |
| 263 | monkeypatch.setattr(platform, "is_cuda_alike", lambda: True) |
| 264 | monkeypatch.setattr(platform, "has_device_capability", lambda *a, **k: True) |
| 265 | if is_rocm: |
| 266 | import vllm.platforms.rocm as rocm_platform |
| 267 | |
| 268 | monkeypatch.setattr(rocm_platform, "on_gfx1x", lambda: on_gfx1x) |
| 269 | |
| 270 | capture = _LaunchCapture() |
| 271 | monkeypatch.setattr(prefill_ops, "_fwd_kernel", capture) |
| 272 | |
| 273 | def meta(*shape): |
| 274 | return torch.empty(shape, dtype=dtype, device="meta") |
| 275 | |
| 276 | seq_lens = torch.empty(2, dtype=torch.int32, device="meta") |
| 277 | context_attention_fwd( |
| 278 | meta(256, 8, 128), |
| 279 | meta(256, 2, 128), |
| 280 | meta(256, 2, 128), |
| 281 | meta(256, 8, 128), |
| 282 | seq_lens, |
| 283 | seq_lens, |
| 284 | 128, |
| 285 | ) |
| 286 | return capture |
| 287 | |
| 288 | |
| 289 | @pytest.mark.parametrize( |
| 290 | ("is_rocm", "on_gfx1x", "dtype", "expected_block_n"), |
| 291 | [ |
| 292 | pytest.param(True, True, torch.bfloat16, 32, id="rdna"), |
| 293 | # on_gfx1x() is what excludes gfx10xx, CDNA and gfx1250. |
| 294 | pytest.param(True, False, torch.bfloat16, 128, id="rocm-not-rdna"), |
| 295 | pytest.param(False, False, torch.bfloat16, 128, id="not-rocm"), |
| 296 | # get_block_size already returns 32 here, so min() must not widen it. |
| 297 | pytest.param(True, True, torch.float32, 32, id="rdna-float32"), |
| 298 | ], |
| 299 | ) |
| 300 | def test_kv_tile_width_is_gated_by_platform( |
| 301 | monkeypatch, is_rocm: bool, on_gfx1x: bool, dtype, expected_block_n: int |
| 302 | ) -> None: |
| 303 | capture = _capture_tile_config( |
| 304 | monkeypatch, is_rocm=is_rocm, on_gfx1x=on_gfx1x, dtype=dtype |
| 305 | ) |
| 306 | assert capture.kwargs["BLOCK_N"] == expected_block_n |
| 307 | |
| 308 | |
| 309 | def test_rdna_narrows_the_kv_tile_and_nothing_else(monkeypatch) -> None: |
| 310 | with monkeypatch.context() as m: |
| 311 | tuned = _capture_tile_config(m, is_rocm=True, on_gfx1x=True) |
| 312 | with monkeypatch.context() as m: |
| 313 | stock = _capture_tile_config(m, is_rocm=True, on_gfx1x=False) |
| 314 | |
| 315 | assert tuned.kwargs.pop("BLOCK_N") != stock.kwargs.pop("BLOCK_N") |
| 316 | assert tuned.kwargs == stock.kwargs |
| 317 | assert tuned.grid == stock.grid |