@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 98aa16ff41

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 14, 05:30 UTC

[v1] Add cross-attention KV cache support for encoder-decoder models (#23664)

Signed-off-by: Russell Bryant <rbryant@redhat.com>

Russell Bryant committed Aug 26, 2025 at 14:49 UTC 98aa16ff41353e3e6c8a3c2f4e933a888dbce1cb

6 files changed +153 -14

vllm/multimodal/registry.py

+19

@@ -372,3 +372,22 @@ class MultiModalRegistry:

372 )

373

374 return dummy_data

375 +

376 + def get_encdec_max_encoder_len(self, model_config: "ModelConfig") -> int:

377 + """

378 + Get the maximum length of the encoder input for encoder-decoder models.

379 + """

380 + if not model_config.is_encoder_decoder:

381 + return 0

382 + max_tokens = self.\

383 + get_max_tokens_per_item_by_nonzero_modality(model_config)

384 + if not max_tokens:

385 + # TODO - this function assumes encoder-decoder models are

386 + # multimodal. This will need to change when adding support for more

387 + # than whisper.

388 + return 0

389 + assert len(max_tokens) == 1, "Encoder-decoder models are expected \

390 + to implement the multimodal interface with at most one modality."

391 +

392 + first_modality = next(iter(max_tokens))

393 + return max_tokens[first_modality]

vllm/v1/core/kv_cache_coordinator.py

+25 -9

@@ -6,7 +6,7 @@ from typing import Optional

6 from vllm.v1.core.block_pool import BlockPool

7 from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

8 from vllm.v1.core.single_type_kv_cache_manager import (

9 - FullAttentionManager, get_manager_for_kv_cache_spec)

9 + CrossAttentionManager, FullAttentionManager, get_manager_for_kv_cache_spec)

10 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

11 KVCacheSpec)

12 from vllm.v1.request import Request

@@ -42,9 +42,10 @@ class KVCacheCoordinator(ABC):

42 ) for i, kv_cache_group in enumerate(

43 self.kv_cache_config.kv_cache_groups))

45 - def get_num_blocks_to_allocate(

46 - self, request_id: str, num_tokens: int,

47 - new_computed_blocks: tuple[list[KVCacheBlock], ...]) -> int:

45 + def get_num_blocks_to_allocate(self, request_id: str, num_tokens: int,

46 + new_computed_blocks: tuple[

47 + list[KVCacheBlock], ...],

48 + num_encoder_tokens: int) -> int:

49 """

50 Get the number of blocks needed to be allocated for the request.

@@ -54,14 +55,22 @@ class KVCacheCoordinator(ABC):

55 tokens that are already allocated).

56 new_computed_blocks: The new computed blocks just hitting the

57 prefix caching.

58 + num_encoder_tokens: The number of encoder tokens for allocating

59 + blocks for cross-attention.

61 Returns:

62 The number of blocks.

63 """

64 num_blocks_to_allocate = 0

65 for i, manager in enumerate(self.single_type_managers):

63 - num_blocks_to_allocate += manager.get_num_blocks_to_allocate(

64 - request_id, num_tokens, new_computed_blocks[i])

66 + if isinstance(manager, CrossAttentionManager):

67 + # For cross-attention, we issue a single static allocation

68 + # of blocks based on the number of encoder input tokens.

69 + num_blocks_to_allocate += manager.get_num_blocks_to_allocate(

70 + request_id, num_encoder_tokens, [])

71 + else:

72 + num_blocks_to_allocate += manager.get_num_blocks_to_allocate(

73 + request_id, num_tokens, new_computed_blocks[i])

74 return num_blocks_to_allocate

76 def save_new_computed_blocks(

@@ -79,8 +88,11 @@ class KVCacheCoordinator(ABC):

88 manager.save_new_computed_blocks(request_id,

89 new_computed_blocks[i])

82 - def allocate_new_blocks(self, request_id: str,

83 - num_tokens: int) -> tuple[list[KVCacheBlock], ...]:

91 + def allocate_new_blocks(

92 + self,

93 + request_id: str,

94 + num_tokens: int,

95 + num_encoder_tokens: int = 0) -> tuple[list[KVCacheBlock], ...]:

96 """

97 Allocate new blocks for the request to give it at least `num_tokens`

98 token slots.

@@ -89,12 +101,16 @@ class KVCacheCoordinator(ABC):

101 request_id: The request ID.

102 num_tokens: The total number of tokens that need a slot (including

103 tokens that are already allocated).

104 + num_encoder_tokens: The number of encoder tokens for allocating

105 + blocks for cross-attention.

106

107 Returns:

108 The new allocated blocks.

109 """

110 return tuple(

97 - manager.allocate_new_blocks(request_id, num_tokens)

111 + manager.allocate_new_blocks(

112 + request_id, num_encoder_tokens if isinstance(

113 + manager, CrossAttentionManager) else num_tokens)

114 for manager in self.single_type_managers)

115

116 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:

vllm/v1/core/kv_cache_manager.py

+4 -2

@@ -187,6 +187,7 @@ class KVCacheManager:

187 new_computed_blocks: Optional[KVCacheBlocks] = None,

188 num_lookahead_tokens: int = 0,

189 delay_cache_blocks: bool = False,

190 + num_encoder_tokens: int = 0,

191 ) -> Optional[KVCacheBlocks]:

192 """Add slots for a request with new tokens to append.

193

@@ -253,6 +254,7 @@ class KVCacheManager:

254 request_id=request.request_id,

255 num_tokens=num_tokens_need_slot,

256 new_computed_blocks=new_computed_block_list,

257 + num_encoder_tokens=num_encoder_tokens,

258 )

259

260 if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():

@@ -273,7 +275,7 @@ class KVCacheManager:

275 new_computed_block_list)

276

277 new_blocks = self.coordinator.allocate_new_blocks(

276 - request.request_id, num_tokens_need_slot)

278 + request.request_id, num_tokens_need_slot, num_encoder_tokens)

279

280 # P/D: delay caching blocks if we have to recv from

281 # remote. Update state for locally cached blocks.

@@ -292,7 +294,7 @@ class KVCacheManager:

294

295 def free(self, request: Request) -> None:

296 """Free the blocks allocated for the request.

295 - We free the blocks in reverse order so that he tail blocks are evicted

297 + We free the blocks in reverse order so that the tail blocks are evicted

298 first when caching is enabled.

299

300 Args:

vllm/v1/core/sched/scheduler.py

+36 -1

@@ -58,6 +58,7 @@ class Scheduler(SchedulerInterface):

58 self.parallel_config = vllm_config.parallel_config

59 self.log_stats = log_stats

60 self.structured_output_manager = structured_output_manager

61 + self.is_encoder_decoder = vllm_config.model_config.is_encoder_decoder

63 # include_finished_set controls whether a separate set of finished

64 # request ids should be included in the EngineCoreOutputs returned

@@ -83,6 +84,9 @@ class Scheduler(SchedulerInterface):

84 assert len(self.kv_cache_config.kv_cache_groups) == 1, (

85 "Multiple KV cache groups are not currently supported "

86 "with KV connectors")

87 + assert not self.is_encoder_decoder, (

88 + "Encoder-decoder models are not currently supported "

89 + "with KV connectors")

90 self.connector = KVConnectorFactory.create_connector(

91 config=self.vllm_config, role=KVConnectorRole.SCHEDULER)

@@ -431,6 +435,22 @@ class Scheduler(SchedulerInterface):

435 == 0 else

436 self.num_lookahead_tokens)

437

438 + # Determine if we need to allocate cross-attention blocks.

439 + if self.is_encoder_decoder and request.has_encoder_inputs:

440 + # TODO(russellb): For Whisper, we know that the input is

441 + # always padded to the maximum length. If we support other

442 + # encoder-decoder models, this will need to be updated if we

443 + # want to only allocate what is needed.

444 + assert ("whisper"

445 + in self.vllm_config.model_config.model.lower()), (

446 + "Whisper is the only supported "

447 + "encoder-decoder model.")

448 + num_encoder_tokens = MULTIMODAL_REGISTRY.\

449 + get_encdec_max_encoder_len(

450 + self.vllm_config.model_config)

451 + else:

452 + num_encoder_tokens = 0

453 +

454 new_blocks = self.kv_cache_manager.allocate_slots(

455 request,

456 num_new_tokens + num_external_computed_tokens,

@@ -438,6 +458,7 @@ class Scheduler(SchedulerInterface):

458 new_computed_blocks,

459 num_lookahead_tokens=effective_lookahead_tokens,

460 delay_cache_blocks=load_kv_async,

461 + num_encoder_tokens=num_encoder_tokens,

462 )

463

464 if new_blocks is None:

@@ -703,7 +724,21 @@ class Scheduler(SchedulerInterface):

724 # The encoder input is not needed in this step.

725 break

726

706 - if start_pos + num_encoder_tokens <= num_computed_tokens:

727 + if self.is_encoder_decoder and num_computed_tokens > 0:

728 + assert start_pos == 0, (

729 + "Encoder input should be processed at the beginning of "

730 + "the sequence when encoder-decoder models are used.")

731 + # Encoder input has already been computed

732 + # The calculation here is a bit different. We don't turn encoder

733 + # output into tokens that get processed by the decoder and

734 + # reflected in num_computed_tokens. Instead, start_pos reflects

735 + # the position where we need to ensure we calculate encoder

736 + # inputs. This should always be 0 to ensure we calculate encoder

737 + # inputs before running the decoder. Once we've calculated some

738 + # decoder tokens (num_computed_tokens > 0), then we know we

739 + # already calculated encoder inputs and can skip here.

740 + continue

741 + elif start_pos + num_encoder_tokens <= num_computed_tokens:

742 # The encoder input is already computed and stored

743 # in the decoder's KV cache.

744 continue

vllm/v1/core/single_type_kv_cache_manager.py

+54 -2

@@ -8,8 +8,9 @@ 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 (ChunkedLocalAttentionSpec,

11 - FullAttentionSpec, KVCacheSpec,

12 - MambaSpec, SlidingWindowSpec)

11 + CrossAttentionSpec, FullAttentionSpec,

12 + KVCacheSpec, MambaSpec,

13 + SlidingWindowSpec)

14 from vllm.v1.request import Request

@@ -552,11 +553,62 @@ class MambaManager(SingleTypeKVCacheManager):

553 return new_blocks

554

555

556 +class CrossAttentionManager(SingleTypeKVCacheManager):

557 + """Manager for cross-attention KV cache in encoder-decoder models."""

558 +

559 + def save_new_computed_blocks(

560 + self, request_id: str,

561 + new_computed_blocks: list[KVCacheBlock]) -> None:

562 + # We do not cache blocks for cross-attention to be shared between

563 + # requests, so `new_computed_blocks` should always be empty.

564 + assert len(new_computed_blocks) == 0

565 +

566 + def cache_blocks(self, request: Request, num_tokens: int) -> None:

567 + # We do not cache blocks for cross-attention to be shared between

568 + # requests, so this method is not relevant.

569 + raise ValueError("Should not be called as prefix caching is disabled.")

570 +

571 + def get_num_common_prefix_blocks(self, request_id: str,

572 + num_running_requests: int) -> int:

573 + # Cross-attention blocks contain request-specific encoder states

574 + # and are not shared between different requests

575 + return 0

576 +

577 + @classmethod

578 + def find_longest_cache_hit(

579 + cls,

580 + block_hashes: list[BlockHash],

581 + max_length: int,

582 + kv_cache_group_ids: list[int],

583 + block_pool: BlockPool,

584 + kv_cache_spec: KVCacheSpec,

585 + use_eagle: bool,

586 + ) -> tuple[list[KVCacheBlock], ...]:

587 + assert isinstance(kv_cache_spec, CrossAttentionSpec), (

588 + "CrossAttentionManager can only be used for cross-attention groups"

589 + )

590 + # Cross-attention does not benefit from prefix caching since:

591 + # 1. Encoder states are unique per request (different audio/image

592 + # inputs)

593 + # 2. Encoder states are computed once per request, not incrementally

594 + # 3. No reusable prefix exists between different multimodal inputs

595 + # Return empty blocks to indicate no cache hits

596 + raise NotImplementedError(

597 + "CrossAttentionManager does not support caching")

598 +

599 + def remove_skipped_blocks(self, request_id: str,

600 + num_computed_tokens: int) -> None:

601 + # Cross-attention blocks represent encoder states which are needed

602 + # for the entire decoding process, so no blocks should be skipped

603 + pass

604 +

605 +

606 spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

607 FullAttentionSpec: FullAttentionManager,

608 SlidingWindowSpec: SlidingWindowManager,

609 ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,

610 MambaSpec: MambaManager,

611 + CrossAttentionSpec: CrossAttentionManager,

612 }

613

614

vllm/v1/kv_cache_interface.py

+15

@@ -11,6 +11,7 @@ from typing_extensions import Self

12 from vllm.config import VllmConfig

13 from vllm.logger import init_logger

14 +from vllm.multimodal import MULTIMODAL_REGISTRY

15 from vllm.utils import cdiv, get_dtype_size

17 logger = init_logger(__name__)

@@ -211,6 +212,20 @@ class EncoderOnlyAttentionSpec(AttentionSpec):

212 return 0

213

214

215 +@dataclass(frozen=True)

216 +class CrossAttentionSpec(AttentionSpec):

217 + """

218 + KV cache spec for cross-attention layers in encoder-decoder models.

219 + """

220 +

221 + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

222 + # For cross-attention, we need to cache encoder states

223 + # Get encoder length (e.g., 1500 for Whisper).

224 + max_encoder_len = MULTIMODAL_REGISTRY.\

225 + get_encdec_max_encoder_len(vllm_config.model_config)

226 + return cdiv(max_encoder_len, self.block_size) * self.page_size_bytes

227 +

228 +

229 @dataclass

230 class KVCacheTensor:

231 """