| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Integration tests for FlexAttention backend vs default backend.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | from packaging import version |
| 8 | |
| 9 | from tests.utils import set_random_seed |
| 10 | from tests.v1.attention.utils import ( |
| 11 | BatchSpec, |
| 12 | create_common_attn_metadata, |
| 13 | create_standard_kv_cache_spec, |
| 14 | create_vllm_config, |
| 15 | ) |
| 16 | from vllm.config import AttentionConfig |
| 17 | from vllm.model_executor.layers.attention import Attention |
| 18 | from vllm.v1.attention.backends.flex_attention import ( |
| 19 | BlockSparsityHint, |
| 20 | FlexAttentionMetadataBuilder, |
| 21 | physical_to_logical_mapping, |
| 22 | ) |
| 23 | |
| 24 | from ..models.utils import check_embeddings_close, check_logprobs_close |
| 25 | |
| 26 | TORCH_VERSION = version.parse(torch.__version__) |
| 27 | MINIMUM_TORCH_VERSION = version.parse("2.7.0") |
| 28 | DIRECT_BUILD_VERSION = version.parse("2.9.dev0") |
| 29 | |
| 30 | |
| 31 | @pytest.mark.parametrize( |
| 32 | ("supports_small_blocks", "uses_paged_kv", "expected"), |
| 33 | [ |
| 34 | (True, True, (16, 16)), |
| 35 | (True, False, (128, 128)), |
| 36 | (False, True, (128, 128)), |
| 37 | (False, False, (128, 128)), |
| 38 | ], |
| 39 | ) |
| 40 | def test_flex_attention_default_block_sizes( |
| 41 | supports_small_blocks: bool, |
| 42 | uses_paged_kv: bool, |
| 43 | expected: tuple[int, int], |
| 44 | ): |
| 45 | block_sizes = FlexAttentionMetadataBuilder._get_block_sizes( |
| 46 | AttentionConfig(), |
| 47 | supports_small_blocks=supports_small_blocks, |
| 48 | cache_block_size=16, |
| 49 | uses_paged_kv=uses_paged_kv, |
| 50 | ) |
| 51 | assert block_sizes == expected |
| 52 | |
| 53 | |
| 54 | def test_flex_attention_explicit_block_sizes_override_encoder_defaults(): |
| 55 | block_sizes = FlexAttentionMetadataBuilder._get_block_sizes( |
| 56 | AttentionConfig( |
| 57 | flex_attn_q_block_size=64, |
| 58 | flex_attn_kv_block_size=32, |
| 59 | ), |
| 60 | supports_small_blocks=True, |
| 61 | cache_block_size=16, |
| 62 | uses_paged_kv=False, |
| 63 | ) |
| 64 | assert block_sizes == (64, 32) |
| 65 | |
| 66 | |
| 67 | @pytest.mark.skipif( |
| 68 | not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION, |
| 69 | reason="CUDA not available or PyTorch version < 2.7", |
| 70 | ) |
| 71 | def test_flex_attention_full_cudagraphs(vllm_runner): |
| 72 | """Test the numerics for flex attention full cudagraphs support.""" |
| 73 | model_name = "Qwen/Qwen2.5-1.5B-Instruct" |
| 74 | seed = 42 |
| 75 | max_tokens = 24 |
| 76 | num_logprobs = 5 |
| 77 | prompts = [ |
| 78 | "Hello, my name is", |
| 79 | "The president of the United States is", |
| 80 | "The capital of France is", |
| 81 | ] |
| 82 | |
| 83 | # Run with flex attention eager |
| 84 | set_random_seed(seed) |
| 85 | with vllm_runner( |
| 86 | model_name, |
| 87 | runner="generate", |
| 88 | tensor_parallel_size=1, |
| 89 | num_gpu_blocks_override=128, |
| 90 | enforce_eager=True, |
| 91 | attention_config={"backend": "FLEX_ATTENTION"}, |
| 92 | ) as llm_flex: |
| 93 | output_eager = llm_flex.generate_greedy_logprobs( |
| 94 | prompts, max_tokens, num_logprobs |
| 95 | ) |
| 96 | |
| 97 | # Run with flex attention compiled |
| 98 | set_random_seed(seed) |
| 99 | with vllm_runner( |
| 100 | model_name, |
| 101 | runner="generate", |
| 102 | tensor_parallel_size=1, |
| 103 | num_gpu_blocks_override=128, |
| 104 | enforce_eager=False, |
| 105 | gpu_memory_utilization=0.85, |
| 106 | attention_config={"backend": "FLEX_ATTENTION"}, |
| 107 | ) as llm_default: |
| 108 | output_compile = llm_default.generate_greedy_logprobs( |
| 109 | prompts, max_tokens, num_logprobs |
| 110 | ) |
| 111 | |
| 112 | check_logprobs_close( |
| 113 | outputs_0_lst=output_eager, |
| 114 | outputs_1_lst=output_compile, |
| 115 | name_0="eager", |
| 116 | name_1="compile", |
| 117 | ) |
| 118 | |
| 119 | |
| 120 | def windowed_causal_mask_mod(b, h, q_idx, kv_idx): |
| 121 | return (kv_idx <= q_idx) & (q_idx - kv_idx < 4) |
| 122 | |
| 123 | |
| 124 | @pytest.mark.skipif( |
| 125 | not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION, |
| 126 | reason="CUDA not available or PyTorch version < 2.7", |
| 127 | ) |
| 128 | def test_flex_attention_custom_mask_full_cudagraphs(vllm_runner, monkeypatch): |
| 129 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 130 | monkeypatch.setattr( |
| 131 | Attention, |
| 132 | "logical_mask_mod", |
| 133 | staticmethod(windowed_causal_mask_mod), |
| 134 | raising=False, |
| 135 | ) |
| 136 | |
| 137 | model_name = "Qwen/Qwen2.5-1.5B-Instruct" |
| 138 | seed = 42 |
| 139 | max_tokens = 24 |
| 140 | num_logprobs = 5 |
| 141 | prompts = [ |
| 142 | "Hello, my name is", |
| 143 | "The president of the United States is", |
| 144 | "The capital of France is", |
| 145 | ] |
| 146 | |
| 147 | set_random_seed(seed) |
| 148 | with vllm_runner( |
| 149 | model_name, |
| 150 | runner="generate", |
| 151 | tensor_parallel_size=1, |
| 152 | num_gpu_blocks_override=128, |
| 153 | enforce_eager=True, |
| 154 | attention_config={"backend": "FLEX_ATTENTION"}, |
| 155 | ) as llm_eager: |
| 156 | output_eager = llm_eager.generate_greedy_logprobs( |
| 157 | prompts, max_tokens, num_logprobs |
| 158 | ) |
| 159 | |
| 160 | set_random_seed(seed) |
| 161 | with vllm_runner( |
| 162 | model_name, |
| 163 | runner="generate", |
| 164 | tensor_parallel_size=1, |
| 165 | num_gpu_blocks_override=128, |
| 166 | enforce_eager=False, |
| 167 | gpu_memory_utilization=0.85, |
| 168 | compilation_config={ |
| 169 | "cudagraph_mode": "FULL", |
| 170 | "cudagraph_capture_sizes": [4], |
| 171 | }, |
| 172 | attention_config={"backend": "FLEX_ATTENTION"}, |
| 173 | ) as llm_cudagraph: |
| 174 | output_cudagraph = llm_cudagraph.generate_greedy_logprobs( |
| 175 | prompts, max_tokens, num_logprobs |
| 176 | ) |
| 177 | |
| 178 | check_logprobs_close( |
| 179 | outputs_0_lst=output_eager, |
| 180 | outputs_1_lst=output_cudagraph, |
| 181 | name_0="eager", |
| 182 | name_1="cudagraph", |
| 183 | ) |
| 184 | |
| 185 | |
| 186 | @pytest.mark.skipif( |
| 187 | not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION, |
| 188 | reason="CUDA not available or PyTorch version < 2.7", |
| 189 | ) |
| 190 | def test_flex_attention_vs_default_backend(vllm_runner): |
| 191 | """Test that FlexAttention produces the same outputs as the default backend. |
| 192 | |
| 193 | This test compares the outputs from the FlexAttention backend with |
| 194 | the default backend, ensuring they are similar when using the same seed. |
| 195 | """ |
| 196 | model_name = "Qwen/Qwen2.5-1.5B-Instruct" |
| 197 | seed = 42 |
| 198 | max_tokens = 24 |
| 199 | num_logprobs = 5 |
| 200 | prompts = [ |
| 201 | "Hello, my name is", |
| 202 | "The president of the United States is", |
| 203 | "The capital of France is", |
| 204 | ] |
| 205 | |
| 206 | # Run with flex attention |
| 207 | set_random_seed(seed) |
| 208 | with vllm_runner( |
| 209 | model_name, |
| 210 | runner="generate", |
| 211 | tensor_parallel_size=1, |
| 212 | num_gpu_blocks_override=128, |
| 213 | enforce_eager=True, |
| 214 | attention_config={"backend": "FLEX_ATTENTION"}, |
| 215 | ) as llm_flex: |
| 216 | output_flex = llm_flex.generate_greedy_logprobs( |
| 217 | prompts, max_tokens, num_logprobs |
| 218 | ) |
| 219 | |
| 220 | # Run with default backend |
| 221 | set_random_seed(seed) |
| 222 | with vllm_runner( |
| 223 | model_name, |
| 224 | runner="generate", |
| 225 | tensor_parallel_size=1, |
| 226 | num_gpu_blocks_override=128, |
| 227 | enforce_eager=True, |
| 228 | gpu_memory_utilization=0.85, |
| 229 | ) as llm_default: |
| 230 | output_default = llm_default.generate_greedy_logprobs( |
| 231 | prompts, max_tokens, num_logprobs |
| 232 | ) |
| 233 | |
| 234 | check_logprobs_close( |
| 235 | outputs_0_lst=output_flex, |
| 236 | outputs_1_lst=output_default, |
| 237 | name_0="flex", |
| 238 | name_1="default", |
| 239 | ) |
| 240 | |
| 241 | |
| 242 | @pytest.mark.skipif( |
| 243 | not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION, |
| 244 | reason="CUDA not available or PyTorch version < 2.7", |
| 245 | ) |
| 246 | def test_encoder_flex_attention_vs_default_backend(vllm_runner): |
| 247 | """Test that FlexAttention produces the same outputs as the default backend. |
| 248 | |
| 249 | This test compares the outputs from the FlexAttention backend with |
| 250 | the default backend for encoder models. |
| 251 | """ |
| 252 | model_name = "BAAI/bge-base-en-v1.5" |
| 253 | # Exercise packed sequence boundaries inside 128-token FlexAttention |
| 254 | # blocks, including sequences that span more than one block. |
| 255 | prompts = [ |
| 256 | "hello " * 120, |
| 257 | "world " * 130, |
| 258 | "attention " * 254, |
| 259 | ] |
| 260 | |
| 261 | # Run with flex attention |
| 262 | with vllm_runner( |
| 263 | model_name, |
| 264 | runner="pooling", |
| 265 | dtype=torch.bfloat16, |
| 266 | tensor_parallel_size=1, |
| 267 | max_model_len=384, |
| 268 | enforce_eager=True, |
| 269 | attention_config={"backend": "FLEX_ATTENTION"}, |
| 270 | ) as llm_flex: |
| 271 | flex_outputs = llm_flex.embed(prompts) |
| 272 | |
| 273 | # Run with default backend |
| 274 | with vllm_runner( |
| 275 | model_name, |
| 276 | runner="pooling", |
| 277 | dtype=torch.bfloat16, |
| 278 | tensor_parallel_size=1, |
| 279 | max_model_len=384, |
| 280 | enforce_eager=True, |
| 281 | ) as llm_default: |
| 282 | default_outputs = llm_default.embed(prompts) |
| 283 | |
| 284 | check_embeddings_close( |
| 285 | embeddings_0_lst=flex_outputs, |
| 286 | embeddings_1_lst=default_outputs, |
| 287 | name_0="flex", |
| 288 | name_1="default", |
| 289 | tol=1e-2, |
| 290 | ) |
| 291 | |
| 292 | |
| 293 | @pytest.mark.skipif( |
| 294 | not torch.cuda.is_available() or TORCH_VERSION < DIRECT_BUILD_VERSION, |
| 295 | reason="CUDA not available or PyTorch version < 2.7", |
| 296 | ) |
| 297 | def test_block_mask_direct_vs_slow_path(): |
| 298 | """Test that direct path block mask is a superset of slow path. |
| 299 | |
| 300 | The direct path may include extra blocks for performance (over-estimation), |
| 301 | but must include all blocks that the slow path determines are necessary. |
| 302 | """ |
| 303 | device = torch.device("cuda") |
| 304 | |
| 305 | vllm_config = create_vllm_config( |
| 306 | model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024 |
| 307 | ) |
| 308 | kv_cache_spec = create_standard_kv_cache_spec(vllm_config) |
| 309 | |
| 310 | # Use a mixed batch that will create groups spanning multiple sequences |
| 311 | batch_spec = BatchSpec( |
| 312 | seq_lens=[35, 64, 128, 256], query_lens=[33, 5, 32, 64], name="test_mixed_batch" |
| 313 | ) |
| 314 | |
| 315 | common_attn_metadata = create_common_attn_metadata( |
| 316 | batch_spec, vllm_config.cache_config.block_size, device |
| 317 | ) |
| 318 | |
| 319 | builder = FlexAttentionMetadataBuilder(kv_cache_spec, [], vllm_config, device) |
| 320 | |
| 321 | metadata_direct = builder.build( |
| 322 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 323 | ) |
| 324 | builder.direct_build = False |
| 325 | metadata_slow = builder.build( |
| 326 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 327 | ) |
| 328 | |
| 329 | assert metadata_direct.block_mask is not None |
| 330 | assert metadata_slow.block_mask is not None |
| 331 | |
| 332 | # Extract block indices for comparison, B, H are the same |
| 333 | direct_indices = metadata_direct.block_mask.kv_indices[0, 0] |
| 334 | slow_indices = metadata_slow.block_mask.kv_indices[0, 0] |
| 335 | direct_num = metadata_direct.block_mask.kv_num_blocks[0, 0] |
| 336 | slow_num = metadata_slow.block_mask.kv_num_blocks[0, 0] |
| 337 | |
| 338 | # main test: every block needed by slow path must be in direct path |
| 339 | num_groups = direct_num.shape[0] |
| 340 | all_contained = True |
| 341 | missing_details = [] |
| 342 | |
| 343 | for group_idx in range(num_groups): |
| 344 | direct_blocks = set(direct_indices[group_idx, : direct_num[group_idx]].tolist()) |
| 345 | slow_blocks = set(slow_indices[group_idx, : slow_num[group_idx]].tolist()) |
| 346 | |
| 347 | missing_blocks = slow_blocks - direct_blocks |
| 348 | if missing_blocks: |
| 349 | all_contained = False |
| 350 | missing_details.append( |
| 351 | f"Group {group_idx}: missing {sorted(missing_blocks)}" |
| 352 | ) |
| 353 | |
| 354 | assert all_contained, ( |
| 355 | "Direct path is missing blocks required by slow path:\n" |
| 356 | + "\n".join(missing_details) |
| 357 | ) |
| 358 | |
| 359 | |
| 360 | @pytest.mark.skipif( |
| 361 | not torch.cuda.is_available() or TORCH_VERSION < DIRECT_BUILD_VERSION, |
| 362 | reason="CUDA not available or PyTorch version < 2.9", |
| 363 | ) |
| 364 | @pytest.mark.parametrize("direct_build", [True, False]) |
| 365 | @torch.inference_mode() |
| 366 | def test_flex_attention_request_count_changes_reuse_compiled_graph(direct_build): |
| 367 | """Request-count changes preserve attention results and compiled graph reuse.""" |
| 368 | from torch._dynamo.testing import CompileCounterWithBackend |
| 369 | from torch.nn.attention.flex_attention import flex_attention |
| 370 | |
| 371 | from vllm.v1.attention.backends.flex_attention import get_kernel_options |
| 372 | |
| 373 | torch.manual_seed(42) |
| 374 | device = torch.device("cuda") |
| 375 | config = create_vllm_config( |
| 376 | model_name="Qwen/Qwen3-0.6B", |
| 377 | max_model_len=128, |
| 378 | num_gpu_blocks=40, |
| 379 | max_num_seqs=8, |
| 380 | max_num_batched_tokens=128, |
| 381 | ) |
| 382 | builder = FlexAttentionMetadataBuilder( |
| 383 | create_standard_kv_cache_spec(config), [], config, device |
| 384 | ) |
| 385 | builder.direct_build = direct_build |
| 386 | query = torch.randn(1, 2, 32, 64, device=device) |
| 387 | key = torch.randn(1, 2, 640, 64, device=device) |
| 388 | value = torch.randn_like(key) |
| 389 | kernel_options = get_kernel_options(query, 16, 16, direct_build) |
| 390 | counter = CompileCounterWithBackend("inductor") |
| 391 | compiled_attention = torch.compile(flex_attention, backend=counter, fullgraph=True) |
| 392 | |
| 393 | def build_metadata(num_reqs): |
| 394 | # Keep Q/K tensor and block-mask sizes fixed to isolate request counts. |
| 395 | common = create_common_attn_metadata( |
| 396 | BatchSpec(seq_lens=[64] * num_reqs, query_lens=[32 // num_reqs] * num_reqs), |
| 397 | 16, |
| 398 | device, |
| 399 | max_block_idx=40, |
| 400 | ) |
| 401 | common.block_table_tensor.copy_( |
| 402 | torch.arange(1, num_reqs * 4 + 1, device=device).view(num_reqs, 4) |
| 403 | ) |
| 404 | return builder.build(0, common) |
| 405 | |
| 406 | def attend(metadata): |
| 407 | return compiled_attention( |
| 408 | query, |
| 409 | key, |
| 410 | value, |
| 411 | block_mask=metadata.block_mask, |
| 412 | kernel_options=kernel_options, |
| 413 | ) |
| 414 | |
| 415 | def reference(num_reqs): |
| 416 | query_len = 32 // num_reqs |
| 417 | q_idx = torch.arange(32, device=device)[:, None] |
| 418 | kv_idx = torch.arange(640, device=device)[None, :] |
| 419 | request = q_idx // query_len |
| 420 | logical_q = q_idx % query_len + 64 - query_len |
| 421 | logical_kv = kv_idx - (request * 4 + 1) * 16 |
| 422 | mask = (logical_kv >= 0) & (logical_kv <= logical_q) |
| 423 | return torch.nn.functional.scaled_dot_product_attention( |
| 424 | query, key, value, attn_mask=mask |
| 425 | ) |
| 426 | |
| 427 | for num_reqs in (4, 2, 1, 8, 1): |
| 428 | metadata = build_metadata(num_reqs) |
| 429 | torch.testing.assert_close( |
| 430 | attend(metadata), reference(num_reqs), atol=1e-4, rtol=1e-4 |
| 431 | ) |
| 432 | assert counter.frame_count == 1 |
| 433 | |
| 434 | if direct_build: |
| 435 | torch.accelerator.synchronize() |
| 436 | graph = torch.cuda.CUDAGraph() |
| 437 | with torch.cuda.graph(graph): |
| 438 | output = attend(metadata) |
| 439 | for num_reqs in (8, 2, 1): |
| 440 | metadata = build_metadata(num_reqs) |
| 441 | graph.replay() |
| 442 | torch.testing.assert_close( |
| 443 | output, reference(num_reqs), atol=1e-4, rtol=1e-4 |
| 444 | ) |
| 445 | |
| 446 | |
| 447 | def test_physical_to_logical_mapping_handles_reused_blocks(): |
| 448 | """Regression test: reused physical blocks map to the latest logical block. |
| 449 | |
| 450 | For sliding-window / hybrid attention layers, physical KV-cache blocks can be |
| 451 | reused over time. The inverse mapping must therefore select the latest |
| 452 | logical block index for a physical block id. |
| 453 | """ |
| 454 | # Padding should not make physical block 0 look live. |
| 455 | block_table = torch.tensor([[6, 0, 0, 0]], dtype=torch.int32) |
| 456 | seq_lens = torch.tensor([1 * 16], dtype=torch.int32) # only 1 block valid |
| 457 | out = physical_to_logical_mapping( |
| 458 | block_table=block_table, seq_lens=seq_lens, block_size=16, total_blocks=10 |
| 459 | ) |
| 460 | assert out[0, 0].item() == -1 |
| 461 | assert out[0, 6].item() == 0 |
| 462 | |
| 463 | # If a physical block id appears multiple times (block reuse), mapping should |
| 464 | # point to the latest logical block index. |
| 465 | block_table2 = torch.tensor([[2, 2, 5]], dtype=torch.int32) |
| 466 | seq_lens2 = torch.tensor([3 * 16], dtype=torch.int32) |
| 467 | out2 = physical_to_logical_mapping( |
| 468 | block_table=block_table2, seq_lens=seq_lens2, block_size=16, total_blocks=8 |
| 469 | ) |
| 470 | assert out2[0, 2].item() == 1 |
| 471 | |
| 472 | |
| 473 | @pytest.mark.skipif( |
| 474 | not torch.cuda.is_available() or TORCH_VERSION < DIRECT_BUILD_VERSION, |
| 475 | reason="CUDA not available or PyTorch version < 2.9", |
| 476 | ) |
| 477 | def test_block_sparsity_hint_prunes_blocks(): |
| 478 | """Test that BlockSparsityHint prunes KV blocks from the direct build path. |
| 479 | |
| 480 | Uses a hint that only keeps the diagonal (q_block == kv_block) to verify |
| 481 | that off-diagonal blocks are excluded from the resulting BlockMask. |
| 482 | """ |
| 483 | device = torch.device("cuda") |
| 484 | |
| 485 | vllm_config = create_vllm_config( |
| 486 | model_name="facebook/opt-125m", |
| 487 | block_size=16, |
| 488 | max_model_len=1024, |
| 489 | ) |
| 490 | kv_cache_spec = create_standard_kv_cache_spec(vllm_config) |
| 491 | |
| 492 | batch_spec = BatchSpec( |
| 493 | seq_lens=[256], |
| 494 | query_lens=[256], |
| 495 | name="test_sparsity_hint", |
| 496 | ) |
| 497 | |
| 498 | common_attn_metadata = create_common_attn_metadata( |
| 499 | batch_spec, vllm_config.cache_config.block_size, device |
| 500 | ) |
| 501 | |
| 502 | builder = FlexAttentionMetadataBuilder(kv_cache_spec, [], vllm_config, device) |
| 503 | |
| 504 | metadata_no_hint = builder.build( |
| 505 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 506 | ) |
| 507 | metadata_no_hint.block_mask = metadata_no_hint._build_block_mask_direct() |
| 508 | assert metadata_no_hint.block_mask.kv_num_blocks.max().item() > 1 |
| 509 | |
| 510 | def diagonal_hint(q_block_idx, kv_block_idx, block_size): |
| 511 | return q_block_idx == kv_block_idx |
| 512 | |
| 513 | metadata_with_hint = builder.build( |
| 514 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 515 | ) |
| 516 | metadata_with_hint.block_sparsity_hint = BlockSparsityHint( |
| 517 | hint_fn=diagonal_hint, |
| 518 | ) |
| 519 | metadata_with_hint.block_mask = metadata_with_hint._build_block_mask_direct() |
| 520 | assert metadata_with_hint.block_mask.kv_num_blocks.max().item() <= 1 |
| 521 | |
| 522 | |
| 523 | if __name__ == "__main__": |
| 524 | pytest.main([__file__]) |