From 62c8695dfd831f129b2e3021ae6d5c1673c51cce Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 13 Mar 2026 03:16:33 +0000 Subject: [PATCH] Update GPUGenerationModelRunner to clarify KV connector metadata handling during draft model runs - Changed variable name to better reflect its purpose in managing KV connector metadata. - Updated comments for improved clarity regarding the handling of speculative configurations. Signed-off-by: tzhouam --- vllm_omni/worker/gpu_generation_model_runner.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 0dd304ed2..58e4dfd72 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -265,9 +265,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # Run the model. # Use persistent buffers for CUDA graphs. - # When spec decode is enabled, defer connector finalization - # (wait_for_save + clear metadata) until after draft model runs. - defer_kv_connector_finalize = self.speculative_config is not None + # When spec decode is enabled, defer clearing KV connector metadata + # until after draft model runs in sample_tokens. + clear_kv_metadata = self.speculative_config is None with ( set_forward_context( attn_metadata, @@ -282,7 +282,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output( scheduler_output, - defer_finalize=defer_kv_connector_finalize, + clear_metadata=clear_kv_metadata, ) as kv_connector_output, ): outputs = self._run_generation_model(