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
[Attention] Make local attention backend agnostic (#21093)
Lucas Wilkinson committed
Jul 18, 2025 at 00:10 UTC
89cab4d01f83f8def180e723cee30c7ef8c53e86
8 files changed
+94
-242
vllm/v1/attention/backends/flash_attn.py
+10
-74
@@ -25,9 +25,9 @@ if is_flash_attn_varlen_func_available():
25
from vllm.config import VllmConfig, get_layers_from_vllm_config
26
from vllm.logger import init_logger
27
from vllm.utils import cdiv
28
-from vllm.v1.attention.backends.utils import (
29
- AttentionMetadataBuilder, CommonAttentionMetadata, get_kv_cache_layout,
30
- make_local_attention_virtual_batches)
28
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
29
+ CommonAttentionMetadata,
30
+ get_kv_cache_layout)
31
from vllm.v1.kv_cache_interface import AttentionSpec
32
33
logger = init_logger(__name__)
@@ -130,18 +130,6 @@ class FlashAttentionMetadata:
130
prefix_scheduler_metadata: Optional[torch.Tensor] = None
131
max_num_splits: int = 0
132
133
- # for local attention
134
- @dataclass
135
- class LocalAttentionMetadata:
136
- local_query_start_loc: torch.Tensor
137
- local_seqused_k: torch.Tensor
138
- local_block_table: torch.Tensor
139
- local_max_query_len: int
140
- local_max_seq_len: int
141
- local_scheduler_metadata: Optional[torch.Tensor]
142
-
143
- local_attn_metadata: Optional[LocalAttentionMetadata] = None
144
-
133
134
def _get_sliding_window_configs(
135
vllm_config: VllmConfig) -> set[Optional[tuple[int, int]]]:
@@ -221,7 +209,6 @@ class FlashAttentionMetadataBuilder(
209
max_query_len = common_attn_metadata.max_query_len
210
max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
211
query_start_loc = common_attn_metadata.query_start_loc
224
- query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
212
seq_lens = common_attn_metadata.seq_lens
213
seq_lens_cpu = common_attn_metadata.seq_lens_cpu
214
block_table_tensor = common_attn_metadata.block_table_tensor
@@ -266,40 +253,6 @@ class FlashAttentionMetadataBuilder(
253
)
254
return None
255
269
- # for local attention
270
- local_attn_metadata = None
271
- if self.model_config.attention_chunk_size is not None:
272
- seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
273
- virt_block_table_tensor = make_local_attention_virtual_batches(
274
- self.model_config.attention_chunk_size,
275
- query_start_loc_cpu.numpy(),
276
- seq_lens_cpu.numpy(),
277
- block_table_tensor,
278
- self.block_size,
279
- )
280
- local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
281
- self.device, non_blocking=True)
282
- local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
283
- self.device, non_blocking=True)
284
- local_max_query_len = seqlens_q_local_np.max()
285
- local_max_seq_len = virt_k_seqlens_np.max()
286
- local_scheduler_metadata = schedule(
287
- batch_size=local_query_start_loc.shape[0] - 1,
288
- cu_query_lens=local_query_start_loc,
289
- max_query_len=local_max_query_len,
290
- seqlens=local_seqused_k,
291
- max_seq_len=local_max_seq_len,
292
- causal=True)
293
-
294
- local_attn_metadata = FlashAttentionMetadata.LocalAttentionMetadata(
295
- local_query_start_loc=local_query_start_loc,
296
- local_seqused_k=local_seqused_k,
297
- local_block_table=virt_block_table_tensor,
298
- local_max_query_len=local_max_query_len,
299
- local_max_seq_len=local_max_seq_len,
300
- local_scheduler_metadata=local_scheduler_metadata,
301
- )
302
-
256
use_cascade = common_prefix_len > 0
257
258
if use_cascade:
@@ -371,7 +324,6 @@ class FlashAttentionMetadataBuilder(
324
cu_prefix_query_lens=cu_prefix_query_lens,
325
prefix_kv_lens=prefix_kv_lens,
326
suffix_kv_lens=suffix_kv_lens,
374
- local_attn_metadata=local_attn_metadata,
327
prefix_scheduler_metadata=prefix_scheduler_metadata,
328
max_num_splits=max_num_splits,
329
)
@@ -517,27 +469,13 @@ class FlashAttentionImpl(AttentionImpl):
469
layer._q_scale)
470
query = query.reshape((num_tokens, num_heads, head_size))
471
520
- # Compute attention and update output up to `num_actual_tokens`.
521
- use_local_attn = \
522
- (self.use_irope and attn_metadata.local_attn_metadata is not None)
523
-
524
- if not attn_metadata.use_cascade or use_local_attn:
525
- if use_local_attn:
526
- assert attn_metadata.local_attn_metadata is not None
527
- local_metadata = attn_metadata.local_attn_metadata
528
- cu_seqlens_q = local_metadata.local_query_start_loc
529
- seqused_k = local_metadata.local_seqused_k
530
- max_seqlen_q = local_metadata.local_max_query_len
531
- max_seqlen_k = local_metadata.local_max_seq_len
532
- block_table = local_metadata.local_block_table
533
- scheduler_metadata = local_metadata.local_scheduler_metadata
534
- else:
535
- cu_seqlens_q = attn_metadata.query_start_loc
536
- seqused_k = attn_metadata.seq_lens
537
- max_seqlen_q = attn_metadata.max_query_len
538
- max_seqlen_k = attn_metadata.max_seq_len
539
- block_table = attn_metadata.block_table
540
- scheduler_metadata = attn_metadata.scheduler_metadata
472
+ if not attn_metadata.use_cascade:
473
+ cu_seqlens_q = attn_metadata.query_start_loc
474
+ seqused_k = attn_metadata.seq_lens
475
+ max_seqlen_q = attn_metadata.max_query_len
476
+ max_seqlen_k = attn_metadata.max_seq_len
477
+ block_table = attn_metadata.block_table
478
+ scheduler_metadata = attn_metadata.scheduler_metadata
479
480
descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
481
@@ -565,8 +503,6 @@ class FlashAttentionImpl(AttentionImpl):
503
)
504
return output
505
568
- assert not use_local_attn, (
569
- "Cascade attention does not support local attention.")
506
# Cascade attention (rare case).
507
cascade_attention(
508
output[:num_actual_tokens],
vllm/v1/attention/backends/flashinfer.py
+1
-4
@@ -496,10 +496,6 @@ class FlashInferImpl(AttentionImpl):
496
kv_sharing_target_layer_name: Optional[int] = None,
497
use_irope: bool = False,
498
) -> None:
499
- if use_irope:
500
- logger.warning_once(
501
- "Using irope in FlashInfer is not supported yet, it will fall"
502
- " back to global attention for long context.")
499
self.num_heads = num_heads
500
self.head_size = head_size
501
self.scale = float(scale)
@@ -514,6 +510,7 @@ class FlashInferImpl(AttentionImpl):
510
self.kv_cache_dtype = kv_cache_dtype
511
self.logits_soft_cap = logits_soft_cap
512
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
513
+ self.use_irope = use_irope
514
515
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
516
vllm/v1/attention/backends/rocm_aiter_fa.py
+7
-90
@@ -13,8 +13,6 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
13
from vllm.config import VllmConfig
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
-from vllm.v1.attention.backends.flash_attn import (
17
- make_local_attention_virtual_batches)
16
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
17
from vllm.v1.kv_cache_interface import AttentionSpec
18
@@ -201,9 +199,7 @@ class AiterFlashAttentionMetadataBuilder:
199
max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
200
total_tokens = int(common_attn_metadata.seq_lens_cpu.sum())
201
query_start_loc = common_attn_metadata.query_start_loc
204
- query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
202
seq_lens = common_attn_metadata.seq_lens
206
- seq_lens_cpu = common_attn_metadata.seq_lens_cpu
203
block_table_tensor = common_attn_metadata.block_table_tensor
204
slot_mapping = common_attn_metadata.slot_mapping
205
@@ -215,56 +211,6 @@ class AiterFlashAttentionMetadataBuilder:
211
dtype=cu_seq_lens.dtype,
212
out=cu_seq_lens[1:])
213
218
- def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
219
- max_seq_len, causal):
220
- return None
221
-
222
- # for local attention
223
- local_attn_metadata = None
224
- if self.model_config.attention_chunk_size is not None:
225
- seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
226
- virt_block_table_tensor = make_local_attention_virtual_batches(
227
- self.model_config.attention_chunk_size,
228
- query_start_loc_cpu.numpy(),
229
- seq_lens_cpu.numpy(),
230
- block_table_tensor,
231
- self.block_size,
232
- )
233
- local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
234
- self.device, non_blocking=True)
235
- local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
236
- self.device, non_blocking=True)
237
- local_max_query_len = seqlens_q_local_np.max().item()
238
- local_max_seq_len = virt_k_seqlens_np.max().item()
239
- local_scheduler_metadata = schedule(
240
- batch_size=local_query_start_loc.shape[0] - 1,
241
- cu_query_lens=local_query_start_loc,
242
- max_query_len=local_max_query_len,
243
- seqlens=local_seqused_k,
244
- max_seq_len=local_max_seq_len,
245
- causal=True)
246
-
247
- local_cu_seq_lens = torch.zeros(virt_k_seqlens_np.shape[0] + 1,
248
- dtype=torch.int32,
249
- device=self.device)
250
- local_cu_seq_lens[1:] = torch.cumsum(
251
- torch.from_numpy(virt_k_seqlens_np).to(device=self.device,
252
- dtype=torch.int32,
253
- non_blocking=True),
254
- dim=0)
255
-
256
-
257
- local_attn_metadata = \
258
- AiterFlashAttentionMetadata.LocalAttentionMetadata(
259
- local_query_start_loc=local_query_start_loc,
260
- local_seqused_k=local_seqused_k,
261
- local_block_table=virt_block_table_tensor,
262
- local_max_query_len=local_max_query_len,
263
- local_max_seq_len=local_max_seq_len,
264
- local_cu_seq_lens=local_cu_seq_lens,
265
- local_scheduler_metadata=local_scheduler_metadata,
266
- )
267
-
214
use_cascade = common_prefix_len > 0
215
216
cu_prefix_query_lens = None
@@ -286,7 +232,6 @@ class AiterFlashAttentionMetadataBuilder:
232
cu_prefix_query_lens=cu_prefix_query_lens,
233
prefix_kv_lens=prefix_kv_lens,
234
suffix_kv_lens=suffix_kv_lens,
289
- local_attn_metadata=local_attn_metadata,
235
)
236
return attn_metadata
237
@@ -377,19 +322,6 @@ class AiterFlashAttentionMetadata:
322
prefix_kv_lens: Optional[torch.Tensor]
323
suffix_kv_lens: Optional[torch.Tensor]
324
380
- # for local attention
381
- @dataclass
382
- class LocalAttentionMetadata:
383
- local_query_start_loc: torch.Tensor
384
- local_seqused_k: torch.Tensor
385
- local_block_table: torch.Tensor
386
- local_max_query_len: int
387
- local_max_seq_len: int
388
- local_cu_seq_lens: torch.Tensor
389
- local_scheduler_metadata: Optional[torch.Tensor]
390
-
391
- local_attn_metadata: Optional[LocalAttentionMetadata] = None
392
-
325
326
class AiterFlashAttentionImpl(AttentionImpl):
327
@@ -521,25 +453,12 @@ class AiterFlashAttentionImpl(AttentionImpl):
453
layer._q_scale)
454
query = query.reshape((num_tokens, num_heads, head_size))
455
524
- # Compute attention and update output up to `num_actual_tokens`.
525
- use_local_attn = \
526
- (self.use_irope and attn_metadata.local_attn_metadata is not None)
527
-
528
- if not attn_metadata.use_cascade or use_local_attn:
529
- if use_local_attn:
530
- assert attn_metadata.local_attn_metadata is not None
531
- local_metadata = attn_metadata.local_attn_metadata
532
- cu_seqlens_q = local_metadata.local_query_start_loc
533
- seqused_k = local_metadata.local_seqused_k
534
- max_seqlen_q = local_metadata.local_max_query_len
535
- max_seqlen_k = local_metadata.local_max_seq_len
536
- block_table = local_metadata.local_block_table
537
- else:
538
- cu_seqlens_q = attn_metadata.query_start_loc
539
- seqused_k = attn_metadata.seq_lens
540
- max_seqlen_q = attn_metadata.max_query_len
541
- max_seqlen_k = attn_metadata.max_seq_len
542
- block_table = attn_metadata.block_table
456
+ if not attn_metadata.use_cascade:
457
+ cu_seqlens_q = attn_metadata.query_start_loc
458
+ seqused_k = attn_metadata.seq_lens
459
+ max_seqlen_q = attn_metadata.max_query_len
460
+ max_seqlen_k = attn_metadata.max_seq_len
461
+ block_table = attn_metadata.block_table
462
463
if max_seqlen_q > 1:
464
cu_seq_lens = attn_metadata.cu_seq_lens
@@ -557,9 +476,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
476
alibi_slopes=self.alibi_slopes,
477
window_size=self.sliding_window,
478
block_table=block_table,
560
- cu_seqlens_k=(cu_seq_lens if not use_local_attn else
561
- local_metadata.local_cu_seq_lens),
562
- )
479
+ cu_seqlens_k=cu_seq_lens)
480
481
_, num_heads, head_size = query.shape
482
_PARTITION_SIZE_ROCM = 256
vllm/v1/attention/backends/triton_attn.py
+7
-61
@@ -18,9 +18,8 @@ from vllm.config import VllmConfig
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
21
-from vllm.v1.attention.backends.utils import (
22
- AttentionMetadataBuilder, CommonAttentionMetadata,
23
- make_local_attention_virtual_batches)
21
+from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
22
+ CommonAttentionMetadata)
23
from vllm.v1.kv_cache_interface import AttentionSpec
24
25
logger = init_logger(__name__)
@@ -55,18 +54,6 @@ class TritonAttentionMetadata:
54
scheduler_metadata: Optional[torch.Tensor] = None
55
prefix_scheduler_metadata: Optional[torch.Tensor] = None
56
58
- # for local attention
59
- @dataclass
60
- class LocalAttentionMetadata:
61
- local_query_start_loc: torch.Tensor
62
- local_seqused_k: torch.Tensor
63
- local_block_table: torch.Tensor
64
- local_max_query_len: int
65
- local_max_seq_len: int
66
- local_scheduler_metadata: Optional[torch.Tensor]
67
-
68
- local_attn_metadata: Optional[LocalAttentionMetadata] = None
69
-
57
58
class TritonAttentionMetadataBuilder(
59
AttentionMetadataBuilder[TritonAttentionMetadata]):
@@ -111,34 +98,6 @@ class TritonAttentionMetadataBuilder(
98
block_table_tensor = common_attn_metadata.block_table_tensor
99
slot_mapping = common_attn_metadata.slot_mapping
100
114
- # for local attention
115
- local_attn_metadata = None
116
- if self.attention_chunk_size is not None:
117
- seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
118
- virt_block_table_tensor = make_local_attention_virtual_batches(
119
- self.attention_chunk_size,
120
- common_attn_metadata.query_start_loc_cpu.numpy(),
121
- common_attn_metadata.seq_lens_cpu.numpy(),
122
- block_table_tensor,
123
- self.block_size,
124
- )
125
- local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
126
- self.device, non_blocking=True)
127
- local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
128
- self.device, non_blocking=True)
129
- local_max_query_len = seqlens_q_local_np.max().item()
130
- local_max_seq_len = virt_k_seqlens_np.max().item()
131
-
132
- local_attn_metadata = TritonAttentionMetadata \
133
- .LocalAttentionMetadata(
134
- local_query_start_loc=local_query_start_loc,
135
- local_seqused_k=local_seqused_k,
136
- local_block_table=virt_block_table_tensor,
137
- local_max_query_len=local_max_query_len,
138
- local_max_seq_len=local_max_seq_len,
139
- local_scheduler_metadata=None,
140
- )
141
-
101
use_cascade = common_prefix_len > 0
102
103
if use_cascade:
@@ -170,7 +129,6 @@ class TritonAttentionMetadataBuilder(
129
cu_prefix_query_lens=cu_prefix_query_lens,
130
prefix_kv_lens=prefix_kv_lens,
131
suffix_kv_lens=suffix_kv_lens,
173
- local_attn_metadata=local_attn_metadata,
132
prefix_scheduler_metadata=prefix_scheduler_metadata,
133
)
134
return attn_metadata
@@ -384,23 +342,11 @@ class TritonAttentionImpl(AttentionImpl):
342
layer._q_scale)
343
query = query.reshape((num_tokens, num_heads, head_size))
344
387
- use_local_attn = \
388
- (self.use_irope and attn_metadata.local_attn_metadata is not None)
389
-
390
- if use_local_attn:
391
- assert attn_metadata.local_attn_metadata is not None
392
- local_metadata = attn_metadata.local_attn_metadata
393
- cu_seqlens_q = local_metadata.local_query_start_loc
394
- seqused_k = local_metadata.local_seqused_k
395
- max_seqlen_q = local_metadata.local_max_query_len
396
- max_seqlen_k = local_metadata.local_max_seq_len
397
- block_table = local_metadata.local_block_table
398
- else:
399
- cu_seqlens_q = attn_metadata.query_start_loc
400
- seqused_k = attn_metadata.seq_lens
401
- max_seqlen_q = attn_metadata.max_query_len
402
- max_seqlen_k = attn_metadata.max_seq_len
403
- block_table = attn_metadata.block_table
345
+ cu_seqlens_q = attn_metadata.query_start_loc
346
+ seqused_k = attn_metadata.seq_lens
347
+ max_seqlen_q = attn_metadata.max_query_len
348
+ max_seqlen_k = attn_metadata.max_seq_len
349
+ block_table = attn_metadata.block_table
350
351
if use_prefill_decode_attn:
352
# Compute attention and update output up to `num_actual_tokens`.
vllm/v1/attention/backends/utils.py
+24
-6
@@ -272,11 +272,14 @@ def infer_global_hyperparameters(
272
# block_table_local : shape[local_virtual_batches, pages_per_local_batch]
273
def make_local_attention_virtual_batches(
274
attn_chunk_size: int,
275
- query_start_loc_np: np.ndarray,
276
- seq_lens_np: np.ndarray,
277
- block_table: torch.Tensor,
275
+ common_attn_metadata: CommonAttentionMetadata,
276
block_size: int = 0,
279
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, torch.Tensor]:
277
+) -> CommonAttentionMetadata:
278
+ query_start_loc_np = common_attn_metadata.query_start_loc_cpu.numpy()
279
+ seq_lens_np = common_attn_metadata.seq_lens_cpu.numpy()
280
+ block_table = common_attn_metadata.block_table_tensor
281
+ device = common_attn_metadata.query_start_loc.device
282
+
283
q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1]
284
actual_batch_size = seq_lens_np.shape[0]
285
@@ -339,6 +342,7 @@ def make_local_attention_virtual_batches(
342
attn_chunk_size,
343
dtype=np.int32)
344
seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block
345
+ num_computed_tokens_local = seqlens_k_local - seqlens_q_local
346
347
k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - \
348
(rarange * attn_chunk_size + \
@@ -380,8 +384,22 @@ def make_local_attention_virtual_batches(
384
block_table_local = block_table[batch_indices, block_indices]\
385
.view(virtual_batches, -1)
386
383
- return seqlens_q_local, cu_seqlens_q_local, seqlens_k_local, \
384
- block_table_local
387
+ query_start_loc_cpu = torch.from_numpy(cu_seqlens_q_local)
388
+ seq_lens_cpu = torch.from_numpy(seqlens_k_local)
389
+
390
+ return CommonAttentionMetadata(
391
+ query_start_loc_cpu=query_start_loc_cpu,
392
+ query_start_loc=query_start_loc_cpu.to(device=device,
393
+ non_blocking=True),
394
+ seq_lens_cpu=seq_lens_cpu,
395
+ seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),
396
+ num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local),
397
+ num_reqs=len(seq_lens_cpu),
398
+ num_actual_tokens=common_attn_metadata.num_actual_tokens,
399
+ max_query_len=seqlens_q_local.max(),
400
+ block_table_tensor=block_table_local,
401
+ slot_mapping=common_attn_metadata.slot_mapping,
402
+ )
403
404
405
def split_decodes_and_prefills(
vllm/v1/core/single_type_kv_cache_manager.py
+7
-3
@@ -7,7 +7,8 @@ from typing import Callable
7
from vllm.utils import cdiv
8
from vllm.v1.core.block_pool import BlockPool
9
from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
10
-from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheSpec,
10
+from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
11
+ FullAttentionSpec, KVCacheSpec,
12
MambaSpec, SlidingWindowSpec)
13
from vllm.v1.request import Request
14
@@ -256,8 +257,10 @@ class FullAttentionManager(SingleTypeKVCacheManager):
257
kv_cache_spec: KVCacheSpec,
258
use_eagle: bool,
259
) -> tuple[list[KVCacheBlock], ...]:
259
- assert isinstance(kv_cache_spec, FullAttentionSpec), (
260
- "FullAttentionManager can only be used for full attention groups")
260
+ assert isinstance(
261
+ kv_cache_spec, (FullAttentionSpec, ChunkedLocalAttentionSpec)
262
+ ), "FullAttentionManager can only be used for full attention " \
263
+ "and chunked local attention groups"
264
computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
265
[] for _ in range(len(kv_cache_group_ids)))
266
max_num_blocks = max_length // kv_cache_spec.block_size
@@ -432,6 +435,7 @@ class MambaManager(SingleTypeKVCacheManager):
435
436
spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
437
FullAttentionSpec: FullAttentionManager,
438
+ ChunkedLocalAttentionSpec: FullAttentionManager,
439
SlidingWindowSpec: SlidingWindowManager,
440
MambaSpec: MambaManager,
441
}
vllm/v1/kv_cache_interface.py
+15
@@ -125,6 +125,21 @@ class FullAttentionSpec(AttentionSpec):
125
return merged_spec
126
127
128
+@dataclass
129
+class ChunkedLocalAttentionSpec(AttentionSpec):
130
+ attention_chunk_size: int
131
+
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
+
136
+ @property
137
+ def type_id(self) -> str:
138
+ return (
139
+ f"local_attention_{self.attention_chunk_size}_{self.block_size}_{self.page_size_bytes}"
140
+ ) # noqa
141
+
142
+
143
@dataclass
144
class SlidingWindowSpec(AttentionSpec):
145
sliding_window: int
vllm/v1/worker/gpu_model_runner.py
+23
-4
@@ -44,11 +44,14 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
44
GiB_bytes, LazyLoader, check_use_alibi, get_dtype_size,
45
is_pin_memory_available, round_up)
46
from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
47
-from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
48
- CommonAttentionMetadata)
47
+from vllm.v1.attention.backends.utils import (
48
+ AttentionMetadataBuilder, CommonAttentionMetadata,
49
+ make_local_attention_virtual_batches)
50
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
50
-from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
51
- KVCacheConfig, KVCacheSpec, MambaSpec,
51
+from vllm.v1.kv_cache_interface import (AttentionSpec,
52
+ ChunkedLocalAttentionSpec,
53
+ FullAttentionSpec, KVCacheConfig,
54
+ KVCacheSpec, MambaSpec,
55
SlidingWindowSpec)
56
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
57
ModelRunnerOutput)
@@ -705,6 +708,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
708
spec_decode_common_attn_metadata is None:
709
spec_decode_common_attn_metadata = common_attn_metadata
710
711
+ if isinstance(kv_cache_group_spec.kv_cache_spec,
712
+ ChunkedLocalAttentionSpec):
713
+ common_attn_metadata = make_local_attention_virtual_batches(
714
+ kv_cache_group_spec.kv_cache_spec.attention_chunk_size,
715
+ common_attn_metadata, self.cache_config.block_size)
716
+
717
# Prepare for cascade attention if enabled & beneficial.
718
common_prefix_len = 0
719
builder = self.attn_metadata_builders[kv_cache_group_id]
@@ -2589,6 +2598,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2598
2599
# TODO: Support other attention modules, e.g., cross-attention
2600
if attn_module.attn_type == AttentionType.DECODER:
2601
+ use_local_attention = (self.attention_chunk_size is not None
2602
+ and attn_module.impl.use_irope)
2603
if attn_module.sliding_window is not None:
2604
kv_cache_spec[layer_name] = SlidingWindowSpec(
2605
block_size=block_size,
@@ -2597,6 +2608,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2608
dtype=self.kv_cache_dtype,
2609
sliding_window=attn_module.sliding_window,
2610
use_mla=use_mla)
2611
+ elif use_local_attention:
2612
+ kv_cache_spec[layer_name] = (ChunkedLocalAttentionSpec(
2613
+ block_size=block_size,
2614
+ num_kv_heads=attn_module.num_kv_heads,
2615
+ head_size=attn_module.head_size,
2616
+ dtype=self.kv_cache_dtype,
2617
+ attention_chunk_size=self.attention_chunk_size,
2618
+ use_mla=use_mla))
2619
else:
2620
kv_cache_spec[layer_name] = FullAttentionSpec(
2621
block_size=block_size,