@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f80b66f548

[Model Runner V2][Spec Decode] Add KV cache support for multi-layer MTP (#50062)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Giancarlo Delfin committed Aug 13, 2026 at 16:22 UTC f80b66f548d855a104c9b2a0527e6c0b1a31750c

10 files changed +196 -32

tests/config/test_speculative_draft_hf_overrides.py

+4 -2

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

88 @pytest.mark.cpu_test

89 -def test_inkling_override_exposes_only_first_mtp_depth():

89 +def test_inkling_override_exposes_all_mtp_depths():

90 text_config = _make_hf_config(

91 architectures=["InklingForCausalLM"],

92 model_type="inkling_model",

@@ -107,7 +107,9 @@ def test_inkling_override_exposes_only_first_mtp_depth():

107 assert out is text_config

108 assert out.model_type == "inkling_mtp"

109 assert out.architectures == ["InklingMTPModel"]

110 - assert out.n_predict == 1

110 + # Multi-module MTP: every checkpoint depth is exposed (module i drafts

111 + # speculative token i), no longer clamped to the first depth.

112 + assert out.n_predict == 8

113 assert out.num_nextn_predict_layers == 8

114 assert out.chain_hidden_post_norm is False

115 assert out.local_layer_ids == [0, 2, 4]

tests/v1/core/test_scheduler.py

+3 -1

@@ -3701,6 +3701,7 @@ def test_mamba_align_eagle_schedules_encoder_at_boundary():

3701 )

3702 scheduler.need_mamba_block_aligned_split = True

3703 scheduler.use_eagle = True

3704 + scheduler.num_prefill_lookahead = 1

3705 scheduler.max_num_encoder_input_tokens = 2048

3706 scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=2048)

3707

@@ -5341,8 +5342,9 @@ def test_free_encoder_inputs_defers_for_eagle_lookahead():

5342 worker-side token-embedding fallback is only a backstop."""

5343 scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")

5344 # create_scheduler only builds ngram spec configs; force the eagle path that

5344 - # _free_encoder_inputs keys off (self.use_eagle).

5345 + # _free_encoder_inputs keys off (its read-ahead deferral).

5346 scheduler.use_eagle = True

5347 + scheduler.num_prefill_lookahead = 1

5348 mm_positions = [[PlaceholderRange(offset=50, length=100)]]

5349 request = create_requests(

5350 num_requests=1,

vllm/config/speculative.py

+1 -10

@@ -630,8 +630,7 @@ class SpeculativeConfig:

630 hf_config.model_type = "inkling_mtp"

631 hf_config.update(

632 {

633 - # Inkling currently exposes only the first checkpoint depth.

634 - "n_predict": 1,

633 + "n_predict": checkpoint_depths,

634 "num_nextn_predict_layers": checkpoint_depths,

635 "chain_hidden_post_norm": mtp_config.get(

636 "chain_hidden_post_norm", False

@@ -1083,14 +1082,6 @@ class SpeculativeConfig:

1082 "`num_speculative_tokens` was not provided"

1083 )

1084

1086 - if (

1087 - self.draft_model_config.hf_config.model_type == "inkling_mtp"

1088 - and self.num_speculative_tokens != 1

1089 - ):

1090 - raise ValueError(

1091 - "Inkling MTP currently supports exactly one speculative token"

1092 - )

1093 -

1085 if self.dspark_draft_topk is not None and self.method != "dspark":

1086 raise ValueError("dspark_draft_topk is only supported by DSpark")

1087

vllm/v1/core/kv_cache_coordinator.py

+53 -3

@@ -80,6 +80,7 @@ class KVCacheCoordinator(ABC):

80 scheduler_block_size: int,

81 hash_block_size: int,

82 metrics_collector: KVCacheMetricsCollector | None = None,

83 + num_prefill_lookahead: int = 0,

84 ):

85 self.kv_cache_config = kv_cache_config

86 self.max_model_len = max_model_len

@@ -91,6 +92,7 @@ class KVCacheCoordinator(ABC):

92 for g in kv_cache_config.kv_cache_groups

93 )

94 self.scheduler_block_size = scheduler_block_size

95 + self.num_reprefillable_tokens = max(0, num_prefill_lookahead - 1)

97 self.block_pool = BlockPool(

98 num_gpu_blocks=kv_cache_config.num_blocks,

@@ -108,6 +110,28 @@ class KVCacheCoordinator(ABC):

110 if use_eagle and not self.eagle_group_ids:

111 self.eagle_group_ids = set(range(len(kv_cache_config.kv_cache_groups)))

112

113 + # During chunked prefill with EAGLE, the single next prefill lookahead

114 + # token past the chunk boundary is combined with the final hidden state

115 + # and written to the KV cache. Therefore, the final chunk token must be

116 + # excluded from prefix cache hits to prevent requests from acquiring the

117 + # KV cache slot polluted with the next prefill token, which may or may not

118 + # be present after the matching prefix. The last-block drop handles this

119 + # edge case. During multi-module MTP, the issue generalizes to a prefill

120 + # lookahead of num_speculative_tokens, so the dropped tail must be large

121 + # enough to contain them. Hits land on scheduler-block boundaries (see

122 + # `_cache_hit_alignment_tokens`), so the excluded tail is

123 + # scheduler_block_size, not the group's own block size.

124 + if (

125 + enable_caching

126 + and self.eagle_group_ids

127 + and scheduler_block_size < num_prefill_lookahead

128 + ):

129 + raise ValueError(

130 + f"Multi-module MTP with prefix caching requires scheduler_block_size"

131 + f" (={scheduler_block_size}) >= num_speculative_tokens"

132 + f" (={num_prefill_lookahead})."

133 + )

134 +

135 self.single_type_managers = tuple(

136 get_manager_for_kv_cache_spec(

137 kv_cache_spec=kv_cache_group.kv_cache_spec,

@@ -286,9 +310,14 @@ class KVCacheCoordinator(ABC):

310 (including tokens that are already cached).

311 """

312 for manager in self.single_type_managers:

313 + # Only cache tokens with finalized KV. The last num_reprefillable_tokens

314 + # tokens can be re-prefilled during multi-module MTP.

315 + num_tokens_to_cache = max(

316 + 0, num_computed_tokens - self.num_reprefillable_tokens

317 + )

318 manager.cache_blocks(

319 request,

291 - num_computed_tokens,

320 + num_tokens_to_cache,

321 retention_interval=self.retention_interval,

322 )

323

@@ -407,6 +436,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):

436 scheduler_block_size: int,

437 hash_block_size: int,

438 metrics_collector: KVCacheMetricsCollector | None = None,

439 + num_prefill_lookahead: int = 0,

440 ):

441 super().__init__(

442 kv_cache_config,

@@ -420,6 +450,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):

450 scheduler_block_size=scheduler_block_size,

451 hash_block_size=hash_block_size,

452 metrics_collector=metrics_collector,

453 + num_prefill_lookahead=num_prefill_lookahead,

454 )

455 self.num_single_type_manager = len(self.single_type_managers)

456

@@ -457,6 +488,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

488 scheduler_block_size: int,

489 hash_block_size: int,

490 metrics_collector: KVCacheMetricsCollector | None = None,

491 + num_prefill_lookahead: int = 0,

492 ):

493 super().__init__(

494 kv_cache_config,

@@ -470,6 +502,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):

502 scheduler_block_size=scheduler_block_size,

503 hash_block_size=hash_block_size,

504 metrics_collector=metrics_collector,

505 + num_prefill_lookahead=num_prefill_lookahead,

506 )

507 self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec

508 self.block_size = self.kv_cache_spec.block_size

@@ -542,6 +575,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

575 scheduler_block_size: int,

576 hash_block_size: int,

577 metrics_collector: KVCacheMetricsCollector | None = None,

578 + num_prefill_lookahead: int = 0,

579 ):

580 super().__init__(

581 kv_cache_config,

@@ -555,6 +589,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

589 scheduler_block_size=scheduler_block_size,

590 hash_block_size=hash_block_size,

591 metrics_collector=metrics_collector,

592 + num_prefill_lookahead=num_prefill_lookahead,

593 )

594 # hash_block_size: the block size used to compute block hashes.

595 # The actual block size usually equals hash_block_size, but in cases where

@@ -688,9 +723,20 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

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

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

725 if manager.use_eagle and aligned_num_computed_tokens > 0:

726 + # Only cache tokens with finalized KV. The last

727 + # num_reprefillable_tokens tokens can be re-prefilled during

728 + # multi-module MTP.

729 + num_finalized_computed_tokens = max(

730 + 0, num_computed_tokens - self.num_reprefillable_tokens

731 + )

732 + aligned_num_finalized_computed_tokens = (

733 + num_finalized_computed_tokens

734 + // self.scheduler_block_size

735 + * self.scheduler_block_size

736 + )

737 num_tokens_to_cache = min(

692 - num_computed_tokens,

693 - aligned_num_computed_tokens + manager.block_size,

738 + num_finalized_computed_tokens,

739 + aligned_num_finalized_computed_tokens + manager.block_size,

740 )

741 # The manager already knows the fine hit granularity

742 # (``scheduler_block_size``); retention is passed separately so it

@@ -880,6 +926,7 @@ def get_kv_cache_coordinator(

926 scheduler_block_size: int,

927 hash_block_size: int,

928 metrics_collector: KVCacheMetricsCollector | None = None,

929 + num_prefill_lookahead: int = 0,

930 ) -> KVCacheCoordinator:

931 if not enable_caching:

932 return KVCacheCoordinatorNoPrefixCache(

@@ -893,6 +940,7 @@ def get_kv_cache_coordinator(

940 scheduler_block_size=scheduler_block_size,

941 hash_block_size=hash_block_size,

942 metrics_collector=metrics_collector,

943 + num_prefill_lookahead=num_prefill_lookahead,

944 )

945 if len(kv_cache_config.kv_cache_groups) == 1:

946 return UnitaryKVCacheCoordinator(

@@ -907,6 +955,7 @@ def get_kv_cache_coordinator(

955 scheduler_block_size=scheduler_block_size,

956 hash_block_size=hash_block_size,

957 metrics_collector=metrics_collector,

958 + num_prefill_lookahead=num_prefill_lookahead,

959 )

960 return HybridKVCacheCoordinator(

961 kv_cache_config,

@@ -920,4 +969,5 @@ def get_kv_cache_coordinator(

969 scheduler_block_size=scheduler_block_size,

970 hash_block_size=hash_block_size,

971 metrics_collector=metrics_collector,

972 + num_prefill_lookahead=num_prefill_lookahead,

973 )

vllm/v1/core/kv_cache_manager.py

@@ -125,6 +125,7 @@ class KVCacheManager:

125 max_in_flight_tokens: int | None = None,

126 enable_caching: bool = True,

127 use_eagle: bool = False,

128 + num_prefill_lookahead: int = 0,

129 log_stats: bool = False,

130 enable_kv_cache_events: bool = False,

131 dcp_world_size: int = 1,

@@ -161,6 +162,7 @@ class KVCacheManager:

162 scheduler_block_size=scheduler_block_size,

163 hash_block_size=hash_block_size,

164 metrics_collector=self.metrics_collector,

165 + num_prefill_lookahead=num_prefill_lookahead,

166 )

167 self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)

168 self.block_pool = self.coordinator.block_pool

vllm/v1/core/kv_cache_utils.py

+18

@@ -1468,6 +1468,9 @@ def _promote_local_kv_cache_specs(

1468 promoted_specs[layer_name] = replace_as(

1469 spec,

1470 target_cls,

1471 + # Promoted specs allocate blocks for all tokens and never free

1472 + # below the window, so the trailing-edge extension is moot.

1473 + drop=("extra_retained_tokens",),

1474 block_size=block_size,

1475 page_size_padded=promoted_page_size_padded(spec, block_size),

1476 )

@@ -2128,6 +2131,21 @@ def get_kv_cache_configs(

2131 # Check if the KV cache specs are registered correctly.

2132 # This is to prevent that some layers are initialized with unregistered specs.

2133 KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs)

2134 +

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

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

2137 + extra_retained_tokens = (

2138 + vllm_config.speculative_config.num_speculative_tokens - 1

2139 + if vllm_config.speculative_config is not None

2140 + and vllm_config.speculative_config.use_multi_module_mtp()

2141 + else 0

2142 + )

2143 + for layer_name, layer_spec in merged_kv_cache_specs.items():

2144 + if isinstance(layer_spec, SlidingWindowSpec):

2145 + merged_kv_cache_specs[layer_name] = replace(

2146 + layer_spec, extra_retained_tokens=extra_retained_tokens

2147 + )

2148 +

2149 # Get global KV cache groups. This also handles spec unification for

2150 # hybrid models when disable_hybrid_kv_cache_manager is enabled.

2151 # After this call, merged_kv_cache_specs may be modified in-place.

vllm/v1/core/sched/scheduler.py

+58 -8

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

247 self.use_eagle = False

248 self.num_spec_tokens = vllm_config.num_speculative_tokens

249 self.num_lookahead_tokens = vllm_config.num_lookahead_tokens

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

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

252 + # num_spec_tokens ahead at chunked-prefill boundaries. Determines the

253 + # encoder scheduling shift, the deferred encoder free, the KV cache

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

255 + # reserve between a chunk boundary and the prefill end.

256 + self.num_prefill_lookahead = 0

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

258 if speculative_config is not None:

259 if speculative_config.num_speculative_tokens_per_batch_size:

@@ -256,6 +263,12 @@ class Scheduler(SchedulerInterface):

263 vllm_num_speculative_tokens=self.num_spec_tokens,

264 )

265 self.use_eagle = speculative_config.use_eagle()

266 + if self.use_eagle:

267 + self.num_prefill_lookahead = (

268 + self.num_spec_tokens

269 + if speculative_config.use_multi_module_mtp()

270 + else 1

271 + )

272

273 # Create the KV cache manager.

274 if hash_block_size is None:

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

280 max_in_flight_tokens=vllm_config.max_in_flight_tokens,

281 enable_caching=self.cache_config.enable_prefix_caching,

282 use_eagle=self.use_eagle,

283 + num_prefill_lookahead=self.num_prefill_lookahead,

284 log_stats=self.log_stats,

285 enable_kv_cache_events=self.enable_kv_cache_events,

286 dcp_world_size=self.dcp_world_size,

@@ -438,6 +452,27 @@ class Scheduler(SchedulerInterface):

452 )

453 return blocks, num_local, shared_prefix_boundary, False

454

455 + def _reserve_prefill_lookahead(

456 + self,

457 + request: Request,

458 + num_computed_tokens: int,

459 + num_new_tokens: int,

460 + ) -> int:

461 + """Never end a prefill chunk within num_prefill_lookahead of the

462 + prefill end.

463 +

464 + At a chunked-prefill boundary, the multi-module MTP drafter consumes

465 + the next num_prefill_lookahead known prefill tokens as draft inputs. A

466 + boundary closer to the end than that would make it fall back to

467 + sampled drafts, permanently polluting the trailing modules' KV caches.

468 + Either finish the prefill or leave at least num_prefill_lookahead for

469 + the next chunk. No-op for eagle-family drafters (lookahead 1).

470 + """

471 + remaining = request.num_tokens - num_computed_tokens - num_new_tokens

472 + if 0 < remaining < self.num_prefill_lookahead:

473 + num_new_tokens -= self.num_prefill_lookahead - remaining

474 + return max(num_new_tokens, 0)

475 +

476 def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:

477 self.current_step += 1

478 # NOTE(woosuk) on the scheduling algorithm:

@@ -561,9 +596,15 @@ class Scheduler(SchedulerInterface):

596 request.num_computed_tokens,

597 num_new_tokens,

598 encoder_compute_budget,

564 - shift_computed_tokens=1 if self.use_eagle else 0,

599 + shift_computed_tokens=self.num_prefill_lookahead,

600 )

601

602 + # Multi-module MTP: avoid ending a prefill chunk within

603 + # num_prefill_lookahead of the prefill end.

604 + num_new_tokens = self._reserve_prefill_lookahead(

605 + request, request.num_computed_tokens, num_new_tokens

606 + )

607 +

608 if num_new_tokens == 0:

609 # The request cannot be scheduled because one of the following

610 # reasons:

@@ -576,6 +617,8 @@ class Scheduler(SchedulerInterface):

617 # 3. The encoder cache is exhausted.

618 # 4. Insufficient budget for a block-aligned chunk in hybrid

619 # models with mamba cache mode \"align\".

620 + # 5. Insufficient budget to keep a multi-module MTP prefill

621 + # chunk out of the prefill-lookahead window.

622 # NOTE(woosuk): Here, by doing `continue` instead of `break`,

623 # we do not strictly follow the FCFS scheduling policy and

624 # allow the lower-priority requests to be scheduled.

@@ -946,11 +989,18 @@ class Scheduler(SchedulerInterface):

989 num_computed_tokens,

990 num_new_tokens,

991 encoder_compute_budget,

949 - shift_computed_tokens=1 if self.use_eagle else 0,

992 + shift_computed_tokens=self.num_prefill_lookahead,

993 )

951 - if num_new_tokens == 0:

952 - # The request cannot be scheduled.

953 - break

994 +

995 + # Multi-module MTP: avoid ending a prefill chunk within

996 + # num_prefill_lookahead of the prefill end.

997 + num_new_tokens = self._reserve_prefill_lookahead(

998 + request, num_computed_tokens, num_new_tokens

999 + )

1000 +

1001 + if num_new_tokens == 0:

1002 + # The request cannot be scheduled.

1003 + break

1004

1005 # During async KV load, no forward pass is run yet.

1006 # Allocate speculative lookahead slots later to avoid

@@ -2157,9 +2207,9 @@ class Scheduler(SchedulerInterface):

2207 return

2208

2209 # Defer the free by the drafter's look-ahead so an entry stays

2160 - # referenced until the drafter's +1 read has also passed it, mirroring

2161 - # the shift the encoder scheduling path applies.

2162 - spec_lookahead = 1 if self.use_eagle else 0

2210 + # referenced until the drafter's read-ahead has also passed it,

2211 + # mirroring the shift the encoder scheduling path applies.

2212 + spec_lookahead = self.num_prefill_lookahead

2213

2214 # Here, we use list(set) to avoid modifying the set while iterating

2215 # over it.

vllm/v1/core/single_type_kv_cache_manager.py

+19 -1

@@ -561,6 +561,11 @@ class SingleTypeKVCacheManager(ABC):

561 return an empty list.

562 If eagle is enabled, drop the last matched block to force recompute the

563 last block to get the required hidden states for eagle drafting head.

564 + For multi-module MTP, this recompute also rewrites the dropped block's

565 + draft-layer KVs, which depend on up to num_speculative_tokens - 1

566 + tokens past the matched prefix (i.e. on the cache writer's

567 + continuation, which the block hash does not cover); the coordinator

568 + asserts the block size covers that window.

569 Need to be customized for each attention type.

570

571 Args:

@@ -877,6 +882,10 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

882 def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None:

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

884 self.sliding_window = kv_cache_spec.sliding_window

885 + # Extra trailing tokens to retain below the window (never attended) so a

886 + # multi-module MTP store-side lag can still reconstruct the window from

887 + # cached blocks.

888 + self.extra_retained_tokens = kv_cache_spec.extra_retained_tokens

889

890 @classmethod

891 def _contiguous_blocks_for_hit(

@@ -1072,13 +1081,22 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

1081 attention computation since they are outside the sliding window.

1082 Thus, get_num_skipped_tokens(7) == 4.

1083

1084 + The trailing edge of the window is extended by ``extra_retained_tokens``

1085 + so that those extra trailing tokens' blocks are retained (but not

1086 + attended). This is needed for multi-module spec decoding which can

1087 + re-prefill the last num_spec_prefill_tokens - 1 tokens from the end

1088 + of the sequence, and thus needs to delay freeing/caching of blocks.

1089 +

1090 Args:

1091 num_computed_tokens: The number of tokens that have been computed.

1092

1093 Returns:

1094 The number of tokens that will be skipped for attention computation.

1095 """

1081 - return max(0, num_computed_tokens - self.sliding_window + 1)

1096 + return max(

1097 + 0,

1098 + num_computed_tokens - self.sliding_window + 1 - self.extra_retained_tokens,

1099 + )

1100

1101 def get_num_common_prefix_blocks(self, running_request_id: str) -> int:

1102 """

vllm/v1/kv_cache_interface.py

+32 -7

@@ -5,6 +5,7 @@ from __future__ import annotations

6 import copy

7 from collections import Counter

8 +from collections.abc import Collection

9 from dataclasses import dataclass, fields, replace

10 from enum import Enum, IntEnum

11 from math import prod

@@ -88,14 +89,24 @@ def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:

89 return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE

91 -def replace_as(spec: KVCacheSpec, target_cls: type[_SpecT], **changes) -> _SpecT:

92 +def replace_as(

93 + spec: KVCacheSpec,

94 + target_cls: type[_SpecT],

95 + *,

96 + drop: Collection[str] = (),

97 + **changes,

98 +) -> _SpecT:

99 """``dataclasses.replace``, but rebuilding *spec* as *target_cls*

100 e.g. ``SlidingWindowSpec`` -> ``FullAttentionSpec``

101

95 - Every field of *spec* must exist on *target_cls*; fields only *target_cls* has keep

96 - their default values.

102 + Every field of *spec* must exist on *target_cls* unless named in *drop*;

103 + fields only *target_cls* has keep their default values.

104 """

98 - kwargs = {f.name: getattr(spec, f.name) for f in fields(spec) if f.init}

105 + kwargs = {

106 + f.name: getattr(spec, f.name)

107 + for f in fields(spec)

108 + if f.init and f.name not in drop

109 + }

110 kwargs.update(changes)

111 return target_cls(**kwargs)

112

@@ -573,6 +584,12 @@ class ChunkedLocalAttentionSpec(AttentionSpec):

584 class SlidingWindowSpec(AttentionSpec):

585 sliding_window: int

586 head_size_v: int = None # type: ignore[assignment]

587 + # The trailing edge of the window is extended by ``extra_retained_tokens``

588 + # so that those extra trailing tokens' blocks are retained (but not

589 + # attended). This is needed for multi-module spec decoding which can

590 + # re-prefill the last num_spec_prefill_tokens - 1 tokens from the end

591 + # of the sequence, and thus needs to delay freeing/caching of blocks.

592 + extra_retained_tokens: int = 0

593

594 def __post_init__(self):

595 if self.head_size_v is None:

@@ -614,8 +631,13 @@ class SlidingWindowSpec(AttentionSpec):

631 """

632 # During chunked prefill, we hold KV for the last `sliding_window-1`

633 # computed tokens plus the in-flight tokens (frees happen on the

617 - # processed-token basis); never more than `max_model_len`.

618 - num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len)

634 + # processed-token basis); never more than `max_model_len`. An additional

635 + # `extra_retained_tokens` trailing tokens are kept alive below the

636 + # window for multi-module spec decoding, and must be accounted here too.

637 + num_tokens = min(

638 + self.sliding_window - 1 + self.extra_retained_tokens + max_in_flight_tokens,

639 + max_model_len,

640 + )

641 # +1 because the sliding window may not start from the beginning of

642 # the block. E.g. block size 4 and num_token 4 needs two blocks

643 # [XXCD][EF] to store the 6-token window [CDEF].

@@ -686,16 +708,18 @@ class SlidingWindowMLASpec(SlidingWindowSpec):

708 model_version_set = set(spec.model_version for spec in specs)

709 sliding_window_set = set(spec.sliding_window for spec in specs)

710 block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs)

711 + extra_retained_set = set(spec.extra_retained_tokens for spec in specs)

712 assert (

713 len(cache_dtype_str_set) == 1

714 and len(compress_ratio_set) == 1

715 and len(model_version_set) == 1

716 and len(sliding_window_set) == 1

717 and len(block_stride_set) == 1

718 + and len(extra_retained_set) == 1

719 ), (

720 "All attention layers in the same KV cache group must use the same "

721 "quantization method, compress ratio, model version, sliding "

698 - "window size, and KV block stride indexing."

722 + "window size, KV block stride indexing, and retained token count."

723 )

724 return cls(

725 block_size=specs[0].block_size,

@@ -705,6 +729,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):

729 page_size_padded=specs[0].page_size_padded,

730 indexes_kv_by_block_stride=block_stride_set.pop(),

731 sliding_window=sliding_window_set.pop(),

732 + extra_retained_tokens=extra_retained_set.pop(),

733 cache_dtype_str=cache_dtype_str_set.pop(),

734 compress_ratio=compress_ratio_set.pop(),

735 model_version=model_version_set.pop(),

vllm/v1/worker/gpu/spec_decode/__init__.py

@@ -26,6 +26,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):

26 )

28 return Gemma4Speculator(vllm_config, device)

29 + elif speculative_config.use_multi_module_mtp():

30 + from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import (

31 + MultiModuleMTPSpeculator,

32 + )

33 +

34 + return MultiModuleMTPSpeculator(vllm_config, device)

35 elif speculative_config.method == "mtp":

36 from vllm.v1.worker.gpu.spec_decode.mtp.speculator import MTPSpeculator