From af8cdb1aa7fa228540d347c5a817a6ed5aaa6b32 Mon Sep 17 00:00:00 2001 From: Greg Weyer Date: Sat, 15 Aug 2026 22:59:52 -0400 Subject: [PATCH 1/5] [Bugfix] Reload speculative draft weights after Level 2 sleep wake Level 2 sleep discards speculative draft parameters, while the documented disk-backed reload path restored only the target model. Reload the configured draft checkpoint in GPUModelRunner.reload_weights before the target so speculative decoding recovers and the target remains authoritative for shared parameter storage. Model Runner V2 keeps its existing one-line delegate into that method. Keep caller-supplied iterator reloads unchanged and cover the disk, no-draft, iterator, and V2-delegate paths with focused unit tests. Assisted-by: OpenAI Codex Assisted-by: Grok Signed-off-by: Greg Weyer --- tests/v1/worker/test_gpu_model_runner.py | 136 ++++++++++++++++++++++- vllm/v1/worker/gpu_model_runner.py | 32 ++++++ 2 files changed, 167 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index bb77f90ced2a..5e55c3e096b6 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -4,7 +4,7 @@ import gc from contextlib import nullcontext from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, call, patch import numpy as np import pytest @@ -60,6 +60,7 @@ from vllm.v1.worker.gpu.lora_utils import LoraState from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras +from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2 from vllm.v1.worker.gpu_input_batch import InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.utils import select_common_block_size @@ -1827,3 +1828,136 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks(): with pytest.raises(ValueError, match="max_num_seqs"): runner.initialize_kv_cache(kv_cache_config) + + +class TestReloadDraftWeights: + """Disk-backed reload restores draft parameters before the target.""" + + def _make_runner(self, cls=GPUModelRunner): + runner = object.__new__(cls) + runner.load_config = Mock() + runner.load_config.load_format = "safetensors" + runner.lora_config = None + runner.model_config = Mock() + runner.model_config.quantization = None + runner.speculative_config = SimpleNamespace( + draft_model_config=Mock(), + draft_load_config=None, + ) + runner.reset_lora_state = Mock() + runner.reset_encoder_cache = Mock() + runner.reset_mm_cache = Mock() + return runner + + def _assert_disk_reload_restores_draft(self, runner): + draft_model = Mock() + target_model = Mock() + target_model.named_parameters.return_value = [] + runner.get_draft_model = Mock(return_value=draft_model) + runner.get_model = Mock(return_value=target_model) + reload_order = [] + draft_weights = iter([("draft", torch.zeros(1))]) + target_weights = iter([("target", torch.zeros(1))]) + + def load_draft(_weights): + reload_order.append("draft") + + def load_target(_weights): + reload_order.append("target") + return set() + + draft_model.load_weights.side_effect = load_draft + target_model.load_weights.side_effect = load_target + + def get_loader(_load_config): + loader = Mock() + + def get_all_weights(_config, model): + return draft_weights if model is draft_model else target_weights + + loader.get_all_weights.side_effect = get_all_weights + return loader + + with ( + patch.object( + gpu_model_runner_module, "get_model_loader", side_effect=get_loader + ) as get_loader_mock, + patch.object( + gpu_model_runner_module, "initialize_layerwise_reload" + ) as initialize_reload, + patch.object( + gpu_model_runner_module, "finalize_layerwise_reload" + ) as finalize_reload, + ): + runner.reload_weights() + + assert reload_order == ["draft", "target"] + assert get_loader_mock.call_args_list == [ + call(runner.load_config), + call(runner.load_config), + ] + draft_model.load_weights.assert_called_once_with(draft_weights) + target_model.load_weights.assert_called_once_with(target_weights) + assert initialize_reload.call_args_list == [ + call(draft_model), + call(target_model), + ] + assert finalize_reload.call_args_list == [ + call(draft_model, runner.speculative_config.draft_model_config), + call(target_model, runner.model_config), + ] + + def test_disk_reload_restores_draft_model(self): + self._assert_disk_reload_restores_draft(self._make_runner()) + + def test_v2_disk_reload_delegates_to_v1_helper(self): + self._assert_disk_reload_restores_draft(self._make_runner(GPUModelRunnerV2)) + + def test_disk_reload_without_draft_only_reloads_target(self): + runner = self._make_runner() + target_model = Mock() + target_model.named_parameters.return_value = [] + target_model.load_weights.return_value = set() + runner.get_draft_model = Mock(return_value=None) + runner.get_model = Mock(return_value=target_model) + target_weights = iter([("target", torch.zeros(1))]) + model_loader = Mock() + model_loader.get_all_weights.return_value = target_weights + + with ( + patch.object( + gpu_model_runner_module, + "get_model_loader", + return_value=model_loader, + ) as get_loader, + patch.object(gpu_model_runner_module, "initialize_layerwise_reload"), + patch.object(gpu_model_runner_module, "finalize_layerwise_reload"), + ): + runner.reload_weights() + + get_loader.assert_called_once_with(runner.load_config) + model_loader.get_all_weights.assert_called_once_with( + runner.model_config, target_model + ) + target_model.load_weights.assert_called_once_with(target_weights) + runner.get_draft_model.assert_called_once_with() + + def test_iterator_reload_does_not_reload_draft_model(self): + runner = self._make_runner() + target_model = Mock() + target_model.named_parameters.return_value = [] + target_model.load_weights.return_value = set() + runner.get_draft_model = Mock() + runner.get_model = Mock(return_value=target_model) + weights = iter([("weight", torch.zeros(1))]) + + with ( + patch.object(gpu_model_runner_module, "get_model_loader") as get_loader, + patch.object(gpu_model_runner_module, "initialize_layerwise_reload"), + patch.object(gpu_model_runner_module, "finalize_layerwise_reload"), + ): + runner.reload_weights(weights_iterator=weights) + + get_loader.assert_not_called() + runner.get_draft_model.assert_not_called() + target_model.load_weights.assert_called_once_with(weights) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 4ff50f82b515..89a1bf3c5caa 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5551,6 +5551,30 @@ def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None: return None + def _reload_draft_weights_from_disk(self) -> None: + get_draft_model = getattr(self, "get_draft_model", None) + if not callable(get_draft_model): + return + draft_model = get_draft_model() + speculative_config = getattr(self, "speculative_config", None) + if draft_model is None or speculative_config is None: + return + + draft_model_config = speculative_config.draft_model_config + assert draft_model_config is not None + load_config = speculative_config.draft_load_config or self.load_config + model_loader = get_model_loader(load_config) + if not hasattr(model_loader, "get_all_weights"): + raise NotImplementedError( + f"Draft model reloading with `{load_config.load_format}` format" + ) + + initialize_layerwise_reload(draft_model) + draft_model.load_weights( + model_loader.get_all_weights(draft_model_config, draft_model) + ) + finalize_layerwise_reload(draft_model, draft_model_config) + def reload_weights( self, weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None, @@ -5585,6 +5609,14 @@ def reload_weights( # load weights from disk if none are provided if weights_iterator is None: + # Speculative drafts are a separate module. Level 2 sleep discards + # their parameters with the target's; wake_up only restores draft + # buffers. Reload the configured draft checkpoint first so the + # target remains authoritative for any shared storage. + # Call unbound: Model Runner V2 delegates this method with a V2 + # `self` and does not inherit the helper. + GPUModelRunner._reload_draft_weights_from_disk(self) + model_loader = get_model_loader(self.load_config) if not hasattr(model_loader, "get_all_weights"): raise NotImplementedError( From c122415981537c5faebc4a83c9bbeaba957f35ac Mon Sep 17 00:00:00 2001 From: Greg Weyer Date: Wed, 19 Aug 2026 09:11:59 -0400 Subject: [PATCH 2/5] Trim draft-reload comment per review Condense the reload_weights draft-reload comment block to three lines; content unchanged (L2-sleep rationale, target-authoritative ordering, unbound-call note for MRV2 delegation). Assisted-by: Aether (Hermes Agent) Signed-off-by: Greg Weyer --- vllm/v1/worker/gpu_model_runner.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 89a1bf3c5caa..8078f6f5f675 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5609,12 +5609,8 @@ def reload_weights( # load weights from disk if none are provided if weights_iterator is None: - # Speculative drafts are a separate module. Level 2 sleep discards - # their parameters with the target's; wake_up only restores draft - # buffers. Reload the configured draft checkpoint first so the - # target remains authoritative for any shared storage. - # Call unbound: Model Runner V2 delegates this method with a V2 - # `self` and does not inherit the helper. + # Reload draft weights (dropped by L2 sleep) before the target so it + # stays authoritative for shared storage. Unbound: MRV2 passes a V2 self. GPUModelRunner._reload_draft_weights_from_disk(self) model_loader = get_model_loader(self.load_config) From 2d30a7e5b68da24f4632c8233e639f0caee21be6 Mon Sep 17 00:00:00 2001 From: Greg Weyer Date: Fri, 21 Aug 2026 13:50:03 -0400 Subject: [PATCH 3/5] Skip target-owned draft aliases on Level 2 weight reload Gemma4 MTP replaces draft embed_tokens with the backbone-dim target module, so replaying the assistant checkpoint into that alias shape-mismatches. Unbind those identities for the draft load, refill a shape-mismatched tied lm_head, and leave same-width distinct heads untouched. Assisted-by: Grok Assisted-by: OpenAI Codex Signed-off-by: Greg Weyer Co-authored-by: Cursor --- tests/v1/worker/test_gpu_model_runner.py | 252 +++++++++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 107 +++++++++- 2 files changed, 352 insertions(+), 7 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 5e55c3e096b6..8f2bfa84b277 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -9,6 +9,7 @@ import numpy as np import pytest import torch +import torch.nn as nn import vllm.v1.worker.gpu_model_runner as gpu_model_runner_module from vllm.config import ( @@ -1961,3 +1962,254 @@ def test_iterator_reload_does_not_reload_draft_model(self): get_loader.assert_not_called() runner.get_draft_model.assert_not_called() target_model.load_weights.assert_called_once_with(weights) + + def test_disk_reload_skips_dim_mismatched_aliased_embeddings(self): + """Skip shared embed load; still fill the tied draft-dim lm_head.""" + + class _TiedHead(nn.Module): + def __init__(self, weight: nn.Parameter): + super().__init__() + self.weight = weight + + class _TinyLM(nn.Module): + def __init__(self, embed_dim: int): + super().__init__() + self.embed_tokens = nn.Embedding(8, embed_dim) + self.lm_head = _TiedHead(self.embed_tokens.weight) + self.proj = nn.Linear(4, 4, bias=False) + + def load_weights(self, weights): + from vllm.model_executor.models.utils import AutoWeightsLoader + + return AutoWeightsLoader(self).load_weights(weights) + + target_model = _TinyLM(embed_dim=8) + draft_model = _TinyLM(embed_dim=2) + original_embed = target_model.embed_tokens.weight.detach().clone() + original_lm_head = draft_model.lm_head.weight + draft_model.embed_tokens = target_model.embed_tokens + + runner = self._make_runner() + runner.get_draft_model = Mock(return_value=draft_model) + runner.get_model = Mock(return_value=target_model) + draft_placeholder = torch.full((8, 2), 7.0) + draft_proj = torch.eye(4) + target_weights = iter([("proj.weight", torch.eye(4))]) + seen_during_draft_init: list[nn.Module] = [] + + def get_loader(_load_config): + loader = Mock() + + def get_all_weights(_config, model): + if model is draft_model: + return iter( + [ + ("embed_tokens.weight", draft_placeholder), + ("proj.weight", draft_proj), + ] + ) + return target_weights + + loader.get_all_weights.side_effect = get_all_weights + return loader + + def capture_init(model): + if model is not draft_model: + return + seen_during_draft_init.extend( + module + for _name, module in model.named_modules(remove_duplicate=False) + if _name + ) + + with ( + patch.object( + gpu_model_runner_module, "get_model_loader", side_effect=get_loader + ), + patch.object( + gpu_model_runner_module, + "initialize_layerwise_reload", + side_effect=capture_init, + ), + patch.object(gpu_model_runner_module, "finalize_layerwise_reload"), + ): + runner.reload_weights() + + assert draft_model.embed_tokens is target_model.embed_tokens + assert torch.equal(target_model.embed_tokens.weight, original_embed) + assert draft_model.lm_head.weight is original_lm_head + assert torch.equal(draft_model.lm_head.weight, draft_placeholder) + assert torch.equal(draft_model.proj.weight, draft_proj) + assert target_model.embed_tokens not in seen_during_draft_init + + def test_disk_reload_keeps_untied_head_off_shared_embed(self): + """Same-width shared embed is discarded. + A distinct lm_head loads from its own checkpoint tensor, not embed_tokens. + """ + + class _Head(nn.Module): + def __init__(self, weight: nn.Parameter): + super().__init__() + self.weight = weight + + class _TinyLM(nn.Module): + def __init__(self, embed_dim: int): + super().__init__() + self.embed_tokens = nn.Embedding(8, embed_dim) + self.lm_head = _Head(nn.Parameter(torch.zeros(8, embed_dim))) + self.proj = nn.Linear(4, 4, bias=False) + + def load_weights(self, weights): + from vllm.model_executor.models.utils import AutoWeightsLoader + + return AutoWeightsLoader(self).load_weights(weights) + + target_model = _TinyLM(embed_dim=8) + draft_model = _TinyLM(embed_dim=8) + original_embed = target_model.embed_tokens.weight.detach().clone() + original_head = draft_model.lm_head.weight.detach().clone() + draft_model.embed_tokens = target_model.embed_tokens + + runner = self._make_runner() + runner.get_draft_model = Mock(return_value=draft_model) + runner.get_model = Mock(return_value=target_model) + embed_placeholder = torch.full((8, 8), 7.0) + head_from_ckpt = torch.full((8, 8), 3.0) + draft_proj = torch.eye(4) + target_weights = iter([("proj.weight", torch.eye(4))]) + + def get_loader(_load_config): + loader = Mock() + + def get_all_weights(_config, model): + if model is draft_model: + return iter( + [ + ("embed_tokens.weight", embed_placeholder), + ("lm_head.weight", head_from_ckpt), + ("proj.weight", draft_proj), + ] + ) + return target_weights + + loader.get_all_weights.side_effect = get_all_weights + return loader + + with ( + patch.object( + gpu_model_runner_module, "get_model_loader", side_effect=get_loader + ), + patch.object(gpu_model_runner_module, "initialize_layerwise_reload"), + patch.object(gpu_model_runner_module, "finalize_layerwise_reload"), + ): + runner.reload_weights() + + assert draft_model.embed_tokens is target_model.embed_tokens + assert torch.equal(target_model.embed_tokens.weight, original_embed) + assert not torch.equal(draft_model.lm_head.weight, original_head) + assert torch.equal(draft_model.lm_head.weight, head_from_ckpt) + assert torch.equal(draft_model.proj.weight, draft_proj) + + def test_disk_reload_detaches_repeated_lm_head_aliases(self): + """MTP shares one target lm_head at draft.lm_head and each shared_head.""" + + class _SharedHead(nn.Module): + def __init__(self): + super().__init__() + self.head = nn.Linear(2, 8, bias=False) + + class _Layer(nn.Module): + def __init__(self): + super().__init__() + self.shared_head = _SharedHead() + + class _Inner(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Layer(), _Layer()]) + + class _Draft(nn.Module): + def __init__(self): + super().__init__() + self.lm_head = nn.Linear(2, 8, bias=False) + self.model = _Inner() + self.proj = nn.Linear(4, 4, bias=False) + + def load_weights(self, weights): + from vllm.model_executor.models.utils import AutoWeightsLoader + + return AutoWeightsLoader(self).load_weights(weights) + + class _Target(nn.Module): + def __init__(self): + super().__init__() + self.lm_head = nn.Linear(8, 8, bias=False) + self.proj = nn.Linear(4, 4, bias=False) + + def load_weights(self, weights): + from vllm.model_executor.models.utils import AutoWeightsLoader + + return AutoWeightsLoader(self).load_weights(weights) + + target_model = _Target() + draft_model = _Draft() + original_head = target_model.lm_head.weight.detach().clone() + draft_model.lm_head = target_model.lm_head + for layer in draft_model.model.layers: + layer.shared_head.head = target_model.lm_head + + runner = self._make_runner() + runner.get_draft_model = Mock(return_value=draft_model) + runner.get_model = Mock(return_value=target_model) + placeholder = torch.full((8, 2), 7.0) + draft_proj = torch.eye(4) + target_weights = iter([("proj.weight", torch.eye(4))]) + seen_target_head: list[bool] = [] + + def get_loader(_load_config): + loader = Mock() + + def get_all_weights(_config, model): + if model is draft_model: + return iter( + [ + ("lm_head.weight", placeholder), + ("model.layers.0.shared_head.head.weight", placeholder), + ("model.layers.1.shared_head.head.weight", placeholder), + ("proj.weight", draft_proj), + ] + ) + return target_weights + + loader.get_all_weights.side_effect = get_all_weights + return loader + + def capture_init(model): + if model is not draft_model: + return + seen_target_head.append( + any( + module is target_model.lm_head + for _name, module in model.named_modules(remove_duplicate=False) + ) + ) + + with ( + patch.object( + gpu_model_runner_module, "get_model_loader", side_effect=get_loader + ), + patch.object( + gpu_model_runner_module, + "initialize_layerwise_reload", + side_effect=capture_init, + ), + patch.object(gpu_model_runner_module, "finalize_layerwise_reload"), + ): + runner.reload_weights() + + assert seen_target_head == [False] + assert draft_model.lm_head is target_model.lm_head + assert draft_model.model.layers[0].shared_head.head is target_model.lm_head + assert draft_model.model.layers[1].shared_head.head is target_model.lm_head + assert torch.equal(target_model.lm_head.weight, original_head) + assert torch.equal(draft_model.proj.weight, draft_proj) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 8078f6f5f675..af453f9f5d0c 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -495,6 +495,97 @@ class ExecuteModelState(NamedTuple): slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None +class _DiscardSharedDraftWeights(nn.Module): + """Drop tensors for draft submodules that are the live target object.""" + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + for _name, _weight in weights: + pass + return set() + + +class _TiedDraftWeightSink(nn.Module): + """Fill draft lm_head from tied embed tensors; do not write the target embed. + Uses load_weights so AutoWeightsLoader can dispatch without dual-registering. + """ + + def __init__(self, src: nn.Module): + super().__init__() + self._dst = src + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded: set[str] = set() + dst = self._dst + for name, tensor in weights: + leaf = name.split(".")[-1] + param = getattr(dst, leaf, None) + if not isinstance(param, nn.Parameter): + continue + loader = getattr(param, "weight_loader", None) + if callable(loader): + loader(param, tensor) + else: + param.data.copy_(tensor) + loaded.add(name) + return loaded + + +def _replacement_for_detached_draft_module( + attr: str, module: nn.Module, draft_model: nn.Module +) -> nn.Module: + if attr == "embed_tokens": + lm_head = getattr(draft_model, "lm_head", None) + detached_weight = getattr(module, "weight", None) + head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None + if ( + isinstance(detached_weight, nn.Parameter) + and isinstance(head_weight, nn.Parameter) + and head_weight is not detached_weight + and head_weight.shape != detached_weight.shape + ): + return _TiedDraftWeightSink(lm_head) + return _DiscardSharedDraftWeights() + + +@contextmanager +def _temporarily_detach_target_owned_draft_modules( + draft_model: object, target_model: object +) -> Iterator[None]: + """Unbind target-owned draft aliases for the duration of draft load_weights. + Dim-mismatched shared embeds (Gemma4 MTP after embed replace) mismatch otherwise. + """ + if not isinstance(draft_model, nn.Module) or not isinstance( + target_model, nn.Module + ): + yield + return + + target_module_ids = {id(module) for module in target_model.modules()} + attachments: list[tuple[nn.Module, str, nn.Module]] = [] + # remove_duplicate=False: MTP aliases the same target lm_head at + # draft.lm_head and every layer.shared_head.head. + for name, module in draft_model.named_modules(remove_duplicate=False): + if not name or id(module) not in target_module_ids: + continue + parent_name, sep, attr = name.rpartition(".") + parent = draft_model if not sep else draft_model.get_submodule(parent_name) + if id(parent) in target_module_ids: + continue + attachments.append((parent, attr, module)) + + for parent, attr, module in attachments: + setattr( + parent, + attr, + _replacement_for_detached_draft_module(attr, module, draft_model), + ) + try: + yield + finally: + for parent, attr, module in attachments: + setattr(parent, attr, module) + + class GPUModelRunner( LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin ): @@ -5569,11 +5660,13 @@ def _reload_draft_weights_from_disk(self) -> None: f"Draft model reloading with `{load_config.load_format}` format" ) - initialize_layerwise_reload(draft_model) - draft_model.load_weights( - model_loader.get_all_weights(draft_model_config, draft_model) - ) - finalize_layerwise_reload(draft_model, draft_model_config) + target_model = self.get_model() + with _temporarily_detach_target_owned_draft_modules(draft_model, target_model): + initialize_layerwise_reload(draft_model) + draft_model.load_weights( + model_loader.get_all_weights(draft_model_config, draft_model) + ) + finalize_layerwise_reload(draft_model, draft_model_config) def reload_weights( self, @@ -5609,8 +5702,8 @@ def reload_weights( # load weights from disk if none are provided if weights_iterator is None: - # Reload draft weights (dropped by L2 sleep) before the target so it - # stays authoritative for shared storage. Unbound: MRV2 passes a V2 self. + # Draft first (L2 dropped its params). Target-owned aliases are + # detached during that load. Unbound: MRV2 passes a V2 self. GPUModelRunner._reload_draft_weights_from_disk(self) model_loader = get_model_loader(self.load_config) From fee28d09725f9def060d3e9e93591f1e7290e34f Mon Sep 17 00:00:00 2001 From: Greg Weyer Date: Wed, 2 Sep 2026 14:46:26 -0400 Subject: [PATCH 4/5] Restrict tied-embed sink to hidden-axis mismatch The tied-placeholder layout (Gemma4 MTP) keeps the embedding's vocab width and differs only on the hidden axis. A draft lm_head that is narrower on the vocab axis (EAGLE3 low-rank heads, draft vocab 32000 vs the shared target embedding) is a real projection with its own checkpoint entry, not a tied placeholder, so its embed_tokens tensors must be discarded rather than redirected into the head. Verified live: on EAGLE3-LLaMA3.1-Instruct-8B + Llama-3.1-8B-Instruct (TP=2, num_speculative_tokens=3, L2 sleep) the old condition engaged the sink on the vocab-axis mismatch; the checkpoint ships no embed tensors so it consumed nothing and acceptance was preserved either way. The guard removes that false engagement; the Gemma4 tied case is unchanged and the new unit test fails without the guard. Signed-off-by: Greg Weyer --- tests/v1/worker/test_gpu_model_runner.py | 67 ++++++++++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 5 ++ 2 files changed, 72 insertions(+) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 8f2bfa84b277..1745ea7b8ab9 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -2110,6 +2110,73 @@ def get_all_weights(_config, model): assert torch.equal(draft_model.lm_head.weight, head_from_ckpt) assert torch.equal(draft_model.proj.weight, draft_proj) + def test_disk_reload_discards_vocab_mismatched_shared_embed(self): + """A head narrower on the vocab axis (EAGLE3 low-rank layout) is a real + projection, not a tied placeholder: embed tensors are discarded and the + head still loads from its own checkpoint entry. + """ + + class _Head(nn.Module): + def __init__(self, weight: nn.Parameter): + super().__init__() + self.weight = weight + + class _TinyLM(nn.Module): + def __init__(self, embed_dim: int, head_vocab: int): + super().__init__() + self.embed_tokens = nn.Embedding(8, embed_dim) + self.lm_head = _Head(nn.Parameter(torch.zeros(head_vocab, embed_dim))) + self.proj = nn.Linear(4, 4, bias=False) + + def load_weights(self, weights): + from vllm.model_executor.models.utils import AutoWeightsLoader + + return AutoWeightsLoader(self).load_weights(weights) + + target_model = _TinyLM(embed_dim=8, head_vocab=8) + draft_model = _TinyLM(embed_dim=8, head_vocab=4) + original_embed = target_model.embed_tokens.weight.detach().clone() + draft_model.embed_tokens = target_model.embed_tokens + + runner = self._make_runner() + runner.get_draft_model = Mock(return_value=draft_model) + runner.get_model = Mock(return_value=target_model) + embed_placeholder = torch.full((8, 8), 7.0) + head_from_ckpt = torch.full((4, 8), 3.0) + draft_proj = torch.eye(4) + target_weights = iter([("proj.weight", torch.eye(4))]) + + def get_loader(_load_config): + loader = Mock() + + def get_all_weights(_config, model): + if model is draft_model: + return iter( + [ + ("embed_tokens.weight", embed_placeholder), + ("lm_head.weight", head_from_ckpt), + ("proj.weight", draft_proj), + ] + ) + return target_weights + + loader.get_all_weights.side_effect = get_all_weights + return loader + + with ( + patch.object( + gpu_model_runner_module, "get_model_loader", side_effect=get_loader + ), + patch.object(gpu_model_runner_module, "initialize_layerwise_reload"), + patch.object(gpu_model_runner_module, "finalize_layerwise_reload"), + ): + runner.reload_weights() + + assert draft_model.embed_tokens is target_model.embed_tokens + assert torch.equal(target_model.embed_tokens.weight, original_embed) + assert torch.equal(draft_model.lm_head.weight, head_from_ckpt) + assert torch.equal(draft_model.proj.weight, draft_proj) + def test_disk_reload_detaches_repeated_lm_head_aliases(self): """MTP shares one target lm_head at draft.lm_head and each shared_head.""" diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index af453f9f5d0c..fd7d8daaf261 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -537,10 +537,15 @@ def _replacement_for_detached_draft_module( lm_head = getattr(draft_model, "lm_head", None) detached_weight = getattr(module, "weight", None) head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None + # A tied placeholder head (Gemma4 MTP) keeps the embedding's vocab + # width and differs only on the hidden axis; a head that differs on + # the vocab axis (EAGLE3 low-rank heads) is a real projection with + # its own checkpoint entry, so its embed tensors are discarded. if ( isinstance(detached_weight, nn.Parameter) and isinstance(head_weight, nn.Parameter) and head_weight is not detached_weight + and head_weight.shape[0] == detached_weight.shape[0] and head_weight.shape != detached_weight.shape ): return _TiedDraftWeightSink(lm_head) From 22d6effaf268a2a5f9a824b311e064789e84a00b Mon Sep 17 00:00:00 2001 From: Greg Weyer Date: Wed, 2 Sep 2026 15:07:40 -0400 Subject: [PATCH 5/5] Rebuild draft fused state after layerwise finalization The dflash/dspark drafters rebuild fused buffers (a stacked copy of layer weights plus a cached rotary-embedding reference) at the end of load_weights. Under the disk draft reload that rebuild now runs inside the layerwise lifecycle, so it captures the temporarily meta-derived rotary buffer; finalization then restores the real buffer on the module tree and the cached reference is left on meta, crashing the first fused RoPE dispatch after the wake. Re-run the drafter's own rebuild after finalize_layerwise_reload so the cached references are re-captured from the restored tree. No-op for drafters without _build_fused_kv_buffers. Verified live on Qwen3.8-27B-FP8 TP=2: DSpark (K=7) previously crashed post-wake and now serves with acceptance restored across repeated L2 cycles. EAGLE3 and MTP drafts do not define the builder and are unaffected. Signed-off-by: Greg Weyer --- tests/v1/worker/test_gpu_model_runner.py | 71 ++++++++++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 15 +++++ 2 files changed, 86 insertions(+) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 1745ea7b8ab9..d4829d60f289 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -2177,6 +2177,77 @@ def get_all_weights(_config, model): assert torch.equal(draft_model.lm_head.weight, head_from_ckpt) assert torch.equal(draft_model.proj.weight, draft_proj) + def test_disk_reload_rebuilds_fused_state_after_finalization(self): + """Drafters whose load_weights rebuilds fused buffers mid-lifecycle + (dflash/dspark) capture the temporarily meta-derived rotary buffer; + the runner must re-run that rebuild after finalize restores the tree. + """ + + class _Rotary(nn.Module): + def __init__(self): + super().__init__() + self.cache = torch.zeros(4) + + class _Inner(nn.Module): + def __init__(self): + super().__init__() + self.rotary_emb = _Rotary() + self.captured = None + self.rebuild_calls = 0 + + def _build_fused_kv_buffers(self) -> None: + self.rebuild_calls += 1 + self.captured = self.rotary_emb.cache + + def load_weights(self, weights): + self._build_fused_kv_buffers() + return set() + + class _Draft(nn.Module): + def __init__(self): + super().__init__() + self.model = _Inner() + + def load_weights(self, weights): + return self.model.load_weights(weights) + + draft_model = _Draft() + target_model = Mock() + target_model.named_parameters.return_value = [] + target_model.load_weights.return_value = set() + runner = self._make_runner() + runner.get_draft_model = Mock(return_value=draft_model) + runner.get_model = Mock(return_value=target_model) + model_loader = Mock() + model_loader.get_all_weights.return_value = iter([]) + restored = _Rotary() + + def finalize(model, _config): + # Emulate the real finalize: the module tree ends up holding a + # different, healthy buffer object than the meta one load_weights + # saw mid-lifecycle. + model.model.rotary_emb = restored + + with ( + patch.object( + gpu_model_runner_module, + "get_model_loader", + return_value=model_loader, + ), + patch.object(gpu_model_runner_module, "initialize_layerwise_reload"), + patch.object( + gpu_model_runner_module, + "finalize_layerwise_reload", + side_effect=finalize, + ), + ): + runner.reload_weights() + + # Mid-lifecycle rebuild from load_weights, then the post-finalize + # re-run: the cached reference must be the restored tree object. + assert draft_model.model.rebuild_calls == 2 + assert draft_model.model.captured is restored.cache + def test_disk_reload_detaches_repeated_lm_head_aliases(self): """MTP shares one target lm_head at draft.lm_head and each shared_head.""" diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index fd7d8daaf261..de07496474d1 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5673,6 +5673,21 @@ def _reload_draft_weights_from_disk(self) -> None: ) finalize_layerwise_reload(draft_model, draft_model_config) + # The dflash/dspark drafters rebuild fused buffers (a stacked copy of + # layer weights plus a cached rotary embedding reference) at the end of + # load_weights. Inside the layerwise lifecycle that rebuild captures + # the temporarily meta-derived rotary buffer; finalization then + # restores the real buffer on the module tree, leaving the cached + # reference on meta. Re-running the drafter's own rebuild after + # finalization re-captures the restored objects. No-op elsewhere. + rebuild_fused_state = getattr( + getattr(draft_model, "model", None) or draft_model, + "_build_fused_kv_buffers", + None, + ) + if callable(rebuild_fused_state): + rebuild_fused_state() + def reload_weights( self, weights_iterator: Iterable[tuple[str, torch.Tensor]] | None = None,