481839ad9e · tduncan/Nvidia-Optimization-Stack-VLLM

tests/test_config.py

+17

@@ -2152,6 +2152,23 @@ def test_draft_sample_method_probabilistic_is_accepted():

2152 assert speculative_config.draft_sample_method == "probabilistic"

2153

2154

2155 +@pytest.mark.parametrize("disable_eagle_block_drop", [False, True])

2156 +def test_eagle_block_drop_can_be_disabled_without_disabling_eagle(

2157 + disable_eagle_block_drop: bool,

2158 +):

2159 + # Start from an ngram config to avoid loading model metadata: these predicates

2160 + # depend only on the speculative method and the new switch.

2161 + speculative_config = SpeculativeConfig(

2162 + method="ngram",

2163 + num_speculative_tokens=3,

2164 + disable_eagle_block_drop=disable_eagle_block_drop,

2165 + )

2166 + speculative_config.method = "eagle3"

2167 +

2168 + assert speculative_config.use_eagle()

2169 + assert speculative_config.use_eagle_block_drop() is not disable_eagle_block_drop

2170 +

2171 +

2172 def test_draft_sample_method_gumbel_is_rejected():

2173 with pytest.raises(ValidationError):

2174 SpeculativeConfig(

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

+3 -3

@@ -170,7 +170,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):

170 cache_config=SimpleNamespace(block_size=block_size),

171 max_num_scheduled_tokens=8192,

172 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

173 - use_eagle=False,

173 + use_eagle_block_drop=False,

174 hash_block_size=hash_block_size,

175 dcp_world_size=dcp_world_size,

176 scheduler_block_size=scheduler_block_size,

@@ -218,7 +218,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():

218 cache_config=SimpleNamespace(block_size=block_size),

219 max_num_scheduled_tokens=token_budget,

220 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

221 - use_eagle=False,

221 + use_eagle_block_drop=False,

222 hash_block_size=32,

223 mamba_partial_cache_hit=False,

224 mamba_has_prefill_checkpoint_blocks=False,

@@ -257,7 +257,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():

257 scheduler_config=SimpleNamespace(

258 long_prefill_token_threshold=long_prefill_threshold

259 ),

260 - use_eagle=False,

260 + use_eagle_block_drop=False,

261 hash_block_size=32,

262 mamba_partial_cache_hit=False,

263 mamba_has_prefill_checkpoint_blocks=False,

tests/v1/core/test_kv_cache_utils.py

+1

@@ -3296,6 +3296,7 @@ def _spec_decode_grouping_config(method="dspark", model_type=None):

3296 speculative_config=SimpleNamespace(

3297 method=method,

3298 use_eagle=lambda: True,

3299 + use_eagle_block_drop=lambda: True,

3300 ),

3301 )

3302

tests/v1/core/test_mamba_align_chunk_split.py

+14 -1

@@ -80,13 +80,16 @@ def _split(

80 request: Request,

81 num_new_tokens: int,

82 use_eagle: bool = True,

83 + use_eagle_block_drop: bool | None = None,

84 partial_hit: bool = False,

85 num_prefill_checkpoint_blocks: int = 0,

86 ) -> int:

87 """Call the real `Scheduler._mamba_block_aligned_split` on a stub self."""

88 + if use_eagle_block_drop is None:

89 + use_eagle_block_drop = use_eagle

90 stub = SimpleNamespace(

91 cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE),

89 - use_eagle=use_eagle,

92 + use_eagle_block_drop=use_eagle_block_drop,

93 max_num_scheduled_tokens=16384,

94 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

95 # `prefix_match_unit` finer than the block size (#46384).

@@ -128,6 +131,16 @@ def test_internal_checkpoint_split(

131 assert all(not block.is_null for block in blocks) # checkpoint + running state

132

133

134 +def test_disabling_eagle_block_drop_keeps_the_trailing_cache_boundary() -> None:

135 + (request,) = create_requests(1, num_tokens=3602, block_size=ATTN_BLOCK_SIZE)

136 +

137 + with_drop = _split(request, request.num_tokens, use_eagle_block_drop=True)

138 + without_drop = _split(request, request.num_tokens, use_eagle_block_drop=False)

139 +

140 + assert with_drop == MAMBA_BLOCK_SIZE

141 + assert without_drop == 2 * MAMBA_BLOCK_SIZE

142 +

143 +

144 def _run_chunked_prefill(

145 manager: KVCacheManager, request: Request, budgets: list[int]

146 ) -> dict[int, int]:

tests/v1/core/test_prefix_caching.py

+1 -1

@@ -1163,7 +1163,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():

1163 cache_config=SimpleNamespace(block_size=block_size),

1164 max_num_scheduled_tokens=3 * block_size,

1165 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),

1166 - use_eagle=False,

1166 + use_eagle_block_drop=False,

1167 hash_block_size=block_size,

1168 mamba_partial_cache_hit=False,

1169 mamba_has_prefill_checkpoint_blocks=False,

tests/v1/core/test_single_type_kv_cache_manager.py

+41

@@ -399,6 +399,47 @@ def test_sliding_window_possible_cached_prefix():

399 )

400

401

402 +def test_sliding_window_cache_hit_with_finer_hash_alignment():

403 + """Sliding-window lookup uses full blocks with finer hybrid-cache hashes."""

404 + hash_block_size = 2

405 + block_size = 4

406 + sliding_window_spec = SlidingWindowSpec(

407 + block_size=block_size,

408 + num_kv_heads=1,

409 + head_size=1,

410 + dtype=torch.float32,

411 + sliding_window=8,

412 + )

413 + block_pool = BlockPool(

414 + num_gpu_blocks=100,

415 + enable_caching=True,

416 + hash_block_size=hash_block_size,

417 + )

418 + manager = get_sliding_window_manager(sliding_window_spec, block_pool)

419 + block_hashes = [BlockHash(str(i).encode()) for i in range(4)]

420 +

421 + # Each physical block uses the last chained hash in its block-size view.

422 + for block_hash, block in zip(

423 + (block_hashes[1], block_hashes[3]), block_pool.blocks[10:12]

424 + ):

425 + block_pool.cached_block_hash_to_block.insert(

426 + make_block_hash_with_group_id(block_hash, 0), block

427 + )

428 +

429 + computed_blocks, hit_length = manager.find_longest_cache_hit(

430 + block_hashes=block_hashes,

431 + max_length=8,

432 + kv_cache_group_ids=[0],

433 + block_pool=block_pool,

434 + kv_cache_spec=sliding_window_spec,

435 + drop_eagle_block=False,

436 + alignment_tokens=hash_block_size,

437 + )

438 +

439 + assert hit_length == 8

440 + assert computed_blocks[0] == block_pool.blocks[10:12]

441 +

442 +

443 def test_chunked_local_attention_remove_skipped_blocks():

444 attention_spec = ChunkedLocalAttentionSpec(

445 block_size=2,

vllm/config/speculative.py

+9

@@ -436,6 +436,11 @@ class SpeculativeConfig:

436 speculative input batches can contain sequences of different lengths,

437 which may only be supported by certain attention backends. This currently

438 only affects the EAGLE method of speculation."""

439 + disable_eagle_block_drop: bool = False

440 + """Disable dropping the trailing prefix-cache block for EAGLE-like

441 + speculative methods. This is an experimental option for measuring the

442 + acceptance-rate impact of reusing that block. It does not disable the

443 + speculative drafter itself."""

444 use_local_argmax_reduction: bool = False

445 """Use vocab-parallel local argmax instead of all-gathering full logits

446 for draft token generation. Reduces communication from O(vocab_size) to

@@ -1846,6 +1851,10 @@ class SpeculativeConfig:

1851 # TODO(ben): Refactor this so the naming is clearer

1852 return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark")

1853

1854 + def use_eagle_block_drop(self) -> bool:

1855 + """Whether volatile trailing cache blocks should be discarded."""

1856 + return self.use_eagle() and not self.disable_eagle_block_drop

1857 +

1858 def use_dflash(self) -> bool:

1859 return self.method == "dflash"

1860

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

+5 -4

@@ -1593,16 +1593,17 @@ class MooncakeStoreWorker:

1593 for group in kv_cache_config.transfer_groups

1594 ]

1595 spec_cfg = getattr(vllm_config, "speculative_config", None)

1596 - use_eagle = bool(

1597 - spec_cfg.use_eagle()

1598 - if spec_cfg is not None and callable(getattr(spec_cfg, "use_eagle", None))

1596 + use_eagle_block_drop = bool(

1597 + spec_cfg.use_eagle_block_drop()

1598 + if spec_cfg is not None

1599 + and callable(getattr(spec_cfg, "use_eagle_block_drop", None))

1600 else False

1601 )

1602 self.coord = MooncakeStoreCoordinator(

1603 self._kv_cache_groups,

1604 scheduler_block_size=self.block_size,

1605 hash_block_size=self.hash_block_size,

1605 - use_eagle=use_eagle,

1606 + use_eagle=use_eagle_block_drop,

1607 retention_interval=kv_cache_config.prefix_cache_retention_interval,

1608 dcp_world_size=self.dcp_size,

1609 )

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+3 -3

@@ -223,11 +223,11 @@ class SchedulerOffloadConfig(NamedTuple):

223 if g.is_eagle_group

224 }

225

226 - use_eagle = (

226 + use_eagle_block_drop = (

227 vllm_config.speculative_config is not None

228 - and vllm_config.speculative_config.use_eagle()

228 + and vllm_config.speculative_config.use_eagle_block_drop()

229 )

230 - if use_eagle and not eagle_groups:

230 + if use_eagle_block_drop and not eagle_groups:

231 eagle_groups = set(range(len(kv_cache_config.kv_cache_groups)))

232

233 if eagle_groups:

vllm/v1/core/kv_cache_utils.py

+1 -1

@@ -1841,7 +1841,7 @@ def _annotate_eagle_groups(

1841 use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group.

1842 """

1843 spec_config = vllm_config.speculative_config

1844 - if spec_config is None or not spec_config.use_eagle():

1844 + if spec_config is None or not spec_config.use_eagle_block_drop():

1845 return

1846

1847 for group in kv_cache_groups:

vllm/v1/core/sched/scheduler.py

+10 -2

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

255 )

256 speculative_config = vllm_config.speculative_config

257 self.use_eagle = False

258 + self.use_eagle_block_drop = False

259 self.num_spec_tokens = vllm_config.num_speculative_tokens

260 self.num_lookahead_tokens = vllm_config.num_lookahead_tokens

261 # Positions past the computed tokens that the drafter reads mid-prefill.

@@ -279,6 +280,13 @@ class Scheduler(SchedulerInterface):

280 if speculative_config.use_multi_module_mtp()

281 else 1

282 )

283 + self.use_eagle_block_drop = speculative_config.use_eagle_block_drop()

284 + if self.use_eagle and not self.use_eagle_block_drop:

285 + logger.warning(

286 + "EAGLE trailing prefix-cache block dropping is disabled. "

287 + "This is experimental and may affect speculative-token "

288 + "acceptance rates."

289 + )

290

291 # Create the KV cache manager.

292 if hash_block_size is None:

@@ -289,7 +297,7 @@ class Scheduler(SchedulerInterface):

297 max_model_len=self.max_model_len,

298 max_in_flight_tokens=vllm_config.max_in_flight_tokens,

299 enable_caching=self.cache_config.enable_prefix_caching,

292 - use_eagle=self.use_eagle,

300 + use_eagle=self.use_eagle_block_drop,

301 num_prefill_lookahead=self.num_prefill_lookahead,

302 log_stats=self.log_stats,

303 enable_kv_cache_events=self.enable_kv_cache_events,

@@ -414,7 +422,7 @@ class Scheduler(SchedulerInterface):

422 # Eagle, FullAttn prunes the last matching block, so back off one

423 # block to avoid a Mamba cache miss.

424 last_cache_position = request.num_tokens - request.num_tokens % block_size

417 - if self.use_eagle:

425 + if self.use_eagle_block_drop:

426 last_cache_position = max(last_cache_position - block_size, 0)

427

428 end = start + num_new_tokens

vllm/v1/core/single_type_kv_cache_manager.py

+3 -4

@@ -925,10 +925,9 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

925 )

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

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

928 - # Fine-grained partial hits are not supported for sliding window now

929 - assert alignment_tokens % kv_cache_spec.block_size == 0, (

930 - "SlidingWindowManager does not support fine-grained (partial) cache hits"

931 - )

928 + # Sliding-window cache hits must stay at the group's physical block

929 + # granularity. resolve_block_hashes() converts finer-grained hashes to

930 + # that view when the hybrid-cache alignment is smaller than block_size.

931 block_hashes = resolve_block_hashes(

932 block_hashes,

933 block_pool.hash_block_size,

vllm/v1/simple_kv_offload/manager.py

+4 -2

@@ -125,12 +125,14 @@ class SimpleCPUOffloadScheduler:

125 )

126

127 spec_config = vllm_config.speculative_config

128 - use_eagle = spec_config is not None and spec_config.use_eagle()

128 + use_eagle_block_drop = (

129 + spec_config is not None and spec_config.use_eagle_block_drop()

130 + )

131 self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(

132 kv_cache_config=self.cpu_kv_cache_config,

133 max_model_len=vllm_config.model_config.max_model_len,

134 max_in_flight_tokens=vllm_config.max_in_flight_tokens,

133 - use_eagle=use_eagle,

135 + use_eagle=use_eagle_block_drop,

136 enable_caching=True,

137 enable_kv_cache_events=self.enable_kv_cache_events,

138 dcp_world_size=dcp_world_size,

[Feature][Spec] Support disabling trailing prefix-cache block dropping (#53388)