A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 17:45 UTC
expose flex block size for batch invariant mode (#41252)
Signed-off-by: Angel Li <liangel@meta.com>
liangel-02 committed
May 13, 2026 at 17:11 UTC
6b5c389ee326195534859a480cd2ac714d2ac14e
4 files changed
+114
-8
tests/v1/determinism/test_batch_invariance.py
+11
-1
@@ -147,8 +147,14 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
147
"backend",
148
BACKENDS,
149
)
150
+@pytest.mark.parametrize(
151
+ "block_m,block_n",
152
+ [(16, 16), (8, 16)],
153
+)
154
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
155
backend,
156
+ block_m,
157
+ block_n,
158
):
159
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
160
random.seed(seed)
@@ -172,7 +178,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
178
max_model_len=8192,
179
dtype="auto", # not everything is supported
180
gpu_memory_utilization=0.9,
175
- attention_config={"backend": backend},
181
+ attention_config={
182
+ "backend": backend,
183
+ "flex_attn_block_m": block_m,
184
+ "flex_attn_block_n": block_n,
185
+ },
186
)
187
188
# Use more realistic prompts for better token generation
vllm/config/attention.py
+21
@@ -67,6 +67,27 @@ class AttentionConfig:
67
use_non_causal: bool = False
68
"""Whether to use non-causal (bidirectional) attention."""
69
70
+ flex_attn_block_m: int | None = None
71
+ """Triton kernel BLOCK_M tile size for flex attention.
72
+ Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1,
73
+ defaults to 16."""
74
+
75
+ flex_attn_block_n: int | None = None
76
+ """Triton kernel BLOCK_N tile size for flex attention.
77
+ Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1,
78
+ defaults to 16."""
79
+
80
+ flex_attn_q_block_size: int | None = None
81
+ """Logical Q block size for the flex attention block mask.
82
+ Must be a power of 2 and divisible by flex_attn_block_m.
83
+ If None, uses the default (16 on PyTorch >= 2.9, 128 otherwise)."""
84
+
85
+ flex_attn_kv_block_size: int | None = None
86
+ """Logical KV block size for the flex attention block mask.
87
+ Must be a power of 2 and divisible by flex_attn_block_n.
88
+ If None, uses the default (kv_cache_block_size on PyTorch >= 2.9,
89
+ 128 otherwise)."""
90
+
91
def compute_hash(self) -> str:
92
"""
93
Provide a hash that uniquely identifies all the configs
vllm/model_executor/layers/attention/attention.py
+23
@@ -341,6 +341,29 @@ class Attention(nn.Module, AttentionLayerBase):
341
f"but got {self.attn_backend.get_name()}."
342
)
343
344
+ if self.attn_backend.get_name() == "FLEX_ATTENTION":
345
+ block_m = vllm_config.attention_config.flex_attn_block_m
346
+ block_n = vllm_config.attention_config.flex_attn_block_n
347
+
348
+ if envs.VLLM_BATCH_INVARIANT and cache_config is not None:
349
+ if block_m is not None and block_m > cache_config.block_size:
350
+ raise ValueError(
351
+ f"flex_attn_block_m ({block_m}) must be "
352
+ f"<= cache block size ({cache_config.block_size}) for "
353
+ f"batch invariance"
354
+ )
355
+ if block_n is not None and block_n > cache_config.block_size:
356
+ raise ValueError(
357
+ f"flex_attn_block_n ({block_n}) must be "
358
+ f"<= cache block size ({cache_config.block_size}) for "
359
+ f"batch invariance"
360
+ )
361
+
362
+ if block_m is not None:
363
+ extra_impl_args.setdefault("block_m", block_m)
364
+ if block_n is not None:
365
+ extra_impl_args.setdefault("block_n", block_n)
366
+
367
impl_cls = self.attn_backend.get_impl_cls()
368
self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an AttentionImpl subclass
369
num_heads,
vllm/v1/attention/backends/flex_attention.py
+59
-7
@@ -769,8 +769,15 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
769
self.kv_cache_spec = kv_cache_spec
770
supports_small_blocks = is_torch_equal_or_newer("2.9.0.dev0")
771
self.direct_build: bool = supports_small_blocks
772
- self.q_block_size: int = 16 if supports_small_blocks else 128
773
- self.kv_block_size: int = self.block_size if supports_small_blocks else 128
772
+
773
+ self.q_block_size, self.kv_block_size = self._get_block_sizes(
774
+ vllm_config.attention_config,
775
+ supports_small_blocks,
776
+ self.block_size,
777
+ )
778
+
779
+ if self.direct_build and self.kv_block_size != self.block_size:
780
+ self.direct_build = False
781
782
self.max_model_len = self.model_config.max_model_len
783
max_num_seqs = vllm_config.scheduler_config.max_num_seqs
@@ -792,6 +799,39 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
799
self.persistent_physical_to_logical = None
800
self.persistent_kv_indices = None
801
802
+ @staticmethod
803
+ def _get_block_sizes(
804
+ attn_cfg,
805
+ supports_small_blocks: bool,
806
+ cache_block_size: int,
807
+ ) -> tuple[int, int]:
808
+ q_block_size = 16 if supports_small_blocks else 128
809
+ kv_block_size = cache_block_size if supports_small_blocks else 128
810
+
811
+ q_block_size = attn_cfg.flex_attn_q_block_size or q_block_size
812
+ if (q_block_size & (q_block_size - 1)) != 0 or (
813
+ attn_cfg.flex_attn_block_m is not None
814
+ and q_block_size % attn_cfg.flex_attn_block_m != 0
815
+ ):
816
+ raise ValueError(
817
+ f"flex_attn_q_block_size must be a power of 2 "
818
+ f"and divisible by flex_attn_block_m, got "
819
+ f"{q_block_size}, {attn_cfg.flex_attn_block_m}"
820
+ )
821
+
822
+ kv_block_size = attn_cfg.flex_attn_kv_block_size or kv_block_size
823
+ if (kv_block_size & (kv_block_size - 1)) != 0 or (
824
+ attn_cfg.flex_attn_block_n is not None
825
+ and kv_block_size % attn_cfg.flex_attn_block_n != 0
826
+ ):
827
+ raise ValueError(
828
+ f"flex_attn_kv_block_size must be a power of 2 "
829
+ f"and divisible by flex_attn_block_n, got "
830
+ f"{kv_block_size}, {attn_cfg.flex_attn_block_n}"
831
+ )
832
+
833
+ return q_block_size, kv_block_size
834
+
835
def build_for_cudagraph_capture(
836
self, common_attn_metadata: CommonAttentionMetadata
837
) -> FlexAttentionMetadata:
@@ -946,6 +986,8 @@ class FlexAttentionImpl(AttentionImpl):
986
logits_soft_cap: float | None = None,
987
attn_type: AttentionType = AttentionType.DECODER,
988
kv_sharing_target_layer_name: str | None = None,
989
+ block_m: int | None = None,
990
+ block_n: int | None = None,
991
**kwargs,
992
) -> None:
993
self.num_heads = num_heads
@@ -986,6 +1028,14 @@ class FlexAttentionImpl(AttentionImpl):
1028
"FlexAttention does not support quantized kv-cache. Yet"
1029
)
1030
1031
+ self.block_m = 16 if envs.VLLM_BATCH_INVARIANT else None
1032
+ self.block_n = 16 if envs.VLLM_BATCH_INVARIANT else None
1033
+
1034
+ if block_m is not None:
1035
+ self.block_m = block_m
1036
+ if block_n is not None:
1037
+ self.block_n = block_n
1038
+
1039
@staticmethod
1040
def view_as_4d(tensor: torch.Tensor) -> torch.Tensor:
1041
"""View a 3d tensor as 4D."""
@@ -1131,6 +1181,13 @@ class FlexAttentionImpl(AttentionImpl):
1181
kernel_options = get_kernel_options(
1182
query, block_m, block_n, attn_metadata.direct_build
1183
)
1184
+
1185
+ if self.block_m is not None:
1186
+ kernel_options["BLOCK_M"] = self.block_m
1187
+ if self.block_n is not None:
1188
+ kernel_options["BLOCK_N"] = self.block_n
1189
+ if envs.VLLM_BATCH_INVARIANT:
1190
+ kernel_options["IS_DIVISIBLE"] = False
1191
out = flex_attention_compiled(
1192
query,
1193
key_tensor,
@@ -1170,11 +1227,6 @@ def get_kernel_options(
1227
return block_size
1228
return candidate
1229
1173
- if envs.VLLM_BATCH_INVARIANT:
1174
- kernel_options["BLOCK_M"] = 16
1175
- kernel_options["BLOCK_N"] = 16
1176
- kernel_options["IS_DIVISIBLE"] = False
1177
- return kernel_options
1230
if use_direct_build:
1231
kernel_options["BLOCK_M"] = block_m
1232
kernel_options["BLOCK_N"] = block_n