| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm.platforms import current_platform |
| 8 | from vllm.utils.torch_utils import set_random_seed |
| 9 | |
| 10 | try: |
| 11 | import flashinfer |
| 12 | except ImportError: |
| 13 | if current_platform.is_rocm(): |
| 14 | pytest.skip( |
| 15 | "flashinfer is not supported for vLLM on ROCm.", allow_module_level=True |
| 16 | ) |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | NUM_HEADS = [(32, 8), (6, 1)] |
| 21 | HEAD_SIZES = [128, 256] |
| 22 | BLOCK_SIZES = [16, 32] |
| 23 | DTYPES = [torch.bfloat16] |
| 24 | NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation. |
| 25 | SOFT_CAPS = [None, 30.0] |
| 26 | SLIDING_WINDOWS = [None, 64] |
| 27 | |
| 28 | |
| 29 | def ref_paged_attn( |
| 30 | query: torch.Tensor, |
| 31 | key_cache: torch.Tensor, |
| 32 | value_cache: torch.Tensor, |
| 33 | query_lens: list[int], |
| 34 | kv_lens: list[int], |
| 35 | block_tables: torch.Tensor, |
| 36 | scale: float, |
| 37 | sliding_window: int | None = None, |
| 38 | soft_cap: float | None = None, |
| 39 | ) -> torch.Tensor: |
| 40 | num_seqs = len(query_lens) |
| 41 | block_tables = block_tables.cpu().numpy() |
| 42 | _, block_size, num_kv_heads, head_size = key_cache.shape |
| 43 | |
| 44 | outputs: list[torch.Tensor] = [] |
| 45 | start_idx = 0 |
| 46 | for i in range(num_seqs): |
| 47 | query_len = query_lens[i] |
| 48 | kv_len = kv_lens[i] |
| 49 | q = query[start_idx : start_idx + query_len] |
| 50 | q *= scale |
| 51 | |
| 52 | num_kv_blocks = (kv_len + block_size - 1) // block_size |
| 53 | block_indices = block_tables[i, :num_kv_blocks] |
| 54 | |
| 55 | k = key_cache[block_indices].view(-1, num_kv_heads, head_size) |
| 56 | k = k[:kv_len] |
| 57 | v = value_cache[block_indices].view(-1, num_kv_heads, head_size) |
| 58 | v = v[:kv_len] |
| 59 | |
| 60 | if q.shape[1] != k.shape[1]: |
| 61 | k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1) |
| 62 | v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1) |
| 63 | attn = torch.einsum("qhd,khd->hqk", q, k).float() |
| 64 | empty_mask = torch.ones(query_len, kv_len) |
| 65 | mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool() |
| 66 | if sliding_window is not None: |
| 67 | sliding_window_mask = ( |
| 68 | torch.triu( |
| 69 | empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1 |
| 70 | ) |
| 71 | .bool() |
| 72 | .logical_not() |
| 73 | ) |
| 74 | mask |= sliding_window_mask |
| 75 | if soft_cap is not None: |
| 76 | attn = soft_cap * torch.tanh(attn / soft_cap) |
| 77 | attn.masked_fill_(mask, float("-inf")) |
| 78 | attn = torch.softmax(attn, dim=-1).to(v.dtype) |
| 79 | out = torch.einsum("hqk,khd->qhd", attn, v) |
| 80 | |
| 81 | outputs.append(out) |
| 82 | start_idx += query_len |
| 83 | |
| 84 | return torch.cat(outputs, dim=0) |
| 85 | |
| 86 | |
| 87 | def _make_paged_kv_metadata( |
| 88 | kv_lens: list[int], |
| 89 | block_size: int, |
| 90 | num_blocks: int, |
| 91 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 92 | """Build paged-KV metadata tensors for fast_plan_decode tests. |
| 93 | |
| 94 | Returns: |
| 95 | kv_indptr – CPU int32, shape [num_seqs + 1] |
| 96 | kv_indices – CUDA int32, shape [total_blocks] |
| 97 | kv_last_page_lens – CPU int32, shape [num_seqs] |
| 98 | block_tables – CUDA int32, shape [num_seqs, max_blocks_per_seq] |
| 99 | |
| 100 | """ |
| 101 | num_seqs = len(kv_lens) |
| 102 | max_blocks = (max(kv_lens) + block_size - 1) // block_size |
| 103 | block_tables = torch.randint( |
| 104 | 0, num_blocks, (num_seqs, max_blocks), dtype=torch.int32, device="cuda" |
| 105 | ) |
| 106 | |
| 107 | indptr_list = [0] |
| 108 | indices_list: list[int] = [] |
| 109 | last_lens_list: list[int] = [] |
| 110 | for i, seq_len in enumerate(kv_lens): |
| 111 | n = (seq_len + block_size - 1) // block_size |
| 112 | indices_list.extend(block_tables[i, :n].cpu().tolist()) |
| 113 | indptr_list.append(indptr_list[-1] + n) |
| 114 | last_lens_list.append(seq_len % block_size or block_size) |
| 115 | |
| 116 | return ( |
| 117 | torch.tensor(indptr_list, dtype=torch.int32, device="cpu"), |
| 118 | torch.tensor(indices_list, dtype=torch.int32, device="cuda"), |
| 119 | torch.tensor(last_lens_list, dtype=torch.int32, device="cpu"), |
| 120 | block_tables, |
| 121 | ) |
| 122 | |
| 123 | |
| 124 | def _make_cg_decode_wrapper( |
| 125 | num_seqs: int, |
| 126 | kv_indices_buffer: torch.Tensor, |
| 127 | workspace_buffer: torch.Tensor, |
| 128 | use_tensor_cores: bool = True, |
| 129 | ) -> "flashinfer.BatchDecodeWithPagedKVCacheWrapper": |
| 130 | """Create a cudagraph-enabled BatchDecodeWithPagedKVCacheWrapper. |
| 131 | |
| 132 | *kv_indices_buffer* is shared with the caller so that fast_plan_decode |
| 133 | can avoid the device-to-device index copy on subsequent (cudagraph) calls. |
| 134 | """ |
| 135 | return flashinfer.BatchDecodeWithPagedKVCacheWrapper( |
| 136 | workspace_buffer, |
| 137 | "NHD", |
| 138 | use_cuda_graph=True, |
| 139 | paged_kv_indptr_buffer=torch.zeros( |
| 140 | num_seqs + 1, dtype=torch.int32, device="cuda" |
| 141 | ), |
| 142 | paged_kv_indices_buffer=kv_indices_buffer, |
| 143 | paged_kv_last_page_len_buffer=torch.zeros( |
| 144 | num_seqs, dtype=torch.int32, device="cuda" |
| 145 | ), |
| 146 | use_tensor_cores=use_tensor_cores, |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | def test_fast_decode_plan_importable() -> None: |
| 151 | """fast_decode_plan must be importable from flashinfer.decode. |
| 152 | |
| 153 | This is a forward-compatibility smoke test: if FlashInfer reorganises its |
| 154 | public API the import will fail before any other test does. |
| 155 | """ |
| 156 | from flashinfer.decode import fast_decode_plan # noqa: F401 |
| 157 | |
| 158 | assert callable(fast_decode_plan) |
| 159 | |
| 160 | |
| 161 | @pytest.mark.parametrize("dtype", DTYPES) |
| 162 | @torch.inference_mode |
| 163 | def test_fast_plan_decode_warmup_uses_full_plan(dtype: torch.dtype) -> None: |
| 164 | """On the first call fast_plan_decode must route through self.plan() and |
| 165 | flip vllm_first_call to False on the wrapper object.""" |
| 166 | from unittest.mock import patch |
| 167 | |
| 168 | from vllm.v1.attention.backends.flashinfer import fast_plan_decode |
| 169 | |
| 170 | torch.set_default_device("cuda") |
| 171 | set_random_seed(0) |
| 172 | |
| 173 | kv_lens = [128, 64] |
| 174 | block_size = 16 |
| 175 | num_seqs = len(kv_lens) |
| 176 | num_query_heads, num_kv_heads = 8, 2 |
| 177 | head_size = 128 |
| 178 | |
| 179 | kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata( |
| 180 | kv_lens, block_size, NUM_BLOCKS |
| 181 | ) |
| 182 | |
| 183 | workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8) |
| 184 | wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices.clone(), workspace) |
| 185 | |
| 186 | assert getattr(wrapper, "vllm_first_call", True) is True |
| 187 | |
| 188 | with patch.object(wrapper, "plan", wraps=wrapper.plan) as mock_plan: |
| 189 | fast_plan_decode( |
| 190 | wrapper, |
| 191 | indptr_cpu=kv_indptr, |
| 192 | indices=kv_indices, |
| 193 | last_page_len_cpu=kv_last_page_lens, |
| 194 | seq_lens_cpu=torch.tensor(kv_lens, dtype=torch.int32, device="cpu"), |
| 195 | num_qo_heads=num_query_heads, |
| 196 | num_kv_heads=num_kv_heads, |
| 197 | head_dim=head_size, |
| 198 | page_size=block_size, |
| 199 | q_data_type=dtype, |
| 200 | kv_data_type=dtype, |
| 201 | ) |
| 202 | mock_plan.assert_called_once() |
| 203 | |
| 204 | assert wrapper.vllm_first_call is False, ( |
| 205 | "vllm_first_call should be False after the first fast_plan_decode call" |
| 206 | ) |
| 207 | |
| 208 | |
| 209 | @pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]]) |
| 210 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 211 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 212 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 213 | @pytest.mark.parametrize("dtype", DTYPES) |
| 214 | @torch.inference_mode |
| 215 | def test_fast_plan_decode_matches_full_plan( |
| 216 | kv_lens: list[int], |
| 217 | num_heads: tuple[int, int], |
| 218 | head_size: int, |
| 219 | block_size: int, |
| 220 | dtype: torch.dtype, |
| 221 | ) -> None: |
| 222 | """fast_plan_decode's cudagraph path (delegating to FlashInfer's |
| 223 | fast_decode_plan) must produce attention output numerically identical to |
| 224 | a standard plan() call. |
| 225 | |
| 226 | Both the warmup call (self.plan) and the subsequent fast call |
| 227 | (fast_decode_plan) are verified against the same reference. |
| 228 | """ |
| 229 | from vllm.v1.attention.backends.flashinfer import fast_plan_decode |
| 230 | |
| 231 | torch.set_default_device("cuda") |
| 232 | set_random_seed(0) |
| 233 | num_seqs = len(kv_lens) |
| 234 | num_query_heads, num_kv_heads = num_heads |
| 235 | |
| 236 | query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype) |
| 237 | key_value_cache = torch.randn( |
| 238 | NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype |
| 239 | ) |
| 240 | |
| 241 | kv_indptr, kv_indices, kv_last_page_lens, _ = _make_paged_kv_metadata( |
| 242 | kv_lens, block_size, NUM_BLOCKS |
| 243 | ) |
| 244 | |
| 245 | # Reference output via the standard plan() |
| 246 | workspace_ref = torch.empty(128 * 1024 * 1024, dtype=torch.int8) |
| 247 | ref_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( |
| 248 | workspace_ref, "NHD", use_tensor_cores=True |
| 249 | ) |
| 250 | ref_wrapper.plan( |
| 251 | kv_indptr, |
| 252 | kv_indices, |
| 253 | kv_last_page_lens, |
| 254 | num_query_heads, |
| 255 | num_kv_heads, |
| 256 | head_size, |
| 257 | block_size, |
| 258 | "NONE", |
| 259 | q_data_type=dtype, |
| 260 | kv_data_type=dtype, |
| 261 | ) |
| 262 | ref_output = ref_wrapper.run(query, key_value_cache) |
| 263 | |
| 264 | # CUDAGraph wrapper exercised through fast_plan_decode |
| 265 | kv_indices_buf = kv_indices.clone() |
| 266 | workspace_cg = torch.empty(128 * 1024 * 1024, dtype=torch.int8) |
| 267 | cg_wrapper = _make_cg_decode_wrapper(num_seqs, kv_indices_buf, workspace_cg) |
| 268 | |
| 269 | plan_kwargs: dict = dict( |
| 270 | indptr_cpu=kv_indptr, |
| 271 | indices=kv_indices_buf, |
| 272 | last_page_len_cpu=kv_last_page_lens, |
| 273 | seq_lens_cpu=torch.tensor(kv_lens, dtype=torch.int32, device="cpu"), |
| 274 | num_qo_heads=num_query_heads, |
| 275 | num_kv_heads=num_kv_heads, |
| 276 | head_dim=head_size, |
| 277 | page_size=block_size, |
| 278 | q_data_type=dtype, |
| 279 | kv_data_type=dtype, |
| 280 | ) |
| 281 | |
| 282 | # First call – warmup path (routes through self.plan) |
| 283 | fast_plan_decode(cg_wrapper, **plan_kwargs) |
| 284 | warmup_output = cg_wrapper.run(query, key_value_cache) |
| 285 | torch.testing.assert_close(warmup_output, ref_output, atol=1e-2, rtol=1e-2) |
| 286 | |
| 287 | # Second call – fast path (routes through fast_decode_plan from FlashInfer) |
| 288 | fast_plan_decode(cg_wrapper, **plan_kwargs) |
| 289 | fast_output = cg_wrapper.run(query, key_value_cache) |
| 290 | torch.testing.assert_close(fast_output, ref_output, atol=1e-2, rtol=1e-2) |
| 291 | |
| 292 | |
| 293 | @pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]]) |
| 294 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 295 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 296 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 297 | @pytest.mark.parametrize("dtype", DTYPES) |
| 298 | @pytest.mark.parametrize("soft_cap", SOFT_CAPS) |
| 299 | @pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) |
| 300 | @torch.inference_mode |
| 301 | def test_flashinfer_decode_with_paged_kv( |
| 302 | kv_lens: list[int], |
| 303 | num_heads: tuple[int, int], |
| 304 | head_size: int, |
| 305 | dtype: torch.dtype, |
| 306 | block_size: int, |
| 307 | soft_cap: float | None, |
| 308 | sliding_window: int | None, |
| 309 | ) -> None: |
| 310 | torch.set_default_device("cuda") |
| 311 | set_random_seed(0) |
| 312 | num_seqs = len(kv_lens) |
| 313 | num_query_heads = num_heads[0] |
| 314 | num_kv_heads = num_heads[1] |
| 315 | assert num_query_heads % num_kv_heads == 0 |
| 316 | max_kv_len = max(kv_lens) |
| 317 | scale = head_size**-0.5 |
| 318 | |
| 319 | query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype) |
| 320 | |
| 321 | key_value_cache = torch.randn( |
| 322 | NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype |
| 323 | ) |
| 324 | key_cache = key_value_cache[:, 0, :, :, :].squeeze(1) |
| 325 | value_cache = key_value_cache[:, 1, :, :, :].squeeze(1) |
| 326 | |
| 327 | max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size |
| 328 | block_tables = torch.randint( |
| 329 | 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 |
| 330 | ) |
| 331 | |
| 332 | kv_indptr = [0] |
| 333 | kv_indices = [] |
| 334 | kv_last_page_lens = [] |
| 335 | for i in range(num_seqs): |
| 336 | seq_len = kv_lens[i] |
| 337 | assert seq_len > 0 |
| 338 | num_blocks = (seq_len + block_size - 1) // block_size |
| 339 | kv_indices.extend(block_tables[i, :num_blocks]) |
| 340 | kv_indptr.append(kv_indptr[-1] + num_blocks) |
| 341 | kv_last_page_len = seq_len % block_size |
| 342 | if kv_last_page_len == 0: |
| 343 | kv_last_page_len = block_size |
| 344 | kv_last_page_lens.append(kv_last_page_len) |
| 345 | |
| 346 | kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32) |
| 347 | kv_indices = torch.tensor(kv_indices, dtype=torch.int32) |
| 348 | kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) |
| 349 | |
| 350 | workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8) |
| 351 | wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( |
| 352 | workspace_buffer, "NHD", use_tensor_cores=True |
| 353 | ) |
| 354 | wrapper.plan( |
| 355 | kv_indptr, |
| 356 | kv_indices, |
| 357 | kv_last_page_lens, |
| 358 | num_query_heads, |
| 359 | num_kv_heads, |
| 360 | head_size, |
| 361 | block_size, |
| 362 | "NONE", |
| 363 | window_left=sliding_window - 1 if sliding_window is not None else -1, |
| 364 | q_data_type=dtype, |
| 365 | kv_data_type=dtype, |
| 366 | logits_soft_cap=soft_cap, |
| 367 | ) |
| 368 | |
| 369 | output = wrapper.run(query, key_value_cache) |
| 370 | |
| 371 | ref_output = ref_paged_attn( |
| 372 | query=query, |
| 373 | key_cache=key_cache, |
| 374 | value_cache=value_cache, |
| 375 | query_lens=[1] * num_seqs, |
| 376 | kv_lens=kv_lens, |
| 377 | block_tables=block_tables, |
| 378 | scale=scale, |
| 379 | soft_cap=soft_cap, |
| 380 | sliding_window=sliding_window, |
| 381 | ) |
| 382 | ( |
| 383 | torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2), |
| 384 | f"{torch.max(torch.abs(output - ref_output))}", |
| 385 | ) |
| 386 | |
| 387 | |
| 388 | @pytest.mark.parametrize("seq_lens", [[(1, 1328), (5, 18), (129, 463)]]) |
| 389 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 390 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 391 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 392 | @pytest.mark.parametrize("dtype", DTYPES) |
| 393 | @pytest.mark.parametrize("soft_cap", SOFT_CAPS) |
| 394 | @pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) |
| 395 | @torch.inference_mode |
| 396 | def test_flashinfer_prefill_with_paged_kv( |
| 397 | seq_lens: list[tuple[int, int]], |
| 398 | num_heads: tuple[int, int], |
| 399 | head_size: int, |
| 400 | dtype: torch.dtype, |
| 401 | block_size: int, |
| 402 | soft_cap: float | None, |
| 403 | sliding_window: int | None, |
| 404 | ) -> None: |
| 405 | torch.set_default_device("cuda") |
| 406 | set_random_seed(0) |
| 407 | num_seqs = len(seq_lens) |
| 408 | query_lens = [x[0] for x in seq_lens] |
| 409 | kv_lens = [x[1] for x in seq_lens] |
| 410 | num_query_heads = num_heads[0] |
| 411 | num_kv_heads = num_heads[1] |
| 412 | assert num_query_heads % num_kv_heads == 0 |
| 413 | max_kv_len = max(kv_lens) |
| 414 | scale = head_size**-0.5 |
| 415 | |
| 416 | query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype) |
| 417 | key_value_cache = torch.randn( |
| 418 | NUM_BLOCKS, 2, block_size, num_kv_heads, head_size, dtype=dtype |
| 419 | ) |
| 420 | key_cache = key_value_cache[:, 0, :, :, :].squeeze(1) |
| 421 | value_cache = key_value_cache[:, 1, :, :, :].squeeze(1) |
| 422 | |
| 423 | # Normalize the scale of the key and value caches to mitigate |
| 424 | # numerical instability. |
| 425 | key_cache /= head_size**0.5 |
| 426 | value_cache /= head_size**0.5 |
| 427 | |
| 428 | max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size |
| 429 | block_tables = torch.randint( |
| 430 | 0, NUM_BLOCKS, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 |
| 431 | ) |
| 432 | |
| 433 | qo_indptr = [0] |
| 434 | kv_indptr = [0] |
| 435 | kv_indices = [] |
| 436 | kv_last_page_lens = [] |
| 437 | for i in range(num_seqs): |
| 438 | seq_len = kv_lens[i] |
| 439 | assert seq_len > 0 |
| 440 | num_blocks = (seq_len + block_size - 1) // block_size |
| 441 | kv_indices.extend(block_tables[i, :num_blocks]) |
| 442 | kv_indptr.append(kv_indptr[-1] + num_blocks) |
| 443 | kv_last_page_len = seq_len % block_size |
| 444 | if kv_last_page_len == 0: |
| 445 | kv_last_page_len = block_size |
| 446 | kv_last_page_lens.append(kv_last_page_len) |
| 447 | qo_indptr.append(qo_indptr[-1] + query_lens[i]) |
| 448 | |
| 449 | qo_indptr = torch.tensor(qo_indptr, dtype=torch.int32) |
| 450 | kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32) |
| 451 | kv_indices = torch.tensor(kv_indices, dtype=torch.int32) |
| 452 | kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) |
| 453 | |
| 454 | workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8) |
| 455 | wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, "NHD") |
| 456 | wrapper.plan( |
| 457 | qo_indptr, |
| 458 | kv_indptr, |
| 459 | kv_indices, |
| 460 | kv_last_page_lens, |
| 461 | num_query_heads, |
| 462 | num_kv_heads, |
| 463 | head_size, |
| 464 | block_size, |
| 465 | window_left=sliding_window - 1 if sliding_window is not None else -1, |
| 466 | q_data_type=dtype, |
| 467 | kv_data_type=dtype, |
| 468 | logits_soft_cap=soft_cap, |
| 469 | ) |
| 470 | |
| 471 | output = wrapper.run( |
| 472 | query, |
| 473 | key_value_cache, |
| 474 | ) |
| 475 | |
| 476 | ref_output = ref_paged_attn( |
| 477 | query=query, |
| 478 | key_cache=key_cache, |
| 479 | value_cache=value_cache, |
| 480 | query_lens=query_lens, |
| 481 | kv_lens=kv_lens, |
| 482 | block_tables=block_tables, |
| 483 | scale=scale, |
| 484 | soft_cap=soft_cap, |
| 485 | sliding_window=sliding_window, |
| 486 | ) |
| 487 | ( |
| 488 | torch.testing.assert_close(output, ref_output, atol=5e-2, rtol=1e-2), |
| 489 | f"{torch.max(torch.abs(output - ref_output))}", |
| 490 | ) |
| 491 | |
| 492 | |
| 493 | @pytest.mark.parametrize("seq_lens", [[(1, 132), (5, 18)]]) |
| 494 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 495 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 496 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 497 | @pytest.mark.parametrize("dtype", DTYPES) |
| 498 | @pytest.mark.parametrize("soft_cap", SOFT_CAPS) |
| 499 | def test_flashinfer_prefill_with_paged_fp8_kv( |
| 500 | seq_lens: list[tuple[int, int]], |
| 501 | num_heads: tuple[int, int], |
| 502 | head_size: int, |
| 503 | dtype: torch.dtype, |
| 504 | block_size: int, |
| 505 | soft_cap: float | None, |
| 506 | ) -> None: |
| 507 | pytest.skip("TODO: fix the accuracy issue") |
| 508 | torch.set_default_device("cuda") |
| 509 | set_random_seed(0) |
| 510 | num_seqs = len(seq_lens) |
| 511 | query_lens = [x[0] for x in seq_lens] |
| 512 | kv_lens = [x[1] for x in seq_lens] |
| 513 | num_query_heads = num_heads[0] |
| 514 | num_kv_heads = num_heads[1] |
| 515 | assert num_query_heads % num_kv_heads == 0 |
| 516 | max_kv_len = max(kv_lens) |
| 517 | scale = head_size**-0.5 |
| 518 | |
| 519 | kv_cache_dtype = torch.float8_e4m3fn |
| 520 | |
| 521 | query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype) |
| 522 | NUM_BLOCKS_FP8 = 2048 |
| 523 | key_value_cache = torch.randn( |
| 524 | NUM_BLOCKS_FP8, 2, block_size, num_kv_heads, head_size, dtype=dtype |
| 525 | ) |
| 526 | key_cache, value_cache = torch.chunk(key_value_cache, 2, dim=1) |
| 527 | key_cache /= head_size**0.5 |
| 528 | value_cache /= head_size**0.5 |
| 529 | |
| 530 | k_scale = key_cache.amax().item() / 448.0 |
| 531 | v_scale = value_cache.amax().item() / 448.0 |
| 532 | |
| 533 | kv_cache_fp8 = torch.cat([key_cache / k_scale, value_cache / v_scale], dim=1).to( |
| 534 | kv_cache_dtype |
| 535 | ) |
| 536 | |
| 537 | assert kv_cache_fp8.shape == key_value_cache.shape |
| 538 | max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size |
| 539 | block_tables = torch.randint( |
| 540 | 0, NUM_BLOCKS_FP8, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 |
| 541 | ) |
| 542 | |
| 543 | qo_indptr = [0] |
| 544 | kv_indptr = [0] |
| 545 | kv_indices = [] |
| 546 | kv_last_page_lens = [] |
| 547 | for i in range(num_seqs): |
| 548 | seq_len = kv_lens[i] |
| 549 | assert seq_len > 0 |
| 550 | num_blocks = (seq_len + block_size - 1) // block_size |
| 551 | kv_indices.extend(block_tables[i, :num_blocks]) |
| 552 | kv_indptr.append(kv_indptr[-1] + num_blocks) |
| 553 | kv_last_page_len = seq_len % block_size |
| 554 | if kv_last_page_len == 0: |
| 555 | kv_last_page_len = block_size |
| 556 | kv_last_page_lens.append(kv_last_page_len) |
| 557 | qo_indptr.append(qo_indptr[-1] + query_lens[i]) |
| 558 | |
| 559 | qo_indptr = torch.tensor(qo_indptr, dtype=torch.int32) |
| 560 | kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32) |
| 561 | kv_indices = torch.tensor(kv_indices, dtype=torch.int32) |
| 562 | kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) |
| 563 | |
| 564 | workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8) |
| 565 | wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace_buffer, "NHD") |
| 566 | wrapper.plan( |
| 567 | qo_indptr, |
| 568 | kv_indptr, |
| 569 | kv_indices, |
| 570 | kv_last_page_lens, |
| 571 | num_query_heads, |
| 572 | num_kv_heads, |
| 573 | head_size, |
| 574 | block_size, |
| 575 | q_data_type=dtype, |
| 576 | kv_data_type=kv_cache_dtype, |
| 577 | logits_soft_cap=soft_cap, |
| 578 | ) |
| 579 | |
| 580 | output = wrapper.run(query, kv_cache_fp8, k_scale=k_scale, v_scale=v_scale) |
| 581 | |
| 582 | ref_output = ref_paged_attn( |
| 583 | query=query, |
| 584 | key_cache=key_cache.squeeze(1), |
| 585 | value_cache=value_cache.squeeze(1), |
| 586 | query_lens=query_lens, |
| 587 | kv_lens=kv_lens, |
| 588 | block_tables=block_tables, |
| 589 | scale=scale, |
| 590 | soft_cap=soft_cap, |
| 591 | ) |
| 592 | del query |
| 593 | del block_tables |
| 594 | # verify prefill fp8 |
| 595 | ( |
| 596 | torch.testing.assert_close(output, ref_output, atol=5e-2, rtol=1e-2), |
| 597 | f"{torch.max(torch.abs(output - ref_output))}", |
| 598 | ) |
| 599 | |
| 600 | |
| 601 | @pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]]) |
| 602 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 603 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 604 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 605 | @pytest.mark.parametrize("dtype", DTYPES) |
| 606 | @pytest.mark.parametrize("soft_cap", SOFT_CAPS) |
| 607 | @pytest.mark.skip(reason="TODO: fix the accuracy issue") |
| 608 | @torch.inference_mode |
| 609 | def test_flashinfer_decode_with_paged_fp8_kv( |
| 610 | kv_lens: list[int], |
| 611 | num_heads: tuple[int, int], |
| 612 | head_size: int, |
| 613 | dtype: torch.dtype, |
| 614 | block_size: int, |
| 615 | soft_cap: float | None, |
| 616 | ) -> None: |
| 617 | # test doesn't work for num_heads = (16,16) |
| 618 | torch.set_default_device("cuda") |
| 619 | set_random_seed(0) |
| 620 | num_seqs = len(kv_lens) |
| 621 | num_query_heads = num_heads[0] |
| 622 | num_kv_heads = num_heads[1] |
| 623 | assert num_query_heads % num_kv_heads == 0 |
| 624 | max_kv_len = max(kv_lens) |
| 625 | scale = head_size**-0.5 |
| 626 | use_tensor_cores = True |
| 627 | kv_cache_dtype = torch.float8_e4m3fn |
| 628 | |
| 629 | query = torch.randn(num_seqs, num_query_heads, head_size, dtype=dtype) |
| 630 | NUM_BLOCKS_FP8 = 2048 |
| 631 | key_value_cache = torch.randn( |
| 632 | NUM_BLOCKS_FP8, 2, block_size, num_kv_heads, head_size, dtype=dtype |
| 633 | ) |
| 634 | key_cache, value_cache = torch.chunk(key_value_cache, 2, dim=1) |
| 635 | key_cache /= head_size**0.5 |
| 636 | value_cache /= head_size**0.5 |
| 637 | |
| 638 | k_scale = key_cache.amax().item() / 448.0 |
| 639 | v_scale = value_cache.amax().item() / 448.0 |
| 640 | |
| 641 | key_cache_fp8 = (key_cache / k_scale).to(kv_cache_dtype) |
| 642 | value_cache_fp8 = (value_cache / v_scale).to(kv_cache_dtype) |
| 643 | assert key_cache_fp8.shape[1] == 1 and value_cache_fp8.shape[1] == 1 |
| 644 | kv_cache_fp8 = torch.cat([key_cache_fp8, value_cache_fp8], dim=1) |
| 645 | |
| 646 | max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size |
| 647 | block_tables = torch.randint( |
| 648 | 0, NUM_BLOCKS_FP8, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 |
| 649 | ) |
| 650 | |
| 651 | kv_indptr = [0] |
| 652 | kv_indices = [] |
| 653 | kv_last_page_lens = [] |
| 654 | for i in range(num_seqs): |
| 655 | seq_len = kv_lens[i] |
| 656 | assert seq_len > 0 |
| 657 | num_blocks = (seq_len + block_size - 1) // block_size |
| 658 | kv_indices.extend(block_tables[i, :num_blocks]) |
| 659 | kv_indptr.append(kv_indptr[-1] + num_blocks) |
| 660 | kv_last_page_len = seq_len % block_size |
| 661 | if kv_last_page_len == 0: |
| 662 | kv_last_page_len = block_size |
| 663 | kv_last_page_lens.append(kv_last_page_len) |
| 664 | |
| 665 | kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32) |
| 666 | kv_indices = torch.tensor(kv_indices, dtype=torch.int32) |
| 667 | kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32) |
| 668 | |
| 669 | workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8) |
| 670 | wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( |
| 671 | workspace_buffer, "NHD", use_tensor_cores=use_tensor_cores |
| 672 | ) |
| 673 | wrapper.plan( |
| 674 | kv_indptr, |
| 675 | kv_indices, |
| 676 | kv_last_page_lens, |
| 677 | num_query_heads, |
| 678 | num_kv_heads, |
| 679 | head_size, |
| 680 | block_size, |
| 681 | "NONE", |
| 682 | q_data_type=dtype, |
| 683 | kv_data_type=kv_cache_dtype, |
| 684 | logits_soft_cap=soft_cap, |
| 685 | ) |
| 686 | output = wrapper.run(query, kv_cache_fp8, k_scale=k_scale, v_scale=v_scale) |
| 687 | key_cache = key_value_cache[:, 0, :, :, :].squeeze(1) |
| 688 | value_cache = key_value_cache[:, 1, :, :, :].squeeze(1) |
| 689 | |
| 690 | ref_output = ref_paged_attn( |
| 691 | query=query, |
| 692 | key_cache=key_cache, |
| 693 | value_cache=value_cache, |
| 694 | query_lens=[1] * num_seqs, |
| 695 | kv_lens=kv_lens, |
| 696 | block_tables=block_tables, |
| 697 | scale=scale, |
| 698 | soft_cap=soft_cap, |
| 699 | ) |
| 700 | # Temporary fix: Increasing the tolerance. Seems like a flashinfer issue |
| 701 | ( |
| 702 | torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=1e-2), |
| 703 | f"{torch.max(torch.abs(output - ref_output))}", |
| 704 | ) |