A high-throughput and memory-efficient inference and serving engine for LLMs
[FlexAttention] make custom mask mods fully cudagraphable (#45232)
Signed-off-by: Angel Li <liangel@meta.com>
liangel-02 committed
Jun 16, 2026 at 23:53 UTC
20a5f8b43ba5e1f691cc98f143d09c0593e1828a
2 files changed
+98
-1
tests/kernels/test_flex_attention.py
+67
@@ -13,6 +13,7 @@ from tests.v1.attention.utils import (
13
create_standard_kv_cache_spec,
14
create_vllm_config,
15
)
16
+from vllm.model_executor.layers.attention import Attention
17
from vllm.v1.attention.backends.flex_attention import (
18
BlockSparsityHint,
19
FlexAttentionMetadataBuilder,
@@ -79,6 +80,72 @@ def test_flex_attention_full_cudagraphs(vllm_runner):
80
)
81
82
83
+def windowed_causal_mask_mod(b, h, q_idx, kv_idx):
84
+ return (kv_idx <= q_idx) & (q_idx - kv_idx < 4)
85
+
86
+
87
+@pytest.mark.skipif(
88
+ not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
89
+ reason="CUDA not available or PyTorch version < 2.7",
90
+)
91
+def test_flex_attention_custom_mask_full_cudagraphs(vllm_runner, monkeypatch):
92
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
93
+ monkeypatch.setattr(
94
+ Attention,
95
+ "logical_mask_mod",
96
+ staticmethod(windowed_causal_mask_mod),
97
+ raising=False,
98
+ )
99
+
100
+ model_name = "Qwen/Qwen2.5-1.5B-Instruct"
101
+ seed = 42
102
+ max_tokens = 24
103
+ num_logprobs = 5
104
+ prompts = [
105
+ "Hello, my name is",
106
+ "The president of the United States is",
107
+ "The capital of France is",
108
+ ]
109
+
110
+ set_random_seed(seed)
111
+ with vllm_runner(
112
+ model_name,
113
+ runner="generate",
114
+ tensor_parallel_size=1,
115
+ num_gpu_blocks_override=128,
116
+ enforce_eager=True,
117
+ attention_config={"backend": "FLEX_ATTENTION"},
118
+ ) as llm_eager:
119
+ output_eager = llm_eager.generate_greedy_logprobs(
120
+ prompts, max_tokens, num_logprobs
121
+ )
122
+
123
+ set_random_seed(seed)
124
+ with vllm_runner(
125
+ model_name,
126
+ runner="generate",
127
+ tensor_parallel_size=1,
128
+ num_gpu_blocks_override=128,
129
+ enforce_eager=False,
130
+ gpu_memory_utilization=0.85,
131
+ compilation_config={
132
+ "cudagraph_mode": "FULL",
133
+ "cudagraph_capture_sizes": [4],
134
+ },
135
+ attention_config={"backend": "FLEX_ATTENTION"},
136
+ ) as llm_cudagraph:
137
+ output_cudagraph = llm_cudagraph.generate_greedy_logprobs(
138
+ prompts, max_tokens, num_logprobs
139
+ )
140
+
141
+ check_logprobs_close(
142
+ outputs_0_lst=output_eager,
143
+ outputs_1_lst=output_cudagraph,
144
+ name_0="eager",
145
+ name_1="cudagraph",
146
+ )
147
+
148
+
149
@pytest.mark.skipif(
150
not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
151
reason="CUDA not available or PyTorch version < 2.7",
vllm/v1/attention/backends/flex_attention.py
+31
-1
@@ -22,9 +22,10 @@ from torch.nn.attention.flex_attention import (
22
)
23
24
import vllm.envs as envs
25
-from vllm.config import VllmConfig
25
+from vllm.config import VllmConfig, get_layers_from_vllm_config
26
from vllm.config.cache import CacheDType
27
from vllm.logger import init_logger
28
+from vllm.model_executor.layers.attention import Attention
29
from vllm.platforms import current_platform
30
from vllm.utils.math_utils import cdiv
31
from vllm.utils.torch_utils import is_quantized_kv_cache, is_torch_equal_or_newer
@@ -807,6 +808,13 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
808
self.persistent_physical_to_logical = None
809
self.persistent_kv_indices = None
810
811
+ self.custom_logical_mask_mod: _mask_mod_signature | None = None
812
+ if self._uses_full_cudagraphs():
813
+ layers = get_layers_from_vllm_config(
814
+ vllm_config, Attention, self.layer_names
815
+ )
816
+ self.custom_logical_mask_mod = self._maybe_get_custom_mask_mod(layers)
817
+
818
@staticmethod
819
def _get_block_sizes(
820
attn_cfg,
@@ -853,6 +861,21 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
861
common_prefix_len=0, common_attn_metadata=common_attn_metadata
862
)
863
864
+ def _maybe_get_custom_mask_mod(self, layers) -> _mask_mod_signature | None:
865
+ mask_mods = {
866
+ getattr(layer, "logical_mask_mod", None) for layer in layers.values()
867
+ }
868
+ if len(mask_mods) > 1:
869
+ raise ValueError(
870
+ f"Found differing mask mods {mask_mods}, "
871
+ "cannot use alternating mask mods w/ full CUDA graphs"
872
+ )
873
+ return next(iter(mask_mods), None)
874
+
875
+ def _uses_full_cudagraphs(self) -> bool:
876
+ mode = self.vllm_config.compilation_config.cudagraph_mode
877
+ return mode is not None and mode.has_full_cudagraphs()
878
+
879
def build(
880
self,
881
common_prefix_len: int,
@@ -924,9 +947,16 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
947
else causal_mask_mod
948
)
949
950
+ sliding_window = None
951
+ if self._uses_full_cudagraphs():
952
+ if self.custom_logical_mask_mod is not None:
953
+ logical_mask_mod = self.custom_logical_mask_mod
954
+ sliding_window = getattr(self.kv_cache_spec, "sliding_window", None)
955
+
956
out = FlexAttentionMetadata(
957
causal=common_attn_metadata.causal,
958
logical_mask_mod=logical_mask_mod,
959
+ sliding_window=sliding_window,
960
num_actual_tokens=num_actual_tokens,
961
max_query_len=max_query_len,
962
query_start_loc=query_start_loc,