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 11, 06:15 UTC
[Fix][FlexAttention] return max logical block index to handle reused blocks (#30915)
Signed-off-by: Yifan Qiao <yifanqiao@berkeley.edu>
Yifan Qiao committed
Dec 17, 2025 at 22:42 UTC
11a89cf95caaec8dec13fab1e8e3d64c9a852a08
2 files changed
+42
-4
tests/kernels/test_flex_attention.py
+30
-1
@@ -15,7 +15,10 @@ from tests.v1.attention.utils import (
15
create_standard_kv_cache_spec,
16
create_vllm_config,
17
)
18
-from vllm.v1.attention.backends.flex_attention import FlexAttentionMetadataBuilder
18
+from vllm.v1.attention.backends.flex_attention import (
19
+ FlexAttentionMetadataBuilder,
20
+ physical_to_logical_mapping,
21
+)
22
23
from ..models.utils import check_embeddings_close, check_logprobs_close
24
@@ -205,5 +208,31 @@ def test_block_mask_direct_vs_slow_path():
208
)
209
210
211
+def test_physical_to_logical_mapping_handles_reused_blocks():
212
+ """Regression test: reused physical blocks map to the latest logical block.
213
+
214
+ For sliding-window / hybrid attention layers, physical KV-cache blocks can be
215
+ reused over time. The inverse mapping must therefore select the latest
216
+ logical block index for a physical block id.
217
+ """
218
+ # Padding should not make physical block 0 look live.
219
+ block_table = torch.tensor([[6, 0, 0, 0]], dtype=torch.int32)
220
+ seq_lens = torch.tensor([1 * 16], dtype=torch.int32) # only 1 block valid
221
+ out = physical_to_logical_mapping(
222
+ block_table=block_table, seq_lens=seq_lens, block_size=16, total_blocks=10
223
+ )
224
+ assert out[0, 0].item() == -1
225
+ assert out[0, 6].item() == 0
226
+
227
+ # If a physical block id appears multiple times (block reuse), mapping should
228
+ # point to the latest logical block index.
229
+ block_table2 = torch.tensor([[2, 2, 5]], dtype=torch.int32)
230
+ seq_lens2 = torch.tensor([3 * 16], dtype=torch.int32)
231
+ out2 = physical_to_logical_mapping(
232
+ block_table=block_table2, seq_lens=seq_lens2, block_size=16, total_blocks=8
233
+ )
234
+ assert out2[0, 2].item() == 1
235
+
236
+
237
if __name__ == "__main__":
238
pytest.main([__file__])
vllm/v1/attention/backends/flex_attention.py
+12
-3
@@ -160,7 +160,7 @@ def physical_to_logical_mapping(
160
└───────────────────────────────────────────┘
161
162
If multiple logical blocks map to the same physical block,
163
- this function returns the first (minimum) logical block index.
163
+ this function returns the latest (maximum) logical block index.
164
165
If a physical block is not mapped to by any logical block,
166
its value in the result will be -1.
@@ -183,6 +183,15 @@ def physical_to_logical_mapping(
183
To prevent this, we use seq_lens and block_size to mask out unused
184
entries, ensuring only valid block references are processed.
185
186
+ IMPORTANT: Reused physical blocks (sliding-window / hybrid attention)
187
+ ────────────────────────────────────────────────────────────────────
188
+ For some attention types, physical cache blocks can be reused over time.
189
+ This can cause the same physical block id to appear multiple times in a row
190
+ of `block_table` at different logical block indices. In that case, only the
191
+ latest logical block index corresponds to the current contents of that
192
+ physical block. Therefore, the inverse mapping must pick the maximum logical
193
+ block index for each physical block id.
194
+
195
Args:
196
block_table: Tensor of shape [max_reqs, max_num_blocks]
197
mapping logical blocks to physical locations. May contain
@@ -217,8 +226,8 @@ def physical_to_logical_mapping(
226
mask, torch.arange(max_num_blocks, device=device)[None, :], 0
227
)
228
220
- physical_to_logical.scatter_(
221
- -1, valid_block_table.to(torch.int64), valid_logical_indices
229
+ physical_to_logical.scatter_reduce_(
230
+ -1, valid_block_table.to(torch.int64), valid_logical_indices, reduce="amax"
231
)
232
# NB - Seems like block 0 is always empty so we reset it manually
233
physical_to_logical[:, 0] = -1