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 15, 07:15 UTC
Add FlexAttention to V1 (#16078)
Signed-off-by: drisspg <drisspguessous@gmail.com>
Driss Guessous committed
Jun 7, 2025 at 00:58 UTC
cf02f9b283a6a85f872e87b57dbeff7b767715e5
5 files changed
+575
tests/kernels/test_flex_attention.py
new
+93
@@ -0,0 +1,93 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Integration tests for FlexAttention backend vs default backend"""
3
+
4
+import random
5
+
6
+import numpy as np
7
+import pytest
8
+import torch
9
+from packaging import version
10
+
11
+from vllm import LLM, SamplingParams
12
+
13
+TORCH_VERSION = version.parse(torch.__version__)
14
+MINIMUM_TORCH_VERSION = version.parse("2.7.0")
15
+
16
+
17
+def set_seed(seed):
18
+ """Set seeds for reproducibility"""
19
+ random.seed(seed)
20
+ np.random.seed(seed)
21
+ torch.manual_seed(seed)
22
+ if torch.cuda.is_available():
23
+ torch.cuda.manual_seed_all(seed)
24
+
25
+
26
+@pytest.mark.skipif(
27
+ not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
28
+ reason="CUDA not available or PyTorch version < 2.7",
29
+)
30
+def test_flex_attention_vs_default_backend(monkeypatch):
31
+ """Test that FlexAttention produces the same outputs as the default backend.
32
+
33
+ This test compares the outputs from the FlexAttention backend with
34
+ the default backend, ensuring they are identical when using the same seed.
35
+ """
36
+ model_name = "Qwen/Qwen2.5-1.5B-Instruct"
37
+ seed = 42
38
+ max_tokens = 32
39
+ prompts = [
40
+ "Hello, my name is",
41
+ "The president of the United States is",
42
+ "The capital of France is",
43
+ ]
44
+
45
+ sampling_params = SamplingParams(temperature=0.0,
46
+ top_p=1.0,
47
+ seed=seed,
48
+ max_tokens=max_tokens)
49
+
50
+ # Run with flex attention
51
+ with monkeypatch.context() as m:
52
+ m.setenv("VLLM_USE_V1", "1")
53
+ m.setenv("VLLM_ATTENTION_BACKEND", "FLEX_ATTENTION")
54
+ m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
55
+
56
+ set_seed(seed)
57
+
58
+ llm_flex = LLM(
59
+ model_name,
60
+ tensor_parallel_size=1,
61
+ num_gpu_blocks_override=128,
62
+ enforce_eager=True,
63
+ )
64
+ output_flex = llm_flex.generate(prompts, sampling_params)
65
+
66
+ # Run with default backend
67
+ with monkeypatch.context() as m:
68
+ m.setenv("VLLM_USE_V1", "1")
69
+ m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
70
+ set_seed(seed)
71
+ llm_default = LLM(
72
+ model_name,
73
+ tensor_parallel_size=1,
74
+ num_gpu_blocks_override=128,
75
+ enforce_eager=True,
76
+ )
77
+ output_default = llm_default.generate(prompts, sampling_params)
78
+
79
+ # Compare outputs from both backends
80
+ for i, (flex_result,
81
+ default_result) in enumerate(zip(output_flex, output_default)):
82
+ prompt = prompts[i]
83
+ flex_text = flex_result.outputs[0].text
84
+ default_text = default_result.outputs[0].text
85
+
86
+ assert flex_text == default_text, (
87
+ f"FlexAttention output doesn't match default for: {prompt!r}\n"
88
+ f"FlexAttention: {flex_text!r}\n"
89
+ f"Default: {default_text!r}")
90
+
91
+
92
+if __name__ == "__main__":
93
+ pytest.main([__file__])
vllm/engine/arg_utils.py
+1
@@ -1409,6 +1409,7 @@ class EngineArgs:
1409
"FLASHINFER_VLLM_V1",
1410
"ROCM_AITER_MLA",
1411
"TORCH_SDPA_VLLM_V1",
1412
+ "FLEX_ATTENTION",
1413
]
1414
if (envs.is_set("VLLM_ATTENTION_BACKEND")
1415
and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):
vllm/platforms/cuda.py
+3
@@ -226,6 +226,9 @@ class CudaPlatformBase(Platform):
226
if selected_backend == _Backend.FLASHINFER:
227
logger.info_once("Using FlashInfer backend on V1 engine.")
228
return "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
229
+ if selected_backend == _Backend.FLEX_ATTENTION:
230
+ logger.info("Using FlexAttenion backend on V1 engine.")
231
+ return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501
232
if selected_backend == _Backend.TRITON_ATTN_VLLM_V1:
233
logger.info_once("Using Triton backend on V1 engine.")
234
return ("vllm.v1.attention.backends."
vllm/platforms/interface.py
+1
@@ -60,6 +60,7 @@ class _Backend(enum.Enum):
60
BLOCK_SPARSE_FLASH_ATTN = enum.auto()
61
DUAL_CHUNK_FLASH_ATTN = enum.auto()
62
NO_ATTENTION = enum.auto()
63
+ FLEX_ATTENTION = enum.auto()
64
65
66
class PlatformEnum(enum.Enum):
vllm/v1/attention/backends/flex_attention.py
new
+477
@@ -0,0 +1,477 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Attention layer with FlashAttention."""
3
+
4
+from dataclasses import dataclass
5
+from typing import TYPE_CHECKING, Any, Optional
6
+
7
+import torch
8
+from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
9
+ _score_mod_signature,
10
+ create_block_mask,
11
+ flex_attention)
12
+
13
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
14
+ AttentionMetadata, AttentionType,
15
+ is_quantized_kv_cache)
16
+from vllm.logger import init_logger
17
+from vllm.platforms import current_platform
18
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
19
+from vllm.v1.kv_cache_interface import AttentionSpec
20
+from vllm.v1.worker.block_table import BlockTable
21
+
22
+if current_platform.is_cuda():
23
+ pass
24
+
25
+logger = init_logger(__name__)
26
+
27
+if TYPE_CHECKING:
28
+ from vllm.v1.core.sched.output import SchedulerOutput
29
+ from vllm.v1.worker.gpu_input_batch import InputBatch
30
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
31
+
32
+create_block_mask_compiled = torch.compile(create_block_mask,
33
+ fullgraph=True,
34
+ mode="reduce-overhead")
35
+flex_attention_compiled = torch.compile(flex_attention, fullgraph=True)
36
+
37
+
38
+def _offsets_to_doc_ids_tensor(offsets: torch.Tensor) -> torch.Tensor:
39
+ device = offsets.device
40
+ counts = offsets[1:] - offsets[:-1]
41
+ return torch.repeat_interleave(
42
+ torch.arange(len(counts), device=device, dtype=torch.int32), counts)
43
+
44
+
45
+class FlexAttentionBackend(AttentionBackend):
46
+ accept_output_buffer: bool = True
47
+
48
+ @staticmethod
49
+ def get_supported_head_sizes() -> list[int]:
50
+ return [16, 32, 64, 96, 128, 160, 192, 224, 256]
51
+
52
+ @staticmethod
53
+ def get_name() -> str:
54
+ return "FLEX_ATTENTION"
55
+
56
+ @staticmethod
57
+ def get_impl_cls() -> type["FlexAttentionImpl"]:
58
+ return FlexAttentionImpl
59
+
60
+ @staticmethod
61
+ def get_metadata_cls() -> type["AttentionMetadata"]:
62
+ return FlexAttentionMetadata
63
+
64
+ @staticmethod
65
+ def get_kv_cache_shape(
66
+ num_blocks: int,
67
+ block_size: int,
68
+ num_kv_heads: int,
69
+ head_size: int,
70
+ ) -> tuple[int, ...]:
71
+ return (2, num_blocks, block_size, num_kv_heads, head_size)
72
+
73
+ @staticmethod
74
+ def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]:
75
+ return FlexAttentionMetadataBuilder
76
+
77
+ @staticmethod
78
+ def use_cascade_attention(*args, **kwargs) -> bool:
79
+ return False
80
+
81
+
82
+# @torch.compile(fullgraph=True, mode="reduce-overhead")
83
+def physical_to_logical_mapping(
84
+ block_table: torch.Tensor,
85
+ total_blocks: Optional[int] = None) -> torch.Tensor:
86
+ """
87
+ Creates an inverse mapping from physical block locations to logical indices.
88
+
89
+ The original block_table maps from logical blocks to physical locations:
90
+
91
+ Logical to Physical (Original block_table):
92
+ ┌───────────────────────────────────────────┐
93
+ │ Request 0: │
94
+ │ │
95
+ │ Logical Blocks: 0 1 2 3 4 5 6 7 │
96
+ │ │ │ │ │ │ │ │ │ │
97
+ │ v v v v v v v v │
98
+ │ Physical Blocks: 3 5 1 7 4 2 0 6 │
99
+ └───────────────────────────────────────────┘
100
+
101
+ This function creates the inverse mapping:
102
+
103
+ Physical to Logical (Inverse mapping):
104
+ ┌───────────────────────────────────────────┐
105
+ │ Request 0: │
106
+ │ │
107
+ │ Physical Blocks: 0 1 2 3 4 5 6 7 │
108
+ │ │ │ │ │ │ │ │ │ │
109
+ │ v v v v v v v v │
110
+ │ Logical Blocks: 6 2 5 0 4 1 7 3 │
111
+ └───────────────────────────────────────────┘
112
+
113
+ If multiple logical blocks map to the same physical block,
114
+ this function returns the first (minimum) logical block index.
115
+
116
+ If a physical block is not mapped to by any logical block,
117
+ its value in the result will be -1.
118
+
119
+
120
+ Args:
121
+ block_table: Tensor of shape [max_reqs, max_num_blocks]
122
+ mapping logical blocks to physical locations
123
+
124
+ Returns:
125
+ A tensor of shape [max_reqs, max_physical_block]
126
+ """
127
+ max_reqs, max_num_blocks = block_table.shape
128
+ device = block_table.device
129
+
130
+ physical_to_logical = torch.full((max_reqs, total_blocks),
131
+ -1,
132
+ dtype=torch.long,
133
+ device=device)
134
+
135
+ logical_indices = (torch.arange(max_num_blocks,
136
+ device=device).unsqueeze(0).expand(
137
+ max_reqs, -1))
138
+
139
+ physical_to_logical.scatter_(-1, block_table.to(torch.int64),
140
+ logical_indices)
141
+ # TODO Confirm - Seems like block 0 is always empty so we reset it manually
142
+ physical_to_logical[:, 0] = -1
143
+ return physical_to_logical
144
+
145
+
146
+def causal_mask_mod(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor,
147
+ kv_idx: torch.Tensor):
148
+ return q_idx >= kv_idx
149
+
150
+
151
+@dataclass
152
+class FlexAttentionMetadata:
153
+ num_actual_tokens: int # Number of tokens excluding padding.
154
+ max_query_len: int
155
+ query_start_loc: torch.Tensor
156
+ max_seq_len: int
157
+ seq_lens: torch.Tensor
158
+ block_table: torch.Tensor
159
+ slot_mapping: torch.Tensor
160
+
161
+ use_cascade: bool
162
+ common_prefix_len: int
163
+ cu_prefix_query_lens: Optional[torch.Tensor]
164
+ prefix_kv_lens: Optional[torch.Tensor]
165
+ suffix_kv_lens: Optional[torch.Tensor]
166
+
167
+ # Block info
168
+ total_cache_tokens: int
169
+ block_size: int
170
+ max_possible_sequence_length: int
171
+ num_reqs: int
172
+ physical_to_logical: torch.Tensor
173
+ decode_offset: torch.Tensor
174
+
175
+ # For logging.
176
+ num_input_tokens: int = 0 # Number of tokens including padding.
177
+
178
+ # Flex Metadata
179
+ num_blocks = 0
180
+ block_mask: Optional[BlockMask] = None
181
+ score_mod: Optional[_score_mod_signature] = None
182
+ mask_mod: Optional[_mask_mod_signature] = None
183
+ logical_mask_mod: _mask_mod_signature = causal_mask_mod
184
+
185
+ def get_mask_mod(self) -> _mask_mod_signature:
186
+ """Creates the mask_mod function for FlexAttention.
187
+
188
+ This function creates the combined mask mod function that handles:
189
+ 1. The paged attention block mapping
190
+ 2. The mapping from packed query sequences to logical query entries
191
+
192
+ It also by defaults adds the decoding offset to the query indices.
193
+ With this info we create the "logical" indices that are passed to
194
+ mask_mod functions. This allows mask mod functions to be agnostic to
195
+ layout of the query and key/value tensors.
196
+
197
+ TODO is_within_lower_bound: do sequences start on block_boundaries?
198
+ """
199
+ # Create a lookup mapping from query indices -> request number
200
+ request_lookup = _offsets_to_doc_ids_tensor(self.query_start_loc)
201
+
202
+ def final_mask_mod(
203
+ b: torch.Tensor,
204
+ h: torch.Tensor,
205
+ q_idx: torch.Tensor,
206
+ physical_kv_idx: torch.Tensor,
207
+ ) -> torch.Tensor:
208
+ # Map query indices to corresponding request indices
209
+ q_req = request_lookup[q_idx]
210
+
211
+ # Convert physical KV indices to logical indices
212
+ physical_kv_block = physical_kv_idx // self.block_size
213
+ physical_kv_offset = physical_kv_idx % self.block_size
214
+ logical_block_idx = self.physical_to_logical[q_req,
215
+ physical_kv_block]
216
+ logical_kv_idx = logical_block_idx * self.block_size + physical_kv_offset # noqa: E501
217
+
218
+ # Determine valid kv indices
219
+ live_block = logical_block_idx >= 0
220
+ within_upper_bound = logical_kv_idx < self.seq_lens[q_req]
221
+ within_lower_bound = logical_kv_idx >= 0
222
+
223
+ is_valid = live_block & within_upper_bound & within_lower_bound
224
+
225
+ # Convert physical query indices to logical indices
226
+ local_q_idx = q_idx - self.query_start_loc[q_req]
227
+ logical_q_idx = local_q_idx + self.decode_offset[q_req]
228
+
229
+ # Apply mask modification only for valid indices
230
+ return torch.where(
231
+ is_valid,
232
+ self.logical_mask_mod(b, h, logical_q_idx, logical_kv_idx),
233
+ False,
234
+ )
235
+
236
+ return final_mask_mod
237
+
238
+ def build_block_mask(self) -> BlockMask:
239
+ assert self.mask_mod is not None
240
+ return create_block_mask_compiled(
241
+ self.mask_mod,
242
+ None,
243
+ None,
244
+ self.num_actual_tokens,
245
+ self.total_cache_tokens,
246
+ )
247
+
248
+ def __post_init__(self):
249
+ assert self.use_cascade is False, "Not implemented yet."
250
+ assert self.common_prefix_len == 0, "Not implemented yet."
251
+ assert self.cu_prefix_query_lens is None, "Not implemented yet."
252
+ assert self.prefix_kv_lens is None, "Not implemented yet."
253
+ assert self.suffix_kv_lens is None, "Not implemented yet."
254
+ self.num_blocks = self.total_cache_tokens // self.block_size
255
+ self.mask_mod = self.get_mask_mod()
256
+ self.block_mask = self.build_block_mask()
257
+
258
+
259
+class FlexAttentionMetadataBuilder:
260
+
261
+ def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
262
+ block_table: BlockTable):
263
+ model_config = runner.model_config
264
+
265
+ self.runner = runner
266
+ self.num_heads_q = model_config.get_num_attention_heads(
267
+ runner.parallel_config)
268
+ self.num_heads_kv = model_config.get_num_kv_heads(
269
+ runner.parallel_config)
270
+ self.headdim = model_config.get_head_size()
271
+ self.block_size = kv_cache_spec.block_size
272
+ self.kv_cache_spec = kv_cache_spec
273
+ self.block_table = block_table
274
+
275
+ def reorder_batch(self, input_batch: "InputBatch",
276
+ scheduler_output: "SchedulerOutput") -> bool:
277
+ return False
278
+
279
+ def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
280
+ common_prefix_len: int,
281
+ common_attn_metadata: CommonAttentionMetadata):
282
+ max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
283
+ query_start_loc = common_attn_metadata.query_start_loc
284
+ seq_lens = common_attn_metadata.seq_lens
285
+
286
+ block_table = self.block_table
287
+ block_table_tensor = block_table.get_device_tensor()[:num_reqs]
288
+ block_table.slot_mapping[:num_actual_tokens].copy_(
289
+ block_table.slot_mapping_cpu[:num_actual_tokens],
290
+ non_blocking=True)
291
+ slot_mapping = block_table.slot_mapping[:num_actual_tokens]
292
+
293
+ use_cascade = common_prefix_len > 0
294
+ cu_prefix_query_lens = None
295
+ prefix_kv_lens = None
296
+ suffix_kv_lens = None
297
+ if use_cascade:
298
+ raise NotImplementedError("Not yet my friend")
299
+
300
+ block_size = self.kv_cache_spec.block_size
301
+ max_possible_seq_len = self.runner.model_config.max_model_len
302
+ total_cache_tokens = (self.runner.cache_config.num_gpu_blocks *
303
+ block_size)
304
+
305
+ inverse_block_table = physical_to_logical_mapping(
306
+ block_table_tensor, self.runner.cache_config.num_gpu_blocks)
307
+
308
+ # Get the original offset tensor
309
+ offset_tensor = torch.tensor(
310
+ self.runner.input_batch.num_computed_tokens_cpu[:num_reqs]).to(
311
+ self.runner.device, non_blocking=True)
312
+
313
+ out = FlexAttentionMetadata(
314
+ num_actual_tokens=num_actual_tokens,
315
+ max_query_len=max_query_len,
316
+ query_start_loc=query_start_loc,
317
+ max_seq_len=max_seq_len,
318
+ seq_lens=seq_lens,
319
+ block_table=block_table_tensor,
320
+ slot_mapping=slot_mapping,
321
+ use_cascade=use_cascade,
322
+ common_prefix_len=common_prefix_len,
323
+ cu_prefix_query_lens=cu_prefix_query_lens,
324
+ prefix_kv_lens=prefix_kv_lens,
325
+ suffix_kv_lens=suffix_kv_lens,
326
+ block_size=block_size,
327
+ max_possible_sequence_length=max_possible_seq_len,
328
+ num_reqs=num_reqs,
329
+ physical_to_logical=inverse_block_table,
330
+ total_cache_tokens=total_cache_tokens,
331
+ decode_offset=offset_tensor,
332
+ )
333
+ return out
334
+
335
+ def use_cascade_attention(self, *args, **kwargs) -> bool:
336
+ return False
337
+
338
+
339
+class FlexAttentionImpl(AttentionImpl):
340
+ sliding_window: Optional[tuple[int, int]]
341
+ alibi_slopes: Optional[torch.Tensor]
342
+ logits_soft_cap: Optional[float]
343
+
344
+ def __init__(
345
+ self,
346
+ num_heads: int,
347
+ head_size: int,
348
+ scale: float,
349
+ num_kv_heads: int,
350
+ alibi_slopes: Optional[list[float]],
351
+ sliding_window: Optional[int],
352
+ kv_cache_dtype: str,
353
+ blocksparse_params: Optional[dict[str, Any]] = None,
354
+ logits_soft_cap: Optional[float] = None,
355
+ attn_type: AttentionType = AttentionType.DECODER,
356
+ kv_sharing_target_layer_name: Optional[str] = None,
357
+ ) -> None:
358
+ if blocksparse_params is not None:
359
+ # TODO we should support this :think
360
+ raise ValueError(
361
+ "FlashAttention does not support block-sparse attention.")
362
+ self.num_heads = num_heads
363
+ self.head_size = head_size
364
+ self.scale = float(scale)
365
+ self.num_kv_heads = num_kv_heads
366
+
367
+ if alibi_slopes is not None:
368
+ raise NotImplementedError(
369
+ "FlexAttention does not support alibi slopes yet.")
370
+ else:
371
+ self.alibi_slopes = None
372
+ if sliding_window is not None:
373
+ raise NotImplementedError(
374
+ "FlexAttention does not support sliding window yet.")
375
+ else:
376
+ self.sliding_window = (-1, -1)
377
+ self.kv_cache_dtype = kv_cache_dtype
378
+ self.logits_soft_cap = logits_soft_cap
379
+ if self.logits_soft_cap is not None:
380
+ raise NotImplementedError(
381
+ "FlexAttention does not support logits soft cap yet.")
382
+
383
+ assert self.num_heads % self.num_kv_heads == 0
384
+ self.num_queries_per_kv = self.num_heads // self.num_kv_heads
385
+
386
+ if kv_sharing_target_layer_name is not None:
387
+ raise NotImplementedError(
388
+ "FlexAttention does not support kv sharing yet.")
389
+
390
+ support_head_sizes = FlexAttentionBackend.get_supported_head_sizes()
391
+ if head_size not in support_head_sizes:
392
+ raise ValueError(
393
+ f"Head size {head_size} is not supported by FlashAttention. "
394
+ f"Supported head sizes are: {support_head_sizes}. "
395
+ "Set VLLM_USE_V1=0 to use another attention backend.")
396
+ if is_quantized_kv_cache(self.kv_cache_dtype):
397
+ raise NotImplementedError(
398
+ "FlexAttention does not support quantized kv-cache. Yet")
399
+
400
+ @staticmethod
401
+ def view_as_4d(tensor: torch.Tensor) -> torch.Tensor:
402
+ """View a 3d tensor as 4D."""
403
+ if tensor.ndim == 4:
404
+ return tensor
405
+ assert tensor.ndim == 3
406
+ return tensor[None, :, :, :]
407
+
408
+ def forward(
409
+ self,
410
+ layer: torch.nn.Module,
411
+ query: torch.Tensor,
412
+ key: torch.Tensor,
413
+ value: torch.Tensor,
414
+ kv_cache: torch.Tensor,
415
+ attn_metadata: FlexAttentionMetadata,
416
+ output: Optional[torch.Tensor] = None,
417
+ ) -> torch.Tensor:
418
+ """Forward pass with FLexAttention.
419
+
420
+ Args:
421
+ query: shape = [num_tokens, num_heads, head_size]
422
+ key: shape = [num_tokens, num_kv_heads, head_size]
423
+ value: shape = [num_tokens, num_kv_heads, head_size]
424
+ kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
425
+ attn_metadata: Metadata for attention.
426
+ Returns:
427
+ shape = [num_tokens, num_heads * head_size]
428
+ """
429
+ assert output is not None, "Output tensor must be provided."
430
+ enable_gqa = self.num_kv_heads != self.num_heads
431
+
432
+ if attn_metadata is None:
433
+ # Profiling run.
434
+ return output
435
+ # query = self.view_as_4d(query).permute(0, 2, 1, 3)
436
+ # return torch.empty_like(query)
437
+
438
+ num_actual_tokens = attn_metadata.num_actual_tokens
439
+
440
+ key_cache, value_cache = kv_cache.unbind(0)
441
+
442
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
443
+ key,
444
+ value,
445
+ key_cache,
446
+ value_cache,
447
+ attn_metadata.slot_mapping,
448
+ self.kv_cache_dtype,
449
+ layer._k_scale,
450
+ layer._v_scale,
451
+ )
452
+
453
+ # View out the block_size dim
454
+ key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
455
+ value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size)
456
+ query, key_cache, value_cache = map(
457
+ lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
458
+ (query, key_cache, value_cache),
459
+ )
460
+ query = query[:, :, :num_actual_tokens, :]
461
+ # Doesn't work for now -> constraint violation
462
+ # torch._dynamo.try_mark_dynamic(query, 2)
463
+ out = flex_attention_compiled(
464
+ query,
465
+ key_cache,
466
+ value_cache,
467
+ attn_metadata.score_mod,
468
+ attn_metadata.block_mask,
469
+ self.scale,
470
+ enable_gqa=enable_gqa,
471
+ kernel_options={"FORCE_USE_FLEX_ATTENTION": True},
472
+ )
473
+
474
+ # Flex doesn't have an out variant today, rely on epilogue fusion
475
+ out = out.permute(0, 2, 1, 3).squeeze(0)
476
+ output[:num_actual_tokens, :, :].copy_(out)
477
+ return output