@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9fec0e13d5

Fork 0

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, 08:45 UTC

[Attention] Cache attention metadata builds across hybrid KV-cache groups (#29627)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Stanislaw Wozniak <stw@zurich.ibm.com>

Lucas Wilkinson committed Dec 16, 2025 at 17:10 UTC 9fec0e13d512b6b9082e40297582d8052f434610

7 files changed +105 -13

tests/v1/attention/test_chunked_local_attention.py

+1 -1

@@ -172,7 +172,7 @@ def test_local_attention_virtual_batches(test_data: LocalAttentionTestData):

172 )

173

174 # Call the function

175 - result = make_local_attention_virtual_batches(

175 + result, _ = make_local_attention_virtual_batches(

176 attn_chunk_size, common_attn_metadata, block_size

177 )

178

vllm/attention/layers/chunked_local_attention.py

+12 -4

@@ -4,7 +4,7 @@ import functools

5 import torch

7 -from vllm.attention.backends.abstract import AttentionBackend, AttentionMetadata

7 +from vllm.attention.backends.abstract import AttentionBackend

8 from vllm.attention.layer import Attention

9 from vllm.attention.selector import get_attn_backend

10 from vllm.config import CacheConfig

@@ -51,11 +51,19 @@ def create_chunked_local_attention_backend(

51 common_prefix_len: int,

52 common_attn_metadata: CommonAttentionMetadata,

53 fast_build: bool = False,

54 - ) -> AttentionMetadata:

55 - common_attn_metadata = make_local_attention_virtual_batches(

54 + ):

55 + cm, make_virtual_batches_block_table = make_local_attention_virtual_batches(

56 attention_chunk_size, common_attn_metadata, block_size

57 )

58 - return super().build(common_prefix_len, common_attn_metadata, fast_build)

58 + metadata = super().build(common_prefix_len, cm, fast_build)

59 + metadata.make_virtual_batches_block_table = make_virtual_batches_block_table

60 + return metadata

61 +

62 + def update_block_table(

63 + self, metadata, blk_table: torch.Tensor, slot_mapping: torch.Tensor

64 + ):

65 + blk_table = metadata.make_virtual_batches_block_table(blk_table)

66 + return super().update_block_table(metadata, blk_table, slot_mapping)

68 attn_backend = subclass_attention_backend(

69 name_prefix=prefix,

vllm/envs.py

+2 -2

@@ -207,7 +207,7 @@ if TYPE_CHECKING:

207 VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL: bool = False

208 VLLM_ENABLE_CUDAGRAPH_GC: bool = False

209 VLLM_LOOPBACK_IP: str = ""

210 - VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = False

210 + VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True

211 VLLM_ENABLE_RESPONSES_API_STORE: bool = False

212 VLLM_USE_TRTLLM_ATTENTION: str | None = None

213 VLLM_NVFP4_GEMM_BACKEND: str | None = None

@@ -1430,7 +1430,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1430 # kv-cache memory usage and enable longer contexts)

1431 # TODO(lucas): Remove this flag once latency regression is resolved.

1432 "VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE": lambda: bool(

1433 - int(os.getenv("VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE", "0"))

1433 + int(os.getenv("VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE", "1"))

1434 ),

1435 # Enables support for the "store" option in the OpenAI Responses API.

1436 # When set to 1, vLLM's OpenAI server will retain the input and output

vllm/v1/attention/backends/flash_attn.py

+13

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Attention layer with FlashAttention."""

5 +import copy

6 from dataclasses import dataclass

7 from typing import ClassVar

@@ -250,6 +251,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

251 if get_flash_attn_version() == 3

252 else AttentionCGSupport.UNIFORM_BATCH

253 )

254 + supports_update_block_table: bool = True

255

256 def __init__(

257 self,

@@ -493,6 +495,17 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

495 )

496 return attn_metadata

497

498 + def update_block_table(

499 + self,

500 + metadata: FlashAttentionMetadata,

501 + blk_table: torch.Tensor,

502 + slot_mapping: torch.Tensor,

503 + ) -> FlashAttentionMetadata:

504 + new_metadata = copy.copy(metadata)

505 + new_metadata.block_table = blk_table

506 + new_metadata.slot_mapping = slot_mapping

507 + return new_metadata

508 +

509 def use_cascade_attention(self, *args, **kwargs) -> bool:

510 return use_cascade_attention(*args, **kwargs)

511

vllm/v1/attention/backends/mamba2_attn.py

+27

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import copy

4 import itertools

5 from dataclasses import dataclass

@@ -134,6 +135,8 @@ class Mamba2AttentionMetadata:

135 class Mamba2AttentionMetadataBuilder(

136 BaseMambaAttentionMetadataBuilder[Mamba2AttentionMetadata]

137 ):

138 + supports_update_block_table: bool = True

139 +

140 def __init__(

141 self,

142 kv_cache_spec: AttentionSpec,

@@ -346,3 +349,27 @@ class Mamba2AttentionMetadataBuilder(

349 num_computed_tokens_p=num_computed_tokens_p,

350 )

351 return attn_metadata

352 +

353 + def update_block_table(

354 + self,

355 + metadata: Mamba2AttentionMetadata,

356 + blk_table: torch.Tensor,

357 + slot_mapping: torch.Tensor,

358 + ) -> Mamba2AttentionMetadata:

359 + new_metadata = copy.copy(metadata)

360 + prefix_caching = self.vllm_config.cache_config.enable_prefix_caching

361 + state_indices_t = blk_table if prefix_caching else blk_table[:, 0]

362 + num_reqs = blk_table.shape[0]

363 +

364 + # For CUDA graphs, copy to persistent buffer

365 + if (

366 + metadata.num_prefills == 0

367 + and num_reqs <= self.decode_cudagraph_max_bs

368 + and self.compilation_config.cudagraph_mode.has_full_cudagraphs()

369 + ):

370 + persistent_state_indices_t = self.state_indices_tensor[:num_reqs]

371 + persistent_state_indices_t.copy_(state_indices_t, non_blocking=True)

372 + state_indices_t = persistent_state_indices_t

373 +

374 + new_metadata.state_indices_tensor = state_indices_t

375 + return new_metadata

vllm/v1/attention/backends/utils.py

+27 -5

@@ -4,6 +4,7 @@ import abc

4 import enum

5 import functools

6 from abc import abstractmethod

7 +from collections.abc import Callable

8 from dataclasses import dataclass, field, fields, make_dataclass

9 from typing import (

10 TYPE_CHECKING,

@@ -317,6 +318,9 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

318 # If not, set this to None. Otherwise set it to the query

319 # length that will be pulled into the front of the batch.

320 reorder_batch_threshold: int | None = None

321 + # Does this backend/builder support updating the block table in existing

322 + # metadata

323 + supports_update_block_table: bool = False

324

325 @abstractmethod

326 def __init__(

@@ -387,6 +391,21 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):

391 """

392 raise NotImplementedError

393

394 + def update_block_table(

395 + self,

396 + metadata: M,

397 + blk_table: torch.Tensor,

398 + slot_mapping: torch.Tensor,

399 + ) -> M:

400 + """

401 + Update the block table for the attention metadata.

402 + Faster when theres multiple kv-cache groups that create virtually the

403 + same metadata but just with different block tables.

404 +

405 + Only needs to be implemented if supports_update_block_table is True.

406 + """

407 + raise NotImplementedError

408 +

409 def build_for_cudagraph_capture(

410 self, common_attn_metadata: CommonAttentionMetadata

411 ) -> M:

@@ -603,7 +622,7 @@ def make_local_attention_virtual_batches(

622 attn_chunk_size: int,

623 common_attn_metadata: CommonAttentionMetadata,

624 block_size: int = 0,

606 -) -> CommonAttentionMetadata:

625 +) -> tuple[CommonAttentionMetadata, Callable[[torch.Tensor], torch.Tensor]]:

626 query_start_loc_np = common_attn_metadata.query_start_loc_cpu.numpy()

627 seq_lens_np = common_attn_metadata.seq_lens_cpu.numpy()

628 block_table = common_attn_metadata.block_table_tensor

@@ -715,9 +734,12 @@ def make_local_attention_virtual_batches(

734 # tensor first, which recovers perf.

735 batch_indices_torch = torch.from_numpy(batch_indices)

736 block_indices_torch = torch.from_numpy(block_indices)

718 - block_table_local = block_table[batch_indices_torch, block_indices_torch].view(

719 - virtual_batches, -1

720 - )

737 +

738 + # Save as a lambda so we can return this for update_block_table

739 + make_block_table = lambda block_table: block_table[

740 + batch_indices_torch, block_indices_torch

741 + ].view(virtual_batches, -1)

742 + block_table_local = make_block_table(block_table)

743

744 query_start_loc_cpu = torch.from_numpy(cu_seqlens_q_local)

745 seq_lens_cpu = torch.from_numpy(seqlens_k_local)

@@ -736,7 +758,7 @@ def make_local_attention_virtual_batches(

758 causal=True,

759 _seq_lens_cpu=seq_lens_cpu,

760 _num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local),

739 - )

761 + ), make_block_table

762

763

764 def make_kv_sharing_fast_prefill_common_attn_metadata(

vllm/v1/worker/gpu_model_runner.py

+23 -1

@@ -1630,6 +1630,15 @@ class GPUModelRunner(

1630 logits_indices

1631 )

1632

1633 + # Cache attention metadata builds across hybrid KV-cache groups

1634 + # The only thing that changes between different hybrid KV-cache groups when the

1635 + # same metadata builder and KVCacheSpec is the same is the block table, so we

1636 + # can cache the attention metadata builds and just update the block table using

1637 + # `builder.update_block_table` if the builder supports it.

1638 + cached_attn_metadata: dict[

1639 + tuple[KVCacheSpec, type[AttentionMetadataBuilder]], AttentionMetadata

1640 + ] = {}

1641 +

1642 def _build_attn_group_metadata(

1643 kv_cache_gid: int,

1644 attn_gid: int,

@@ -1637,13 +1646,15 @@ class GPUModelRunner(

1646 ubid: int | None = None,

1647 ) -> None:

1648 attn_group = self.attn_groups[kv_cache_gid][attn_gid]

1649 + builder = attn_group.get_metadata_builder(ubid or 0)

1650 + cache_key = (kv_cache_groups[kv_cache_gid].kv_cache_spec, type(builder))

1651 +

1652 cascade_attn_prefix_len = (

1653 cascade_attn_prefix_lens[kv_cache_gid][attn_gid]

1654 if cascade_attn_prefix_lens

1655 else 0

1656 )

1657

1646 - builder = attn_group.get_metadata_builder(ubid or 0)

1658 extra_attn_metadata_args = {}

1659 if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):

1660 assert ubid is None, "UBatching not supported with GDN yet"

@@ -1658,12 +1669,23 @@ class GPUModelRunner(

1669 attn_metadata_i = builder.build_for_cudagraph_capture(

1670 common_attn_metadata

1671 )

1672 + elif (

1673 + cache_key in cached_attn_metadata

1674 + and builder.supports_update_block_table

1675 + ):

1676 + attn_metadata_i = builder.update_block_table(

1677 + cached_attn_metadata[cache_key],

1678 + common_attn_metadata.block_table_tensor,

1679 + common_attn_metadata.slot_mapping,

1680 + )

1681 else:

1682 attn_metadata_i = builder.build(

1683 common_prefix_len=cascade_attn_prefix_len,

1684 common_attn_metadata=common_attn_metadata,

1685 **extra_attn_metadata_args,

1686 )

1687 + if builder.supports_update_block_table:

1688 + cached_attn_metadata[cache_key] = attn_metadata_i

1689

1690 if ubid is None:

1691 assert isinstance(attn_metadata, dict)