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

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 10, 10:15 UTC

[Feat][Core] Support multiple KV cache groups in Hybrid KV Coordinator (#31707)

Signed-off-by: Yifan Qiao <yifanqiao@berkeley.edu>

Yifan Qiao committed Jan 9, 2026 at 10:53 UTC cd4a95e3aa0478d441964fa9afc062febffd65bf

2 files changed +324 -127

tests/v1/core/test_prefix_caching.py

+235 -4

@@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import (

35 FullAttentionSpec,

36 KVCacheConfig,

37 KVCacheGroupSpec,

38 + MambaSpec,

39 SlidingWindowSpec,

40 )

@@ -106,8 +107,23 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:

107

108

109 def make_kv_cache_config_hybrid_model(

109 - block_size: int, num_blocks: int

110 + block_size: int, num_blocks: int, second_spec_type: str = "sliding_window"

111 ) -> KVCacheConfig:

112 + if second_spec_type == "sliding_window":

113 + second_spec = SlidingWindowSpec(

114 + block_size=block_size,

115 + num_kv_heads=1,

116 + head_size=1,

117 + dtype=torch.float32,

118 + sliding_window=2 * block_size,

119 + )

120 + elif second_spec_type == "mamba":

121 + second_spec = MambaSpec(

122 + block_size=block_size,

123 + shapes=(1, 1),

124 + dtypes=(torch.float32,),

125 + )

126 +

127 return KVCacheConfig(

128 num_blocks=num_blocks,

129 kv_cache_tensors=[],

@@ -123,16 +139,49 @@ def make_kv_cache_config_hybrid_model(

139 ),

140 KVCacheGroupSpec(

141 ["layer2"],

126 - SlidingWindowSpec(

142 + second_spec,

143 + ),

144 + KVCacheGroupSpec(

145 + ["layer3"],

146 + second_spec,

147 + ),

148 + ],

149 + )

150 +

151 +

152 +def make_kv_cache_config_three_types(

153 + block_size: int, num_blocks: int, third_spec_type: str = "mamba"

154 +) -> KVCacheConfig:

155 + if third_spec_type == "mamba":

156 + third_spec = MambaSpec(

157 + block_size=block_size,

158 + shapes=(1, 1),

159 + dtypes=(torch.float32,),

160 + )

161 + elif third_spec_type == "sliding_window":

162 + third_spec = SlidingWindowSpec(

163 + block_size=block_size,

164 + num_kv_heads=1,

165 + head_size=1,

166 + dtype=torch.float32,

167 + sliding_window=4 * block_size,

168 + )

169 +

170 + return KVCacheConfig(

171 + num_blocks=num_blocks,

172 + kv_cache_tensors=[],

173 + kv_cache_groups=[

174 + KVCacheGroupSpec(

175 + ["layer1"],

176 + FullAttentionSpec(

177 block_size=block_size,

178 num_kv_heads=1,

179 head_size=1,

180 dtype=torch.float32,

131 - sliding_window=2 * block_size,

181 ),

182 ),

183 KVCacheGroupSpec(

135 - ["layer3"],

184 + ["layer2"],

185 SlidingWindowSpec(

186 block_size=block_size,

187 num_kv_heads=1,

@@ -141,6 +190,10 @@ def make_kv_cache_config_hybrid_model(

190 sliding_window=2 * block_size,

191 ),

192 ),

193 + KVCacheGroupSpec(

194 + ["layer3"],

195 + third_spec,

196 + ),

197 ],

198 )

199

@@ -424,6 +477,184 @@ def test_prefill_hybrid_model():

477 )

478

479

480 +def _make_hybrid_kv_cache_config(

481 + block_size: int, num_blocks: int, spec_types: list[str]

482 +) -> KVCacheConfig:

483 + """

484 + Create a KVCacheConfig with the specified spec types.

485 +

486 + Args:

487 + block_size: The block size for KV cache.

488 + num_blocks: The number of blocks in the KV cache.

489 + spec_types: List of spec type strings. Supported types:

490 + - "full": FullAttentionSpec

491 + - "sliding_window": SlidingWindowSpec with window=2*block_size

492 + - "sliding_window_large": SlidingWindowSpec with window=4*block_size

493 + - "mamba": MambaSpec

494 + """

495 + spec_map = {

496 + "full": lambda: FullAttentionSpec(

497 + block_size=block_size,

498 + num_kv_heads=1,

499 + head_size=1,

500 + dtype=torch.float32,

501 + ),

502 + "sliding_window": lambda: SlidingWindowSpec(

503 + block_size=block_size,

504 + num_kv_heads=1,

505 + head_size=1,

506 + dtype=torch.float32,

507 + sliding_window=2 * block_size,

508 + ),

509 + "sliding_window_large": lambda: SlidingWindowSpec(

510 + block_size=block_size,

511 + num_kv_heads=1,

512 + head_size=1,

513 + dtype=torch.float32,

514 + sliding_window=4 * block_size,

515 + ),

516 + "mamba": lambda: MambaSpec(

517 + block_size=block_size,

518 + shapes=(1, 1),

519 + dtypes=(torch.float32,),

520 + ),

521 + }

522 +

523 + kv_cache_groups = [

524 + KVCacheGroupSpec([f"layer{i}"], spec_map[spec_type]())

525 + for i, spec_type in enumerate(spec_types)

526 + ]

527 +

528 + return KVCacheConfig(

529 + num_blocks=num_blocks,

530 + kv_cache_tensors=[],

531 + kv_cache_groups=kv_cache_groups,

532 + )

533 +

534 +

535 +# Test cases covering various combinations of KV cache spec types:

536 +# - Varying number of groups (2, 3, or 4)

537 +# - 0, 1, or 2 full attention groups

538 +# - Sliding window with different window sizes

539 +# - Interleaved group IDs (full attn and other types mixed)

540 +# - Mamba spec combinations

541 +_HYBRID_MODEL_TEST_CASES = [

542 + # 2 groups: 1 full + 1 other

543 + pytest.param(["full", "sliding_window"], id="2g-full+sw"),

544 + pytest.param(["full", "mamba"], id="2g-full+mamba"),

545 + # 2 groups: 0 full (all other types)

546 + pytest.param(["sliding_window", "mamba"], id="2g-sw+mamba"),

547 + pytest.param(["sliding_window", "sliding_window_large"], id="2g-sw+sw_large"),

548 + # 3 groups: 1 full + 2 others (same type)

549 + pytest.param(["full", "sliding_window", "sliding_window"], id="3g-full+2sw"),

550 + pytest.param(["full", "mamba", "mamba"], id="3g-full+2mamba"),

551 + # 3 groups: 1 full + 2 others (different types)

552 + pytest.param(["full", "sliding_window", "mamba"], id="3g-full+sw+mamba"),

553 + pytest.param(

554 + ["full", "sliding_window", "sliding_window_large"],

555 + id="3g-full+sw+sw_large",

556 + ),

557 + # 3 groups: 2 full + 1 other

558 + pytest.param(["full", "full", "sliding_window"], id="3g-2full+sw"),

559 + pytest.param(["full", "full", "mamba"], id="3g-2full+mamba"),

560 + # 4 groups: interleaved (full, other, full, other)

561 + pytest.param(

562 + ["full", "sliding_window", "full", "sliding_window_large"],

563 + id="4g-interleaved-full+sw+sw_large",

564 + ),

565 + pytest.param(

566 + ["full", "mamba", "full", "mamba"],

567 + id="4g-interleaved-full+mamba",

568 + ),

569 + # 4 groups: interleaved with different sliding windows

570 + pytest.param(

571 + ["full", "sliding_window", "full", "sliding_window_large"],

572 + id="4g-interleaved-full+sw_mixed",

573 + ),

574 + # 4 groups: 0 full (all other types)

575 + pytest.param(

576 + ["sliding_window", "mamba", "sliding_window_large", "mamba"],

577 + id="4g-sw+mamba+sw_large+mamba",

578 + ),

579 + # 4 groups: 2 full + 2 others (grouped)

580 + pytest.param(

581 + ["full", "full", "sliding_window", "mamba"],

582 + id="4g-2full+sw+mamba",

583 + ),

584 +]

585 +

586 +

587 +@pytest.mark.parametrize("spec_types", _HYBRID_MODEL_TEST_CASES)

588 +def test_prefill_hybrid_model_combinations(spec_types: list[str]):

589 + """

590 + Test prefix caching with hybrid models containing various combinations of

591 + KV cache spec types.

592 +

593 + This unified test covers:

594 + - Various combinations (full attn + other attn types)

595 + - Varying number of groups (2, 3, or 4)

596 + - 0, 1, or 2 full attention groups in the combination

597 + - Two sliding_window attn groups with different window sizes

598 + - Interleaved group IDs (full attn and other types alternating)

599 + - Mamba spec with other attention types

600 + """

601 + block_size = 16

602 + num_groups = len(spec_types)

603 + # Allocate enough blocks for all groups

604 + num_blocks = 10 * num_groups

605 +

606 + kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types)

607 + manager = KVCacheManager(

608 + kv_cache_config,

609 + max_model_len=8192,

610 + enable_caching=True,

611 + hash_block_size=block_size,

612 + )

613 +

614 + hash_fn = sha256

615 +

616 + # Complete 3 blocks (48 tokens)

617 + common_token_ids = [i for i in range(3) for _ in range(block_size)]

618 + unique_token_ids = [3] * 7

619 + all_token_ids = common_token_ids + unique_token_ids

620 +

621 + # First request: no cache hit initially

622 + req0 = make_request("0", all_token_ids, block_size, hash_fn)

623 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

624 +

625 + assert len(req0.block_hashes) == 3

626 + assert not computed_blocks.blocks[0] # No cache hit initially

627 + assert num_computed_tokens == 0

628 +

629 + blocks = manager.allocate_slots(

630 + req0, 55, len(computed_blocks.blocks[0]) * block_size, computed_blocks

631 + )

632 + assert blocks is not None

633 + # Should have blocks for all groups

634 + assert len(blocks.get_block_ids()) == num_groups

635 +

636 + # Second request: should hit cached blocks for common prefix

637 + req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn)

638 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

639 +

640 + # Should hit cached blocks for all groups

641 + assert num_computed_tokens == 3 * block_size

642 + assert len(computed_blocks.blocks) == num_groups

643 +

644 + # Allocate and verify blocks for second request

645 + blocks = manager.allocate_slots(

646 + req1,

647 + len(common_token_ids) + 5 - num_computed_tokens,

648 + num_computed_tokens,

649 + computed_blocks,

650 + )

651 + assert blocks is not None

652 + assert len(blocks.get_block_ids()) == num_groups

653 +

654 + manager.free(req0)

655 + manager.free(req1)

656 +

657 +

658 def test_prefill_plp():

659 """Test prefill with APC and some prompt logprobs (plp) requests.

660

vllm/v1/core/kv_cache_coordinator.py

+89 -123

@@ -14,7 +14,7 @@ from vllm.v1.core.kv_cache_utils import (

14 )

15 from vllm.v1.core.single_type_kv_cache_manager import (

16 CrossAttentionManager,

17 - FullAttentionManager,

17 + SingleTypeKVCacheManager,

18 get_manager_for_kv_cache_spec,

19 )

20 from vllm.v1.kv_cache_interface import (

@@ -354,9 +354,6 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

354 """

355 KV cache coordinator for hybrid models with multiple KV cache types, and

356 thus multiple kv cache groups.

357 - To simplify `find_longest_cache_hit`, it only supports the combination of

358 - two types of KV cache groups, and one of them must be full attention.

359 - May extend to more general cases in the future.

357 """

358

359 def __init__(

@@ -397,70 +394,46 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

394

395 def verify_and_split_kv_cache_groups(self) -> None:

396 """

400 - Verifies that the model has exactly two types of KV cache groups, and

401 - one of them is full attention. Then, split the kv cache groups into full

402 - attention groups and other groups.

397 + Groups KV cache groups by their spec type for efficient batch processing

398 + during cache hit lookup.

399 """

404 - full_attention_spec: FullAttentionSpec | None = None

405 - other_spec: KVCacheSpec | None = None

406 - self.full_attention_group_ids: list[int] = []

407 - self.other_group_ids: list[int] = []

400 + attention_groups: list[

401 + tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]]

402 + ] = []

403 +

404 for i, g in enumerate(self.kv_cache_config.kv_cache_groups):

409 - if isinstance(g.kv_cache_spec, FullAttentionSpec):

410 - if full_attention_spec is None:

411 - full_attention_spec = g.kv_cache_spec

412 - else:

413 - assert full_attention_spec == g.kv_cache_spec, (

414 - "HybridKVCacheCoordinator assumes exactly one type of "

415 - "full attention groups now."

405 + manager_cls = self.single_type_managers[i].__class__

406 + spec = g.kv_cache_spec

407 +

408 + # Try to find an existing group with the same spec

409 + for existing_spec, group_ids, existing_cls in attention_groups:

410 + if existing_spec == spec:

411 + assert manager_cls is existing_cls, (

412 + "Expected same manager class for identical KV cache specs."

413 )

417 - self.full_attention_group_ids.append(i)

414 + group_ids.append(i)

415 + break

416 else:

419 - if other_spec is None:

420 - other_spec = g.kv_cache_spec

421 - else:

422 - assert other_spec == g.kv_cache_spec, (

423 - "HybridKVCacheCoordinator assumes "

424 - "exactly one other type of groups now."

425 - )

426 - self.other_group_ids.append(i)

417 + attention_groups.append((spec, [i], manager_cls))

418

428 - assert full_attention_spec is not None, (

429 - "HybridKVCacheCoordinator assumes exactly one type of full "

430 - "attention groups now."

419 + assert len(attention_groups) > 1, (

420 + "HybridKVCacheCoordinator requires at least two attention groups."

421 )

432 - assert other_spec is not None, (

433 - "HybridKVCacheCoordinator assumes exactly one type of other groups now."

422 +

423 + # Put full attention first: its efficient left-to-right scan provides

424 + # a tighter initial bound, reducing work for subsequent groups.

425 + self.attention_groups = sorted(

426 + attention_groups,

427 + key=lambda x: not isinstance(x[0], FullAttentionSpec),

428 )

429

436 - self.full_attention_manager_cls = FullAttentionManager

437 - self.other_attention_cls = self.single_type_managers[

438 - self.other_group_ids[0]

439 - ].__class__

440 - self.full_attention_spec = full_attention_spec

441 - self.other_spec = other_spec

442 - self.full_attention_block_size = self.full_attention_spec.block_size

443 - self.other_block_size = self.other_spec.block_size

444 - # The LCM of the block sizes of full attention and other attention.

430 + # The LCM of the block sizes of all attention types.

431 # The cache hit length must be a multiple of the LCM of the block sizes

432 # to make sure the cache hit length is a multiple of the block size of

433 # each attention type. Requiring this because we don't support partial

434 # block cache hit yet.

449 - self.lcm_block_size = lcm(self.full_attention_block_size, self.other_block_size)

450 -

451 - if max(self.full_attention_group_ids) < min(self.other_group_ids):

452 - self.full_attn_first = True

453 - elif max(self.other_group_ids) < min(self.full_attention_group_ids):

454 - self.full_attn_first = False

455 - else:

456 - raise ValueError(

457 - "HybridKVCacheCoordinator assumes the full "

458 - "attention group ids and other attention group ids "

459 - "do not interleave, either full attention group ids "

460 - "are before other attention group ids or vice versa."

461 - "This is for simplifying merging hit_blocks_full_attn and "

462 - "hit_blocks_other_attn to hit_blocks."

463 - )

435 + block_sizes = [spec.block_size for spec, _, _ in attention_groups]

436 + self.lcm_block_size = lcm(*block_sizes)

437

438 def find_longest_cache_hit(

439 self,

@@ -468,7 +441,12 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

441 max_cache_hit_length: int,

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

443 """

471 - Find the longest cache hit for the request.

444 + Find the longest cache hit using an iterative fixed-point algorithm.

445 +

446 + Each attention type either accepts the current candidate length or

447 + reduces it. If any type reduces the length, restart checks over all

448 + types. This converges because length monotonically decreases and is

449 + bounded below by 0.

450

451 Args:

452 block_hashes: The block hashes of the request.

@@ -476,75 +454,63 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

454

455 Returns:

456 A tuple containing:

479 - - A list of the cache hit blocks for each single type manager.

457 + - A tuple of the cache hit blocks for each single type manager.

458 - The number of tokens of the longest cache hit.

459 """

482 - # First, find the longest cache hit for full attention.

483 - if self.full_attention_spec.block_size == self.hash_block_size:

484 - # Common case.

485 - full_attention_block_hashes: BlockHashList = block_hashes

486 - else:

487 - # block_size is a multiple of hash_block_size. This happens when different

488 - # KV cache groups have different block sizes. In this case, we need to

489 - # recalculate block_hashes at the granularity of block_size, using the

490 - # original block_hashes (at the granularity of hash_block_size).

491 - full_attention_block_hashes = BlockHashListWithBlockSize(

492 - block_hashes, self.hash_block_size, self.full_attention_spec.block_size

493 - )

494 - hit_blocks_full_attn = self.full_attention_manager_cls.find_longest_cache_hit(

495 - block_hashes=full_attention_block_hashes,

496 - max_length=max_cache_hit_length,

497 - kv_cache_group_ids=self.full_attention_group_ids,

498 - block_pool=self.block_pool,

499 - kv_cache_spec=self.full_attention_spec,

500 - use_eagle=self.use_eagle,

501 - alignment_tokens=self.lcm_block_size,

502 - )

503 - hit_length = len(hit_blocks_full_attn[0]) * self.full_attention_block_size

504 -

505 - # Next, find the cache hit for the other attention WITHIN

506 - # the cache hit of full attention.

507 - if self.other_spec.block_size == self.hash_block_size:

508 - # Common case.

509 - other_block_hashes: BlockHashList = block_hashes

510 - else:

511 - # Similar to the full attention case, here we need to recalculate

512 - # block_hashes at the granularity of block_size, using the original

513 - # block_hashes (at the granularity of hash_block_size).

514 - other_block_hashes = BlockHashListWithBlockSize(

515 - block_hashes, self.hash_block_size, self.other_spec.block_size

460 +

461 + def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:

462 + if kv_cache_spec.block_size == self.hash_block_size:

463 + return block_hashes

464 + return BlockHashListWithBlockSize(

465 + block_hashes, self.hash_block_size, kv_cache_spec.block_size

466 )

517 - hit_blocks_other_attn = self.other_attention_cls.find_longest_cache_hit(

518 - block_hashes=other_block_hashes,

519 - max_length=hit_length,

520 - kv_cache_group_ids=self.other_group_ids,

521 - block_pool=self.block_pool,

522 - kv_cache_spec=self.other_spec,

523 - use_eagle=self.use_eagle,

524 - alignment_tokens=self.lcm_block_size,

525 - )

526 - hit_length = len(hit_blocks_other_attn[0]) * self.other_block_size

527 -

528 - # NOTE: the prefix cache hit length must be a multiple of block_size as

529 - # we don't support partial block cache hit yet. The cache hit length

530 - # of other attention is ensured to be a multiple of the block size of

531 - # full attention layers in current implementation, because hit_length is

532 - # a multiple of other attention's block size, and other attention's

533 - # block size is a multiple of full attention's block size (verified in

534 - # `verify_and_split_kv_cache_groups`).

535 - assert hit_length % self.full_attention_block_size == 0

536 -

537 - # Truncate the full attention cache hit to the length of the

538 - # cache hit of the other attention.

539 - for group_hit_blocks in hit_blocks_full_attn:

540 - del group_hit_blocks[hit_length // self.full_attention_block_size :]

541 -

542 - # Merge the hit blocks of full attention and other attention.

543 - if self.full_attn_first:

544 - hit_blocks = hit_blocks_full_attn + hit_blocks_other_attn

545 - else:

546 - hit_blocks = hit_blocks_other_attn + hit_blocks_full_attn

547 - return hit_blocks, hit_length

467 +

468 + num_groups = len(self.kv_cache_config.kv_cache_groups)

469 + hit_length = max_cache_hit_length

470 + hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups

471 +

472 + while True:

473 + curr_hit_length = hit_length

474 +

475 + for spec, group_ids, manager_cls in self.attention_groups:

476 + is_full_attn = isinstance(spec, FullAttentionSpec)

477 +

478 + # Full attention: reuse cached blocks (downward-closed property)

479 + cached_blocks = hit_blocks_by_group[group_ids[0]]

480 + if is_full_attn and cached_blocks is not None:

481 + # For full attention, we only need to compute the cache hit

482 + # length once. Starting from the second iteration, if the

483 + # curr_hit_length is reduced by other groups, we can simply

484 + # keep the first (curr_hit_length // block_size) blocks from

485 + # the last iteration.

486 + num_blocks = curr_hit_length // spec.block_size

487 + curr_hit_length = num_blocks * spec.block_size

488 + for group_id in group_ids:

489 + blocks = hit_blocks_by_group[group_id]

490 + assert blocks is not None

491 + del blocks[num_blocks:]

492 + else:

493 + hit_blocks = manager_cls.find_longest_cache_hit(

494 + block_hashes=_get_block_hashes(spec),

495 + max_length=curr_hit_length,

496 + kv_cache_group_ids=group_ids,

497 + block_pool=self.block_pool,

498 + kv_cache_spec=spec,

499 + use_eagle=self.use_eagle,

500 + alignment_tokens=self.lcm_block_size,

501 + )

502 + curr_hit_length = len(hit_blocks[0]) * spec.block_size

503 + for group_id, blocks in zip(group_ids, hit_blocks):

504 + hit_blocks_by_group[group_id] = blocks

505 +

506 + if curr_hit_length < hit_length:

507 + hit_length = curr_hit_length

508 + else:

509 + break

510 +

511 + return tuple(

512 + blocks if blocks is not None else [] for blocks in hit_blocks_by_group

513 + ), hit_length

514

515

516 def get_kv_cache_coordinator(