From 18f01befa76865ab8ab01ddba72aad76eaf18ff6 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Sun, 30 Aug 2026 20:28:35 +0800 Subject: [PATCH] [NPU] upgrade to v0.28.0 (#6674) Signed-off-by: Weiming Liao Co-authored-by: Hongsheng Liu Co-authored-by: amy-why-3459 --- docker/Dockerfile.npu | 1 + docker/Dockerfile.npu.a3 | 1 + docker/Dockerfile.npu.ci | 1 + docker/Dockerfile.npu.ci.a3 | 1 + .../npu/worker/npu_ar_model_runner.py | 50 ++++++++--------- .../npu/worker/npu_generation_model_runner.py | 54 +++++++++---------- .../platforms/npu/worker/npu_model_runner.py | 12 ++--- 7 files changed, 55 insertions(+), 65 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 9bb72ff65..3ebacdc98 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \ "stepaudio2-minicpmo>=0.1.1" \ "step-audio2>=1.0.0" +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 1b234d3fe..d341d456d 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \ "stepaudio2-minicpmo>=0.1.1" \ "step-audio2>=1.0.0" +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 42968e467..029a905fd 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -65,6 +65,7 @@ RUN python3 -m pip install -v --no-deps \ --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ --default-timeout=60 +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 index 48568275f..fe3741d04 100644 --- a/docker/Dockerfile.npu.ci.a3 +++ b/docker/Dockerfile.npu.ci.a3 @@ -66,6 +66,7 @@ RUN python3 -m pip install -v --no-deps \ --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ --default-timeout=60 +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] 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 9cbbc4fb5..36a3499cc 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import logging import time from collections.abc import Mapping -from copy import copy, deepcopy +from copy import copy +from dataclasses import replace from typing import Any, NamedTuple import numpy as np @@ -394,9 +395,10 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + 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): - self.ascend_config.scheduler_config.profiling_chunk_config.need_timing = False + profiling_chunk_config.need_timing = False else: self._sync_device() self._execution_start_time = time.perf_counter() @@ -452,21 +454,17 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex set(scheduler_output.finished_req_ids), self.input_batch, ) - # self._draft_token_ids is None when `input_fits_in_drafter=False` - # and there is no draft tokens scheduled. so it need to update the - # spec_decoding info in scheduler_output with async_scheduling. - # use deepcopy to avoid the modification has influence on the - # scheduler_output in engine core process. - # TODO(Ronald1995): deepcopy is expensive when there is a large - # number of requests, optimize it later. - if ( - self.use_async_scheduling - and self.num_spec_tokens - and self._draft_token_ids is None # type: ignore[has-type] - ): - scheduler_output = deepcopy(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- + if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy() + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) + if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata if kv_connector_metadata is not None: @@ -485,7 +483,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex # -------------------------------------- Omni-new ------------------------------------------------- if has_ec_transfer() and get_ec_transfer().is_producer: - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -536,7 +534,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex "it when the requests need prompt logprobs" ) - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) num_reqs = self.input_batch.num_reqs req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] @@ -712,13 +710,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex with record_function_or_nullcontext("EPLB weight D2D"): self.eplb_updator.forward_before() - # Set cudagraph mode to none if calc_kv_scales is true. - # KV scales calculation involves dynamic operations that are incompatible - # with CUDA graph capture. - if self.calculate_kv_scales: # type: ignore[has-type] - cudagraph_mode = CUDAGraphMode.NONE - # Mark KV scales as calculated after the first forward pass - self.calculate_kv_scales = False # type: ignore[has-type] # Encoder-decoder models can only compile the pure decode steps where no # encoder inputs are present. Use eager for the first pass. num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) @@ -919,6 +910,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex def sample_tokens( self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config kv_connector_output = self.kv_connector_output self.kv_connector_output = None pp = get_pp_group() @@ -1294,8 +1286,10 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex model_runner_output.omni_connector_output = self.get_omni_connector_output() # -------------------------------------- Omni-new ------------------------------------------------- - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing and hasattr( - self, "_execution_start_time" + if ( + profiling_chunk_config.enabled + and profiling_chunk_config.need_timing + and hasattr(self, "_execution_start_time") ): self._sync_device() model_runner_output.execution_time_ms = (time.perf_counter() - self._execution_start_time) * 1000.0 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 f16724dff..eeac7d059 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -7,7 +7,8 @@ import gc import logging import time from collections.abc import Mapping -from copy import copy, deepcopy +from copy import copy +from dataclasses import replace import numpy as np import torch @@ -91,25 +92,24 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + 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): - self.ascend_config.scheduler_config.profiling_chunk_config.need_timing = False + profiling_chunk_config.need_timing = False else: self._sync_device() self._execution_start_time = time.perf_counter() if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") - # self._draft_token_ids is None when `input_fits_in_drafter=False` - # and there is no draft tokens scheduled. so it need to update the - # spec_decoding info in scheduler_output with async_scheduling. - # use deepcopy to avoid the modification has influence on the - # scheduler_output in engine core process. - # TODO(Ronald1995): deepcopy is expensive when there is a large - # number of requests, optimize it later. - if ( - self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type] - ): - scheduler_output = deepcopy(scheduler_output) + + if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy() + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata @@ -142,7 +142,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # -------------------------------------- Omni-new ------------------------------------------------- if has_ec_transfer() and get_ec_transfer().is_producer: - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -179,7 +179,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin "it when the requests need prompt logprobs" ) - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) num_reqs = self.input_batch.num_reqs req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] @@ -349,13 +349,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin with record_function_or_nullcontext("EPLB weight D2D"): self.eplb_updator.forward_before() - # Set cudagraph mode to none if calc_kv_scales is true. - # KV scales calculation involves dynamic operations that are incompatible - # with CUDA graph capture. - if self.calculate_kv_scales: # type: ignore[has-type] - cudagraph_mode = CUDAGraphMode.NONE - # Mark KV scales as calculated after the first forward pass - self.calculate_kv_scales = False # type: ignore[has-type] # Encoder-decoder models can only compile the pure decode steps where no # encoder inputs are present. Use eager for the first pass. num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) @@ -431,6 +424,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin def sample_tokens( self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -546,8 +540,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if self.speculative_config is not None: self.finalize_kv_connector() - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing and hasattr( - self, "_execution_start_time" + if ( + profiling_chunk_config.enabled + and profiling_chunk_config.need_timing + and hasattr(self, "_execution_start_time") ): self._sync_device() output.execution_time_ms = (time.perf_counter() - self._execution_start_time) * 1000.0 @@ -856,9 +852,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if get_pp_group().is_first_rank: intermediate_tensors = None else: - # When PP and flashcomm1 are enabled, during dummy_run the estimated space should divide num_tokens by - # tp_size; otherwise, on non-first PP ranks it would effectively perform an extra all-gather, leading - # to incorrect memory estimation and potentially causing OOM. + # When PP and sequence parallelism are enabled, during dummy_run the estimated space should divide + # num_tokens by tp_size; otherwise, on non-first PP ranks it would effectively perform an extra + # all-gather, leading to incorrect memory estimation and potentially causing OOM. intermediate_tokens = num_tokens_padded if enable_sp(): tp_size = get_tensor_model_parallel_world_size() diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 375d36047..8f0369808 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -335,9 +335,9 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): if get_pp_group().is_first_rank: intermediate_tensors = None else: - # When PP and flashcomm1 are enabled, during dummy_run the estimated space should divide num_tokens by - # tp_size; otherwise, on non-first PP ranks it would effectively perform an extra all-gather, leading - # to incorrect memory estimation and potentially causing OOM. + # When PP and sequence parallelism are enabled, during dummy_run the estimated space should divide + # num_tokens by tp_size; otherwise, on non-first PP ranks it would effectively perform an extra + # all-gather, leading to incorrect memory estimation and potentially causing OOM. intermediate_tokens = num_tokens_padded if enable_sp(): tp_size = get_tensor_model_parallel_world_size() @@ -448,7 +448,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): 1. Accepts num_tokens_padded as required by NPUModelRunner 2. Injects omni-specific kwargs (runtime_additional_information) 3. Caches model output for multimodal results - 4. Handles NPU-specific post-forward logic (graph params update, SP all-gather) + 4. Handles NPU-specific post-forward logic (graph params update) """ # Omni-specific: build and inject extra model kwargs model_kwargs_extra = self._build_model_kwargs_extra() @@ -482,10 +482,6 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # Omni-specific: cache model output for later sample_tokens self._omni_last_model_output = model_output - # NPU-specific: all-gather for sequence parallelism - if forward_context.flash_comm_v1_enabled and not isinstance(model_output, IntermediateTensors): - model_output = self._all_gather_hidden_states_and_aux(model_output) - return model_output def _gather_runtime_additional_information(self):