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 14, 05:30 UTC
Updates to Flex + VLLm integration (#21416)
Signed-off-by: drisspg <drisspguessous@gmail.com>
Driss Guessous committed
Aug 25, 2025 at 06:32 UTC
e0329ed4b426af432a8cc0997f964ba1e59cfdc2
3 files changed
+439
-103
tests/kernels/test_flex_attention.py
+87
-23
@@ -9,12 +9,17 @@ import pytest
9
import torch
10
from packaging import version
11
12
-from vllm import SamplingParams
12
+from tests.v1.attention.utils import (BatchSpec, create_common_attn_metadata,
13
+ create_standard_kv_cache_spec,
14
+ create_vllm_config)
15
+from vllm.v1.attention.backends.flex_attention import (
16
+ FlexAttentionMetadataBuilder)
17
14
-from ..models.utils import check_embeddings_close
18
+from ..models.utils import check_embeddings_close, check_logprobs_close
19
20
TORCH_VERSION = version.parse(torch.__version__)
21
MINIMUM_TORCH_VERSION = version.parse("2.7.0")
22
+DIRECT_BUILD_VERSION = version.parse("2.9.dev0")
23
24
25
def set_seed(seed):
@@ -34,22 +39,18 @@ def test_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
39
"""Test that FlexAttention produces the same outputs as the default backend.
40
41
This test compares the outputs from the FlexAttention backend with
37
- the default backend, ensuring they are identical when using the same seed.
42
+ the default backend, ensuring they are similar when using the same seed.
43
"""
44
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
45
seed = 42
46
max_tokens = 24
47
+ num_logprobs = 5
48
prompts = [
49
"Hello, my name is",
50
"The president of the United States is",
51
"The capital of France is",
52
]
53
48
- sampling_params = SamplingParams(temperature=0.0,
49
- top_p=1.0,
50
- seed=seed,
51
- max_tokens=max_tokens)
52
-
54
# Run with flex attention
55
with monkeypatch.context() as m:
56
m.setenv("VLLM_USE_V1", "1")
@@ -61,7 +62,8 @@ def test_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
62
tensor_parallel_size=1,
63
num_gpu_blocks_override=128,
64
enforce_eager=True) as llm_flex:
64
- output_flex = llm_flex.generate(prompts, sampling_params)
65
+ output_flex = llm_flex.generate_greedy_logprobs(
66
+ prompts, max_tokens, num_logprobs)
67
68
# Run with default backend
69
with monkeypatch.context() as m:
@@ -71,20 +73,17 @@ def test_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
73
runner="generate",
74
tensor_parallel_size=1,
75
num_gpu_blocks_override=128,
74
- enforce_eager=True) as llm_default:
75
- output_default = llm_default.generate(prompts, sampling_params)
76
-
77
- # Compare outputs from both backends
78
- for i, (flex_result,
79
- default_result) in enumerate(zip(output_flex, output_default)):
80
- prompt = prompts[i]
81
- flex_text = flex_result[1][0]
82
- default_text = default_result[1][0]
83
-
84
- assert flex_text == default_text, (
85
- f"FlexAttention output doesn't match default for: {prompt!r}\n"
86
- f"FlexAttention: {flex_text!r}\n"
87
- f"Default: {default_text!r}")
76
+ enforce_eager=True,
77
+ gpu_memory_utilization=0.85) as llm_default:
78
+ output_default = llm_default.generate_greedy_logprobs(
79
+ prompts, max_tokens, num_logprobs)
80
+
81
+ check_logprobs_close(
82
+ outputs_0_lst=output_flex,
83
+ outputs_1_lst=output_default,
84
+ name_0="flex",
85
+ name_1="default",
86
+ )
87
88
89
@pytest.mark.skipif(
@@ -136,5 +135,70 @@ def test_encoder_flex_attention_vs_default_backend(vllm_runner, monkeypatch):
135
)
136
137
138
+@pytest.mark.skipif(
139
+ not torch.cuda.is_available() or TORCH_VERSION < DIRECT_BUILD_VERSION,
140
+ reason="CUDA not available or PyTorch version < 2.7",
141
+)
142
+def test_block_mask_direct_vs_slow_path():
143
+ """Test that direct path block mask is a superset of slow path.
144
+
145
+ The direct path may include extra blocks for performance (over-estimation),
146
+ but must include all blocks that the slow path determines are necessary.
147
+ """
148
+ device = torch.device("cuda")
149
+
150
+ vllm_config = create_vllm_config(model_name="meta-llama/Meta-Llama-3-8B",
151
+ block_size=16,
152
+ max_model_len=1024)
153
+ kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
154
+
155
+ # Use a mixed batch that will create groups spanning multiple sequences
156
+ batch_spec = BatchSpec(seq_lens=[35, 64, 128, 256],
157
+ query_lens=[33, 5, 32, 64],
158
+ name="test_mixed_batch")
159
+
160
+ common_attn_metadata = create_common_attn_metadata(
161
+ batch_spec, vllm_config.cache_config.block_size, device)
162
+
163
+ builder = FlexAttentionMetadataBuilder(kv_cache_spec, [], vllm_config,
164
+ device)
165
+
166
+ metadata_direct = builder.build(common_prefix_len=0,
167
+ common_attn_metadata=common_attn_metadata)
168
+ builder.direct_build = False
169
+ metadata_slow = builder.build(common_prefix_len=0,
170
+ common_attn_metadata=common_attn_metadata)
171
+
172
+ assert metadata_direct.block_mask is not None
173
+ assert metadata_slow.block_mask is not None
174
+
175
+ # Extract block indices for comparison, B, H are the same
176
+ direct_indices = metadata_direct.block_mask.kv_indices[0, 0]
177
+ slow_indices = metadata_slow.block_mask.kv_indices[0, 0]
178
+ direct_num = metadata_direct.block_mask.kv_num_blocks[0, 0]
179
+ slow_num = metadata_slow.block_mask.kv_num_blocks[0, 0]
180
+
181
+ # main test: every block needed by slow path must be in direct path
182
+ num_groups = direct_num.shape[0]
183
+ all_contained = True
184
+ missing_details = []
185
+
186
+ for group_idx in range(num_groups):
187
+ direct_blocks = set(
188
+ direct_indices[group_idx, :direct_num[group_idx]].tolist())
189
+ slow_blocks = set(
190
+ slow_indices[group_idx, :slow_num[group_idx]].tolist())
191
+
192
+ missing_blocks = slow_blocks - direct_blocks
193
+ if missing_blocks:
194
+ all_contained = False
195
+ missing_details.append(
196
+ f"Group {group_idx}: missing {sorted(missing_blocks)}")
197
+
198
+ assert all_contained, (
199
+ "Direct path is missing blocks required by slow path:\n" +
200
+ "\n".join(missing_details))
201
+
202
+
203
if __name__ == "__main__":
204
pytest.main([__file__])
tests/v1/attention/test_attention_backends.py
+18
-12
@@ -10,14 +10,15 @@ from tests.v1.attention.utils import (BatchSpec, _Backend,
10
create_standard_kv_cache_spec,
11
create_vllm_config,
12
get_attention_backend)
13
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
13
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv, is_torch_equal_or_newer
14
from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
15
set_kv_cache_layout)
16
from vllm.v1.kv_cache_interface import FullAttentionSpec
17
18
BACKENDS_TO_TEST = [
19
_Backend.FLASH_ATTN_VLLM_V1, _Backend.FLASHINFER_VLLM_V1,
20
- _Backend.FLEX_ATTENTION, _Backend.TRITON_ATTN_VLLM_V1, _Backend.TREE_ATTN
20
+ _Backend.FLEX_ATTENTION, _Backend.TRITON_ATTN_VLLM_V1, _Backend.TREE_ATTN,
21
+ "FLEX_ATTENTION_SLOW"
22
]
23
24
# Remove flashinfer from the list if it's not available
@@ -97,7 +98,7 @@ def create_and_prepopulate_kv_cache(
98
common_attn_metadata: CommonAttentionMetadata,
99
randomize_blocks: bool = True) -> torch.Tensor:
100
"""Create and prepopulate a KV cache with context data.
100
-
101
+
102
Args:
103
k_contexts: List of key context tensors for each sequence
104
v_contexts: List of value context tensors for each sequence
@@ -109,9 +110,9 @@ def create_and_prepopulate_kv_cache(
110
device: Device to create the cache on
111
num_blocks: Total number of blocks in the cache
112
block_table: Block table tensor to populate
112
- randomize_blocks: Whether to randomly permute blocks
113
+ randomize_blocks: Whether to randomly permute blocks
114
or use sequential order
114
-
115
+
116
Returns:
117
Tuple of (kv_cache, updated_block_table)
118
"""
@@ -206,10 +207,18 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,
207
kv_cache: torch.Tensor) -> torch.Tensor:
208
"""Run attention computation using the specified backend's AttentionImpl."""
209
209
- builder_cls, impl_cls = get_attention_backend(backend)
210
+ # Handle special case for FLEX_ATTENTION_SLOW
211
+ actual_backend = backend
212
+
213
+ use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
214
+ if backend == "FLEX_ATTENTION_SLOW":
215
+ actual_backend = _Backend.FLEX_ATTENTION
216
+ use_direct_block_mask = False
217
+
218
+ builder_cls, impl_cls = get_attention_backend(actual_backend)
219
220
# Mock flashinfer's get_per_layer_parameters if needed
212
- if backend == _Backend.FLASHINFER_VLLM_V1:
221
+ if actual_backend == _Backend.FLASHINFER_VLLM_V1:
222
import unittest.mock
223
224
from vllm.v1.attention.backends.utils import PerLayerParameters
@@ -239,6 +248,8 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,
248
else:
249
# Build metadata
250
builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
251
+ if actual_backend == _Backend.FLEX_ATTENTION:
252
+ builder.direct_build = use_direct_block_mask
253
attn_metadata = builder.build(
254
common_prefix_len=0,
255
common_attn_metadata=common_attn_metadata,
@@ -453,11 +464,6 @@ def test_backend_correctness(batch_spec_name: str, model: str):
464
rtol = 1e-2
465
atol = 5e-3
466
456
- if backend_name == _Backend.FLEX_ATTENTION:
457
- atol = 5e-1 # TODO: figure out why flex_attention has such large
458
- # numerical differences for medium_decode, medium_prefill,
459
- # mixed_medium
460
-
467
max_diff = torch.max(torch.abs(backend_output - sdpa_output)).item()
468
max_rel_diff = torch.max(
469
torch.abs(backend_output - sdpa_output) /
vllm/v1/attention/backends/flex_attention.py
+334
-68
@@ -1,11 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Attention layer with FlashAttention."""
4
-from collections import defaultdict
3
+"""Attention layer with FlexAttention."""
4
+
5
from dataclasses import dataclass
6
-from typing import Optional
6
+from typing import TYPE_CHECKING, Optional, Union
7
8
import torch
9
+import torch._dynamo.decorators
10
+import torch.nn.functional as F
11
from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
12
_score_mod_signature,
13
create_block_mask,
@@ -16,13 +18,17 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
18
is_quantized_kv_cache)
19
from vllm.config import VllmConfig
20
from vllm.logger import init_logger
19
-from vllm.platforms import current_platform
21
+from vllm.utils import cdiv, is_torch_equal_or_newer
22
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
23
CommonAttentionMetadata)
24
from vllm.v1.kv_cache_interface import AttentionSpec
25
26
logger = init_logger(__name__)
27
28
+if TYPE_CHECKING:
29
+ from vllm.v1.core.sched.output import SchedulerOutput
30
+ from vllm.v1.worker.gpu_input_batch import InputBatch
31
+
32
create_block_mask_compiled = torch.compile(create_block_mask,
33
fullgraph=True,
34
mode="reduce-overhead")
@@ -36,6 +42,23 @@ def _offsets_to_doc_ids_tensor(offsets: torch.Tensor) -> torch.Tensor:
42
torch.arange(len(counts), device=device, dtype=torch.int32), counts)
43
44
45
+def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int):
46
+ difference = (multiple - (x.shape[dim] % multiple)) % multiple
47
+ if difference == 0:
48
+ return x
49
+
50
+ dim = dim if dim >= 0 else x.ndim + dim
51
+ pad_list = []
52
+
53
+ for i in range(x.ndim - 1, dim - 1, -1):
54
+ if i == dim:
55
+ pad_list.extend([0, difference])
56
+ else:
57
+ pad_list.extend([0, 0])
58
+
59
+ return F.pad(x, pad_list, mode="constant", value=0)
60
+
61
+
62
class FlexAttentionBackend(AttentionBackend):
63
accept_output_buffer: bool = True
64
@@ -77,10 +100,10 @@ class FlexAttentionBackend(AttentionBackend):
100
return False
101
102
80
-# @torch.compile(fullgraph=True, mode="reduce-overhead")
81
-def physical_to_logical_mapping(
82
- block_table: torch.Tensor,
83
- total_blocks: Optional[int] = None) -> torch.Tensor:
103
+#@torch.compile(fullgraph=True, mode="reduce-overhead")
104
+def physical_to_logical_mapping(block_table: torch.Tensor,
105
+ seq_lens: torch.Tensor, block_size: int,
106
+ total_blocks: int) -> torch.Tensor:
107
"""
108
Creates an inverse mapping from physical block locations to logical indices.
109
@@ -114,13 +137,38 @@ def physical_to_logical_mapping(
137
If a physical block is not mapped to by any logical block,
138
its value in the result will be -1.
139
140
+ IMPORTANT: Garbage Value Protection
141
+ ────────────────────────────────────
142
+ The block_table tensor may contain garbage values in unused positions
143
+ (beyond the actual sequence length). For example, if a sequence only
144
+ needs 3 blocks but the table has space for 8:
145
+
146
+ block_table[0] = [10, 25, 7, 999, 1234, 888, ...]
147
+ ^^^^^^^^^^^^^^^^^^^^
148
+ garbage values
149
+
150
+ These garbage values can cause issues because:
151
+ 1. They may map to valid physical blocks by coincidence
152
+ 2. The scatter_ operation will assign them logical indices
153
+ 3. Later attention computations may incorrectly access these blocks
154
+
155
+ To prevent this, we use seq_lens and block_size to mask out unused
156
+ entries, ensuring only valid block references are processed.
157
158
Args:
159
block_table: Tensor of shape [max_reqs, max_num_blocks]
120
- mapping logical blocks to physical locations
160
+ mapping logical blocks to physical locations. May contain
161
+ garbage values in unused positions.
162
+ seq_lens: Tensor of sequence lengths for each request. Used to
163
+ determine how many blocks are actually needed per sequence.
164
+ block_size: Size of each block in tokens. Used with seq_lens to
165
+ compute the number of valid blocks per sequence.
166
+ total_blocks: Total number of physical blocks available
167
168
Returns:
123
- A tensor of shape [max_reqs, max_physical_block]
169
+ A tensor of shape [max_reqs, total_blocks] where each entry
170
+ physical_to_logical[req_id, physical_block] contains the logical
171
+ block index for that physical block, or -1 if unused.
172
"""
173
max_reqs, max_num_blocks = block_table.shape
174
device = block_table.device
@@ -130,17 +178,76 @@ def physical_to_logical_mapping(
178
dtype=torch.long,
179
device=device)
180
133
- logical_indices = (torch.arange(max_num_blocks,
134
- device=device).unsqueeze(0).expand(
135
- max_reqs, -1))
181
+ # Only process valid blocks to avoid garbage values
182
+ num_blocks_per_seq = cdiv(seq_lens, block_size)
183
+ mask = torch.arange(max_num_blocks,
184
+ device=device)[None, :] < num_blocks_per_seq[:, None]
185
137
- physical_to_logical.scatter_(-1, block_table.to(torch.int64),
138
- logical_indices)
139
- # TODO Confirm - Seems like block 0 is always empty so we reset it manually
186
+ valid_block_table = torch.where(mask, block_table, 0)
187
+ valid_logical_indices = torch.where(
188
+ mask,
189
+ torch.arange(max_num_blocks, device=device)[None, :], 0)
190
+
191
+ physical_to_logical.scatter_(-1, valid_block_table.to(torch.int64),
192
+ valid_logical_indices)
193
+ # NB - Seems like block 0 is always empty so we reset it manually
194
physical_to_logical[:, 0] = -1
195
return physical_to_logical
196
197
198
+def unique_static_unsorted(
199
+ x: torch.Tensor,
200
+ *,
201
+ M: int, # maximum positive value (0 is “skip me”)
202
+ dim: int = -1, # axis along which to deduplicate
203
+ ignored_val: int = 0, # value to ignore
204
+ pad_val: int = -1, # sentinel for unused slots
205
+) -> torch.Tensor:
206
+ """
207
+ - Keeps the first occurrence of each non-zero value while preserving order,
208
+ then left-packs those uniques and fills the rest with `pad_val`.
209
+ - Returns (packed, keep_mask) with the *same shape* as `x`.
210
+ - Requires that all values be in the range [0, M]
211
+ - Skips ignored_val
212
+
213
+ Works on CPU or GPU, no Python loops, O(B·N) time / O(B·M) memory.
214
+
215
+ Example:
216
+ x =[3, 1, 0, 1, 2], M=3, ignored_val=0 => [3, 1, 2, -1, -1]
217
+ """
218
+ if not (-1 <= pad_val <= M):
219
+ raise ValueError("`pad_val` must lie in [-1, M]")
220
+
221
+ # ── move `dim` to the end so we can treat tensor as [B, N] ──────────
222
+ dim = dim % x.ndim
223
+ x_perm = x.movedim(dim, -1) # shape [..., N]
224
+ B, N = x_perm.numel() // x_perm.shape[-1], x_perm.shape[-1]
225
+ x_flat = x_perm.reshape(B, N) # [B, N]
226
+
227
+ device = x.device
228
+ idx = torch.arange(N, device=device).expand(B, N) # per-row indices
229
+
230
+ # ── build first-occurrence table for every v ∈ [0, M] ───────────────
231
+ first_idx = torch.full((B, M + 1), N, device=device) # “∞”
232
+ # scatter_reduce_: first_idx[b, v] = min(first_idx[b, v], i) for each i
233
+ first_idx.scatter_reduce_(1, x_flat, idx, reduce="amin")
234
+
235
+ # ── keep mask: first occurrence *and* value ≠ 0 ─────────────────────
236
+ keep = (x_flat != ignored_val) & (idx == first_idx.gather(1, x_flat)
237
+ ) # [B, N]
238
+
239
+ # ── left-pack uniques into a fresh tensor ───────────────────────────
240
+ dest_pos = torch.cumsum(keep.to(torch.long), dim=1) - 1 # where to go
241
+ packed_flat = torch.full_like(x_flat, pad_val)
242
+
243
+ rows, src_cols = torch.nonzero(keep, as_tuple=True)
244
+ packed_flat[rows, dest_pos[rows, src_cols]] = x_flat[rows, src_cols]
245
+
246
+ # ── restore original layout ─────────────────────────────────────────
247
+ packed = packed_flat.reshape(x_perm.shape).movedim(-1, dim)
248
+ return packed
249
+
250
+
251
def causal_mask_mod(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor,
252
kv_idx: torch.Tensor):
253
return q_idx >= kv_idx
@@ -170,6 +277,7 @@ class FlexAttentionMetadata:
277
num_reqs: int
278
physical_to_logical: torch.Tensor
279
decode_offset: torch.Tensor
280
+ num_blocks_per_seq: torch.Tensor
281
282
# For logging.
283
num_input_tokens: int = 0 # Number of tokens including padding.
@@ -179,6 +287,46 @@ class FlexAttentionMetadata:
287
block_mask: Optional[BlockMask] = None
288
score_mod: Optional[_score_mod_signature] = None
289
logical_mask_mod: _mask_mod_signature = causal_mask_mod
290
+ doc_ids: Optional[torch.Tensor] = None
291
+ direct_build: bool = True
292
+ q_block_size: int = 16
293
+ kv_block_size: int = 16
294
+ transformed_score_mod: Optional[_score_mod_signature] = None
295
+
296
+ def _convert_physical_to_logical(
297
+ self,
298
+ request_lookup: torch.Tensor,
299
+ q_idx: torch.Tensor,
300
+ physical_kv_idx: torch.Tensor,
301
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
302
+ """Convert physical indices to logical indices for both query and kv.
303
+
304
+ NB is_within_lower_bound: do sequences start on block_boundaries?
305
+
306
+ Returns:
307
+ tuple of (is_valid, logical_q_idx, logical_kv_idx)
308
+ """
309
+ # Map query indices to corresponding request indices
310
+ q_req = request_lookup[q_idx]
311
+
312
+ # Convert physical KV indices to logical indices
313
+ physical_kv_block = physical_kv_idx // self.block_size
314
+ physical_kv_offset = physical_kv_idx % self.block_size
315
+ logical_block_idx = self.physical_to_logical[q_req, physical_kv_block]
316
+ logical_kv_idx = (logical_block_idx * self.block_size +
317
+ physical_kv_offset)
318
+
319
+ # Determine valid kv indices
320
+ live_block = logical_block_idx >= 0
321
+ within_upper_bound = logical_kv_idx < self.seq_lens[q_req]
322
+ within_lower_bound = logical_kv_idx >= 0
323
+ is_valid = live_block & within_upper_bound & within_lower_bound
324
+
325
+ # Convert physical query indices to logical indices
326
+ local_q_idx = q_idx - self.query_start_loc[q_req]
327
+ logical_q_idx = local_q_idx + self.decode_offset[q_req]
328
+
329
+ return is_valid, logical_q_idx, logical_kv_idx
330
331
def get_causal_mask_mod(self) -> _mask_mod_signature:
332
"""Creates the mask_mod function for FlexAttention.
@@ -191,11 +339,8 @@ class FlexAttentionMetadata:
339
With this info we create the "logical" indices that are passed to
340
mask_mod functions. This allows mask mod functions to be agnostic to
341
layout of the query and key/value tensors.
194
-
195
- TODO is_within_lower_bound: do sequences start on block_boundaries?
342
"""
197
- # Create a lookup mapping from query indices -> request number
198
- request_lookup = _offsets_to_doc_ids_tensor(self.query_start_loc)
343
+ assert self.doc_ids is not None
344
345
def final_mask_mod(
346
b: torch.Tensor,
@@ -203,27 +348,9 @@ class FlexAttentionMetadata:
348
q_idx: torch.Tensor,
349
physical_kv_idx: torch.Tensor,
350
) -> torch.Tensor:
206
- # Map query indices to corresponding request indices
207
- q_req = request_lookup[q_idx]
208
-
209
- # Convert physical KV indices to logical indices
210
- physical_kv_block = physical_kv_idx // self.block_size
211
- physical_kv_offset = physical_kv_idx % self.block_size
212
- logical_block_idx = self.physical_to_logical[q_req,
213
- physical_kv_block]
214
- logical_kv_idx = logical_block_idx * self.block_size + physical_kv_offset # noqa: E501
215
-
216
- # Determine valid kv indices
217
- live_block = logical_block_idx >= 0
218
- within_upper_bound = logical_kv_idx < self.seq_lens[q_req]
219
- within_lower_bound = logical_kv_idx >= 0
220
-
221
- is_valid = live_block & within_upper_bound & within_lower_bound
222
-
223
- # Convert physical query indices to logical indices
224
- local_q_idx = q_idx - self.query_start_loc[q_req]
225
- logical_q_idx = local_q_idx + self.decode_offset[q_req]
226
-
351
+ (is_valid, logical_q_idx,
352
+ logical_kv_idx) = self._convert_physical_to_logical(
353
+ self.doc_ids, q_idx, physical_kv_idx)
354
# Apply mask modification only for valid indices
355
return torch.where(
356
is_valid,
@@ -236,7 +363,7 @@ class FlexAttentionMetadata:
363
def get_bidirectional_mask_mod(self) -> _mask_mod_signature:
364
"""Creates the encoder mask_mod function for FlexAttention.
365
239
- Since the encoder bidirectional attention doesn't run with
366
+ Since the encoder bidirectional attention doesn't run with
367
KV cache, this function creates a mask based on the
368
packed query sequences.
369
"""
@@ -253,6 +380,97 @@ class FlexAttentionMetadata:
380
381
return final_mask_mod
382
383
+ def get_transformed_score_mod(self) -> Optional[_score_mod_signature]:
384
+ """Creates the transformed score_mod function for FlexAttention.
385
+
386
+ This function wraps the user's score_mod to handle physical-to-logical
387
+ index conversion, similar to how get_mask_mod works for mask functions.
388
+ """
389
+ if self.score_mod is None:
390
+ return None
391
+
392
+ # Create a lookup mapping from query indices -> request number
393
+ request_lookup = _offsets_to_doc_ids_tensor(self.query_start_loc)
394
+ user_score_mod = self.score_mod
395
+
396
+ def transformed_score_mod(
397
+ score: torch.Tensor,
398
+ b: torch.Tensor,
399
+ h: torch.Tensor,
400
+ q_idx: torch.Tensor,
401
+ physical_kv_idx: torch.Tensor,
402
+ ) -> torch.Tensor:
403
+ (is_valid, logical_q_idx,
404
+ logical_kv_idx) = self._convert_physical_to_logical(
405
+ request_lookup, q_idx, physical_kv_idx)
406
+
407
+ return torch.where(
408
+ is_valid,
409
+ user_score_mod(score,
410
+ b,
411
+ h,
412
+ logical_q_idx,
413
+ logical_kv_idx,
414
+ physical_q=q_idx), -float('inf'))
415
+
416
+ return transformed_score_mod
417
+
418
+ def _build_block_mask_direct(self) -> BlockMask:
419
+ """Direct block mask construction for standard causal attention.
420
+
421
+ This method constructs the block mask directly using
422
+ BlockMask.from_kv_blocks which is much more efficient than the
423
+ generic create_block_mask approach.
424
+
425
+ The direct path works as follows:
426
+ 1. For each query token, fetch blocks from block_table using max_seq_len
427
+ (this fetches more blocks than needed for shorter sequences)
428
+ 2. Group query tokens into chunks of q_block_size
429
+ 3. For each group, deduplicate the blocks using unique_static_unsorted
430
+ 4. Create BlockMask using the deduplicated block indices
431
+
432
+ Over-estimation occurs when a group of q_block_size tokens contains
433
+ multiple sequence IDs (doc_ids). In this case, we fetch ALL blocks for
434
+ each sequence represented in the group, even though individual query
435
+ tokens may only need a subset of those blocks based on causal masking
436
+ and their position.
437
+
438
+ """
439
+ page_to_block_ratio = self.kv_block_size // self.block_size
440
+ if page_to_block_ratio != 1:
441
+ raise ValueError(
442
+ f"FlexAttention currently requires the cache block size "
443
+ f"({self.block_size}) to be equal to the kv_block_size "
444
+ f"({self.kv_block_size}). Please check your model's "
445
+ f"configuration.")
446
+
447
+ used_pages = self.block_table[
448
+ self.doc_ids, :cdiv(self.max_seq_len, self.block_size)]
449
+ used_pages_padded = pad_to_multiple(used_pages,
450
+ multiple=self.q_block_size,
451
+ dim=0)
452
+ used_pages_padded = used_pages_padded.reshape(
453
+ used_pages_padded.shape[0] // self.q_block_size, -1)
454
+ used_pages_padded = used_pages_padded // page_to_block_ratio
455
+ kv_indices = unique_static_unsorted((used_pages_padded.long()),
456
+ M=self.num_blocks).to(torch.int32)
457
+
458
+ kv_num_blocks = (kv_indices >= 0).sum(dim=-1).to(torch.int32)
459
+ block_mask_kwargs = {
460
+ "seq_lengths": (self.num_actual_tokens, self.total_cache_tokens),
461
+ "kv_num_blocks": kv_num_blocks[None, None],
462
+ "kv_indices": kv_indices[None, None],
463
+ "full_kv_num_blocks": None,
464
+ "full_kv_indices": None,
465
+ "BLOCK_SIZE": (self.q_block_size, self.kv_block_size),
466
+ "mask_mod": self.mask_mod,
467
+ }
468
+
469
+ # compute_q_blocks parameter is available in PyTorch 2.9+
470
+ if is_torch_equal_or_newer("2.9.0.dev0"):
471
+ block_mask_kwargs["compute_q_blocks"] = False
472
+ return BlockMask.from_kv_blocks(**block_mask_kwargs)
473
+
474
def build_block_mask(self) -> BlockMask:
475
if self.causal:
476
mask_mod = self.get_causal_mask_mod()
@@ -267,6 +485,7 @@ class FlexAttentionMetadata:
485
self.num_actual_tokens,
486
kv_len,
487
device=self.block_table.device,
488
+ BLOCK_SIZE=(self.q_block_size, self.kv_block_size),
489
)
490
491
def __post_init__(self):
@@ -275,8 +494,21 @@ class FlexAttentionMetadata:
494
assert self.cu_prefix_query_lens is None, "Not implemented yet."
495
assert self.prefix_kv_lens is None, "Not implemented yet."
496
assert self.suffix_kv_lens is None, "Not implemented yet."
497
+ # Create a lookup mapping from query indices -> request number
498
+ self.doc_ids = _offsets_to_doc_ids_tensor(self.query_start_loc)
499
self.num_blocks = self.total_cache_tokens // self.block_size
279
- self.block_mask = self.build_block_mask()
500
+
501
+ if self.causal:
502
+ self.mask_mod = self.get_causal_mask_mod()
503
+ else:
504
+ self.mask_mod = self.get_bidirectional_mask_mod()
505
+
506
+ self.transformed_score_mod = self.get_transformed_score_mod()
507
+
508
+ if self.direct_build and self.causal:
509
+ self.block_mask = self._build_block_mask_direct()
510
+ else:
511
+ self.block_mask = self.build_block_mask()
512
513
514
class FlexAttentionMetadataBuilder(
@@ -287,15 +519,24 @@ class FlexAttentionMetadataBuilder(
519
self.model_config = vllm_config.model_config
520
self.parallel_config = vllm_config.parallel_config
521
self.cache_config = vllm_config.cache_config
522
+ self.device = device
523
524
self.num_heads_q = self.model_config.get_num_attention_heads(
292
- vllm_config.parallel_config)
525
+ self.parallel_config)
526
self.num_heads_kv = self.model_config.get_num_kv_heads(
294
- vllm_config.parallel_config)
527
+ self.parallel_config)
528
self.headdim = self.model_config.get_head_size()
529
self.block_size = kv_cache_spec.block_size
530
self.kv_cache_spec = kv_cache_spec
298
- self.device = device
531
+ self.direct_build: bool = is_torch_equal_or_newer("2.9.0.dev0")
532
+ self.q_block_size: int = 16 if is_torch_equal_or_newer(
533
+ "2.9.0.dev0") else 128
534
+ self.kv_block_size: int = 16 if is_torch_equal_or_newer(
535
+ "2.9.0.dev0") else 128
536
+
537
+ def reorder_batch(self, input_batch: "InputBatch",
538
+ scheduler_output: "SchedulerOutput") -> bool:
539
+ return False
540
541
def build(self,
542
common_prefix_len: int,
@@ -310,6 +551,7 @@ class FlexAttentionMetadataBuilder(
551
seq_lens = common_attn_metadata.seq_lens
552
block_table_tensor = common_attn_metadata.block_table_tensor
553
slot_mapping = common_attn_metadata.slot_mapping
554
+ num_blocks_per_seq = cdiv(seq_lens, self.block_size)
555
556
use_cascade = common_prefix_len > 0
557
cu_prefix_query_lens = None
@@ -320,12 +562,15 @@ class FlexAttentionMetadataBuilder(
562
563
block_size = self.kv_cache_spec.block_size
564
max_possible_seq_len = self.model_config.max_model_len
323
- total_cache_tokens = self.cache_config.num_gpu_blocks * block_size
565
+ num_gpu_blocks = self.cache_config.num_gpu_blocks
566
+
567
+ assert num_gpu_blocks is not None, \
568
+ "FlexAttention requires num_gpu_blocks to be set"
569
+ total_cache_tokens = (num_gpu_blocks * block_size)
570
571
inverse_block_table = physical_to_logical_mapping(
326
- block_table_tensor, self.cache_config.num_gpu_blocks)
572
+ block_table_tensor, seq_lens, block_size, num_gpu_blocks)
573
328
- # Get the original offset tensor
574
offset_tensor = common_attn_metadata.num_computed_tokens_cpu.to(
575
self.device, non_blocking=True)
576
@@ -349,9 +594,16 @@ class FlexAttentionMetadataBuilder(
594
physical_to_logical=inverse_block_table,
595
total_cache_tokens=total_cache_tokens,
596
decode_offset=offset_tensor,
597
+ num_blocks_per_seq=num_blocks_per_seq,
598
+ direct_build=self.direct_build,
599
+ q_block_size=self.q_block_size,
600
+ kv_block_size=self.kv_block_size,
601
)
602
return out
603
604
+ def use_cascade_attention(self, *args, **kwargs) -> bool:
605
+ return False
606
+
607
608
class FlexAttentionImpl(AttentionImpl):
609
sliding_window: Optional[tuple[int, int]]
@@ -370,6 +622,7 @@ class FlexAttentionImpl(AttentionImpl):
622
logits_soft_cap: Optional[float] = None,
623
attn_type: AttentionType = AttentionType.DECODER,
624
kv_sharing_target_layer_name: Optional[str] = None,
625
+ **kwargs,
626
) -> None:
627
self.num_heads = num_heads
628
self.head_size = head_size
@@ -398,6 +651,7 @@ class FlexAttentionImpl(AttentionImpl):
651
raise NotImplementedError(
652
"FlexAttention does not support logits soft cap yet.")
653
654
+ assert self.num_heads % self.num_kv_heads == 0
655
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
656
657
if kv_sharing_target_layer_name is not None:
@@ -405,7 +659,6 @@ class FlexAttentionImpl(AttentionImpl):
659
"FlexAttention does not support kv sharing yet.")
660
661
FlexAttentionBackend.validate_head_size(head_size)
408
-
662
if is_quantized_kv_cache(self.kv_cache_dtype):
663
raise NotImplementedError(
664
"FlexAttention does not support quantized kv-cache. Yet")
@@ -493,35 +746,48 @@ class FlexAttentionImpl(AttentionImpl):
746
# Doesn't work for now -> constraint violation
747
# torch._dynamo.try_mark_dynamic(query, 2)
748
496
- # default M=64, N=64 may run out of shared memory on some GPUs
497
- # TODO: Explicit configs for each GPU?
498
- # Not sure how to calculate the shared memory requirement
499
- extra_kernel_options = defaultdict[str, int](lambda: 64)
500
- if query.dtype == torch.float32:
501
- extra_kernel_options["BLOCK_M"] //= 2
502
- extra_kernel_options["BLOCK_N"] //= 2
503
- if current_platform.is_cuda():
504
- device_props = torch.cuda.get_device_properties()
505
- max_shared_memory = device_props.shared_memory_per_block_optin
506
- if max_shared_memory < 144 * 1024:
507
- extra_kernel_options["BLOCK_M"] //= 2
508
- extra_kernel_options["BLOCK_N"] //= 2
749
+ assert attn_metadata.block_mask is not None
750
+ block_m, block_n = attn_metadata.block_mask.BLOCK_SIZE
751
752
+ kernel_options = get_kernel_options(query, block_m, block_n,
753
+ attn_metadata.direct_build)
754
out = flex_attention_compiled(
755
query,
756
key_tensor,
757
value_tensor,
514
- attn_metadata.score_mod,
758
+ attn_metadata.transformed_score_mod,
759
attn_metadata.block_mask,
760
self.scale,
761
enable_gqa=enable_gqa,
518
- kernel_options={
519
- "FORCE_USE_FLEX_ATTENTION": True,
520
- **extra_kernel_options
521
- },
762
+ kernel_options=kernel_options,
763
)
764
765
# Flex doesn't have an out variant today, rely on epilogue fusion
766
out = out.permute(0, 2, 1, 3).squeeze(0)
767
output[:num_actual_tokens, :, :].copy_(out)
768
return output
769
+
770
+
771
+def get_kernel_options(query, block_m, block_n,
772
+ use_direct_build: bool) -> dict[str, Union[int, bool]]:
773
+ kernel_options: dict[str, Union[int, bool]] = {
774
+ "FORCE_USE_FLEX_ATTENTION": True,
775
+ }
776
+ if use_direct_build:
777
+ kernel_options["BLOCK_M"] = block_m
778
+ kernel_options["BLOCK_N"] = block_n
779
+ return kernel_options
780
+ else:
781
+ kernel_options["BLOCK_M"] = 64
782
+ kernel_options["BLOCK_N"] = 64
783
+ if query.dtype == torch.float32:
784
+ kernel_options["BLOCK_M"] = 32
785
+ kernel_options["BLOCK_N"] = 32
786
+ # if current_platform.is_cuda():
787
+ if torch.cuda.is_available():
788
+ device_props = torch.cuda.get_device_properties()
789
+ max_shared_memory = device_props.shared_memory_per_block_optin
790
+ if max_shared_memory < 144 * 1024:
791
+ kernel_options["BLOCK_M"] = 32
792
+ kernel_options["BLOCK_N"] = 32
793
+ return kernel_options