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

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 2, 03:30 UTC

[Bugfix] Cache the EAGLE/MTP lookahead block in the SWA prefix-cache mask (#44082)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>

Yifan Qiao committed Jun 2, 2026 at 12:21 UTC e9e08c49b966a7d36bd37745590f658e76a2082d

5 files changed +338 -82

tests/v1/core/test_prefix_caching.py

+196 -1

@@ -2466,6 +2466,201 @@ def test_eagle_with_sliding_window():

2466 assert num_tokens == 0

2467

2468

2469 +def test_eagle_swa_alignment_caches_extra_block():

2470 + """Regression: SWA + EAGLE with `sliding_window <= alignment_tokens`.

2471 +

2472 + When the cache-hit alignment (lcm of per-group block sizes) is larger than

2473 + the SWA window, the SWA mask only kept the last block of each aligned

2474 + segment. EAGLE/MTP lookup needs ``tail + 1`` contiguous cached blocks and

2475 + that +1 block lives at the next segment's first position, which was left

2476 + uncached. The fix caches that extra block when ``use_eagle=True``.

2477 + """

2478 + block_size = 8

2479 + # Full group uses 4 * block_size, so lcm/alignment is 4 * block_size.

2480 + # SWA group has sliding_window = block_size (i.e., tail = 1 block).

2481 + # Without the fix, the second cached block needed for the EAGLE 2-block

2482 + # match never exists -> EAGLE cache hit fails entirely.

2483 + kv_cache_config = KVCacheConfig(

2484 + num_blocks=100,

2485 + kv_cache_tensors=[],

2486 + kv_cache_groups=[

2487 + KVCacheGroupSpec(

2488 + ["full"],

2489 + FullAttentionSpec(

2490 + block_size=4 * block_size,

2491 + num_kv_heads=1,

2492 + head_size=1,

2493 + dtype=torch.float16,

2494 + ),

2495 + ),

2496 + KVCacheGroupSpec(

2497 + ["swa_mtp"],

2498 + SlidingWindowSpec(

2499 + block_size=block_size,

2500 + num_kv_heads=1,

2501 + head_size=1,

2502 + dtype=torch.float32,

2503 + sliding_window=block_size,

2504 + ),

2505 + is_eagle_group=True,

2506 + ),

2507 + ],

2508 + )

2509 + manager = make_kv_cache_manager(

2510 + kv_cache_config=kv_cache_config,

2511 + max_model_len=8192,

2512 + enable_caching=True,

2513 + hash_block_size=block_size,

2514 + use_eagle=True,

2515 + )

2516 +

2517 + # Prime the cache with a long prompt (16 swa blocks = 4 aligned segments).

2518 + token_ids = [i for i in range(16) for _ in range(block_size)]

2519 + req0 = make_request("0", token_ids, block_size, sha256)

2520 + computed_blocks, _ = manager.get_computed_blocks(req0)

2521 + blocks = manager.allocate_slots(

2522 + req0,

2523 + len(token_ids),

2524 + len(computed_blocks.blocks[0]) * block_size,

2525 + computed_blocks,

2526 + )

2527 + assert blocks is not None

2528 + manager.free(req0)

2529 +

2530 + # Second request with identical prompt should find an EAGLE cache hit.

2531 + # Without the fix, ``num_computed_tokens`` is 0; with the fix, it lands at

2532 + # an alignment boundary (multiple of 32 tokens, minus the EAGLE drop).

2533 + req1 = make_request("1", token_ids, block_size, sha256)

2534 + _, num_computed_tokens = manager.get_computed_blocks(req1)

2535 + assert num_computed_tokens > 0, (

2536 + "EAGLE + SWA with sliding_window <= alignment failed to find any "

2537 + "cache hit; the +1 block past each segment boundary must be cached."

2538 + )

2539 + # Each aligned segment contributes 4 * block_size = 32 tokens; EAGLE drops

2540 + # the last block (block_size tokens) from the hit.

2541 + assert num_computed_tokens % (4 * block_size) == 0

2542 +

2543 +

2544 +def test_eagle_swa_boundary_caches_post_boundary_block():

2545 + """EAGLE + SWA must cache the first block after an alignment boundary.

2546 +

2547 + A 40-token computed prefix with 8-token SWA blocks and 32-token hybrid

2548 + alignment needs SWA blocks 3 and 4 cached to reuse a 32-token prefix:

2549 + block 3 is the segment tail, and block 4 is the EAGLE lookahead block

2550 + that gets dropped after lookup.

2551 + """

2552 + block_size = 8

2553 + kv_cache_config = KVCacheConfig(

2554 + num_blocks=100,

2555 + kv_cache_tensors=[],

2556 + kv_cache_groups=[

2557 + KVCacheGroupSpec(

2558 + ["full"],

2559 + FullAttentionSpec(

2560 + block_size=4 * block_size,

2561 + num_kv_heads=1,

2562 + head_size=1,

2563 + dtype=torch.float16,

2564 + ),

2565 + ),

2566 + KVCacheGroupSpec(

2567 + ["swa_mtp"],

2568 + SlidingWindowSpec(

2569 + block_size=block_size,

2570 + num_kv_heads=1,

2571 + head_size=1,

2572 + dtype=torch.float32,

2573 + sliding_window=block_size,

2574 + ),

2575 + is_eagle_group=True,

2576 + ),

2577 + ],

2578 + )

2579 + manager = make_kv_cache_manager(

2580 + kv_cache_config=kv_cache_config,

2581 + max_model_len=8192,

2582 + enable_caching=True,

2583 + hash_block_size=block_size,

2584 + use_eagle=True,

2585 + )

2586 +

2587 + token_ids = [i for i in range(5) for _ in range(block_size)]

2588 + req0 = make_request("0", token_ids, block_size, sha256)

2589 + computed_blocks, _ = manager.get_computed_blocks(req0)

2590 + blocks = manager.allocate_slots(

2591 + req0,

2592 + len(token_ids),

2593 + len(computed_blocks.blocks[0]) * block_size,

2594 + computed_blocks,

2595 + )

2596 + assert blocks is not None

2597 +

2598 + pool = manager.block_pool

2599 + assert pool.get_cached_block(req0.block_hashes[3], kv_cache_group_ids=[1])

2600 + assert pool.get_cached_block(req0.block_hashes[4], kv_cache_group_ids=[1])

2601 + manager.free(req0)

2602 +

2603 + req1 = make_request("1", token_ids + [999], block_size, sha256)

2604 + _, num_computed_tokens = manager.get_computed_blocks(req1)

2605 + assert num_computed_tokens == 4 * block_size

2606 +

2607 +

2608 +def test_eagle_grouped_swa_siblings_use_same_cache_mask():

2609 + """Grouped SWA siblings must cache the EAGLE lookahead block together."""

2610 + block_size = 8

2611 + swa_spec = SlidingWindowSpec(

2612 + block_size=block_size,

2613 + num_kv_heads=1,

2614 + head_size=1,

2615 + dtype=torch.float32,

2616 + sliding_window=block_size,

2617 + )

2618 + kv_cache_config = KVCacheConfig(

2619 + num_blocks=100,

2620 + kv_cache_tensors=[],

2621 + kv_cache_groups=[

2622 + KVCacheGroupSpec(

2623 + ["full"],

2624 + FullAttentionSpec(

2625 + block_size=4 * block_size,

2626 + num_kv_heads=1,

2627 + head_size=1,

2628 + dtype=torch.float16,

2629 + ),

2630 + ),

2631 + KVCacheGroupSpec(["swa_main"], swa_spec),

2632 + KVCacheGroupSpec(["swa_mtp"], swa_spec, is_eagle_group=True),

2633 + ],

2634 + )

2635 + manager = make_kv_cache_manager(

2636 + kv_cache_config=kv_cache_config,

2637 + max_model_len=8192,

2638 + enable_caching=True,

2639 + hash_block_size=block_size,

2640 + use_eagle=True,

2641 + )

2642 +

2643 + token_ids = [i for i in range(9) for _ in range(block_size)]

2644 + req0 = make_request("0", token_ids, block_size, sha256)

2645 + computed_blocks, _ = manager.get_computed_blocks(req0)

2646 + blocks = manager.allocate_slots(

2647 + req0,

2648 + len(token_ids),

2649 + len(computed_blocks.blocks[0]) * block_size,

2650 + computed_blocks,

2651 + )

2652 + assert blocks is not None

2653 +

2654 + pool = manager.block_pool

2655 + assert pool.get_cached_block(req0.block_hashes[4], kv_cache_group_ids=[1, 2])

2656 + assert pool.get_cached_block(req0.block_hashes[8], kv_cache_group_ids=[1, 2])

2657 + manager.free(req0)

2658 +

2659 + req1 = make_request("1", token_ids + [999], block_size, sha256)

2660 + _, num_computed_tokens = manager.get_computed_blocks(req1)

2661 + assert num_computed_tokens == 8 * block_size

2662 +

2663 +

2664 def test_different_block_size():

2665 block_size = 16

2666 # full attention and sliding window attention layers have the same page size:

@@ -2614,7 +2809,7 @@ def test_hybrid_cache_blocks_swa_tail_window_only():

2809

2810

2811 def test_hybrid_cache_blocks_clamped_to_lcm():

2617 - """HybridKVCacheCoordinator.cache_blocks() clamps to lcm_block_size.

2812 + """HybridKVCacheCoordinator.cache_blocks() clamps to scheduler_block_size.

2813 Chunks past the last lcm-aligned boundary can never participate in a

2814 cache hit (find_longest_cache_hit always returns lcm-aligned hits), so

2815 caching them only pollutes the prefix-cache hash map and keeps blocks

tests/v1/core/test_single_type_kv_cache_manager.py

+2 -2

@@ -86,7 +86,7 @@ def test_chunked_local_attention_possible_cached_prefix():

86 kv_cache_group_ids=[0],

87 block_pool=block_pool,

88 kv_cache_spec=chunked_local_attention_spec,

89 - use_eagle=False,

89 + drop_eagle_block=False,

90 alignment_tokens=block_size,

91 )[0]

92 assert len(computed_blocks) == expect_length

@@ -157,7 +157,7 @@ def test_sliding_window_possible_cached_prefix():

157 kv_cache_group_ids=[0],

158 block_pool=block_pool,

159 kv_cache_spec=sliding_window_spec,

160 - use_eagle=False,

160 + drop_eagle_block=False,

161 alignment_tokens=block_size,

162 )[0]

163 assert len(computed_blocks) == expect_length

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

+5 -5

@@ -233,7 +233,7 @@ class MooncakeStoreCoordinator:

233 kv_cache_group_ids=group_ids,

234 block_pool=cast(BlockPool, cached_block_pool),

235 kv_cache_spec=spec,

236 - use_eagle=(0 in eagle_indices),

236 + drop_eagle_block=(0 in eagle_indices),

237 alignment_tokens=spec.block_size,

238 )

239 num_groups = len(self.kv_cache_groups)

@@ -262,9 +262,9 @@ class MooncakeStoreCoordinator:

262 )

263 continue

264

265 - use_eagle = idx in eagle_indices and idx not in eagle_verified

265 + drop_eagle_block = idx in eagle_indices and idx not in eagle_verified

266 _max_length = curr_hit_length

267 - if use_eagle:

267 + if drop_eagle_block:

268 _max_length = min(curr_hit_length + spec.block_size, max_length)

269 hashes = self.block_hashes_for_spec(block_hashes, spec)

270 hit_blocks = manager_cls.find_longest_cache_hit(

@@ -273,11 +273,11 @@ class MooncakeStoreCoordinator:

273 kv_cache_group_ids=group_ids,

274 block_pool=cast(BlockPool, cached_block_pool),

275 kv_cache_spec=spec,

276 - use_eagle=use_eagle,

276 + drop_eagle_block=drop_eagle_block,

277 alignment_tokens=self.lcm_block_size,

278 )

279 _new_hit_length = len(hit_blocks[0]) * spec.block_size

280 - if use_eagle:

280 + if drop_eagle_block:

281 eagle_verified.add(idx)

282 elif _new_hit_length < curr_hit_length:

283 eagle_verified.clear()

vllm/v1/core/kv_cache_coordinator.py

+61 -38

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

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

3 from abc import ABC, abstractmethod

4 from collections.abc import Sequence

5 +from typing import NamedTuple

7 from vllm.v1.core.block_pool import BlockPool

8 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector

@@ -381,6 +382,8 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

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

383 "UnitaryKVCacheCoordinator assumes only one kv cache group"

384 )

385 + # Single group; useless but just set ``use_eagle`` for consistency regardless.

386 + self.single_type_managers[0].use_eagle = 0 in self.eagle_group_ids

387

388 def find_longest_cache_hit(

389 self,

@@ -393,7 +396,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

396 kv_cache_group_ids=[0],

397 block_pool=self.block_pool,

398 kv_cache_spec=self.kv_cache_spec,

396 - use_eagle=0 in self.eagle_group_ids,

399 + drop_eagle_block=0 in self.eagle_group_ids,

400 alignment_tokens=self.block_size,

401 dcp_world_size=self.dcp_world_size,

402 pcp_world_size=self.pcp_world_size,

@@ -401,6 +404,21 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

404 return hit_blocks, len(hit_blocks[0]) * self.block_size

405

406

407 +class SpecGroup(NamedTuple):

408 + """KV cache groups that share one spec, batched together for a single

409 + cache-hit lookup.

410 +

411 + ``use_eagle`` is True iff any member group is an EAGLE/MTP group. Members

412 + sharing a spec are cached and looked up jointly, so the EAGLE last-block drop

413 + is necessarily decided for the whole spec group.

414 + """

415 +

416 + spec: KVCacheSpec

417 + group_ids: list[int]

418 + manager_cls: type[SingleTypeKVCacheManager]

419 + use_eagle: bool

420 +

421 +

422 class HybridKVCacheCoordinator(KVCacheCoordinator):

423 """

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

@@ -452,58 +470,63 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

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

471 during cache hit lookup.

472 """

455 - attention_groups: list[

456 - tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]]

457 - ] = []

458 -

473 + self.attention_groups: list[SpecGroup] = []

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

475 manager_cls = self.single_type_managers[i].__class__

476 spec = g.kv_cache_spec

477 + use_eagle = i in self.eagle_group_ids

478

479 # Try to find an existing group with the same spec

464 - for existing_spec, group_ids, existing_cls in attention_groups:

465 - if existing_spec == spec:

466 - assert manager_cls is existing_cls, (

480 + for idx, group in enumerate(self.attention_groups):

481 + if group.spec == spec:

482 + assert manager_cls is group.manager_cls, (

483 "Expected same manager class for identical KV cache specs."

484 )

469 - group_ids.append(i)

485 + group.group_ids.append(i)

486 + if use_eagle and not group.use_eagle:

487 + self.attention_groups[idx] = group._replace(use_eagle=True)

488 break

489 else:

472 - attention_groups.append((spec, [i], manager_cls))

490 + self.attention_groups.append(

491 + SpecGroup(spec, [i], manager_cls, use_eagle)

492 + )

493

474 - assert len(attention_groups) > 1, (

494 + assert len(self.attention_groups) > 1, (

495 "HybridKVCacheCoordinator requires at least two attention groups."

496 )

497

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

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

480 - self.attention_groups = sorted(

481 - attention_groups,

482 - key=lambda x: not isinstance(x[0], FullAttentionSpec),

500 + self.attention_groups.sort(

501 + key=lambda g: not isinstance(g.spec, FullAttentionSpec)

502 )

503

485 - # Attention-group indices (into ``self.attention_groups``) that

486 - # contain at least one EAGLE/MTP KV cache group.

487 - self.eagle_attn_group_indices: set[int] = {

488 - i

489 - for i, (_, group_ids, _) in enumerate(self.attention_groups)

490 - if any(gid in self.eagle_group_ids for gid in group_ids)

491 - }

504 + # Propagate the eagle bit to each manager (default to ``use_eagle=False``).

505 + for group in self.attention_groups:

506 + if group.use_eagle:

507 + for gid in group.group_ids:

508 + self.single_type_managers[gid].use_eagle = True

509

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

511 # Cache hits in this coordinator are always a multiple of

512 # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``).

496 - # Within an aligned region, SWA groups only consult a subset of blocks

513 + # Within an aligned region, SWA groups may only consult a subset of blocks

514 # per ``scheduler_block_size``-segment so the unused blocks also stay

515 # out of the prefix-cache hash map.

499 - num_computed_tokens = (

516 + aligned_num_computed_tokens = (

517 num_computed_tokens // self.scheduler_block_size * self.scheduler_block_size

518 )

519 for manager in self.single_type_managers:

520 + num_tokens_to_cache = aligned_num_computed_tokens

521 + # EAGLE groups match one block past each aligned boundary and drop

522 + # it, so make that lookahead block eligible to be cached.

523 + if manager.use_eagle and aligned_num_computed_tokens > 0:

524 + num_tokens_to_cache = min(

525 + num_computed_tokens,

526 + aligned_num_computed_tokens + manager.block_size,

527 + )

528 manager.cache_blocks(

504 - request,

505 - num_computed_tokens,

506 - alignment_tokens=self.scheduler_block_size,

529 + request, num_tokens_to_cache, alignment_tokens=self.scheduler_block_size

530 )

531

532 def find_longest_cache_hit(

@@ -543,7 +566,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

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

567 # Full attn is always first if it exists.

568 is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(

546 - self.attention_groups[0][0], FullAttentionSpec

569 + self.attention_groups[0].spec, FullAttentionSpec

570 )

571

572 # Attention-group indices whose EAGLE drop is verified at the current

@@ -554,7 +577,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

577 while True:

578 curr_hit_length = hit_length

579

557 - for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups):

580 + for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate(

581 + self.attention_groups

582 + ):

583 cached_blocks = hit_blocks_by_group[group_ids[0]]

584 if isinstance(spec, FullAttentionSpec) and cached_blocks is not None:

585 # Full attention is downward-closed: we only need to look

@@ -565,12 +590,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

590 )

591 continue

592

568 - use_eagle = (

569 - idx in self.eagle_attn_group_indices and idx not in eagle_verified

570 - )

593 + drop_eagle_block = use_eagle and idx not in eagle_verified

594

595 _max_length = curr_hit_length

573 - if use_eagle:

596 + if drop_eagle_block:

597 # Eagle needs to match one more block and then pop the last.

598 _max_length = min(

599 curr_hit_length + spec.block_size, max_cache_hit_length

@@ -581,11 +604,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

604 kv_cache_group_ids=group_ids,

605 block_pool=self.block_pool,

606 kv_cache_spec=spec,

584 - use_eagle=use_eagle,

607 + drop_eagle_block=drop_eagle_block,

608 alignment_tokens=self.scheduler_block_size,

609 )

610 _new_hit_length = len(hit_blocks[0]) * spec.block_size

588 - if use_eagle:

611 + if drop_eagle_block:

612 eagle_verified.add(idx)

613 elif _new_hit_length < curr_hit_length:

614 # length shrunk; invalidate previous eagle verifications

@@ -601,10 +624,10 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

624 break

625

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

604 - spec, group_ids, _ = self.attention_groups[0]

605 - if isinstance(spec, FullAttentionSpec):

606 - num_blocks = hit_length // spec.block_size

607 - for group_id in group_ids:

627 + first_group = self.attention_groups[0]

628 + if isinstance(first_group.spec, FullAttentionSpec):

629 + num_blocks = hit_length // first_group.spec.block_size

630 + for group_id in first_group.group_ids:

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

632 del blks[num_blocks:]

633

vllm/v1/core/single_type_kv_cache_manager.py

+74 -36

@@ -86,6 +86,12 @@ class SingleTypeKVCacheManager(ABC):

86 self.kv_cache_group_id = kv_cache_group_id

87 self._null_block = block_pool.null_block

89 + # Whether this group's prefix-cache hits drop the EAGLE/MTP lookahead

90 + # block. Only consulted by managers whose hit logic is sparse within an

91 + # aligned segment (SWA). Initialized lazily by the coordinator after

92 + # determining the attention groups.

93 + self.use_eagle = False

94 +

95 @classmethod

96 def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]):

97 return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks)

@@ -317,8 +323,12 @@ class SingleTypeKVCacheManager(ABC):

323 if alignment_tokens is None or alignment_tokens <= self.block_size:

324 block_mask = None

325 else:

320 - block_mask = self._cache_block_mask(

321 - num_cached_blocks, num_full_blocks, alignment_tokens

326 + block_mask = self.reachable_block_mask(

327 + num_cached_blocks,

328 + num_full_blocks,

329 + alignment_tokens,

330 + self.kv_cache_spec,

331 + self.use_eagle,

332 )

333 self.block_pool.cache_full_blocks(

334 request=request,

@@ -332,11 +342,14 @@ class SingleTypeKVCacheManager(ABC):

342

343 self.num_cached_block[request.request_id] = num_full_blocks

344

335 - def _cache_block_mask(

336 - self,

337 - num_cached_blocks: int,

338 - num_full_blocks: int,

345 + @classmethod

346 + def reachable_block_mask(

347 + cls,

348 + start_block: int,

349 + num_blocks: int,

350 alignment_tokens: int,

351 + kv_cache_spec: KVCacheSpec,

352 + use_eagle: bool,

353 ) -> list[bool] | None:

354 """Per-block mask for ``cache_full_blocks``. ``None`` means cache

355 every (non-null) block — the default for full attention.

@@ -389,7 +402,7 @@ class SingleTypeKVCacheManager(ABC):

402 kv_cache_group_ids: list[int],

403 block_pool: BlockPool,

404 kv_cache_spec: KVCacheSpec,

392 - use_eagle: bool,

405 + drop_eagle_block: bool,

406 alignment_tokens: int,

407 dcp_world_size: int = 1,

408 pcp_world_size: int = 1,

@@ -409,7 +422,10 @@ class SingleTypeKVCacheManager(ABC):

422 kv_cache_group_ids: The ids of the kv cache groups.

423 block_pool: The block pool.

424 kv_cache_spec: The kv cache spec.

412 - use_eagle: Whether to use eagle.

425 + drop_eagle_block: Whether to drop the last matched block for EAGLE/MTP.

426 + Always False for non-EAGLE/MTP groups, but can be False for EAGLE/MTP

427 + groups too if the last block is already dropped (e.g., in a

428 + convergence loop in `find_longest_cache_hit`).

429 alignment_tokens: The returned cache hit length (in tokens) should

430 be a multiple of this value (in tokens). By default, it should

431 be set to the block_size.

@@ -499,7 +515,7 @@ class FullAttentionManager(SingleTypeKVCacheManager):

515 kv_cache_group_ids: list[int],

516 block_pool: BlockPool,

517 kv_cache_spec: KVCacheSpec,

502 - use_eagle: bool,

518 + drop_eagle_block: bool,

519 alignment_tokens: int,

520 dcp_world_size: int = 1,

521 pcp_world_size: int = 1,

@@ -528,7 +544,7 @@ class FullAttentionManager(SingleTypeKVCacheManager):

544 computed.append(cached)

545 else:

546 break

531 - if use_eagle and computed_blocks[0]:

547 + if drop_eagle_block and computed_blocks[0]:

548 # Need to drop the last matched block if eagle is enabled.

549 for computed in computed_blocks:

550 computed.pop()

@@ -556,6 +572,19 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

572 super().__init__(kv_cache_spec, **kwargs)

573 self.sliding_window = kv_cache_spec.sliding_window

574

575 + @classmethod

576 + def _contiguous_blocks_for_hit(

577 + cls, window_size: int, block_size: int, use_eagle: bool

578 + ) -> int:

579 + blocks = cdiv(window_size - 1, block_size)

580 + if use_eagle:

581 + # Need to drop the last matched block if eagle is enabled. For

582 + # sliding window layer, we achieve this by increasing the number of

583 + # contiguous blocks needed for prefix cache hit by one and dropping

584 + # the last matched block.

585 + blocks += 1

586 + return blocks

587 +

588 @classmethod

589 def find_longest_cache_hit(

590 cls,

@@ -564,7 +593,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

593 kv_cache_group_ids: list[int],

594 block_pool: BlockPool,

595 kv_cache_spec: KVCacheSpec,

567 - use_eagle: bool,

596 + drop_eagle_block: bool,

597 alignment_tokens: int,

598 dcp_world_size: int = 1,

599 pcp_world_size: int = 1,

@@ -575,17 +604,10 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

604 assert dcp_world_size == 1, "DCP not support sliding window attn now."

605 assert pcp_world_size == 1, "PCP not support sliding window attn now."

606

578 - # The number of contiguous blocks needed for prefix cache hit.

579 - # -1 since the input token itself is also included in the window

580 - sliding_window_contiguous_blocks = cdiv(

581 - kv_cache_spec.sliding_window - 1, kv_cache_spec.block_size

607 + # The number of contiguous blocks needed for a prefix cache hit.

608 + sliding_window_contiguous_blocks = cls._contiguous_blocks_for_hit(

609 + kv_cache_spec.sliding_window, kv_cache_spec.block_size, drop_eagle_block

610 )

583 - if use_eagle:

584 - # Need to drop the last matched block if eagle is enabled. For

585 - # sliding window layer, we achieve this by increasing the number of

586 - # contiguous blocks needed for prefix cache hit by one and dropping

587 - # the last matched block.

588 - sliding_window_contiguous_blocks += 1

611

612 # TODO: reduce i by sliding_window_contiguous_blocks when cache miss, to

613 # optimize the time complexity from O(max_num_blocks) to

@@ -608,7 +630,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

630 # Skip prefix matching check if the block is not aligned with

631 # `alignment_tokens`.

632 if num_contiguous_blocks == 0 and block_size != alignment_tokens:

611 - post_pop_blocks = i if use_eagle else i + 1

633 + post_pop_blocks = i if drop_eagle_block else i + 1

634 if (post_pop_blocks * block_size) % alignment_tokens != 0:

635 continue

636 # Add the cached block to the computed blocks.

@@ -636,7 +658,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

658 ):

659 for computed in computed_blocks:

660 computed.pop()

639 - if use_eagle and computed_blocks[0]:

661 + if drop_eagle_block and computed_blocks[0]:

662 for computed in computed_blocks:

663 computed.pop()

664 # Re-align after eagle pop: the pop may break the alignment

@@ -650,17 +672,33 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

672 computed.pop()

673 return computed_blocks

674

653 - def _cache_block_mask(

654 - self, num_cached_blocks: int, num_full_blocks: int, alignment_tokens: int

675 + @classmethod

676 + def reachable_block_mask(

677 + cls,

678 + start_block: int,

679 + num_blocks: int,

680 + alignment_tokens: int,

681 + kv_cache_spec: KVCacheSpec,

682 + use_eagle: bool,

683 ) -> list[bool] | None:

656 - assert alignment_tokens > self.block_size

657 - per_segment = alignment_tokens // self.block_size

658 - tail = cdiv(self.sliding_window - 1, self.block_size)

659 - if tail >= per_segment:

684 + assert alignment_tokens > kv_cache_spec.block_size

685 + assert isinstance(kv_cache_spec, SlidingWindowSpec)

686 + per_segment = alignment_tokens // kv_cache_spec.block_size

687 + need = cls._contiguous_blocks_for_hit(

688 + window_size=kv_cache_spec.sliding_window,

689 + block_size=kv_cache_spec.block_size,

690 + use_eagle=use_eagle,

691 + )

692 + if need >= per_segment:

693 return None

661 - skip = per_segment - tail

694 + # The matched run's right edge sits on the aligned boundary block when

695 + # EAGLE peeks one block past it (shift=1), otherwise on the last block

696 + # before the boundary (shift=0). A block is reachable iff it falls in

697 + # the ``need``-wide run ending at some boundary's right edge.

698 + shift = 1 if use_eagle else 0

699 return [

663 - i % per_segment >= skip for i in range(num_cached_blocks, num_full_blocks)

700 + i >= shift and (i - shift) % per_segment >= per_segment - need

701 + for i in range(start_block, num_blocks)

702 ]

703

704 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:

@@ -714,7 +752,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):

752 kv_cache_group_ids: list[int],

753 block_pool: BlockPool,

754 kv_cache_spec: KVCacheSpec,

717 - use_eagle: bool,

755 + drop_eagle_block: bool,

756 alignment_tokens: int,

757 dcp_world_size: int = 1,

758 pcp_world_size: int = 1,

@@ -745,7 +783,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):

783 kv_cache_group_ids: The ids of the kv cache groups.

784 block_pool: The block pool.

785 kv_cache_spec: The kv cache spec.

748 - use_eagle: Whether to use eagle.

786 + drop_eagle_block: Whether to drop the last matched block for EAGLE/MTP.

787 dcp_world_size: The world size of decode context parallelism.

788 pcp_world_size: The world size of prefill context parallelism.

789 alignment_tokens: The returned cache hit length (in tokens) should

@@ -758,7 +796,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):

796 "ChunkedLocalAttentionManager can only be used for "

797 "chunked local attention groups"

798 )

761 - assert use_eagle is False, (

799 + assert drop_eagle_block is False, (

800 "Hybrid KV cache is not supported for " + "eagle + chunked local attention."

801 )

802 assert dcp_world_size == 1, "DCP not support chunked local attn now."

@@ -874,7 +912,7 @@ class MambaManager(SingleTypeKVCacheManager):

912 kv_cache_group_ids: list[int],

913 block_pool: BlockPool,

914 kv_cache_spec: KVCacheSpec,

877 - use_eagle: bool,

915 + drop_eagle_block: bool,

916 alignment_tokens: int,

917 dcp_world_size: int = 1,

918 pcp_world_size: int = 1,

@@ -1168,7 +1206,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):

1206 kv_cache_group_ids: list[int],

1207 block_pool: BlockPool,

1208 kv_cache_spec: KVCacheSpec,

1171 - use_eagle: bool,

1209 + drop_eagle_block: bool,

1210 alignment_tokens: int,

1211 dcp_world_size: int = 1,

1212 pcp_world_size: int = 1,