[FlexAttention] Avoid encoder block-mask compile explosion (#50339)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 30, 2026 at 10:24 UTC
5b958907420e3e4dcbccd47ce912218475c885cc
4 files changed
+66
-13
tests/entrypoints/pooling/embed/test_online_long_text.py
+5
-3
@@ -20,7 +20,9 @@ from vllm.platforms import current_platform
20
21
22
def _generate_random_text(word_count: int) -> str:
23
- """Generate random text with approximately the specified word count."""
23
+ """Generate deterministic text with approximately the specified word count."""
24
+ rng = random.Random(word_count)
25
+
26
# Common English words with focus on verbs and nouns for realistic text
27
common_words = [
28
# Essential articles and pronouns (minimal)
@@ -178,7 +180,7 @@ def _generate_random_text(word_count: int) -> str:
180
181
words = []
182
for _ in range(word_count):
181
- words.append(random.choice(common_words))
183
+ words.append(rng.choice(common_words))
184
185
# Add some punctuation for more realistic text
186
text = " ".join(words)
@@ -187,7 +189,7 @@ def _generate_random_text(word_count: int) -> str:
189
result = []
190
for i, word in enumerate(words_list):
191
result.append(word)
190
- if (i + 1) % random.randint(10, 20) == 0 and i < len(words_list) - 1:
192
+ if (i + 1) % rng.randint(10, 20) == 0 and i < len(words_list) - 1:
193
result[-1] += "."
194
195
return " ".join(result)
tests/kernels/test_flex_attention.py
+44
-5
@@ -13,6 +13,7 @@ from tests.v1.attention.utils import (
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,
@@ -27,6 +28,42 @@ 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",
@@ -213,10 +250,12 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner):
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 = [
217
- "Hello, my name is",
218
- "The president of the United States is",
219
- "The capital of France is",
256
+ "hello " * 120,
257
+ "world " * 130,
258
+ "attention " * 254,
259
]
260
261
# Run with flex attention
@@ -225,7 +264,7 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner):
264
runner="pooling",
265
dtype=torch.bfloat16,
266
tensor_parallel_size=1,
228
- max_model_len=100,
267
+ max_model_len=384,
268
enforce_eager=True,
269
attention_config={"backend": "FLEX_ATTENTION"},
270
) as llm_flex:
@@ -237,7 +276,7 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner):
276
runner="pooling",
277
dtype=torch.bfloat16,
278
tensor_parallel_size=1,
240
- max_model_len=100,
279
+ max_model_len=384,
280
enforce_eager=True,
281
) as llm_default:
282
default_outputs = llm_default.embed(prompts)
vllm/config/attention.py
+5
-3
@@ -89,13 +89,15 @@ class AttentionConfig:
89
flex_attn_q_block_size: int | None = None
90
"""Logical Q block size for the flex attention block mask.
91
Must be a power of 2 and divisible by flex_attn_block_m.
92
- If None, uses the default (16 on PyTorch >= 2.9, 128 otherwise)."""
92
+ If None, uses 16 for paged KV attention on PyTorch >= 2.9, and 128
93
+ for encoder-only attention or older PyTorch versions."""
94
95
flex_attn_kv_block_size: int | None = None
96
"""Logical KV block size for the flex attention block mask.
97
Must be a power of 2 and divisible by flex_attn_block_n.
97
- If None, uses the default (kv_cache_block_size on PyTorch >= 2.9,
98
- 128 otherwise)."""
98
+ If None, uses the KV cache block size for paged KV attention on
99
+ PyTorch >= 2.9, and 128 for encoder-only attention or older PyTorch
100
+ versions."""
101
102
def compute_hash(self) -> str:
103
"""
vllm/v1/attention/backends/flex_attention.py
+12
-2
@@ -877,12 +877,14 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
877
self.block_size = kv_cache_spec.block_size
878
self.kv_cache_spec = kv_cache_spec
879
supports_small_blocks = is_torch_equal_or_newer("2.9.0.dev0")
880
+ uses_paged_kv = not isinstance(kv_cache_spec, EncoderOnlyAttentionSpec)
881
self.direct_build: bool = supports_small_blocks
882
883
self.q_block_size, self.kv_block_size = self._get_block_sizes(
884
vllm_config.attention_config,
885
supports_small_blocks,
886
self.block_size,
887
+ uses_paged_kv,
888
)
889
890
if self.direct_build and self.kv_block_size != self.block_size:
@@ -945,9 +947,17 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
947
attn_cfg,
948
supports_small_blocks: bool,
949
cache_block_size: int,
950
+ uses_paged_kv: bool,
951
) -> tuple[int, int]:
949
- q_block_size = 16 if supports_small_blocks else 128
950
- kv_block_size = cache_block_size if supports_small_blocks else 128
952
+ # Small blocks are efficient with the direct mask builder, which maps
953
+ # logical blocks to paged KV cache blocks without materializing and
954
+ # sorting a generic block mask. Encoder-only attention has no paged KV
955
+ # cache and therefore cannot use that path. Keeping its generic path at
956
+ # 16-token blocks creates very large mask-conversion graphs and makes
957
+ # Inductor compilation dominate the first request.
958
+ use_small_blocks = supports_small_blocks and uses_paged_kv
959
+ q_block_size = 16 if use_small_blocks else 128
960
+ kv_block_size = cache_block_size if use_small_blocks else 128
961
962
q_block_size = attn_cfg.flex_attn_q_block_size or q_block_size
963
if (q_block_size & (q_block_size - 1)) != 0 or (