From 4414701161cc8c77b956d35215badf87802dd449 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sun, 4 Oct 2026 13:40:24 +0800 Subject: [PATCH] Use published release images and wire NPU auxiliary outputs Signed-off-by: tzhouam --- docker/Dockerfile.ci | 6 +- docker/Dockerfile.cuda | 3 +- docker/Dockerfile.rocm | 3 +- tests/buildkite/test_rocm_dockerfile.py | 8 +- tests/engine/test_arg_utils_shared_fields.py | 8 +- tests/platforms/npu/test_aux_output.py | 160 ++++++++++++++++++ vllm_omni/engine/arg_utils.py | 4 +- vllm_omni/platforms/npu/worker/aux_output.py | 96 +++++++++++ vllm_omni/platforms/npu/worker/base.py | 8 + .../npu/worker/npu_ar_model_runner.py | 19 +-- .../npu/worker/npu_generation_model_runner.py | 17 +- .../platforms/npu/worker/npu_model_runner.py | 4 +- 12 files changed, 294 insertions(+), 42 deletions(-) create mode 100644 tests/platforms/npu/test_aux_output.py create mode 100644 vllm_omni/platforms/npu/worker/aux_output.py diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index c1e9cdbe2..949fa0970 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,9 +1,7 @@ # Use the official release image with its matching CUDA/PyTorch stack. -# The release images are available in upstream ECR before Docker Hub promotion. ARG VLLM_VERSION=0.31.0 -ARG VLLM_PRECOMPILED_WHEEL_COMMIT=db9527a46873454610df6dbedf79a36d6bf1a7f6 -ARG VLLM_BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo -ARG VLLM_BASE_TAG=${VLLM_PRECOMPILED_WHEEL_COMMIT} +ARG VLLM_BASE_IMAGE=vllm/vllm-openai +ARG VLLM_BASE_TAG=v${VLLM_VERSION} # We use registry-based layer caching for the dependencies, # where the tag is determined based on the hash of the contents diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index cb9af0753..69b598ab2 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,5 +1,4 @@ -# Official vLLM v0.31.0 release image, before Docker Hub tag promotion. -ARG BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:db9527a46873454610df6dbedf79a36d6bf1a7f6 +ARG BASE_IMAGE=vllm/vllm-openai:v0.31.0 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 04dff295a..379262afa 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,8 +1,7 @@ # Keep the ROCm runtime aligned with the vLLM release targeted by # docker/Dockerfile.ci. A stale base image can build successfully while failing # at runtime when vLLM-Omni imports APIs added by the newer vLLM release. -# Official vLLM v0.31.0 release image, before Docker Hub tag promotion. -ARG BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:db9527a46873454610df6dbedf79a36d6bf1a7f6-rocm +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0 # Keep the dependency-heavy portion independent from the source tree. The # registry cache is keyed on these inputs, while the final source copy can be diff --git a/tests/buildkite/test_rocm_dockerfile.py b/tests/buildkite/test_rocm_dockerfile.py index fc42fc478..4c40a7c7b 100644 --- a/tests/buildkite/test_rocm_dockerfile.py +++ b/tests/buildkite/test_rocm_dockerfile.py @@ -38,9 +38,11 @@ def test_rocm_base_tracks_cuda_vllm_release() -> None: image_ref, separator, image_tag = rocm_base.rpartition(":") assert separator, f"expected a tagged ROCm base image, got {rocm_base}" - assert image_ref == cuda_image - assert image_tag == f"{cuda_release}-rocm" - assert cuda_release == _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_PRECOMPILED_WHEEL_COMMIT") + assert cuda_image == "vllm/vllm-openai" + assert image_ref == "vllm/vllm-openai-rocm" + assert image_tag == cuda_release == f"v{_docker_arg(CUDA_CI_DOCKERFILE, 'VLLM_VERSION')}" + assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_IMAGE") == cuda_image + assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_TAG") == "v${VLLM_VERSION}" def test_rocm_defaults_to_prebuilt_base_image() -> None: diff --git a/tests/engine/test_arg_utils_shared_fields.py b/tests/engine/test_arg_utils_shared_fields.py index e5c20ccd5..cdf5f272d 100644 --- a/tests/engine/test_arg_utils_shared_fields.py +++ b/tests/engine/test_arg_utils_shared_fields.py @@ -61,12 +61,14 @@ def test_internal_blacklist_keys_derived_from_orchestrator(): ], ) @pytest.mark.parametrize("use_v2_model_runner", [False, True]) -def test_aux_outputs_require_the_actual_omni_stage_v2_runner(aux_args, use_v2_model_runner, monkeypatch): +@pytest.mark.parametrize("is_npu", [False, True]) +def test_aux_outputs_require_a_supported_omni_stage_runner(aux_args, use_v2_model_runner, is_npu, monkeypatch): # Upstream runner selection reads this environment variable; Omni workers # select their runner from the separate per-stage flag instead. monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - if use_v2_model_runner: - args = OmniEngineArgs(model="unused", use_v2_model_runner=True, **aux_args) + monkeypatch.setattr("vllm_omni.engine.arg_utils.current_omni_platform.is_npu", lambda: is_npu) + if use_v2_model_runner or is_npu: + args = OmniEngineArgs(model="unused", use_v2_model_runner=use_v2_model_runner, **aux_args) assert args.aux_output_config.enabled else: with pytest.raises(ValueError, match="use_v2_model_runner=True for this Omni stage"): diff --git a/tests/platforms/npu/test_aux_output.py b/tests/platforms/npu/test_aux_output.py new file mode 100644 index 000000000..f7b0573cf --- /dev/null +++ b/tests/platforms/npu/test_aux_output.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU contract coverage for NPU routing capture and the native data plane.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.config import AuxOutputConfig +from vllm.distributed.aux_output_connector import worker +from vllm.distributed.aux_output_connector.connector import ( + AuxOutputConnectorMetadata, + AuxOutputSchedulerConnector, + PackedBlockHashes, +) +from vllm.model_executor.layers.fused_moe import routed_experts_capturer +from vllm.v1.kv_cache_interface import KVCacheConfig + +from vllm_omni.platforms.npu.worker.aux_output import NPUAuxOutputMixin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Router: + def _select_experts(self, logical_ids): + return torch.ones_like(logical_ids, dtype=torch.float32), logical_ids + + +class _AscendMoE(torch.nn.Module): + __module__ = "vllm_ascend.ops.fused_moe" + + def __init__(self): + super().__init__() + self.layer_id = 0 + self.router = _Router() + + +class _Runner(NPUAuxOutputMixin): + def _sync_device(self): + self.synchronized += 1 + + +@pytest.fixture +def runner(monkeypatch): + # Exercise the real capturer, block store, worker and scheduler connectors + # with CPU tensors; importing vllm-ascend or owning NPU hardware is unnecessary. + monkeypatch.setattr(routed_experts_capturer, "current_platform", SimpleNamespace(device_type="cpu")) + monkeypatch.setattr(routed_experts_capturer, "get_forward_context", lambda: SimpleNamespace(dp_metadata=None)) + monkeypatch.setattr(worker, "get_tp_group", lambda: SimpleNamespace(is_first_rank=True)) + runner = _Runner() + runner.model = torch.nn.Sequential(_AscendMoE()) + runner.vllm_config = SimpleNamespace( + aux_output_config=AuxOutputConfig(enable_return_routed_experts=True), + model_config=SimpleNamespace( + get_total_num_hidden_layers=lambda: 1, + get_num_experts=lambda: 16, + get_num_experts_per_tok=lambda: 2, + hf_text_config=SimpleNamespace(model_type="test_moe"), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=16, max_num_seqs=4), + parallel_config=SimpleNamespace(data_parallel_rank=0, tensor_parallel_size=1, decode_context_parallel_size=1), + cache_config=SimpleNamespace(block_size=4, prefix_match_unit=None), + max_concurrent_batches=2, + ) + runner.input_batch = SimpleNamespace(req_ids=["r"], num_computed_tokens_cpu=np.array([0])) + runner.query_start_loc = SimpleNamespace(np=np.array([0, 4])) + runner.synchronized = 0 + runner._init_omni_aux_output(KVCacheConfig(8, [], [])) + try: + yield runner + finally: + runner._close_omni_aux_output() + + +def _step(runner, *, emit_start=0, hashes=None, finished=()): + scheduler_output = SimpleNamespace( + aux_output_connector_metadata=AuxOutputConnectorMetadata( + generation=0, + requests={} if finished else {"r": emit_start}, + block_hashes=hashes or {}, + finished_requests=finished, + ), + scheduled_spec_decode_tokens={}, + ) + runner._begin_omni_aux_output_step(scheduler_output) + return scheduler_output + + +def _capture(runner, ids): + logical_ids = torch.tensor(ids, dtype=torch.int32) + weights, returned = runner.model[0].router._select_experts(logical_ids) + assert torch.equal(returned, logical_ids) + assert weights.shape == logical_ids.shape + return logical_ids.numpy()[:, None, :].astype(np.uint8) + + +def test_routing_reaches_scheduler_and_owns_legacy_batch_snapshot(runner): + step = _step(runner) + expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]]) + pending = runner._prepare_omni_aux_output() + # Simulate bookkeeping and the next batch reusing the runner's buffers. + runner.input_batch.num_computed_tokens_cpu[:] = 99 + runner.query_start_loc.np[:] = 99 + _capture(runner, [[7, 8]] * 4) + output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]])) + request = SimpleNamespace(request_id="r", num_tokens=5) + rows = AuxOutputSchedulerConnector().take_output(request, output) + np.testing.assert_array_equal(rows, expected) + assert output["r"].token_start == 0 + assert runner.synchronized == 1 + + +def test_speculative_rejection_drops_unaccepted_routing_suffix(runner): + step = _step(runner) + step.scheduled_spec_decode_tokens = {"r": [7, 8, 9]} + expected = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]]) + pending = runner._prepare_omni_aux_output() + output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5, 6, -1, -1]])) + np.testing.assert_array_equal(output["r"].rows, expected[:2]) + + +def test_partial_prefill_accumulates_rows_and_no_forward_step_tears_down(runner): + step = _step(runner) + prefix = _capture(runner, [[0, 1], [1, 2], [2, 3], [3, 4]]) + pending = runner._prepare_omni_aux_output() + assert runner._finish_omni_aux_output(pending, step, torch.tensor([[5]]), [0]) == {} + + step = _step(runner, hashes={"r": PackedBlockHashes(b"h0", 2)}) + runner.input_batch.num_computed_tokens_cpu[:] = 4 + suffix = _capture(runner, [[4, 5], [5, 6], [6, 7], [7, 8]]) + pending = runner._prepare_omni_aux_output() + output = runner._finish_omni_aux_output(pending, step, torch.tensor([[5]])) + np.testing.assert_array_equal(output["r"].rows, np.concatenate((prefix, suffix))) + _step(runner, finished=("r",)) + assert not runner.aux_output_connector._requests + store = runner.aux_output_connector._store + runner._close_omni_aux_output() + assert not store._thread.is_alive() + assert runner.model[0].capture_fn is None + runner._close_omni_aux_output() # shutdown is idempotent + + +def test_reinitialization_does_not_wrap_router_twice(runner): + selection = runner.model[0].router._select_experts + runner._init_omni_aux_output(KVCacheConfig(8, [], [])) + assert runner.model[0].router._select_experts is selection + step = _step(runner) + expected = _capture(runner, [[0, 1]] * 4) + output = runner._finish_omni_aux_output(runner._prepare_omni_aux_output(), step, torch.tensor([[5]])) + np.testing.assert_array_equal(output["r"].rows, expected) + + +def test_disabled_capture_preserves_normal_runner_path(runner): + runner._close_omni_aux_output() + runner.vllm_config.aux_output_config.enable_return_routed_experts = False + runner._init_omni_aux_output(KVCacheConfig(8, [], [])) + runner._begin_omni_aux_output_step(SimpleNamespace()) + assert runner._prepare_omni_aux_output() is None + assert runner._finish_omni_aux_output(None, SimpleNamespace()) is None diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 358276e30..a46fc49e2 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -270,7 +270,9 @@ class OmniEngineArgs(EngineArgs): ) validate_worker_omni_connector(self.worker_cls, needs_connector) super().__post_init__() - if self.aux_output_config.enabled and not self.use_v2_model_runner: + # The NPU runner implements the auxiliary connector on its legacy + # execution path; CUDA/ROCm Omni stages use the V2 implementation. + if self.aux_output_config.enabled and not self.use_v2_model_runner and not current_omni_platform.is_npu(): raise ValueError("Auxiliary outputs require use_v2_model_runner=True for this Omni stage.") def _ensure_omni_models_registered(self): diff --git a/vllm_omni/platforms/npu/worker/aux_output.py b/vllm_omni/platforms/npu/worker/aux_output.py new file mode 100644 index 000000000..3fc093ae4 --- /dev/null +++ b/vllm_omni/platforms/npu/worker/aux_output.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Auxiliary-output delivery for Ascend's legacy model runner.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np + + +def _install_ascend_capture_sources(model: Any) -> None: + # AscendRoutedExperts invokes the router's private selection method, before + # logical-to-physical remapping and shared-expert expansion. It bypasses + # BaseRouter.select_experts(), where upstream normally invokes capture_fn. + for module in model.modules(): + if not type(module).__module__.startswith("vllm_ascend.") or not hasattr(module, "layer_id"): + continue + router = getattr(module, "router", None) + select_experts = getattr(router, "_select_experts", None) + if not callable(select_experts) or getattr(module, "_omni_npu_capture_source", False): + continue + module.capture_fn = None + + def select_and_capture(*args, _select=select_experts, _source=module, **kwargs): + weights, logical_ids = _select(*args, **kwargs) + if _source.capture_fn is not None: + _source.capture_fn(logical_ids) + return weights, logical_ids + + router._select_experts = select_and_capture + # layer_id + capture_fn implement upstream RoutedExpertsCaptureSource. + module._omni_npu_capture_source = True + + +class NPUAuxOutputMixin: + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.aux_output_connector = None + + def _init_omni_aux_output(self, kv_cache_config) -> None: + if not self.vllm_config.aux_output_config.enabled: + return + from vllm.distributed.aux_output_connector.worker import get_aux_output_connector + + self._close_omni_aux_output() + _install_ascend_capture_sources(self.model) + self.aux_output_connector = get_aux_output_connector(self.model, self.vllm_config, kv_cache_config) + + def _begin_omni_aux_output_step(self, scheduler_output) -> None: + if self.aux_output_connector is not None: + # Apply terminal/hash-only metadata even on no-forward steps. + self.aux_output_connector.begin_step(scheduler_output.aux_output_connector_metadata) + + def _prepare_omni_aux_output(self): + if self.aux_output_connector is None: + return None + num_reqs = len(self.input_batch.req_ids) + # Bookkeeping and the next async batch mutate these legacy buffers. + batch = SimpleNamespace( + req_ids=list(self.input_batch.req_ids), + query_start_loc_np=self.query_start_loc.np[: num_reqs + 1].copy(), + num_computed_tokens_np=self.input_batch.num_computed_tokens_cpu[:num_reqs].copy(), + ) + return self.aux_output_connector.prepare_output(batch) + + def _finish_omni_aux_output(self, pending, scheduler_output, sampled_token_ids=None, invalid_req_indices=()): + if pending is None: + return None + num_reqs = len(pending.request_ids) + num_sampled = np.zeros(num_reqs, dtype=np.int32) + num_rejected = np.zeros(num_reqs, dtype=np.int32) + if sampled_token_ids is not None: + # Instrumentation opts into a synchronized copy. The normal NPU + # async sampling path remains unchanged when capture is disabled. + sampled = sampled_token_ids.detach().to("cpu").numpy() + num_sampled[:] = (sampled >= 0).sum(axis=-1) + for i, request_id in enumerate(pending.request_ids): + draft_count = len(scheduler_output.scheduled_spec_decode_tokens.get(request_id, ())) + if draft_count: + num_rejected[i] = max(0, draft_count + 1 - int(num_sampled[i])) + # Partial prefills execute valid routing rows but emit no token. + num_sampled[list(invalid_req_indices)] = 0 + pending.enqueue_cpu_copy(num_sampled=num_sampled, num_rejected=num_rejected) + self._sync_device() + return pending.process_output() + + def _close_omni_aux_output(self) -> None: + connector = self.aux_output_connector + self.aux_output_connector = None + if connector is not None: + try: + connector.close() + finally: + for module in self.model.modules(): + if getattr(module, "_omni_npu_capture_source", False): + module.capture_fn = None diff --git a/vllm_omni/platforms/npu/worker/base.py b/vllm_omni/platforms/npu/worker/base.py index 5dd929ea1..a79290640 100644 --- a/vllm_omni/platforms/npu/worker/base.py +++ b/vllm_omni/platforms/npu/worker/base.py @@ -59,3 +59,11 @@ class OmniNPUWorkerBase(NPUWorker): self.profiler.start() else: self.profiler.stop() + + def shutdown(self) -> None: + try: + runner = getattr(self, "model_runner", None) + if runner is not None: + runner._close_omni_aux_output() + finally: + super().shutdown() diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index ee5407065..082b087fa 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -315,10 +315,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: - if self.vllm_config.model_config.enable_return_routed_experts: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(capturer, "finalize_pending_copy"): - capturer.finalize_pending_copy() profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config if profiling_chunk_config.enabled and profiling_chunk_config.need_timing: if getattr(scheduler_output, "disable_profiling_timing", False): @@ -392,6 +388,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex get_kv_transfer_group().handle_preemptions(kv_connector_metadata) # -------------------------------------- Omni-new ------------------------------------------------- + self._begin_omni_aux_output_step(scheduler_output) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with record_function_or_nullcontext("prepare input"): with self.synchronize_input_prep(): @@ -791,9 +788,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if deferred_state_corrections_fn: deferred_state_corrections_fn() - if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): - self._omni_routed_experts_d2h(scheduler_output) - return None def _sample( @@ -884,6 +878,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex multimodal_outputs, # Omni-Specific prefix_cache_step_id, ) = self.execute_model_state + pending_aux_output = self._prepare_omni_aux_output() # Clear ephemeral state. self.execute_model_state = None hidden_seq_len = int(hidden_states.shape[0]) @@ -1015,11 +1010,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex draft_by_req_id = dict(zip(draft_req_ids, draft_ids_list)) output_spec_token_ids = [draft_by_req_id.get(req_id, []) for req_id in req_ids_output_copy] - routed_experts_lists = None - if self.model_config.enable_return_routed_experts: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(self.input_batch, "num_tokens_no_spec"): - routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) + aux_output = self._finish_omni_aux_output( + pending_aux_output, scheduler_output, sampler_output.sampled_token_ids, invalid_req_indices + ) # -------------------------------------- Omni-new ------------------------------------------------- engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) @@ -1209,7 +1202,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex num_nans_in_logits=num_nans_in_logits, ) model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids - model_runner_output.routed_experts = routed_experts_lists + model_runner_output.aux_output_connector_output = aux_output model_runner_output.spec_token_ids = output_spec_token_ids with record_function_or_nullcontext("omni_output_builder:get_omni_connector_output"): model_runner_output.omni_connector_output = self.get_omni_connector_output() diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index bc486e4b5..8948f3125 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -89,10 +89,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors | None: - if self.vllm_config.model_config.enable_return_routed_experts: - capturer = self.routed_experts_capturer - if capturer is not None and hasattr(capturer, "finalize_pending_copy"): - capturer.finalize_pending_copy() profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config if profiling_chunk_config.enabled and profiling_chunk_config.need_timing: if getattr(scheduler_output, "disable_profiling_timing", False): @@ -127,6 +123,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if flush_ids: self.flush_full_payload_outputs(flush_ids) + self._begin_omni_aux_output_step(scheduler_output) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with record_function_or_nullcontext("prepare input"): # -------------------------------------- Omni-new ------------------------------------------------- @@ -418,9 +415,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if deferred_state_corrections_fn: deferred_state_corrections_fn() - if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"): - self._omni_routed_experts_d2h(scheduler_output) - return None @torch.inference_mode() @@ -461,6 +455,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin multimodal_outputs_raw, # Omni-Specific _prefix_cache_step_id, # generation stages never save to the prefix cache ) = self.execute_model_state + pending_aux_output = self._prepare_omni_aux_output() # Clear ephemeral state. self.execute_model_state = None @@ -518,11 +513,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin req_state = self.requests.get(rid) if req_state is not None and inter_stage_outputs[i]: self.accumulate_full_payload_output(rid, inter_stage_outputs[i], req_state) - routed_experts_lists = None - if self.vllm_config.model_config.enable_return_routed_experts and hasattr( - self.input_batch, "num_tokens_no_spec" - ): - routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) + aux_output = self._finish_omni_aux_output(pending_aux_output, scheduler_output) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -537,7 +528,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin cudagraph_stats=cudagraph_stats, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) - output.routed_experts = routed_experts_lists + output.aux_output_connector_output = aux_output output.omni_connector_output = self.get_omni_connector_output() # -------------------------------------- Omni-new ------------------------------------------------- diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index da663f626..a490950a5 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -24,6 +24,7 @@ from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE from vllm_omni.core.prefix_cache import stage_prefix_cache_config from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms.npu._310p import is_310p +from vllm_omni.platforms.npu.worker.aux_output import NPUAuxOutputMixin from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner logger = init_logger(__name__) @@ -35,7 +36,7 @@ else: from vllm_ascend.worker.model_runner_v1 import NPUModelRunner -class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): +class OmniNPUModelRunner(NPUAuxOutputMixin, OmniGPUModelRunner, NPUModelRunner): def initialize_kv_cache( self, kv_cache_config, kv_cache_allocation_context: AbstractContextManager | None = None ) -> None: @@ -43,6 +44,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): NPUModelRunner.initialize_kv_cache( self, kv_cache_config, kv_cache_allocation_context=kv_cache_allocation_context ) + self._init_omni_aux_output(kv_cache_config) if getattr(self, "_omni_prefix_cache_cfg", None) is None: # Same gate as the GPU runner (pooling stage, kv_consumer / # kv_both, hybrid kv groups). Read the config back off