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, 01:00 UTC
[Core] Support Local Chunked Attention for Hybrid KV Cache (#19351)
Signed-off-by: Lucia Fang <fanglu@fb.com> Signed-off-by: Lu Fang <fanglu@meta.com> Signed-off-by: Lu Fang <fanglu@fb.com> Co-authored-by: Lu Fang <fanglu@meta.com>
Lucia Fang committed
Jul 19, 2025 at 11:48 UTC
9a9fda1423c96aa8ea62a56e8f1ad88fc080ae2c
9 files changed
+351
-19
tests/v1/core/test_specialized_manager.py
+155
-2
@@ -1,13 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import random
5
+
6
import torch
7
8
from vllm.v1.core.block_pool import BlockPool
9
from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
10
KVCacheBlock)
9
-from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
10
-from vllm.v1.kv_cache_interface import SlidingWindowSpec
11
+from vllm.v1.core.single_type_kv_cache_manager import (
12
+ ChunkedLocalAttentionManager, SlidingWindowManager)
13
+from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
14
+ SlidingWindowSpec)
15
16
17
def get_sliding_window_manager(sliding_window_spec, block_pool):
@@ -17,6 +21,80 @@ def get_sliding_window_manager(sliding_window_spec, block_pool):
21
kv_cache_group_id=0)
22
23
24
+def get_chunked_local_attention_manager(chunked_local_attention_spec,
25
+ block_pool):
26
+ return ChunkedLocalAttentionManager(chunked_local_attention_spec,
27
+ block_pool,
28
+ caching_hash_fn=lambda x: x,
29
+ kv_cache_group_id=0)
30
+
31
+
32
+def test_chunked_local_attention_possible_cached_prefix():
33
+ block_size = 2
34
+ chunked_local_attention_spec = ChunkedLocalAttentionSpec(
35
+ block_size=block_size,
36
+ num_kv_heads=1,
37
+ head_size=1,
38
+ dtype=torch.float32,
39
+ attention_chunk_size=4,
40
+ use_mla=False,
41
+ )
42
+
43
+ block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
44
+ manager = get_chunked_local_attention_manager(chunked_local_attention_spec,
45
+ block_pool)
46
+
47
+ def run_one_case(block_is_cached, tail_token, expect_length):
48
+ block_hash_list = [
49
+ BlockHash(i, ()) for i in range(len(block_is_cached))
50
+ ]
51
+
52
+ block_pool.cached_block_hash_to_block.clear()
53
+
54
+ # Mock the block pool with the cached blocks
55
+ for i, (block_hash,
56
+ is_cached) in enumerate(zip(block_hash_list, block_is_cached)):
57
+ if is_cached:
58
+ block_pool.cached_block_hash_to_block[BlockHashWithGroupId(
59
+ block_hash, 0)] = {
60
+ i: block_pool.blocks[i + 10],
61
+ }
62
+
63
+ computed_blocks = manager.find_longest_cache_hit(
64
+ block_hashes=block_hash_list,
65
+ max_length=len(block_hash_list) * block_size + tail_token,
66
+ kv_cache_group_ids=[0],
67
+ block_pool=block_pool,
68
+ kv_cache_spec=chunked_local_attention_spec,
69
+ use_eagle=False)[0]
70
+ assert len(computed_blocks) == expect_length
71
+
72
+ assert all(block == block_pool.null_block
73
+ for block in computed_blocks[:(expect_length - 1) // 2])
74
+
75
+ run_one_case([True], 0, 1)
76
+ run_one_case([True], 1, 1)
77
+ run_one_case([True, False], 0, 2)
78
+ run_one_case([True, False], 1, 2)
79
+ run_one_case([True, True], 0, 2)
80
+ run_one_case([True, True], 1, 2)
81
+ run_one_case([True, True, False], 0, 2)
82
+ run_one_case([True, True, False], 1, 2)
83
+ run_one_case([True, True, True], 0, 3)
84
+ run_one_case([True, True, True], 1, 3)
85
+ run_one_case([True, True, True, False], 0, 4)
86
+ run_one_case([True, True, True, False], 1, 4)
87
+ run_one_case([random.choice([True, False])] * 8 + [True], 1, 9)
88
+ run_one_case([random.choice([True, False])] * 8 + [False], 1, 8)
89
+ run_one_case([random.choice([True, False])] * 8 + [True, True], 1, 10)
90
+ run_one_case([random.choice([True, False])] * 8 + [True, False], 0, 10)
91
+ run_one_case([random.choice([True, False])] * 8 + [True, False], 1, 10)
92
+ run_one_case([random.choice([True, False])] * 8 + [False, True], 0, 10)
93
+ run_one_case([random.choice([True, False])] * 8 + [False, True], 1, 10)
94
+ run_one_case([random.choice([True, False])] * 8 + [False, False], 0, 10)
95
+ run_one_case([random.choice([True, False])] * 8 + [False, False], 1, 10)
96
+
97
+
98
def test_sliding_window_possible_cached_prefix():
99
block_size = 2
100
sliding_window_spec = SlidingWindowSpec(
@@ -84,6 +162,58 @@ def test_sliding_window_possible_cached_prefix():
162
], 8)
163
164
165
+def test_chunked_local_attention_remove_skipped_blocks():
166
+ attention_spec = ChunkedLocalAttentionSpec(
167
+ block_size=2,
168
+ num_kv_heads=1,
169
+ head_size=1,
170
+ dtype=torch.float32,
171
+ attention_chunk_size=4,
172
+ use_mla=False,
173
+ )
174
+
175
+ block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
176
+
177
+ manager = get_chunked_local_attention_manager(attention_spec, block_pool)
178
+
179
+ null_block_id = block_pool.null_block.block_id
180
+
181
+ def id_to_block_table(ids) -> list[KVCacheBlock]:
182
+ return [
183
+ KVCacheBlock(id_)
184
+ if id_ != null_block_id else block_pool.null_block for id_ in ids
185
+ ]
186
+
187
+ def assert_block_id(block_table: list[KVCacheBlock], ids: list[int]):
188
+ for block, id_ in zip(block_table, ids):
189
+ if id_ == null_block_id:
190
+ assert block == block_pool.null_block
191
+ else:
192
+ assert block.block_id == id_
193
+
194
+ original_block_ids = [
195
+ 1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010
196
+ ]
197
+ block_table = id_to_block_table(original_block_ids)
198
+ manager.req_to_blocks["test"] = block_table
199
+
200
+ manager.remove_skipped_blocks("test", 0)
201
+ assert_block_id(block_table, original_block_ids)
202
+
203
+ # For 4th token (0-indexed), token 0-3 is out of the local attention window.
204
+ manager.remove_skipped_blocks("test", 4)
205
+ assert_block_id(block_table, [null_block_id] * 2)
206
+
207
+ # For 6th token (0-indexed), token 4 - 6 are in local attention window,
208
+ # token 0 - 3 are out, 2 blocks can be removed.
209
+ manager.remove_skipped_blocks("test", 6)
210
+ assert_block_id(block_table, [null_block_id] * 2 + original_block_ids[2:])
211
+ # For 12th token (0-indexed),
212
+ # token 0-11 are out, 6 block can be removed.
213
+ manager.remove_skipped_blocks("test", 12)
214
+ assert_block_id(block_table, [null_block_id] * 6)
215
+
216
+
217
def test_sliding_window_remove_skipped_blocks():
218
sliding_window_spec = SlidingWindowSpec(
219
block_size=2,
@@ -172,3 +302,26 @@ def test_get_num_blocks_to_allocate():
302
cached_blocks_1) == 20
303
assert manager.get_num_blocks_to_allocate("2", 20 * block_size,
304
cached_blocks_2) == 15
305
+
306
+
307
+def test_chunked_local_attention_get_num_blocks_to_allocate():
308
+ block_size = 2
309
+ attention_spec = ChunkedLocalAttentionSpec(
310
+ block_size=block_size,
311
+ num_kv_heads=1,
312
+ head_size=1,
313
+ dtype=torch.float32,
314
+ attention_chunk_size=4, # Placeholder value, not related to test result
315
+ use_mla=False,
316
+ )
317
+
318
+ block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
319
+ manager = get_chunked_local_attention_manager(attention_spec, block_pool)
320
+ cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
321
+ cached_blocks_2 = [block_pool.null_block for _ in range(5)
322
+ ] + [KVCacheBlock(i + 1) for i in range(5)]
323
+
324
+ assert manager.get_num_blocks_to_allocate("1", 20 * block_size,
325
+ cached_blocks_1) == 20
326
+ assert manager.get_num_blocks_to_allocate("2", 20 * block_size,
327
+ cached_blocks_2) == 15
vllm/attention/layer.py
+1
@@ -172,6 +172,7 @@ class Attention(nn.Module):
172
kv_sharing_target_layer_name, **extra_impl_args)
173
self.backend = backend_name_to_enum(attn_backend.get_name())
174
self.dtype = dtype
175
+ self.use_irope = extra_impl_args.get("use_irope", False)
176
177
# For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
178
# torch.compile works by registering the attention as one giant
vllm/config.py
+7
@@ -4722,6 +4722,13 @@ class VllmConfig:
4722
if self.kv_events_config is not None:
4723
# Hybrid KV cache manager is not compatible with KV events.
4724
self.scheduler_config.disable_hybrid_kv_cache_manager = True
4725
+ if self.model_config is not None and \
4726
+ self.model_config.attention_chunk_size is not None and \
4727
+ self.speculative_config is not None and \
4728
+ self.speculative_config.use_eagle():
4729
+ # Hybrid KV cache manager is not yet supported with chunked
4730
+ # local attention + eagle.
4731
+ self.scheduler_config.disable_hybrid_kv_cache_manager = True
4732
4733
def update_sizes_for_sequence_parallelism(self,
4734
possible_sizes: list) -> list:
vllm/v1/attention/backends/flash_attn.py
+2
-1
@@ -538,6 +538,7 @@ def use_cascade_attention(
538
num_kv_heads: int,
539
use_alibi: bool,
540
use_sliding_window: bool,
541
+ use_local_attention: bool,
542
num_sms: int,
543
) -> bool:
544
"""Decide whether to use cascade attention.
@@ -553,7 +554,7 @@ def use_cascade_attention(
554
if common_prefix_len < 256:
555
return False
556
# Cascade attention is currently not supported with these variants.
556
- if use_alibi or use_sliding_window:
557
+ if use_alibi or use_sliding_window or use_local_attention:
558
return False
559
# Too few queries. Probably not worth using cascade attention.
560
# We use an arbitrary threshold of 8 queries. TODO: Tune this threshold.
vllm/v1/attention/backends/utils.py
+1
@@ -120,6 +120,7 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
120
num_kv_heads: int,
121
use_alibi: bool,
122
use_sliding_window: bool,
123
+ use_local_attention: bool,
124
num_sms: int,
125
) -> bool:
126
return False
vllm/v1/core/kv_cache_utils.py
+16
-3
@@ -11,7 +11,8 @@ from typing import Any, Callable, NamedTuple, Optional
11
from vllm.config import VllmConfig
12
from vllm.logger import init_logger
13
from vllm.utils import GiB_bytes, cdiv, sha256_cbor_64bit
14
-from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
14
+from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
15
+ FullAttentionSpec, KVCacheConfig,
16
KVCacheGroupSpec, KVCacheSpec,
17
KVCacheTensor, SlidingWindowSpec)
18
from vllm.v1.metrics.stats import PrefixCacheStats
@@ -976,7 +977,11 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
977
isinstance(spec, FullAttentionSpec) for spec in kv_cache_spec.values())
978
has_sliding_window = any(
979
isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values())
979
- if has_full_attention and has_sliding_window:
980
+ has_chunked_local_attention = any(
981
+ isinstance(spec, ChunkedLocalAttentionSpec)
982
+ for spec in kv_cache_spec.values())
983
+ if has_full_attention and (has_sliding_window
984
+ or has_chunked_local_attention):
985
for layer_name, spec in kv_cache_spec.items():
986
if isinstance(spec, SlidingWindowSpec):
987
kv_cache_spec[layer_name] = FullAttentionSpec(
@@ -987,6 +992,15 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
992
use_mla=spec.use_mla,
993
sliding_window=spec.sliding_window,
994
)
995
+ elif isinstance(spec, ChunkedLocalAttentionSpec):
996
+ kv_cache_spec[layer_name] = FullAttentionSpec(
997
+ block_size=spec.block_size,
998
+ num_kv_heads=spec.num_kv_heads,
999
+ head_size=spec.head_size,
1000
+ dtype=spec.dtype,
1001
+ use_mla=spec.use_mla,
1002
+ attention_chunk_size=spec.attention_chunk_size,
1003
+ )
1004
1005
if is_hybrid(kv_cache_spec):
1006
raise ValueError("Hybrid KV cache manager is disabled but failed to "
@@ -1010,7 +1024,6 @@ def get_kv_cache_config(
1024
The generated KVCacheConfigs
1025
"""
1026
check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
1013
-
1027
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
1028
unify_hybrid_kv_cache_specs(kv_cache_spec)
1029
vllm/v1/core/single_type_kv_cache_manager.py
+124
-1
@@ -394,6 +394,129 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
394
return 0
395
396
397
+class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
398
+
399
+ def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec,
400
+ block_pool: BlockPool, **kwargs) -> None:
401
+ super().__init__(kv_cache_spec, block_pool, **kwargs)
402
+ self.attention_chunk_size = kv_cache_spec.attention_chunk_size
403
+ self._null_block = block_pool.null_block
404
+
405
+ @classmethod
406
+ def find_longest_cache_hit(
407
+ cls,
408
+ block_hashes: list[BlockHash],
409
+ max_length: int,
410
+ kv_cache_group_ids: list[int],
411
+ block_pool: BlockPool,
412
+ kv_cache_spec: KVCacheSpec,
413
+ use_eagle: bool,
414
+ ) -> tuple[list[KVCacheBlock], ...]:
415
+ """
416
+ For chunked local attention, we need to find the longest cache hit
417
+ prefix of the blocks that is not longer than `max_length`. The prefix
418
+ should be a common prefix hit for all the kv cache groups in
419
+ `kv_cache_group_ids`. If no cache hit is found, return an empty list.
420
+ note we mark as computed if the whole block is outside of the local
421
+ window, and set the block as null. Examples:
422
+
423
+ 1. Attention chunk size of 8, block size of 4, max length of 15
424
+ for next token at 15th (zero-indexed), 8th - 14th tokens are in
425
+ the window(needs lookup), 0th - 7th are not in the window,
426
+ so they are already marked as computed. We check the complete
427
+ block3 (8th - 11th tokens), Assume block 3 is hit, we will return
428
+ [null, null, block 3], otherwise, we return [null, null]
429
+
430
+ 2. Attention chunk size of 8, block size of 4, max length of 16
431
+ for next token at 16th (zero-indexed), 0th - 15th tokens are not
432
+ in the window, so they are already marked as computed.
433
+ we return 4 blocks[null, null, null, null]
434
+
435
+ Args:
436
+ block_hashes: The block hashes of the request.
437
+ max_length: The maximum length of the cache hit prefix.
438
+ kv_cache_group_ids: The ids of the kv cache groups.
439
+ block_pool: The block pool.
440
+ kv_cache_spec: The kv cache spec.
441
+ use_eagle: Whether to use eagle.
442
+
443
+ Returns:
444
+ A list of cached blocks
445
+ """
446
+ assert isinstance(kv_cache_spec, ChunkedLocalAttentionSpec), (
447
+ "ChunkedLocalAttentionManager can only be used for " +
448
+ "chunked local attention groups")
449
+ assert use_eagle is False, ("Hybrid KV cache is not supported for " +
450
+ "eagle + chunked local attention.")
451
+ max_num_blocks = max_length // kv_cache_spec.block_size
452
+ if max_length > 0:
453
+ local_attention_start_idx = (max_length //
454
+ kv_cache_spec.attention_chunk_size *
455
+ kv_cache_spec.attention_chunk_size)
456
+ else:
457
+ local_attention_start_idx = 0
458
+ # we marked blocks out of window as computed
459
+ # with null blocks, and blocks inside window based on cache lookup
460
+ # result [null] [null] ... [null] [hit block 1 (1st block contain
461
+ # last window)] [hit block 2] ... [hit block x]
462
+ local_attention_start_block_idx = (local_attention_start_idx //
463
+ kv_cache_spec.block_size)
464
+ computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
465
+ [block_pool.null_block] * local_attention_start_block_idx
466
+ for _ in range(len(kv_cache_group_ids)))
467
+ for i in range(local_attention_start_block_idx, max_num_blocks):
468
+ block_hash = block_hashes[i]
469
+ if cached_block := block_pool.get_cached_block(
470
+ block_hash, kv_cache_group_ids):
471
+ for computed, cached in zip(computed_blocks, cached_block):
472
+ computed.append(cached)
473
+ else:
474
+ break
475
+ return computed_blocks
476
+
477
+ def remove_skipped_blocks(self, request_id: str,
478
+ num_computed_tokens: int) -> None:
479
+ # Remove the blocks that are no longer be in the chunked attention
480
+ # window and skipped during the attention computation.
481
+
482
+ # [chunk 0][chunk 1]local_attention_start_idx ... current
483
+ # we computed previous number of chunks to get the idx of
484
+ # current chunk window starting offset,
485
+ # e.g. for computed 1024 tokens, the 1024th token (0 indexed)
486
+ # is in the second chunk, there are 1 prev chunk, the start idx
487
+ # is 1024. for 1023, it will be 0.
488
+ num_cached_block = self.num_cached_block.get(request_id, 0)
489
+ local_attention_start_idx = (
490
+ num_computed_tokens
491
+ ) // self.attention_chunk_size * self.attention_chunk_size
492
+ first_useful_block_idx = local_attention_start_idx // self.block_size
493
+ if num_cached_block > 0:
494
+ # Make sure we don't delete the last cached block
495
+ first_useful_block_idx = min(first_useful_block_idx,
496
+ num_cached_block - 1)
497
+ # if block size = 128, 0 -> block 0, 1024 (= 128 * 8) ->
498
+ # block 8, 372 (= 128 * 2 + 116) -> block 2
499
+ blocks = self.req_to_blocks[request_id]
500
+ removed_blocks: list[KVCacheBlock] = []
501
+ # we need to keep the last block to get the previous hash key
502
+ for i in range(first_useful_block_idx - 1, -1, -1):
503
+ if blocks[i] == self._null_block:
504
+ # If the block is already a null block, the blocks before it
505
+ # should also have been set to null blocks by the previous calls
506
+ # to this function.
507
+ break
508
+ removed_blocks.append(blocks[i])
509
+ blocks[i] = self._null_block
510
+ self.block_pool.free_blocks(removed_blocks)
511
+
512
+ def get_num_common_prefix_blocks(self, request_id: str,
513
+ num_running_requests: int) -> int:
514
+ """
515
+ cascade attention is not supported by chunked local attention.
516
+ """
517
+ return 0
518
+
519
+
520
class MambaManager(SingleTypeKVCacheManager):
521
522
@classmethod
@@ -435,8 +558,8 @@ class MambaManager(SingleTypeKVCacheManager):
558
559
spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
560
FullAttentionSpec: FullAttentionManager,
438
- ChunkedLocalAttentionSpec: FullAttentionManager,
561
SlidingWindowSpec: SlidingWindowManager,
562
+ ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,
563
MambaSpec: MambaManager,
564
}
565
vllm/v1/kv_cache_interface.py
+37
-12
@@ -87,6 +87,7 @@ class AttentionSpec(KVCacheSpec):
87
@dataclass
88
class FullAttentionSpec(AttentionSpec):
89
sliding_window: Optional[int] = None
90
+ attention_chunk_size: Optional[int] = None
91
"""
92
When hybrid allocator is disabled and the model contains both full
93
attention layers and sliding window attention layers, sliding
@@ -105,6 +106,17 @@ class FullAttentionSpec(AttentionSpec):
106
max_model_len = vllm_config.model_config.max_model_len
107
return cdiv(max_model_len, self.block_size) * self.page_size_bytes
108
109
+ @classmethod
110
+ def merge_window_sizes(cls, window_sizes: set[int]) -> Optional[int]:
111
+ if len(window_sizes) == 0:
112
+ return None
113
+ elif len(window_sizes) == 1:
114
+ return window_sizes.pop()
115
+ else:
116
+ raise ValueError(
117
+ "All attention layers in the same KV cache group must have the "
118
+ "same window size.")
119
+
120
@classmethod
121
def merge(cls, specs: list[Self]) -> Self:
122
"""
@@ -114,14 +126,17 @@ class FullAttentionSpec(AttentionSpec):
126
merged_spec = super().merge(specs)
127
sliding_window = set(spec.sliding_window for spec in specs
128
if spec.sliding_window is not None)
117
- if len(sliding_window) == 0:
118
- merged_spec.sliding_window = None
119
- elif len(sliding_window) == 1:
120
- merged_spec.sliding_window = sliding_window.pop()
121
- else:
122
- raise ValueError(
123
- "All sliding window layers in the same KV cache group "
124
- "must have the same window size.")
129
+ attention_chunk_size = set(spec.attention_chunk_size for spec in specs
130
+ if spec.attention_chunk_size is not None)
131
+
132
+ merged_spec.sliding_window = cls.merge_window_sizes(sliding_window)
133
+ merged_spec.attention_chunk_size = (
134
+ cls.merge_window_sizes(attention_chunk_size))
135
+ assert (
136
+ (merged_spec.sliding_window is not None) +
137
+ (merged_spec.attention_chunk_size is not None) <= 1
138
+ ), ("Model with both sliding window layers and chunked local attention "
139
+ "layers is not supported.")
140
return merged_spec
141
142
@@ -129,16 +144,26 @@ class FullAttentionSpec(AttentionSpec):
144
class ChunkedLocalAttentionSpec(AttentionSpec):
145
attention_chunk_size: int
146
132
- def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
133
- max_model_len = vllm_config.model_config.max_model_len
134
- return cdiv(max_model_len, self.block_size) * self.page_size_bytes
135
-
147
@property
148
def type_id(self) -> str:
149
return (
150
f"local_attention_{self.attention_chunk_size}_{self.block_size}_{self.page_size_bytes}"
151
) # noqa
152
153
+ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
154
+ max_model_len = vllm_config.model_config.max_model_len
155
+ max_num_batched_tokens = (
156
+ vllm_config.scheduler_config.max_num_batched_tokens)
157
+
158
+ # During chunked prefill, we allocate KV cache for at most
159
+ # `self.attention_chunk_size` computed tokens plus the newly scheduled
160
+ # tokens. And we won't allocate KV cache for more than `max_model_len`
161
+ # tokens.
162
+ num_tokens = min(self.attention_chunk_size + max_num_batched_tokens,
163
+ max_model_len)
164
+
165
+ return cdiv(num_tokens, self.block_size) * self.page_size_bytes
166
+
167
168
@dataclass
169
class SlidingWindowSpec(AttentionSpec):
vllm/v1/worker/gpu_model_runner.py
+8
@@ -862,6 +862,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
862
use_sliding_window = (isinstance(kv_cache_spec, SlidingWindowSpec) or
863
(isinstance(kv_cache_spec, FullAttentionSpec)
864
and kv_cache_spec.sliding_window is not None))
865
+ use_local_attention = (
866
+ isinstance(kv_cache_spec, ChunkedLocalAttentionSpec)
867
+ or (isinstance(kv_cache_spec, FullAttentionSpec)
868
+ and kv_cache_spec.attention_chunk_size is not None))
869
assert isinstance(kv_cache_spec, AttentionSpec)
870
use_cascade = attn_metadata_builder.use_cascade_attention(
871
common_prefix_len=common_prefix_len,
@@ -870,6 +874,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
874
num_kv_heads=kv_cache_spec.num_kv_heads,
875
use_alibi=self.use_alibi,
876
use_sliding_window=use_sliding_window,
877
+ use_local_attention=use_local_attention,
878
num_sms=self.num_sms,
879
)
880
return common_prefix_len if use_cascade else 0
@@ -2672,6 +2677,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2677
dtype=self.kv_cache_dtype,
2678
sliding_window=attn_module.sliding_window,
2679
use_mla=use_mla)
2680
+ assert not use_local_attention, (
2681
+ "attention module can not be with ",
2682
+ "both local attention and sliding window")
2683
elif use_local_attention:
2684
kv_cache_spec[layer_name] = (ChunkedLocalAttentionSpec(
2685
block_size=block_size,