From 1ccffce2bc7c5967805f17a109cf0725f5c7c906 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 6 Mar 2026 02:32:26 +0000 Subject: [PATCH] rebase to updated commit Signed-off-by: tzhouam --- docker/Dockerfile.ci | 4 ++-- vllm_omni/worker/gpu_ar_model_runner.py | 14 +++++++++----- vllm_omni/worker/gpu_generation_model_runner.py | 7 +++++-- vllm_omni/worker/gpu_model_runner.py | 7 +++++-- 4 files changed, 21 insertions(+), 11 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index dc32df22e..9d2f49344 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -14,11 +14,11 @@ RUN apt-get update && \ # Install vLLM wheel for the specific commit # URL format: https://wheels.vllm.ai// -# Current vLLM wheel commit: 3dd142adb7f7fa3b0b8345ee046ba36a12d9e530 (releases/v0.17.0) +# Current vLLM wheel commit: 097eb544e9a22810c9b7a59e586b61627b308362 (main) # Note: --prerelease=allow is required because nightly/commit wheels use # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default, # causing it to fall back to the older stable release from PyPI. -RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/e196028dbaa41b8e3d1c27b7dbd026403af76c27 +RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/097eb544e9a22810c9b7a59e586b61627b308362 RUN uv pip install --system --no-cache-dir ".[dev]" RUN uv pip install --system --upgrade \ diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 72ccaa195..1ea56f724 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -24,6 +24,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( @@ -382,13 +383,12 @@ class GPUARModelRunner(OmniGPUModelRunner): self, grammar_output: GrammarOutput | None, ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: - kv_connector_output = self.kv_connector_output - self.kv_connector_output = None - kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None) self.kv_extracted_req_ids = None if self.execute_model_state is None: + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None # Nothing to do (PP non-final rank case), output isn't used. if not kv_connector_output: return None # type: ignore[return-value] @@ -461,12 +461,12 @@ class GPUARModelRunner(OmniGPUModelRunner): <= self.effective_drafter_max_model_len ) use_gpu_toks = ( - spec_config.use_eagle() or spec_config.uses_draft_model() + spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states() ) and not spec_config.disable_padded_drafter_batch if use_gpu_toks: assert isinstance( self.drafter, - EagleProposer | DraftModelProposer, + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer, ) sampled_token_ids = sampler_output.sampled_token_ids if input_fits_in_drafter: @@ -522,6 +522,10 @@ class GPUARModelRunner(OmniGPUModelRunner): with record_function_or_nullcontext("gpu_model_runner: eplb"): self.eplb_step() + # kv_connector_output may be modified during drafting + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None) if num_scheduled_tokens_np is None: diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 71f173424..05785d7aa 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -25,6 +25,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, @@ -736,11 +737,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) if self.speculative_config and ( - self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() ): assert isinstance( self.drafter, - EagleProposer | DraftModelProposer, + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer, ) assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 1d70e0eb3..dc353d8d4 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -15,6 +15,7 @@ from vllm.utils.import_utils import LazyLoader from vllm.utils.math_utils import cdiv from vllm.v1.spec_decode.draft_model import DraftModelProposer from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices @@ -797,11 +798,13 @@ class OmniGPUModelRunner(GPUModelRunner): hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) if self.speculative_config and ( - self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() ): assert isinstance( self.drafter, - EagleProposer | DraftModelProposer, + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer, ) assert self.speculative_config is not None # Eagle currently only supports PIECEWISE cudagraphs.