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 7, 22:45 UTC
[Spec Decode] Defer clearing KV connector metadata for EAGLE3 speculative decode + prefill / decode disagg setup (#34529)
Signed-off-by: qizixi <qizixi@meta.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
qizixi committed
Feb 22, 2026 at 08:08 UTC
b9c2a565ccfd5cd865ebca47cfb7ef0f7a8ba76c
3 files changed
+38
-5
vllm/v1/worker/gpu/kv_connector.py
+11
-2
@@ -77,7 +77,10 @@ class ActiveKVConnector(KVConnector):
77
self.kv_connector.start_load_kv(get_forward_context())
78
79
def post_forward(
80
- self, scheduler_output: "SchedulerOutput", wait_for_save: bool = True
80
+ self,
81
+ scheduler_output: "SchedulerOutput",
82
+ wait_for_save: bool = True,
83
+ clear_metadata: bool = True,
84
) -> KVConnectorOutput | None:
85
if self._disabled:
86
return None
@@ -91,9 +94,15 @@ class ActiveKVConnector(KVConnector):
94
output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors()
95
output.kv_connector_stats = self.kv_connector.get_kv_connector_stats()
96
output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events()
94
- self.kv_connector.clear_connector_metadata()
97
+ if clear_metadata:
98
+ self.kv_connector.clear_connector_metadata()
99
return output
100
101
+ def clear_metadata(self) -> None:
102
+ """Clear the connector metadata. Call this after draft model runs."""
103
+ if not self._disabled:
104
+ self.kv_connector.clear_connector_metadata()
105
+
106
def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
107
if self._disabled:
108
return EMPTY_MODEL_RUNNER_OUTPUT
vllm/v1/worker/gpu_model_runner.py
+12
-1
@@ -3524,6 +3524,9 @@ class GPUModelRunner(
3524
3525
# Run the model.
3526
# Use persistent buffers for CUDA graphs.
3527
+ # When spec decode is enabled, delay clearing connector metadata
3528
+ # until after draft model runs in sample_tokens.
3529
+ clear_kv_metadata = self.speculative_config is None
3530
with (
3531
set_forward_context(
3532
attn_metadata,
@@ -3537,7 +3540,9 @@ class GPUModelRunner(
3540
skip_compiled=has_encoder_input,
3541
),
3542
record_function_or_nullcontext("gpu_model_runner: forward"),
3540
- self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
3543
+ self.maybe_get_kv_connector_output(
3544
+ scheduler_output, clear_metadata=clear_kv_metadata
3545
+ ) as kv_connector_output,
3546
):
3547
model_output = self._model_forward(
3548
input_ids=input_ids,
@@ -3765,6 +3770,12 @@ class GPUModelRunner(
3770
# tokens on the CPU, so they are run after bookkeeping.
3771
propose_draft_token_ids(valid_sampled_token_ids)
3772
3773
+ # Clear KV connector metadata after draft model runs (if spec decode).
3774
+ # This was deferred from target model forward to allow draft model
3775
+ # to also save its KV cache.
3776
+ if self.speculative_config is not None:
3777
+ self.clear_kv_connector_metadata()
3778
+
3779
with record_function_or_nullcontext("gpu_model_runner: eplb"):
3780
self.eplb_step()
3781
vllm/v1/worker/kv_connector_model_runner_mixin.py
+15
-2
@@ -67,9 +67,12 @@ class KVConnectorModelRunnerMixin:
67
@staticmethod
68
def maybe_get_kv_connector_output(
69
scheduler_output: "SchedulerOutput",
70
+ clear_metadata: bool = True,
71
) -> AbstractContextManager[KVConnectorOutput | None]:
72
return (
72
- KVConnectorModelRunnerMixin._get_kv_connector_output(scheduler_output)
73
+ KVConnectorModelRunnerMixin._get_kv_connector_output(
74
+ scheduler_output, clear_metadata=clear_metadata
75
+ )
76
if has_kv_transfer_group()
77
else nullcontext()
78
)
@@ -79,7 +82,9 @@ class KVConnectorModelRunnerMixin:
82
@staticmethod
83
@contextmanager
84
def _get_kv_connector_output(
82
- scheduler_output: "SchedulerOutput", wait_for_save: bool = True
85
+ scheduler_output: "SchedulerOutput",
86
+ wait_for_save: bool = True,
87
+ clear_metadata: bool = True,
88
) -> Generator[KVConnectorOutput, None, None]:
89
output = KVConnectorOutput()
90
@@ -108,6 +113,14 @@ class KVConnectorModelRunnerMixin:
113
output.kv_connector_stats = kv_connector.get_kv_connector_stats()
114
output.kv_cache_events = kv_connector.get_kv_connector_kv_cache_events()
115
116
+ if clear_metadata:
117
+ kv_connector.clear_connector_metadata()
118
+
119
+ @staticmethod
120
+ def clear_kv_connector_metadata() -> None:
121
+ """Clear the KV connector metadata. Call after draft model runs."""
122
+ if has_kv_transfer_group():
123
+ kv_connector = get_kv_transfer_group()
124
kv_connector.clear_connector_metadata()
125
126
@staticmethod