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

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

[Fix] prefix cache hit rate == 0 bug with gpt-oss style models (#33524)

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

Yifan Qiao committed Feb 1, 2026 at 17:59 UTC a01ef3fa51a0312914ca60ecbab2bc45aa43e32c

2 files changed +366 -47

tests/v1/core/test_prefix_caching.py

+343 -40

@@ -107,7 +107,10 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:

107

108

109 def make_kv_cache_config_hybrid_model(

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

110 + block_size: int,

111 + num_blocks: int,

112 + sliding_window_blocks: int,

113 + second_spec_type: str = "sliding_window",

114 ) -> KVCacheConfig:

115 if second_spec_type == "sliding_window":

116 second_spec = SlidingWindowSpec(

@@ -115,7 +118,7 @@ def make_kv_cache_config_hybrid_model(

118 num_kv_heads=1,

119 head_size=1,

120 dtype=torch.float32,

118 - sliding_window=2 * block_size,

121 + sliding_window=sliding_window_blocks * block_size,

122 )

123 elif second_spec_type == "mamba":

124 second_spec = MambaSpec(

@@ -325,7 +328,7 @@ def test_prefill(hash_fn):

328 def test_prefill_hybrid_model():

329 block_size = 16

330 manager = KVCacheManager(

328 - make_kv_cache_config_hybrid_model(block_size, 21),

331 + make_kv_cache_config_hybrid_model(block_size, 21, 2),

332 max_model_len=8192,

333 enable_caching=True,

334 hash_block_size=block_size,

@@ -334,7 +337,8 @@ def test_prefill_hybrid_model():

337 hash_fn = sha256

338

339 # Complete 3 blocks (48 tokens)

337 - common_token_ids = [i for i in range(3) for _ in range(block_size)]

340 + num_full_blocks = 3

341 + common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]

342

343 # Fully cache miss

344 # Incomplete 1 block (7 tokens)

@@ -375,6 +379,7 @@ def test_prefill_hybrid_model():

379 # Cache hit in the common prefix

380 # Incomplete 1 block (5 tokens)

381 unique_token_ids = [3] * 5

382 + all_token_ids = common_token_ids + unique_token_ids

383 req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)

384 computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

385 assert len(req1.block_hashes) == 3

@@ -394,34 +399,13 @@ def test_prefill_hybrid_model():

399 manager.free(req0)

400 manager.free(req1)

401

397 - cached_block_hash_to_block_bak = copy.copy(

398 - manager.block_pool.cached_block_hash_to_block._cache

399 - )

400 -

401 - def test_partial_request_hit(

402 - request_id: str,

403 - hash_to_evict: list[BlockHashWithGroupId],

404 - expect_hit_length: int,

405 - ):

406 - req = make_request(

407 - request_id, common_token_ids + unique_token_ids, block_size, sha256

408 - )

409 - for hash_with_group_id in hash_to_evict:

410 - manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)

411 - computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)

412 - assert len(req.block_hashes) == 3

413 - assert num_computed_tokens == expect_hit_length * block_size

414 - for block_per_group in computed_blocks.blocks:

415 - assert len(block_per_group) == num_computed_tokens // block_size

416 - for hash_with_group_id in hash_to_evict:

417 - manager.block_pool.cached_block_hash_to_block._cache[hash_with_group_id] = (

418 - cached_block_hash_to_block_bak[hash_with_group_id]

419 - )

420 - manager.free(req)

421 -

402 # Evict the blocks outside sliding window, does not affect the hit length.

423 - test_partial_request_hit(

403 + _test_partial_request_hit(

404 + manager,

405 + block_size,

406 + num_full_blocks,

407 "2",

408 + all_token_ids,

409 [

410 make_block_hash_with_group_id(block_hashes[0], 1),

411 make_block_hash_with_group_id(block_hashes[0], 2),

@@ -430,13 +414,23 @@ def test_prefill_hybrid_model():

414 )

415

416 # Evict the first block of full attention, makes total cache miss.

433 - test_partial_request_hit(

434 - "3", [make_block_hash_with_group_id(block_hashes[0], 0)], 0

417 + _test_partial_request_hit(

418 + manager,

419 + block_size,

420 + num_full_blocks,

421 + "3",

422 + all_token_ids,

423 + [make_block_hash_with_group_id(block_hashes[0], 0)],

424 + 0,

425 )

426

427 # Evict the last block of all layers, reduces the hit length to 2.

438 - test_partial_request_hit(

428 + _test_partial_request_hit(

429 + manager,

430 + block_size,

431 + num_full_blocks,

432 "4",

433 + all_token_ids,

434 [

435 make_block_hash_with_group_id(block_hashes[2], 0),

436 make_block_hash_with_group_id(block_hashes[2], 1),

@@ -446,18 +440,36 @@ def test_prefill_hybrid_model():

440 )

441

442 # Evict the last block of full attention, reduces the hit length to 2.

449 - test_partial_request_hit(

450 - "5", [make_block_hash_with_group_id(block_hashes[2], 0)], 2

443 + _test_partial_request_hit(

444 + manager,

445 + block_size,

446 + num_full_blocks,

447 + "5",

448 + all_token_ids,

449 + [make_block_hash_with_group_id(block_hashes[2], 0)],

450 + 2,

451 )

452

453 # Evict the last block of sliding window, reduces the hit length to 2.

454 - test_partial_request_hit(

455 - "6", [make_block_hash_with_group_id(block_hashes[2], 1)], 2

454 + _test_partial_request_hit(

455 + manager,

456 + block_size,

457 + num_full_blocks,

458 + "6",

459 + all_token_ids,

460 + [make_block_hash_with_group_id(block_hashes[2], 1)],

461 + 2,

462 )

463

464 # Evict the last block of sliding window, reduces the hit length to 2.

459 - test_partial_request_hit(

460 - "7", [make_block_hash_with_group_id(block_hashes[2], 2)], 2

465 + _test_partial_request_hit(

466 + manager,

467 + block_size,

468 + num_full_blocks,

469 + "7",

470 + all_token_ids,

471 + [make_block_hash_with_group_id(block_hashes[2], 2)],

472 + 2,

473 )

474

475 # Evict different set of blocks for full attention and sliding window makes

@@ -466,8 +478,12 @@ def test_prefill_hybrid_model():

478 # The cache hit length of sliding window is 2 * block_size.

479 # Then it is cache miss as the two type of layers

480 # have different hit length.

469 - test_partial_request_hit(

481 + _test_partial_request_hit(

482 + manager,

483 + block_size,

484 + num_full_blocks,

485 "8",

486 + all_token_ids,

487 [

488 make_block_hash_with_group_id(block_hashes[2], 0),

489 make_block_hash_with_group_id(block_hashes[0], 1),

@@ -477,6 +493,214 @@ def test_prefill_hybrid_model():

493 )

494

495

496 +def test_prefill_hybrid_model_eagle():

497 + block_size = 16

498 + kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 3)

499 + manager = KVCacheManager(

500 + kv_cache_config,

501 + max_model_len=8192,

502 + enable_caching=True,

503 + hash_block_size=block_size,

504 + use_eagle=True,

505 + )

506 +

507 + hash_fn = sha256

508 +

509 + # Complete 6 blocks (96 tokens)

510 + num_full_blocks = 6

511 + common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]

512 +

513 + # Fully cache miss

514 + # Incomplete 1 block (7 tokens)

515 + unique_token_ids = [6] * 7

516 + all_token_ids = common_token_ids + unique_token_ids

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

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

519 + assert len(req0.block_hashes) == len(all_token_ids) // block_size

520 + assert not computed_blocks.blocks[0]

521 + assert num_computed_tokens == 0

522 + blocks = manager.allocate_slots(

523 + req0, len(all_token_ids), num_computed_tokens, computed_blocks

524 + )

525 + block_ids = (

526 + [1, 2, 3, 4, 5, 6, 7],

527 + [8, 9, 10, 11, 12, 13, 14],

528 + [15, 16, 17, 18, 19, 20, 21],

529 + )

530 + assert blocks is not None and blocks.get_block_ids() == block_ids

531 +

532 + # Check full block metadata

533 + parent_block_hash = None

534 + for i, full_block_ids in enumerate(zip(*(row[:-1] for row in block_ids))):

535 + block_tokens = tuple(all_token_ids[i * block_size : (i + 1) * block_size])

536 + block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_tokens)

537 + for group_id, block_id in enumerate(full_block_ids):

538 + blk_hash = manager.block_pool.blocks[block_id].block_hash

539 + assert blk_hash is not None

540 + assert get_block_hash(blk_hash) == block_hash

541 + assert get_group_id(blk_hash) == group_id

542 + assert manager.block_pool.blocks[block_id].ref_cnt == 1

543 + parent_block_hash = block_hash

544 +

545 + # Check partial block metadata

546 + for partial_block_id in (row[-1] for row in block_ids):

547 + assert manager.block_pool.blocks[partial_block_id].block_hash is None

548 + assert manager.block_pool.blocks[partial_block_id].ref_cnt == 1

549 +

550 + # Cache hit in the common prefix

551 + # Incomplete 1 block (5 tokens)

552 + unique_token_ids = [6] * 5

553 + all_token_ids = common_token_ids + unique_token_ids

554 + req1 = make_request("1", all_token_ids, block_size, hash_fn)

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

556 + assert len(req1.block_hashes) == num_full_blocks

557 + assert computed_blocks.get_block_ids() == (

558 + [1, 2, 3, 4],

559 + [0, 9, 10, 11],

560 + [0, 16, 17, 18],

561 + )

562 + assert num_computed_tokens == 4 * block_size

563 + num_new_tokens = len(all_token_ids) - num_computed_tokens

564 + blocks = manager.allocate_slots(

565 + req1, num_new_tokens, num_computed_tokens, computed_blocks

566 + )

567 + assert blocks is not None and blocks.get_block_ids() == (

568 + [22, 23, 24],

569 + [25, 26, 27],

570 + [28, 29, 30],

571 + )

572 + for block_per_group in computed_blocks.blocks:

573 + for block in block_per_group:

574 + if block != manager.block_pool.null_block:

575 + assert block.ref_cnt == 2

576 +

577 + block_hashes = req1.block_hashes

578 + manager.free(req0)

579 + manager.free(req1)

580 +

581 + # Evict the blocks outside sliding window, does not affect the hit length.

582 + _test_partial_request_hit(

583 + manager,

584 + block_size,

585 + num_full_blocks,

586 + "2",

587 + all_token_ids,

588 + [

589 + make_block_hash_with_group_id(block_hashes[0], 1),

590 + make_block_hash_with_group_id(block_hashes[0], 2),

591 + ],

592 + 4,

593 + )

594 +

595 + # Evict the first block of full attention, makes total cache miss.

596 + _test_partial_request_hit(

597 + manager,

598 + block_size,

599 + num_full_blocks,

600 + "3",

601 + all_token_ids,

602 + [make_block_hash_with_group_id(block_hashes[0], 0)],

603 + 0,

604 + )

605 +

606 + # Evict the last block of all layers, reduces the hit length to 3.

607 + _test_partial_request_hit(

608 + manager,

609 + block_size,

610 + num_full_blocks,

611 + "4",

612 + all_token_ids,

613 + [

614 + make_block_hash_with_group_id(block_hashes[-1], 0),

615 + make_block_hash_with_group_id(block_hashes[-1], 1),

616 + make_block_hash_with_group_id(block_hashes[-1], 2),

617 + ],

618 + 3,

619 + )

620 +

621 + # Evict the last block of full attention, reduces the hit length to 3.

622 + _test_partial_request_hit(

623 + manager,

624 + block_size,

625 + num_full_blocks,

626 + "5",

627 + all_token_ids,

628 + [make_block_hash_with_group_id(block_hashes[-1], 0)],

629 + 3,

630 + )

631 +

632 + # Since the last block of full attention is dropped for eagle, evict

633 + # the second last block of sliding window, reduces the hit length to 3.

634 + _test_partial_request_hit(

635 + manager,

636 + block_size,

637 + num_full_blocks,

638 + "6",

639 + all_token_ids,

640 + [make_block_hash_with_group_id(block_hashes[-2], 1)],

641 + 3,

642 + )

643 +

644 + # Since the last block of full attention is dropped for eagle, evict

645 + # the second last block of sliding window, reduces the hit length to 3.

646 + _test_partial_request_hit(

647 + manager,

648 + block_size,

649 + num_full_blocks,

650 + "7",

651 + all_token_ids,

652 + [make_block_hash_with_group_id(block_hashes[-2], 2)],

653 + 3,

654 + )

655 +

656 + # Evict different set of blocks for full attention and sliding window makes

657 + # total cache miss.

658 + # The cache hit length of full attention is 4 * block_size.

659 + # The cache hit length of sliding window is 3 * block_size.

660 + # Then it is cache miss as the two type of layers

661 + # have different hit length.

662 + _test_partial_request_hit(

663 + manager,

664 + block_size,

665 + num_full_blocks,

666 + "8",

667 + all_token_ids,

668 + [

669 + make_block_hash_with_group_id(block_hashes[-1], 0),

670 + make_block_hash_with_group_id(block_hashes[0], 1),

671 + make_block_hash_with_group_id(block_hashes[0], 2),

672 + ],

673 + 0,

674 + )

675 +

676 +

677 +def _test_partial_request_hit(

678 + manager: KVCacheManager,

679 + block_size: int,

680 + num_full_blocks,

681 + request_id: str,

682 + prompt_token_ids: list[int],

683 + hash_to_evict: list[BlockHashWithGroupId],

684 + expect_hit_length: int,

685 +):

686 + cached_block_hash_to_block_bak = copy.copy(

687 + manager.block_pool.cached_block_hash_to_block._cache

688 + )

689 + req = make_request(request_id, prompt_token_ids, block_size, sha256)

690 + for hash_with_group_id in hash_to_evict:

691 + manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)

692 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)

693 + assert len(req.block_hashes) == num_full_blocks

694 + assert num_computed_tokens == expect_hit_length * block_size

695 + for block_per_group in computed_blocks.blocks:

696 + assert len(block_per_group) == num_computed_tokens // block_size

697 + for hash_with_group_id in hash_to_evict:

698 + manager.block_pool.cached_block_hash_to_block._cache[hash_with_group_id] = (

699 + cached_block_hash_to_block_bak[hash_with_group_id]

700 + )

701 + manager.free(req)

702 +

703 +

704 def _make_hybrid_kv_cache_config(

705 block_size: int, num_blocks: int, spec_types: list[str]

706 ) -> KVCacheConfig:

@@ -655,6 +879,85 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):

879 manager.free(req1)

880

881

882 +# Test cases with eagle enabled: Only test a single simple case for now.

883 +# - 2 groups: 1 full + 1 other

884 +_EAGLE_HYBRID_MODEL_TEST_CASES = [

885 + # 2 groups: 1 full + 1 other

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

887 +]

888 +

889 +

890 +@pytest.mark.parametrize("spec_types,expect_hit_length", _EAGLE_HYBRID_MODEL_TEST_CASES)

891 +def test_prefill_hybrid_model_combinations_eagle(

892 + spec_types: list[str], expect_hit_length: int

893 +):

894 + """

895 + Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE.

896 + More complex hybrid models with EAGLE are not yet supported (see issue #32802).

897 + """

898 + block_size = 16

899 + num_groups = len(spec_types)

900 + # Allocate enough blocks for all groups

901 + num_blocks = 10 * num_groups

902 +

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

904 + manager = KVCacheManager(

905 + kv_cache_config,

906 + max_model_len=8192,

907 + enable_caching=True,

908 + hash_block_size=block_size,

909 + use_eagle=True,

910 + )

911 +

912 + hash_fn = sha256

913 +

914 + # Complete 3 blocks (48 tokens)

915 + num_full_blocks = 4

916 + common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]

917 + unique_token_ids = [4] * 7

918 + all_token_ids = common_token_ids + unique_token_ids

919 +

920 + # First request: no cache hit initially

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

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

923 +

924 + assert len(req0.block_hashes) == num_full_blocks

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

926 + assert num_computed_tokens == 0

927 +

928 + blocks = manager.allocate_slots(

929 + req0, len(all_token_ids), num_computed_tokens, computed_blocks

930 + )

931 + assert blocks is not None

932 + # Should have blocks for all groups

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

934 +

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

936 + all_token_ids = common_token_ids + [6] * 5

937 + req1 = make_request("1", all_token_ids, block_size, hash_fn)

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

939 +

940 + # Should hit cached blocks for all groups

941 + assert num_computed_tokens == expect_hit_length * block_size

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

943 + # Verify each group has the correct number of computed blocks

944 + for block_per_group in computed_blocks.blocks:

945 + assert len(block_per_group) == expect_hit_length

946 +

947 + # Allocate and verify blocks for second request

948 + blocks = manager.allocate_slots(

949 + req1,

950 + len(all_token_ids) - num_computed_tokens,

951 + num_computed_tokens,

952 + computed_blocks,

953 + )

954 + assert blocks is not None

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

956 +

957 + manager.free(req0)

958 + manager.free(req1)

959 +

960 +

961 def test_prefill_plp():

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

963

vllm/v1/core/kv_cache_coordinator.py

+23 -7

@@ -479,6 +479,16 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

479 hit_length = max_cache_hit_length

480 hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups

481

482 + # Simple hybrid (1 full attn + 1 other): one iteration suffices.

483 + # Full attn is always first if it exists. This avoids EAGLE drops

484 + # being applied multiple times to non-full-attn groups.

485 + # FIXME (yifan): However, for complex hybrid models with multiple attn

486 + # groups, we still have the EAGLE spiral block dropping problem. See

487 + # discussion in issue https://github.com/vllm-project/vllm/issues/32802.

488 + is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(

489 + self.attention_groups[0][0], FullAttentionSpec

490 + )

491 +

492 while True:

493 curr_hit_length = hit_length

494

@@ -495,10 +505,6 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

505 # the last iteration.

506 num_blocks = curr_hit_length // spec.block_size

507 curr_hit_length = num_blocks * spec.block_size

498 - for group_id in group_ids:

499 - blocks = hit_blocks_by_group[group_id]

500 - assert blocks is not None

501 - del blocks[num_blocks:]

508 else:

509 hit_blocks = manager_cls.find_longest_cache_hit(

510 block_hashes=_get_block_hashes(spec),

@@ -513,10 +519,20 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

519 for group_id, blocks in zip(group_ids, hit_blocks):

520 hit_blocks_by_group[group_id] = blocks

521

516 - if curr_hit_length < hit_length:

517 - hit_length = curr_hit_length

518 - else:

522 + if curr_hit_length >= hit_length:

523 break

524 + hit_length = curr_hit_length

525 + # Simple hybrid: exit after one iteration

526 + if is_simple_hybrid:

527 + break

528 +

529 + # Truncate full attention blocks to final hit_length (if present)

530 + spec, group_ids, _ = self.attention_groups[0]

531 + if isinstance(spec, FullAttentionSpec):

532 + num_blocks = hit_length // spec.block_size

533 + for group_id in group_ids:

534 + if (blks := hit_blocks_by_group[group_id]) is not None:

535 + del blks[num_blocks:]

536

537 return tuple(

538 blocks if blocks is not None else [] for blocks in hit_blocks_by_group