From 7e95735eb90fd6497a9386f515db4242688a9abb Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 9 Sep 2026 10:10:25 -0500 Subject: [PATCH] [Bugfix] Avoid FlexAttention recompiles when request counts change (#55888) Signed-off-by: Andreas Karatzas --- tests/kernels/test_flex_attention.py | 87 ++++++++++++++++++++ vllm/v1/attention/backends/flex_attention.py | 29 ++++--- 2 files changed, 105 insertions(+), 11 deletions(-) diff --git a/tests/kernels/test_flex_attention.py b/tests/kernels/test_flex_attention.py index a80467ecc3..12d00ec405 100644 --- a/tests/kernels/test_flex_attention.py +++ b/tests/kernels/test_flex_attention.py @@ -357,6 +357,93 @@ def test_block_mask_direct_vs_slow_path(): ) +@pytest.mark.skipif( + not torch.cuda.is_available() or TORCH_VERSION < DIRECT_BUILD_VERSION, + reason="CUDA not available or PyTorch version < 2.9", +) +@pytest.mark.parametrize("direct_build", [True, False]) +@torch.inference_mode() +def test_flex_attention_request_count_changes_reuse_compiled_graph(direct_build): + """Request-count changes preserve attention results and compiled graph reuse.""" + from torch._dynamo.testing import CompileCounterWithBackend + from torch.nn.attention.flex_attention import flex_attention + + from vllm.v1.attention.backends.flex_attention import get_kernel_options + + torch.manual_seed(42) + device = torch.device("cuda") + config = create_vllm_config( + model_name="Qwen/Qwen3-0.6B", + max_model_len=128, + num_gpu_blocks=40, + max_num_seqs=8, + max_num_batched_tokens=128, + ) + builder = FlexAttentionMetadataBuilder( + create_standard_kv_cache_spec(config), [], config, device + ) + builder.direct_build = direct_build + query = torch.randn(1, 2, 32, 64, device=device) + key = torch.randn(1, 2, 640, 64, device=device) + value = torch.randn_like(key) + kernel_options = get_kernel_options(query, 16, 16, direct_build) + counter = CompileCounterWithBackend("inductor") + compiled_attention = torch.compile(flex_attention, backend=counter, fullgraph=True) + + def build_metadata(num_reqs): + # Keep Q/K tensor and block-mask sizes fixed to isolate request counts. + common = create_common_attn_metadata( + BatchSpec(seq_lens=[64] * num_reqs, query_lens=[32 // num_reqs] * num_reqs), + 16, + device, + max_block_idx=40, + ) + common.block_table_tensor.copy_( + torch.arange(1, num_reqs * 4 + 1, device=device).view(num_reqs, 4) + ) + return builder.build(0, common) + + def attend(metadata): + return compiled_attention( + query, + key, + value, + block_mask=metadata.block_mask, + kernel_options=kernel_options, + ) + + def reference(num_reqs): + query_len = 32 // num_reqs + q_idx = torch.arange(32, device=device)[:, None] + kv_idx = torch.arange(640, device=device)[None, :] + request = q_idx // query_len + logical_q = q_idx % query_len + 64 - query_len + logical_kv = kv_idx - (request * 4 + 1) * 16 + mask = (logical_kv >= 0) & (logical_kv <= logical_q) + return torch.nn.functional.scaled_dot_product_attention( + query, key, value, attn_mask=mask + ) + + for num_reqs in (4, 2, 1, 8, 1): + metadata = build_metadata(num_reqs) + torch.testing.assert_close( + attend(metadata), reference(num_reqs), atol=1e-4, rtol=1e-4 + ) + assert counter.frame_count == 1 + + if direct_build: + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = attend(metadata) + for num_reqs in (8, 2, 1): + metadata = build_metadata(num_reqs) + graph.replay() + torch.testing.assert_close( + output, reference(num_reqs), atol=1e-4, rtol=1e-4 + ) + + def test_physical_to_logical_mapping_handles_reused_blocks(): """Regression test: reused physical blocks map to the latest logical block. diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 61c6344373..7529d8d20f 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -909,6 +909,12 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat self.persistent_offset_tensor = torch.empty( max_num_seqs, dtype=torch.int32, device=device ) + self.persistent_query_start_loc = torch.empty( + max_num_seqs + 1, dtype=torch.int32, device=device + ) + self.persistent_seq_lens = torch.empty( + max_num_seqs, dtype=torch.int32, device=device + ) # Persistent buffer for R-SWA per-request prefix lengths so the device # address stays stable across steps (required for CUDA graph replay). self.rswa_window: int | None = self.model_config.rswa_window @@ -1079,19 +1085,18 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat else self.persistent_kv_num_blocks ) - inverse_block_table = copy_to_persistent( - self.persistent_physical_to_logical, inverse_block_table - ) + copy_to_persistent(self.persistent_physical_to_logical, inverse_block_table) offset_tensor = common_attn_metadata.compute_num_computed_tokens() - offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor) + copy_to_persistent(self.persistent_offset_tensor, offset_tensor) + copy_to_persistent(self.persistent_query_start_loc, query_start_loc) + copy_to_persistent(self.persistent_seq_lens, seq_lens) rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens if use_rswa and rswa_prefix_lens is not None: assert self.persistent_rswa_prefix_lens is not None - rswa_prefix_lens = copy_to_persistent( - self.persistent_rswa_prefix_lens, rswa_prefix_lens - ) + copy_to_persistent(self.persistent_rswa_prefix_lens, rswa_prefix_lens) + rswa_prefix_lens = self.persistent_rswa_prefix_lens uses_paged_kv = not isinstance(self.kv_cache_spec, EncoderOnlyAttentionSpec) logical_mask_mod = ( @@ -1112,10 +1117,12 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat sliding_window=sliding_window, num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, - query_start_loc=query_start_loc, + # Mask closures only index active requests, so retain full-capacity + # tensor shapes to avoid recompiling when the batch shrinks to one. + query_start_loc=self.persistent_query_start_loc, query_start_loc_cpu=query_start_loc_cpu, max_seq_len=max_seq_len, - seq_lens=seq_lens, + seq_lens=self.persistent_seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, use_cascade=use_cascade, @@ -1126,9 +1133,9 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat block_size=block_size, max_possible_sequence_length=max_possible_seq_len, num_reqs=num_reqs, - physical_to_logical=inverse_block_table, + physical_to_logical=self.persistent_physical_to_logical, total_cache_tokens=total_cache_tokens, - decode_offset=offset_tensor, + decode_offset=self.persistent_offset_tensor, num_blocks_per_seq=num_blocks_per_seq, uses_paged_kv=uses_paged_kv, # FIXME(Isotr0py): direct build has issue to build bidirectional