[Bugfix] Avoid FlexAttention recompiles when request counts change (#55888)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Sep 9, 2026 at 10:10 UTC
7e95735eb90fd6497a9386f515db4242688a9abb
2 files changed
+105
-11
tests/kernels/test_flex_attention.py
+87
@@ -357,6 +357,93 @@ def test_block_mask_direct_vs_slow_path():
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
vllm/v1/attention/backends/flex_attention.py
+18
-11
@@ -909,6 +909,12 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
909
self.persistent_offset_tensor = torch.empty(
910
max_num_seqs, dtype=torch.int32, device=device
911
)
912
+ self.persistent_query_start_loc = torch.empty(
913
+ max_num_seqs + 1, dtype=torch.int32, device=device
914
+ )
915
+ self.persistent_seq_lens = torch.empty(
916
+ max_num_seqs, dtype=torch.int32, device=device
917
+ )
918
# Persistent buffer for R-SWA per-request prefix lengths so the device
919
# address stays stable across steps (required for CUDA graph replay).
920
self.rswa_window: int | None = self.model_config.rswa_window
@@ -1079,19 +1085,18 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
1085
else self.persistent_kv_num_blocks
1086
)
1087
1082
- inverse_block_table = copy_to_persistent(
1083
- self.persistent_physical_to_logical, inverse_block_table
1084
- )
1088
+ copy_to_persistent(self.persistent_physical_to_logical, inverse_block_table)
1089
1090
offset_tensor = common_attn_metadata.compute_num_computed_tokens()
1087
- offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor)
1091
+ copy_to_persistent(self.persistent_offset_tensor, offset_tensor)
1092
+ copy_to_persistent(self.persistent_query_start_loc, query_start_loc)
1093
+ copy_to_persistent(self.persistent_seq_lens, seq_lens)
1094
1095
rswa_prefix_lens = common_attn_metadata.rswa_prefix_lens
1096
if use_rswa and rswa_prefix_lens is not None:
1097
assert self.persistent_rswa_prefix_lens is not None
1092
- rswa_prefix_lens = copy_to_persistent(
1093
- self.persistent_rswa_prefix_lens, rswa_prefix_lens
1094
- )
1098
+ copy_to_persistent(self.persistent_rswa_prefix_lens, rswa_prefix_lens)
1099
+ rswa_prefix_lens = self.persistent_rswa_prefix_lens
1100
1101
uses_paged_kv = not isinstance(self.kv_cache_spec, EncoderOnlyAttentionSpec)
1102
logical_mask_mod = (
@@ -1112,10 +1117,12 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
1117
sliding_window=sliding_window,
1118
num_actual_tokens=num_actual_tokens,
1119
max_query_len=max_query_len,
1115
- query_start_loc=query_start_loc,
1120
+ # Mask closures only index active requests, so retain full-capacity
1121
+ # tensor shapes to avoid recompiling when the batch shrinks to one.
1122
+ query_start_loc=self.persistent_query_start_loc,
1123
query_start_loc_cpu=query_start_loc_cpu,
1124
max_seq_len=max_seq_len,
1118
- seq_lens=seq_lens,
1125
+ seq_lens=self.persistent_seq_lens,
1126
block_table=block_table_tensor,
1127
slot_mapping=slot_mapping,
1128
use_cascade=use_cascade,
@@ -1126,9 +1133,9 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
1133
block_size=block_size,
1134
max_possible_sequence_length=max_possible_seq_len,
1135
num_reqs=num_reqs,
1129
- physical_to_logical=inverse_block_table,
1136
+ physical_to_logical=self.persistent_physical_to_logical,
1137
total_cache_tokens=total_cache_tokens,
1131
- decode_offset=offset_tensor,
1138
+ decode_offset=self.persistent_offset_tensor,
1139
num_blocks_per_seq=num_blocks_per_seq,
1140
uses_paged_kv=uses_paged_kv,
1141
# FIXME(Isotr0py): direct build has issue to build bidirectional