@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 85c1f58d50

[BugFix][Model Runner V2][Spec Decode] Fix decode instance's multi-layer MTP kv caches during P/D (#55055)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

Giancarlo Delfin committed Sep 15, 2026 at 17:31 UTC 85c1f58d50b80a6668f1f8158f0d8ef7b1d9112a

7 files changed +69 -46

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+1 -1

@@ -138,7 +138,7 @@ class TestPushScheduler:

138

139 with patch.object(

140 sched,

141 - "_truncate_mamba_request_for_prefill",

141 + "_truncate_request_for_prefill",

142 side_effect=AssertionError("must not truncate after cache lookup"),

143 ):

144 assert sched.get_num_new_matched_tokens(request, 0) == (0, False)

tests/v1/kv_connector/unit/utils.py

+6 -1

@@ -521,6 +521,8 @@ def make_nixl_scheduler(

521 Only sets the flags needed by the tests. When *heartbeat=True* the

522 scheduler-side heartbeat bookkeeping fields are also initialised.

523 """

524 + from types import SimpleNamespace

525 +

526 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (

527 NixlConnectorScheduler,

528 )

@@ -532,6 +534,7 @@ def make_nixl_scheduler(

534 block_size=16,

535 mamba_enabled=has_mamba,

536 )

537 + sched.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=0)

538

539 if heartbeat:

540 sched._heartbeat_by_engine = {}

@@ -598,9 +601,11 @@ def make_nixl_push_scheduler(

601 mamba_enabled=has_mamba,

602 )

603

601 - # vllm_config is consulted for parallel_config.tensor_parallel_size.

604 + # vllm_config is consulted for parallel_config.tensor_parallel_size, and by

605 + # `_prefill_backoff` on both the P and D prefill paths.

606 vllm_config = MagicMock()

607 vllm_config.parallel_config.tensor_parallel_size = 1

608 + vllm_config.num_prefill_lookahead_tokens = 0

609 sched.vllm_config = vllm_config

610

611 # Push-specific state.

vllm/config/vllm.py

+23

@@ -650,6 +650,29 @@ class VllmConfig:

650 return self.num_speculative_tokens

651 return 0

652

653 + @property

654 + def num_prefill_lookahead_tokens(self) -> int:

655 + """Prefill tokens past the computed range that the drafter reads.

656 +

657 + Mid-prefill the drafter consumes tokens the target model has not been

658 + scheduled for yet, so every component that has to keep them available

659 + must apply this margin: the scheduler, which never ends a chunk within

660 + it and shifts encoder scheduling by it, and the KV cache manager, which

661 + treats the trailing `this - 1` tokens as re-prefillable rather than

662 + finalized. Consumers must read this property rather than re-deriving

663 + their own per-method lookahead, so those components cannot drift apart.

664 + """

665 + speculative_config = self.speculative_config

666 + if speculative_config is None or not speculative_config.use_eagle():

667 + return 0

668 + if speculative_config.use_multi_module_mtp():

669 + # Each MTP module reads one token further ahead than the one before

670 + # it, so the chain needs num_speculative_tokens of runway at a

671 + # chunked-prefill boundary.

672 + return self.num_speculative_tokens

673 + # Eagle-family drafters read only the immediate next token.

674 + return 1

675 +

676 @property

677 def uniform_decode_query_len(self) -> int:

678 """Query length of every request in a uniform decode batch.

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

+36 -16

@@ -199,8 +199,8 @@ class NixlBaseConnectorScheduler:

199 def on_new_request(self, request: "Request") -> None:

200 """Track a request that may need heartbeats."""

201 params = request.kv_transfer_params

202 - if params is not None and params.get("do_remote_decode") and self._has_mamba:

203 - self._truncate_mamba_request_for_prefill(request)

202 + if params is not None and params.get("do_remote_decode"):

203 + self._truncate_request_for_prefill(request)

204

205 # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are

206 # effectively disabled for Bidirectional KV transfer.

@@ -382,36 +382,56 @@ class NixlBaseConnectorScheduler:

382 (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts)

383 )

384

385 + def _prefill_backoff(self) -> int:

386 + """Trailing prompt tokens the prefiller must not compute; the decoder

387 + recomputes them locally.

388 +

389 + Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module

390 + MTP needs to keep its whole lookahead window off of the prefiller, which

391 + would otherwise embed the unverified drafts in the MTP layer's KV cache. The

392 + decoder would never rebuild them, because the update is sized by the rejection

393 + count, which is zero for the first decode.

394 + """

395 + return max(

396 + 1 if self._has_mamba else 0,

397 + self.vllm_config.num_prefill_lookahead_tokens - 1,

398 + )

399 +

400 def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:

386 - """D-side only. Returns N-1 for Mamba models since the decoder

387 - always recomputes the last token and must start from h(N-1)."""

388 - if self._has_mamba and num_prompt_tokens > 1:

389 - return num_prompt_tokens - 1

401 + """D-side only. The number of prompt tokens to load from the prefiller.

402 + Stops short of the trailing ``_prefill_backoff()`` tokens that the decoder

403 + will recompute locally."""

404 + backoff = self._prefill_backoff()

405 + if backoff and num_prompt_tokens > backoff:

406 + return num_prompt_tokens - backoff

407 return num_prompt_tokens

408

392 - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:

393 - """P-side only: drop the last prompt token so the prefiller computes

394 - h(N-1) instead of h(N). The decoder recomputes the last token to

395 - derive h(N) correctly.

409 + def _truncate_request_for_prefill(self, request: "Request") -> None:

410 + """P-side only: drop the trailing ``_prefill_backoff()`` prompt tokens

411 + so the prefiller stops short of what the decoder recomputes locally.

412 + For Mamba that is the single token needed to yield h(N-1); for

413 + multi-module MTP it is the drafter's whole lookahead window.

414

415 Guarded by ``_p_side_truncated`` to avoid repeated truncation if the

416 request is preempted and rescheduled."""

417 + backoff = self._prefill_backoff()

418 params = request.kv_transfer_params

419 if (

401 - params is not None

420 + backoff

421 + and params is not None

422 # Guard against repeated truncation after preemption/reschedule.

423 and not params.get("_p_side_truncated")

404 - and request.num_prompt_tokens > 1

424 + and request.num_prompt_tokens > backoff

425 ):

426 if request.prompt_token_ids is not None:

407 - request.prompt_token_ids.pop()

427 + del request.prompt_token_ids[-backoff:]

428 elif request.prompt_embeds is not None:

409 - request.prompt_embeds = request.prompt_embeds[:-1]

429 + request.prompt_embeds = request.prompt_embeds[:-backoff]

430 else:

431 return

432

413 - request._all_token_ids.pop()

414 - request.num_prompt_tokens -= 1

433 + del request._all_token_ids[-backoff:]

434 + request.num_prompt_tokens -= backoff

435 request.max_tokens = 1

436 params["_p_side_truncated"] = True

437

vllm/v1/core/kv_cache_utils.py

+1 -6

@@ -2678,12 +2678,7 @@ def get_kv_cache_configs(

2678

2679 # When speculating with more than 1 speculative module (e.g. multi-layered MTP)

2680 # tag every SlidingWindowSpec with how many extra tokens to retain in the window.

2681 - extra_retained_tokens = (

2682 - vllm_config.speculative_config.num_speculative_tokens - 1

2683 - if vllm_config.speculative_config is not None

2684 - and vllm_config.speculative_config.use_multi_module_mtp()

2685 - else 0

2686 - )

2681 + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1)

2682 for layer_name, layer_spec in merged_kv_cache_specs.items():

2683 if isinstance(layer_spec, SlidingWindowSpec):

2684 merged_kv_cache_specs[layer_name] = replace(

vllm/v1/core/sched/scheduler.py

+1 -13

@@ -262,13 +262,7 @@ class Scheduler(SchedulerInterface):

262 self.use_eagle_block_drop = False

263 self.num_spec_tokens = vllm_config.num_speculative_tokens

264 self.num_lookahead_tokens = vllm_config.num_lookahead_tokens

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

266 - # Eagle-family drafters read 1 ahead, but multi-module MTP reads

267 - # num_spec_tokens ahead at chunked-prefill boundaries. Determines the

268 - # encoder scheduling shift, the deferred encoder free, the KV cache

269 - # manager's re-prefillable window (this minus 1), and how many tokens to

270 - # reserve between a chunk boundary and the prefill end.

271 - self.num_prefill_lookahead = 0

265 + self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens

266 self.dynamic_sd_lookup: list[int] | None = None

267 if speculative_config is not None:

268 if speculative_config.num_speculative_tokens_per_batch_size:

@@ -278,12 +272,6 @@ class Scheduler(SchedulerInterface):

272 vllm_num_speculative_tokens=self.num_spec_tokens,

273 )

274 self.use_eagle = speculative_config.use_eagle()

281 - if self.use_eagle:

282 - self.num_prefill_lookahead = (

283 - self.num_spec_tokens

284 - if speculative_config.use_multi_module_mtp()

285 - else 1

286 - )

275 self.use_eagle_block_drop = speculative_config.use_eagle_block_drop()

276 if self.use_eagle and not self.use_eagle_block_drop:

277 logger.warning(

vllm/v1/worker/gpu/model_runner.py

+1 -9

@@ -296,15 +296,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

296 self.is_pooling_model = self.model_config.runner_type == "pooling"

297 self.pooling_runner: PoolingRunner | None = None

298

299 - # Multi-module MTP feeds its modules the next num_speculative_steps prefill

300 - # tokens during chunked prefill. Other speculators only read the immediate

301 - # next one.

302 - num_prefill_lookahead = (

303 - self.num_speculative_steps

304 - if self.speculative_config is not None

305 - and self.speculative_config.use_multi_module_mtp()

306 - else 1

307 - )

299 + num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens)

300

301 self.step_timing = StepTimingCollector()

302