| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.utils.torch_utils import set_random_seed |
| 10 | |
| 11 | try: |
| 12 | from vllm.vllm_flash_attn import ( |
| 13 | fa_version_unsupported_reason, |
| 14 | flash_attn_varlen_func, |
| 15 | is_fa_version_supported, |
| 16 | ) |
| 17 | except ImportError: |
| 18 | if current_platform.is_rocm(): |
| 19 | pytest.skip( |
| 20 | "vllm_flash_attn is not supported for vLLM on ROCm.", |
| 21 | allow_module_level=True, |
| 22 | ) |
| 23 | |
| 24 | |
| 25 | NUM_HEADS = [(4, 4), (8, 2)] |
| 26 | HEAD_SIZES = [40, 72, 80, 128, 256] |
| 27 | BLOCK_SIZES = [16] |
| 28 | DTYPES = [torch.bfloat16] |
| 29 | QDTYPES = [None, torch.float8_e4m3fn] |
| 30 | # one value large enough to test overflow in index calculation. |
| 31 | # one value small enough to test the schema op check |
| 32 | NUM_BLOCKS = [32768, 2048] |
| 33 | SOFT_CAPS = [None] |
| 34 | SLIDING_WINDOWS = [None, 256] |
| 35 | |
| 36 | |
| 37 | def ref_paged_attn( |
| 38 | query: torch.Tensor, |
| 39 | key_cache: torch.Tensor, |
| 40 | value_cache: torch.Tensor, |
| 41 | query_lens: list[int], |
| 42 | kv_lens: list[int], |
| 43 | block_tables: torch.Tensor, |
| 44 | scale: float, |
| 45 | sliding_window: int | None = None, |
| 46 | soft_cap: float | None = None, |
| 47 | ) -> torch.Tensor: |
| 48 | num_seqs = len(query_lens) |
| 49 | block_tables = block_tables.cpu().numpy() |
| 50 | _, block_size, num_kv_heads, head_size = key_cache.shape |
| 51 | |
| 52 | outputs: list[torch.Tensor] = [] |
| 53 | start_idx = 0 |
| 54 | for i in range(num_seqs): |
| 55 | query_len = query_lens[i] |
| 56 | kv_len = kv_lens[i] |
| 57 | q = query[start_idx : start_idx + query_len] |
| 58 | q *= scale |
| 59 | |
| 60 | num_kv_blocks = (kv_len + block_size - 1) // block_size |
| 61 | block_indices = block_tables[i, :num_kv_blocks] |
| 62 | |
| 63 | k = key_cache[block_indices].view(-1, num_kv_heads, head_size) |
| 64 | k = k[:kv_len] |
| 65 | v = value_cache[block_indices].view(-1, num_kv_heads, head_size) |
| 66 | v = v[:kv_len] |
| 67 | |
| 68 | if q.shape[1] != k.shape[1]: |
| 69 | k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1) |
| 70 | v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1) |
| 71 | attn = torch.einsum("qhd,khd->hqk", q, k).float() |
| 72 | empty_mask = torch.ones(query_len, kv_len) |
| 73 | mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool() |
| 74 | if sliding_window is not None: |
| 75 | sliding_window_mask = ( |
| 76 | torch.triu( |
| 77 | empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1 |
| 78 | ) |
| 79 | .bool() |
| 80 | .logical_not() |
| 81 | ) |
| 82 | mask |= sliding_window_mask |
| 83 | if soft_cap is not None: |
| 84 | attn = soft_cap * torch.tanh(attn / soft_cap) |
| 85 | attn.masked_fill_(mask, float("-inf")) |
| 86 | attn = torch.softmax(attn, dim=-1).to(v.dtype) |
| 87 | out = torch.einsum("hqk,khd->qhd", attn, v) |
| 88 | |
| 89 | outputs.append(out) |
| 90 | start_idx += query_len |
| 91 | |
| 92 | return torch.cat(outputs, dim=0) |
| 93 | |
| 94 | |
| 95 | @pytest.mark.parametrize("use_out", [True, False]) |
| 96 | @pytest.mark.parametrize( |
| 97 | "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]] |
| 98 | ) |
| 99 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 100 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 101 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 102 | @pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) |
| 103 | @pytest.mark.parametrize("dtype", DTYPES) |
| 104 | @pytest.mark.parametrize("soft_cap", SOFT_CAPS) |
| 105 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS) |
| 106 | @pytest.mark.parametrize("fa_version", [2, 3]) |
| 107 | @pytest.mark.parametrize("q_dtype", QDTYPES) |
| 108 | @torch.inference_mode() |
| 109 | def test_varlen_with_paged_kv( |
| 110 | use_out: bool, |
| 111 | seq_lens: list[tuple[int, int]], |
| 112 | num_heads: tuple[int, int], |
| 113 | head_size: int, |
| 114 | sliding_window: int | None, |
| 115 | dtype: torch.dtype, |
| 116 | block_size: int, |
| 117 | soft_cap: float | None, |
| 118 | num_blocks: int, |
| 119 | fa_version: int, |
| 120 | q_dtype: torch.dtype | None, |
| 121 | ) -> None: |
| 122 | torch.set_default_device("cuda") |
| 123 | if not is_fa_version_supported(fa_version): |
| 124 | pytest.skip( |
| 125 | f"Flash attention version {fa_version} not supported due " |
| 126 | f'to: "{fa_version_unsupported_reason(fa_version)}"' |
| 127 | ) |
| 128 | if q_dtype is not None and (dtype != torch.bfloat16 or fa_version == 2): |
| 129 | pytest.skip( |
| 130 | "Flash attention with quantized inputs is only " |
| 131 | "supported on version 3 with bfloat16 base type" |
| 132 | ) |
| 133 | set_random_seed(0) |
| 134 | num_seqs = len(seq_lens) |
| 135 | query_lens = [x[0] for x in seq_lens] |
| 136 | kv_lens = [x[1] for x in seq_lens] |
| 137 | num_query_heads = num_heads[0] |
| 138 | num_kv_heads = num_heads[1] |
| 139 | assert num_query_heads % num_kv_heads == 0 |
| 140 | max_query_len = max(query_lens) |
| 141 | max_kv_len = max(kv_lens) |
| 142 | window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1) |
| 143 | scale = head_size**-0.5 |
| 144 | |
| 145 | query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype) |
| 146 | key_cache = torch.randn( |
| 147 | num_blocks, block_size, num_kv_heads, head_size, dtype=dtype |
| 148 | ) |
| 149 | value_cache = torch.randn_like(key_cache) |
| 150 | cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( |
| 151 | dim=0, dtype=torch.int32 |
| 152 | ) |
| 153 | kv_lens = torch.tensor(kv_lens, dtype=torch.int32) |
| 154 | |
| 155 | max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size |
| 156 | block_tables = torch.randint( |
| 157 | 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 |
| 158 | ) |
| 159 | |
| 160 | out = torch.empty_like(query) if use_out else None |
| 161 | |
| 162 | maybe_quantized_query = query |
| 163 | maybe_quantized_key_cache = key_cache |
| 164 | maybe_quantized_value_cache = value_cache |
| 165 | q_descale = None |
| 166 | k_descale = None |
| 167 | v_descale = None |
| 168 | if q_dtype is not None: |
| 169 | # QKV are drawn from N(0, 1): no need for a fp8 scaling factor |
| 170 | maybe_quantized_query = query.to(q_dtype) |
| 171 | maybe_quantized_key_cache = key_cache.to(q_dtype) |
| 172 | maybe_quantized_value_cache = value_cache.to(q_dtype) |
| 173 | |
| 174 | scale_shape = (num_seqs, num_kv_heads) |
| 175 | q_descale = torch.ones(scale_shape, dtype=torch.float32) |
| 176 | k_descale = torch.ones(scale_shape, dtype=torch.float32) |
| 177 | v_descale = torch.ones(scale_shape, dtype=torch.float32) |
| 178 | |
| 179 | output = flash_attn_varlen_func( |
| 180 | q=maybe_quantized_query, |
| 181 | k=maybe_quantized_key_cache, |
| 182 | v=maybe_quantized_value_cache, |
| 183 | out=out, |
| 184 | cu_seqlens_q=cu_query_lens, |
| 185 | seqused_k=kv_lens, |
| 186 | max_seqlen_q=max_query_len, |
| 187 | max_seqlen_k=max_kv_len, |
| 188 | softmax_scale=scale, |
| 189 | causal=True, |
| 190 | window_size=list(window_size), |
| 191 | block_table=block_tables, |
| 192 | softcap=soft_cap if soft_cap is not None else 0, |
| 193 | fa_version=fa_version, |
| 194 | q_descale=q_descale, |
| 195 | k_descale=k_descale, |
| 196 | v_descale=v_descale, |
| 197 | ) |
| 198 | output = output if not use_out else out |
| 199 | |
| 200 | ref_output = ref_paged_attn( |
| 201 | query=query, |
| 202 | key_cache=key_cache, |
| 203 | value_cache=value_cache, |
| 204 | query_lens=query_lens, |
| 205 | kv_lens=kv_lens, |
| 206 | block_tables=block_tables, |
| 207 | scale=scale, |
| 208 | sliding_window=sliding_window, |
| 209 | soft_cap=soft_cap, |
| 210 | ) |
| 211 | atol, rtol = 1.5e-2, 1e-2 |
| 212 | if q_dtype is not None: |
| 213 | atol, rtol = 1.5e-1, 1.5e-1 |
| 214 | ( |
| 215 | torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), |
| 216 | f"{torch.max(torch.abs(output - ref_output))}", |
| 217 | ) |
| 218 | |
| 219 | |
| 220 | @pytest.mark.skipif( |
| 221 | not current_platform.is_device_capability_family(100), |
| 222 | reason="FA4 only has a dedicated head_size=256 kernel on Blackwell", |
| 223 | ) |
| 224 | @pytest.mark.parametrize( |
| 225 | "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]] |
| 226 | ) |
| 227 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 228 | @pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) |
| 229 | @torch.inference_mode() |
| 230 | def test_fa4_hd256_paged_call_shape( |
| 231 | seq_lens: list[tuple[int, int]], |
| 232 | num_heads: tuple[int, int], |
| 233 | sliding_window: int | None, |
| 234 | ) -> None: |
| 235 | """Handles excess block-table columns with page-aligned ``max_seqlen_k``.""" |
| 236 | from vllm.v1.attention.backends.fa_utils import FA4_HD256_PAGE_SIZE |
| 237 | |
| 238 | if not is_fa_version_supported(4): |
| 239 | pytest.skip(f'FA4 unsupported: "{fa_version_unsupported_reason(4)}"') |
| 240 | torch.set_default_device("cuda") |
| 241 | set_random_seed(0) |
| 242 | |
| 243 | head_size = 256 |
| 244 | block_size = FA4_HD256_PAGE_SIZE |
| 245 | num_blocks = 2048 |
| 246 | dtype = torch.bfloat16 |
| 247 | num_query_heads, num_kv_heads = num_heads |
| 248 | query_lens = [x[0] for x in seq_lens] |
| 249 | kv_lens = [x[1] for x in seq_lens] |
| 250 | max_kv_len = max(kv_lens) |
| 251 | window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1) |
| 252 | scale = head_size**-0.5 |
| 253 | |
| 254 | query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype) |
| 255 | key_cache = torch.randn( |
| 256 | num_blocks, block_size, num_kv_heads, head_size, dtype=dtype |
| 257 | ) |
| 258 | value_cache = torch.randn_like(key_cache) |
| 259 | cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( |
| 260 | dim=0, dtype=torch.int32 |
| 261 | ) |
| 262 | num_pages = (max_kv_len + block_size - 1) // block_size |
| 263 | block_tables = torch.randint( |
| 264 | 0, num_blocks, (len(seq_lens), num_pages + 3), dtype=torch.int32 |
| 265 | ) |
| 266 | |
| 267 | output = flash_attn_varlen_func( |
| 268 | q=query, |
| 269 | k=key_cache, |
| 270 | v=value_cache, |
| 271 | cu_seqlens_q=cu_query_lens, |
| 272 | seqused_k=torch.tensor(kv_lens, dtype=torch.int32), |
| 273 | max_seqlen_q=max(query_lens), |
| 274 | max_seqlen_k=num_pages * block_size, |
| 275 | softmax_scale=scale, |
| 276 | causal=True, |
| 277 | window_size=list(window_size), |
| 278 | block_table=block_tables[:, :num_pages], |
| 279 | num_splits=1, |
| 280 | fa_version=4, |
| 281 | ) |
| 282 | ref_output = ref_paged_attn( |
| 283 | query=query, |
| 284 | key_cache=key_cache, |
| 285 | value_cache=value_cache, |
| 286 | query_lens=query_lens, |
| 287 | kv_lens=kv_lens, |
| 288 | block_tables=block_tables, |
| 289 | scale=scale, |
| 290 | sliding_window=sliding_window, |
| 291 | ) |
| 292 | torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2) |