diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 642b89882..2dc94c2b4 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.25.0 +ARG VLLM_ASCEND_TAG=v0.26.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 4fe1d56fe..246467fbe 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.25.0-a3 +ARG VLLM_ASCEND_TAG=v0.26.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 10a62fcac..c56589e9e 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.25.0 +ARG VLLM_ASCEND_TAG=v0.26.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 index 5ce526c96..6a8dd4146 100644 --- a/docker/Dockerfile.npu.ci.a3 +++ b/docker/Dockerfile.npu.ci.a3 @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.25.0-a3 +ARG VLLM_ASCEND_TAG=v0.26.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 7381f5b10..19f67cd1d 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built ```bash # vLLM-Ascend image -# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0 -# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.0-a3 +# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0 +# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 docker run --rm \ --name vllm-omni-npu \ --shm-size=1g \ @@ -32,7 +32,7 @@ docker run --rm \ # Inside the container, install vLLM-Omni from source cd /vllm-workspace -git clone -b v0.25.0 https://github.com/vllm-project/vllm-omni.git +git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git cd vllm-omni pip install -v -e . --no-build-isolation # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . @@ -84,8 +84,8 @@ Supported images as following. Here's an example deployment command that has been verified on 4 x NPUs: ```bash -# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0 -# Atlas A3: quay.io/ascend/vllm-omni:v0.25.0-a3 +# Atlas A2: quay.io/ascend/vllm-omni:v0.26.0 +# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3 docker run --rm \ --name vllm-omni-a3 \ --shm-size=4g \ @@ -110,7 +110,7 @@ docker run --rm \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v ~/.cache:/root/.cache \ -p 8091:8091 \ - -it quay.io/ascend/vllm-omni:v0.25.0-a3 bash + -it quay.io/ascend/vllm-omni:v0.26.0-a3 bash ``` !!! tip diff --git a/vllm_omni/deploy/minicpmo_4_5_batching.yaml b/vllm_omni/deploy/minicpmo_4_5_batching.yaml index 75693feb7..7e6b1fb63 100644 --- a/vllm_omni/deploy/minicpmo_4_5_batching.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_batching.yaml @@ -84,5 +84,6 @@ platforms: cudagraph_mode: PIECEWISE - stage_id: 1 max_num_batched_tokens: 8192 + gpu_memory_utilization: 0.4 compilation_config: cudagraph_mode: PIECEWISE diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 42ff0b418..240f9d576 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -123,7 +123,9 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): cls, selected_backend: str | None, head_size: int, + allow_trtllm_default: bool = True, ) -> str: + # NPU has no TRTLLM backend; accepted for signature parity, ignored. from importlib.util import find_spec if selected_backend is not None: 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 904974829..0ff7144a9 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -3,6 +3,7 @@ from __future__ import annotations +import logging import time from collections.abc import Mapping from copy import copy, deepcopy @@ -28,8 +29,8 @@ from vllm.v1.outputs import ( from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata -from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.attention.utils import AscendCommonAttentionMetadata @@ -43,7 +44,10 @@ from vllm_ascend.worker.model_runner_v1 import graph_capture from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager -from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output +from vllm_omni.distributed.omni_connectors.utils.config import ( + get_stage_connector_role, + stage_sends_async_output, +) from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element @@ -114,6 +118,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) self._async_chunk = getattr(self.model_config, "async_chunk", False) + # -------------------------------------- Omni-new ------------------------------------------------- _OMNI_CONNECTOR_INIT_ARCHS = { "Qwen3OmniMoeForConditionalGeneration", "Qwen2_5OmniForConditionalGeneration", @@ -125,12 +130,17 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): "DyninOmniForConditionalGeneration", "IndexTTS2TalkerForConditionalGeneration", } - if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: + stage_archs = set(getattr(self.model_config, "architectures", None) or ()) + model_arch_override = getattr(self.model_config, "model_arch", None) + if model_arch_override: + stage_archs.add(model_arch_override) + if stage_archs & _OMNI_CONNECTOR_INIT_ARCHS or get_stage_connector_role(self.model_config) is not None: self.init_omni_connectors( model_config=self.model_config, kv_transfer_manager=self.kv_transfer_manager, ) self._downstream_payload_cache: dict[str, bool] = {} + # -------------------------------------- Omni-new ------------------------------------------------- def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size @@ -209,6 +219,35 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): model = getattr(self, "model", None) return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)) + def _deferred_prefix_cache_mm_keys(self) -> set[str]: + """Model-declared multimodal keys whose prefix-cache writes are deferred.""" + model = getattr(self, "model", None) + keys = getattr(model, "deferred_prefix_cache_mm_keys", ()) + return set(keys or ()) + + def _stage_deferred_prefix_cache_mm_outputs( + self, + *, + scheduler_output: SchedulerOutput, + multimodal_outputs: Any, + query_start_loc_cpu: Any, + ) -> None: + """See gpu_ar_model_runner._stage_deferred_prefix_cache_mm_outputs.""" + if self.omni_prefix_cache is None: + return + + deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys() + if not deferred_mm_cache_keys: + return + + self.omni_prefix_cache.stage_deferred_mm_outputs( + query_start_loc=query_start_loc_cpu, + input_batch=self.input_batch, + multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + deferred_mm_cache_keys=deferred_mm_cache_keys, + ) + def _maybe_update_prefix_cache( self, hidden_states: torch.Tensor, @@ -217,19 +256,16 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): num_tokens_padded: int, ): if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: - if multimodal_outputs is not None and not isinstance(multimodal_outputs, Mapping): - logger.warning_once( - "prefix caching expects mm outputs to be a dict, but got %s", - type(multimodal_outputs), - ) - hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None + slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu + slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu() self.omni_prefix_cache.update_omni_tensor_prefix_cache( hidden_states=hs_for_cache, multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, num_tokens_unpadded=num_tokens_unpadded, - slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu, + slot_mapping=slot_mapping_cpu, num_tokens_padded=num_tokens_padded, + skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(), ) def _maybe_get_combined_prefix_cache_tensors( @@ -240,6 +276,11 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) -> tuple[dict[str, torch.Tensor] | None, dict | None]: combined_hidden_states, combined_multimodal_outputs = None, None if self.omni_prefix_cache is not None: + if ( + not self._model_needs_full_prefix_hidden_states() + and not self.omni_prefix_cache.has_prefix_cached_new_req_ids() + ): + return None, None if self._model_needs_full_prefix_hidden_states(): combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states( query_start_loc=self.query_start_loc.cpu, @@ -358,9 +399,12 @@ class NPUARModelRunner(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.profiling_chunk_config.enabled: - self._sync_device() - self._execution_start_time = time.perf_counter() + if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + if getattr(scheduler_output, "disable_profiling_timing", False): + self.ascend_config.scheduler_config.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.") @@ -407,6 +451,12 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests}) if flush_ids: self.flush_full_payload_outputs(flush_ids) + + if self.omni_prefix_cache is not None and scheduler_output.finished_req_ids: + self.omni_prefix_cache.commit_deferred_mm_outputs( + 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. @@ -414,20 +464,11 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): # scheduler_output in engine core process. # TODO(Ronald1995): deepcopy is expensive when there is a large # number of requests, optimize it later. - if (( + if ( self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type] - ) or ( - # NOTE: This branch specifically triggers a deepcopy during the prefill phase - # only for PCP (Parallel Context Processing) + Multi-Modal (MM) scenarios. - # It does not affect other use cases. This is a temporary workaround and - # will be removed once upstream vLLM provides native support for PCP + MM. - self.pcp_size > 1 - and self.supports_mm_inputs - and get_pp_group().is_first_rank - and not self.model_config.is_encoder_decoder - )): + ): scheduler_output = deepcopy(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- @@ -449,6 +490,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): # -------------------------------------- Omni-new ------------------------------------------------- if has_ec_transfer() and get_ec_transfer().is_producer: + self._start_dump_data() with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -458,6 +500,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): kv_ids = self.kv_extracted_req_ids self.kv_extracted_req_ids = None + self._finalize_dump_data() output = make_empty_encoder_model_runner_output(scheduler_output) if kv_ids: output = copy(output) @@ -498,6 +541,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): "it when the requests need prompt logprobs" ) + self._start_dump_data() 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] @@ -514,8 +558,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens - if self.pcp_size > 1: - num_tokens_unpadded = self.pcp_manager.total_num_sampled_tokens_pcp cascade_attn_prefix_lens = None # Disable cascade attention when using microbatching (DBO) if self.cascade_attn_enabled and not self.parallel_config.enable_dbo: @@ -542,14 +584,15 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), ) - logger.debug( - "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " - "should_ubatch: %s, num_tokens_across_dp: %s", - cudagraph_mode, - batch_desc, - should_ubatch, - num_tokens_across_dp, - ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs @@ -561,6 +604,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.parallel_config.num_ubatches, ) + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + pad_attn = cudagraph_mode == CUDAGraphMode.FULL # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877, @@ -574,7 +620,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if deferred_state_corrections_fn: deferred_state_corrections_fn() deferred_state_corrections_fn = None - preprocess_mamba( + mamba_bufs = self._get_mamba_bufs() + preprocess_bufs = mamba_bufs.preprocess + mamba_utils.preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, @@ -583,7 +631,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.requests, self.compilation_config.static_forward_context, self.model.get_mamba_state_copy_func(), - self._get_mamba_copy_bufs(), + preprocess_bufs, ) # preprocess_mamba resets num_accepted_tokens_cpu to 1 # for requests whose state was copied to a new block. @@ -592,13 +640,23 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] self.num_accepted_tokens.copy_to_gpu(num_reqs) + if mamba_bufs.postprocess_align is not None: + mamba_utils.stage_postprocess_inputs_to_gpu( + mamba_bufs.postprocess_align, + scheduler_output, + self.input_batch.req_ids, + num_reqs, + self.requests, + self.mamba_state_idx, + ) + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices if ( cudagraph_mode == CUDAGraphMode.FULL or (enable_sp() and not self.model_config.use_mla) - and self.pcp_size * self.dcp_size == 1 + and self.dcp_size == 1 ): # Currently, Graph Mode and SP will both pad num_tokens, # Another possible condition is num_tokens_padded != num_tokens_unpadded @@ -613,9 +671,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( - num_tokens=num_tokens_unpadded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded, num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, @@ -637,9 +693,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ec_connector_output, ) = self._preprocess( scheduler_output, - num_tokens_padded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens_padded, intermediate_tensors, ) @@ -670,21 +724,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): cudagraph_mode = CUDAGraphMode.NONE # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False # type: ignore[has-type] - # prevent debugger is None - if self.debugger is not None: - dbg_cfg = getattr(self.debugger, "config", None) - dump_level = str(getattr(dbg_cfg, "level", "L1")).upper() if dbg_cfg is not None else "L1" - if dump_level in ("L0", "MIX"): - self.debugger.start(model=self.model) - else: - self.debugger.start() - if self.ascend_config.enable_async_exponential: - self.sampler.do_async_exponential( - b_s=logits_indices.shape[0], - head_dim=self.model_config.get_vocab_size(), - generators=self.input_batch.sampling_metadata.generators, - ) - # 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) @@ -703,8 +742,10 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): batch_descriptor=batch_desc, num_actual_tokens=scheduler_output.total_num_scheduled_tokens, model_instance=self.model, - max_tokens_across_pcp=0 if self.pcp_size == 1 else self.pcp_manager.max_num_tokens_across_pcp, skip_compiled=has_encoder_input, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ), self.maybe_get_kv_connector_output( scheduler_output, @@ -713,6 +754,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ), ) as kv_connector_output, ): + if self.cache_config.mamba_cache_mode == "align": + mamba_utils.do_mamba_copy_block(preprocess_bufs) hidden_states = self._model_forward( num_tokens_padded, input_ids, positions, intermediate_tensors, inputs_embeds, **model_kwargs ) @@ -738,15 +781,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): aux_hidden_states = None if self.use_aux_hidden_state_outputs: hidden_states, aux_hidden_states = hidden_states - if self.pcp_size > 1: - # NOTE we must `slice` hidden_states because pcp_allgather_restore_idx - # ignores the padding from CUDA Graph. - hidden_states = self.pcp_manager.get_restore_hidden_states(hidden_states) - if aux_hidden_states is not None: - aux_hidden_states = [ - self.pcp_manager.get_restore_hidden_states(aux_hidden_states_pcp) - for aux_hidden_states_pcp in aux_hidden_states - ] # -------------------------------------- Omni-new ------------------------------------------------- self._maybe_update_prefix_cache( @@ -764,9 +798,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): assert isinstance(hidden_states, IntermediateTensors) hidden_states.kv_connector_output = kv_connector_output self.kv_connector_output = kv_connector_output - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() + if self.dynamic_eplb: + self.eplb_updator.forward_end(self.eplb_heat_collection_status) return hidden_states if self.is_pooling_model: # Return the pooling output. @@ -774,9 +808,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): hidden_states, num_scheduled_tokens, num_scheduled_tokens_np, kv_connector_output ) output.kv_connector_output = kv_connector_output - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() return output sample_hidden_states = hidden_states[logits_indices] @@ -885,6 +917,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: kv_connector_output = self.kv_connector_output self.kv_connector_output = None + pp = get_pp_group() + use_pp_spec_decode = self.speculative_config is not None and pp.world_size > 1 # -------------------------------------- Omni-new ------------------------------------------------- kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None) @@ -897,11 +931,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if self.execute_model_state is None: # Nothing to do (PP non-final rank case), output isn't used. - # receive sampled token ids from the last PP rank when using - # async scheduling + pipeline parallelism so downstream code - # (e.g., PCP input preparation) can access them. - if self.use_async_scheduling and get_pp_group().world_size > 1: - self._pp_receive_prev_sampled_token_ids_to_input_batch() if not kv_connector_output: return None # noqa # In case of PP with kv transfer, we need to pass through the @@ -992,6 +1021,23 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) self._copy_draft_token_ids_to_cpu(scheduler_output) + output_spec_token_ids = None + use_padded_batch = False + early_pp_padded_drafter = False + if self.speculative_config: + use_padded_batch = ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + or self.speculative_config.uses_extract_hidden_states() + or self.speculative_config.use_ngram_gpu() + ) and not self.speculative_config.disable_padded_drafter_batch + early_pp_padded_drafter = use_pp_spec_decode and not self.use_async_scheduling and use_padded_batch + if early_pp_padded_drafter: + self._draft_token_ids = None + self._draft_token_req_ids = None + with record_function_or_nullcontext("draft_token"): + propose_draft_token_ids(sampler_output.sampled_token_ids) + ( logprobs_lists, valid_sampled_token_ids, @@ -1010,16 +1056,14 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): with record_function_or_nullcontext("draft_token"): if self.speculative_config: - use_padded_batch = ( - self.speculative_config - and (self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model()) - and not self.speculative_config.disable_padded_drafter_batch - ) - if use_padded_batch: + if not early_pp_padded_drafter: + self._draft_token_ids = None + self._draft_token_req_ids = None + if use_padded_batch and not early_pp_padded_drafter: # EAGLE speculative decoding can use the GPU sampled tokens # as inputs, and does not need to wait for bookkeeping to finish. propose_draft_token_ids(sampler_output.sampled_token_ids) - if self.speculative_config and not use_padded_batch: + if not use_padded_batch: # ngram and other speculative decoding methods use the sampled # tokens on the CPU, so they are run after bookkeeping. propose_draft_token_ids(valid_sampled_token_ids) @@ -1030,6 +1074,19 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if self.speculative_config is not None: self.finalize_kv_connector() + draft_token_ids = self._draft_token_ids if use_pp_spec_decode else None + if draft_token_ids is not None: + if isinstance(draft_token_ids, torch.Tensor): + num_draft_reqs = draft_token_ids.shape[0] + draft_ids_list = draft_token_ids[:num_draft_reqs].cpu().tolist() + draft_req_ids = self._draft_token_req_ids + else: + draft_ids_list = draft_token_ids + draft_req_ids = self.input_batch.req_ids + if draft_ids_list and draft_req_ids: + 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 @@ -1073,6 +1130,12 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if callable(query_start_loc_cpu): query_start_loc_cpu = query_start_loc_cpu() + self._stage_deferred_prefix_cache_mm_outputs( + scheduler_output=scheduler_output, + multimodal_outputs=multimodal_outputs, + query_start_loc_cpu=query_start_loc_cpu, + ) + pooler_output: list[dict[str, object]] | None = None if needs_pooler_payload: combined_hidden_states = None @@ -1222,21 +1285,22 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): ) model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids model_runner_output.routed_experts = routed_experts_lists + 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() # -------------------------------------- Omni-new ------------------------------------------------- - if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"): + if self.ascend_config.scheduler_config.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 if self.dynamic_eplb: with record_function_or_nullcontext("EPLB update"): - self.eplb_updator.forward_end() + self.eplb_updator.forward_end(self.eplb_heat_collection_status) - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() if self.need_accepted_tokens: assert self.sampling_done_event is not None @@ -1247,14 +1311,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): global_stream().wait_event(self.sampling_done_event) self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) - # In async scheduling + PP, broadcast sampled token ids from the - # last PP rank so other PP ranks can receive them without going - # through the scheduler/engine IPC path. - if self.use_async_scheduling: - pp = get_pp_group() - if pp.world_size > 1 and pp.is_last_rank: - self._pp_broadcast_prev_sampled_token_ids(sampler_output.sampled_token_ids) - if not self.use_async_scheduling: return model_runner_output async_output = AsyncGPUModelRunnerOutput( 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 f4fc05574..24fa21524 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -4,6 +4,7 @@ from __future__ import annotations import gc +import logging import time from collections.abc import Mapping from copy import copy, deepcopy @@ -21,8 +22,8 @@ from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata -from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_ascend.ascend_forward_context import set_ascend_forward_context @@ -32,6 +33,7 @@ from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE +from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_role from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner @@ -45,6 +47,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._async_chunk = getattr(self.model_config, "async_chunk", False) + # -------------------------------------- Omni-new ------------------------------------------------- + # Mirrors the init allowlist in gpu_generation_model_runner.py. The + # connector-role fallback keeps consumer stages whose arch is not on the + # allowlist (e.g. MiniCPM-o 4.5 stage 2) from starving on input. _OMNI_CONNECTOR_INIT_ARCHS = { "Qwen3OmniMoeForConditionalGeneration", "Qwen2_5OmniForConditionalGeneration", @@ -52,14 +58,20 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin "MiMoAudioModel", "Qwen3TTSTalkerForConditionalGeneration", "Qwen3TTSCode2Wav", + "AudexCode2Wav", + "AudexXCodec1", "CosyVoice3Model", "DyninOmniForConditionalGeneration", "IndexTTS2S2MelDecoder", } - if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: + if ( + getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS + or get_stage_connector_role(self.model_config) is not None + ): self.init_omni_connectors( model_config=self.model_config, ) + # -------------------------------------- Omni-new ------------------------------------------------- def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests @@ -98,9 +110,12 @@ 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.profiling_chunk_config.enabled: - self._sync_device() - self._execution_start_time = time.perf_counter() + if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + if getattr(scheduler_output, "disable_profiling_timing", False): + self.ascend_config.scheduler_config.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` @@ -112,15 +127,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # 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] - ) or ( - # NOTE: This branch specifically triggers a deepcopy during the prefill phase - # only for PCP (Parallel Context Processing) + Multi-Modal (MM) scenarios. - # It does not affect other use cases. This is a temporary workaround and - # will be removed once upstream vLLM provides native support for PCP + MM. - self.pcp_size > 1 - and self.supports_mm_inputs - and get_pp_group().is_first_rank - and not self.model_config.is_encoder_decoder ): scheduler_output = deepcopy(scheduler_output) @@ -149,12 +155,19 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # Update persistent batch states. deferred_state_corrections_fn = self._update_states(scheduler_output) + # -------------------------------------- Omni-new ------------------------------------------------- + if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): + self.model.on_requests_finished(scheduler_output.finished_req_ids) + # -------------------------------------- Omni-new ------------------------------------------------- + if has_ec_transfer() and get_ec_transfer().is_producer: + self._start_dump_data() with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, ) as ec_connector_output: self._execute_mm_encoder(scheduler_output) + self._finalize_dump_data() return self.attach_omni_connector_output( make_empty_encoder_model_runner_output(scheduler_output) ) @@ -185,6 +198,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin "it when the requests need prompt logprobs" ) + self._start_dump_data() 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] @@ -201,8 +215,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens - if self.pcp_size > 1: - num_tokens_unpadded = self.pcp_manager.total_num_sampled_tokens_pcp cascade_attn_prefix_lens = None # Disable cascade attention when using microbatching (DBO) if self.cascade_attn_enabled and not self.parallel_config.enable_dbo: @@ -229,14 +241,15 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), ) - logger.debug( - "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " - "should_ubatch: %s, num_tokens_across_dp: %s", - cudagraph_mode, - batch_desc, - should_ubatch, - num_tokens_across_dp, - ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs @@ -248,6 +261,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin self.parallel_config.num_ubatches, ) + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + pad_attn = cudagraph_mode == CUDAGraphMode.FULL # NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877, @@ -261,7 +277,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if deferred_state_corrections_fn: deferred_state_corrections_fn() deferred_state_corrections_fn = None - preprocess_mamba( + mamba_bufs = self._get_mamba_bufs() + preprocess_bufs = mamba_bufs.preprocess + mamba_utils.preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, @@ -270,7 +288,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin self.requests, self.compilation_config.static_forward_context, self.model.get_mamba_state_copy_func(), - self._get_mamba_copy_bufs(), + preprocess_bufs, ) # preprocess_mamba resets num_accepted_tokens_cpu to 1 # for requests whose state was copied to a new block. @@ -279,13 +297,23 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin self.num_accepted_tokens.np[:num_reqs] = self.input_batch.num_accepted_tokens_cpu[:num_reqs] self.num_accepted_tokens.copy_to_gpu(num_reqs) + if mamba_bufs.postprocess_align is not None: + mamba_utils.stage_postprocess_inputs_to_gpu( + mamba_bufs.postprocess_align, + scheduler_output, + self.input_batch.req_ids, + num_reqs, + self.requests, + self.mamba_state_idx, + ) + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices if ( cudagraph_mode == CUDAGraphMode.FULL or (enable_sp() and not self.model_config.use_mla) - and self.pcp_size * self.dcp_size == 1 + and self.dcp_size == 1 ): # Currently, Graph Mode and SP will both pad num_tokens, # Another possible condition is num_tokens_padded != num_tokens_unpadded @@ -300,9 +328,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ) (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( - num_tokens=num_tokens_unpadded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded, num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, @@ -324,14 +350,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ec_connector_output, ) = self._preprocess( scheduler_output, - num_tokens_padded - if not (self.use_cp and self.pcp_manager.pcp_use_hybrid_attn) - else total_num_scheduled_tokens, + num_tokens_padded, intermediate_tensors, ) + # -------------------------------------- Omni-new ------------------------------------------------- # [Omni] Pass token counts per request for code2wav output slicing model_kwargs["seq_token_counts"] = tokens + if getattr(self.model, "requires_request_ids", False): + model_kwargs["request_ids"] = list(req_ids) + # -------------------------------------- Omni-new ------------------------------------------------- # update global cos, sin update_cos_sin(positions) @@ -347,21 +375,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin cudagraph_mode = CUDAGraphMode.NONE # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False # type: ignore[has-type] - # prevent debugger is None - if self.debugger is not None: - dbg_cfg = getattr(self.debugger, "config", None) - dump_level = str(getattr(dbg_cfg, "level", "L1")).upper() if dbg_cfg is not None else "L1" - if dump_level in ("L0", "MIX"): - self.debugger.start(model=self.model) - else: - self.debugger.start() - if self.ascend_config.enable_async_exponential: - self.sampler.do_async_exponential( - b_s=logits_indices.shape[0], - head_dim=self.model_config.get_vocab_size(), - generators=self.input_batch.sampling_metadata.generators, - ) - # 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) @@ -380,14 +393,18 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin batch_descriptor=batch_desc, num_actual_tokens=scheduler_output.total_num_scheduled_tokens, model_instance=self.model, - max_tokens_across_pcp=0 if self.pcp_size == 1 else self.pcp_manager.max_num_tokens_across_pcp, skip_compiled=has_encoder_input, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ), self.maybe_get_kv_connector_output( scheduler_output, **({"defer_finalize": not clear_kv_metadata}), ) as kv_connector_output, ): + if self.cache_config.mamba_cache_mode == "align": + mamba_utils.do_mamba_copy_block(preprocess_bufs) # -------------------------------------- Omni-new ------------------------------------------------- outputs = self._run_generation_model( num_tokens_padded=num_tokens_padded, @@ -437,11 +454,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if self.execute_model_state is None: # Nothing to do (PP non-final rank case), output isn't used. - # receive sampled token ids from the last PP rank when using - # async scheduling + pipeline parallelism so downstream code - # (e.g., PCP input preparation) can access them. - if self.use_async_scheduling and get_pp_group().world_size > 1: - self._pp_receive_prev_sampled_token_ids_to_input_batch() if not kv_connector_output: return None # noqa # In case of PP with kv transfer, we need to pass through the @@ -552,17 +564,17 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if self.speculative_config is not None: self.finalize_kv_connector() - if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"): + if self.ascend_config.scheduler_config.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 if self.dynamic_eplb: with record_function_or_nullcontext("EPLB update"): - self.eplb_updator.forward_end() + self.eplb_updator.forward_end(self.eplb_heat_collection_status) - if self.debugger is not None: - self.debugger.stop() - self.debugger.step() + self._finalize_dump_data() if not self.use_async_scheduling: return output async_output = AsyncGPUModelRunnerOutput( @@ -708,14 +720,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin force_has_lora=num_active_loras > 0, force_num_active_loras=num_active_loras, ) - if self.use_cp: - self.pcp_manager.init_batch_info( + if self.use_dcp: + self.dcp_manager.init_batch_info( num_scheduled_tokens, num_reqs, + self.input_batch.num_computed_tokens_cpu, + self.input_batch.num_prompt_tokens, ) if self.speculative_config: - self.pcp_manager.query_lens_pcp_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) - self.pcp_manager.query_lens_pcp_full.copy_to_gpu() + self.dcp_manager.query_lens_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) + self.dcp_manager.query_lens_full.copy_to_gpu() if cudagraph_runtime_mode is None: cudagraph_runtime_mode = _cudagraph_mode else: @@ -729,66 +743,82 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher num_tokens_across_dp[:] = num_tokens_padded num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded) + + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + # vllm-ascend does not support ubatch now ubatch_slices, ubatch_slices_padded = None, None attn_metadata: PerLayerAttnMetadata | None = None - # Build attention metadata for dummy_run - if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): - if create_mixed_batch: - raise NotImplementedError( - "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" - ) - self.attn_state = AscendAttentionState.DecodeOnly - if self.speculative_config and self.speculative_config.method == "mtp": - # `AscendAttentionState.SpecDecoding` is only designed for mla - if self.vllm_config.model_config.use_mla: - self.attn_state = AscendAttentionState.SpecDecoding + # Mirrors gpu_model_runner._dummy_run. + with self.synchronize_input_prep(): + # Build attention metadata for dummy_run + if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): + if create_mixed_batch: + raise NotImplementedError( + "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" + ) + self.attn_state = AscendAttentionState.DecodeOnly + if self.speculative_config and self.speculative_config.method == "mtp": + # `AscendAttentionState.SpecDecoding` is only designed for mla + if self.vllm_config.model_config.use_mla: + self.attn_state = AscendAttentionState.SpecDecoding + else: + self.attn_state = AscendAttentionState.ChunkedPrefill + # The reason why we use a fixed seq_len rather than max_query_len is that + # _npu_paged_attention_get_workspace only returns max workspace with specific + # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len + # in inference. This will be removed once npu_fused_infer_attention_score + # outperforms _npu_paged_attention on all cases. + if profile_seq_lens is not None: + seq_lens = profile_seq_lens else: - self.attn_state = AscendAttentionState.ChunkedPrefill - # The reason why we use a fixed seq_len rather than max_query_len is that - # _npu_paged_attention_get_workspace only returns max workspace with specific - # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len - # in inference. This will be removed once npu_fused_infer_attention_score - # outperforms _npu_paged_attention on all cases. - if profile_seq_lens is not None: - seq_lens = profile_seq_lens - else: - seq_lens = ( - SEQ_LEN_WITH_MAX_PA_WORKSPACE - if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) - else max_query_len - ) # type: ignore[assignment] - self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens - self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) - self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + seq_lens = ( + SEQ_LEN_WITH_MAX_PA_WORKSPACE + if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) + else max_query_len + ) # type: ignore[assignment] + self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) - cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) - self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() - if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.gdn_query_start_loc.copy_to_gpu() + cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) + self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() + if self._has_gdn: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.gdn_query_start_loc.copy_to_gpu() - if not profile_cpp: - num_reqs_padded = self._pad_query_start_loc_for_fia( - self.query_start_loc, - num_tokens_padded, - num_reqs_padded, - num_reqs, - cudagraph_runtime_mode, - batch_desc.num_reqs, + if not profile_cpp: + num_reqs_padded = self._pad_query_start_loc_for_fia( + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_runtime_mode, + batch_desc.num_reqs, + ) + + # Dummy graph runs do not go through _prepare_inputs(), but GDN/Mamba + # metadata reads block_table[:num_reqs_padded] below. Sync padded + # rows as well so device-side metadata does not see stale block ids. + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, + for_cudagraph_capture=is_graph_capturing, + num_scheduled_tokens_np=num_scheduled_tokens, ) - - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_tokens_padded=num_tokens_padded, - num_reqs=num_reqs_padded, - max_query_len=max_query_len, - ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, - for_cudagraph_capture=is_graph_capturing, - num_scheduled_tokens_np=num_scheduled_tokens, - ) + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, @@ -827,6 +857,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if hasattr(self.model, "get_dummy_runtime_additional_information"): runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs) model_kwargs["runtime_additional_information"] = runtime_addi + model_kwargs["seq_token_counts"] = [int(num_tokens_padded)] # -------------------------------------- Omni-new ------------------------------------------------- if self.uses_mrope: @@ -885,6 +916,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin aclgraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, model_instance=self.model, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ): outputs = self.model( input_ids=input_ids, @@ -902,7 +936,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # ------------------------------------------------------------------------------------------------- dummy_compute_logits(hidden_states) - if self.drafter: + if self.drafter and not profile_cpp: self.drafter.dummy_run( num_tokens=num_tokens_padded, with_prefill=with_prefill, @@ -915,10 +949,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin is_profile=is_profile, ) if is_profile and self.dynamic_eplb: - target = self.model.language_model if hasattr(self.model, "language_model") else self.model - target.clear_all_moe_loads() - if self.dynamic_eplb: - self.eplb_updator.forward_end() + self.eplb_updator.adaptor.clear_all_moe_loads() + if not is_profile and self.dynamic_eplb: + self.eplb_updator.forward_end(self.eplb_heat_collection_status) + self._finalize_dump_data(dump=False) return hidden_states, None diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index a9a8b529b..bc8e4f0a2 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -20,6 +20,7 @@ from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE +from vllm_omni.core.prefix_cache import OmniTensorPrefixCache from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms.npu._310p import is_310p from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -34,6 +35,37 @@ else: class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): + def initialize_kv_cache(self, kv_cache_config) -> None: + """Create the omni tensor prefix cache. + + The omni prefix cache is used to store the hidden states of the prefix tokens + """ + NPUModelRunner.initialize_kv_cache(self, kv_cache_config) + if self.omni_prefix_cache is None and self.cache_config.enable_prefix_caching: + # Read num_blocks back off self.kv_cache_config: vllm-ascend + # deepcopies the config it was handed, so the value it stored is the + # authoritative one, not our caller's argument. + num_blocks = self.kv_cache_config.num_blocks + self.omni_prefix_cache = OmniTensorPrefixCache( + num_blocks=num_blocks, + block_size=self.cache_config.block_size, + hidden_size=self.model_config.get_hidden_size(), + hs_dtype=self.dtype, + ) + logger.info( + "Initialized omni prefix cache on NPU (num_blocks=%d, block_size=%d, hidden_size=%d). " + "Hidden-state cache is pinned host memory of roughly %.1f GiB; each per-token " + "multimodal output key allocates another tensor of the same block shape.", + num_blocks, + self.cache_config.block_size, + self.model_config.get_hidden_size(), + num_blocks + * self.cache_config.block_size + * self.model_config.get_hidden_size() + * self.dtype.itemsize + / (1024**3), + ) + def load_model(self, *args, **kwargs) -> None: if is_310p(): from vllm_omni.platforms.npu._310p.patch import apply_model_patches @@ -44,8 +76,37 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # in _pad_for_sequence_parallelism during execute_model. # This is a workaround for vllm-ascend not passing vllm_config to enable_sp(). enable_sp(self.vllm_config) + # ---------------------------------------Omni-new---------------------------------------------- + model = getattr(self, "model", None) + override_fn = None + if bool(getattr(model, "supports_sampled_token_ids_cpu_override", False)): + candidate = getattr(model, "consume_sampled_token_ids_cpu_override", None) + if callable(candidate): + override_fn = candidate + self._sampled_token_ids_cpu_override = override_fn + self._omni_query_start_loc_model_kwarg = bool(getattr(model, "supports_omni_query_start_loc", False)) self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() + self._prewarm_attention_capture_workspaces() + # ---------------------------------------Omni-new---------------------------------------------- + + def _update_states(self, scheduler_output) -> Any: + """Mirror NPUModelRunner._update_states and + call OmniGPUModelRunner._update_states to update Omni-specific states. + """ + req_data = scheduler_output.scheduled_cached_reqs + + if self.use_async_scheduling: + for i, req_id in enumerate(req_data.req_ids): + req_state = self.requests.get(req_id) + if req_state is None: + continue + if req_data.num_computed_tokens[i] < req_state.num_computed_tokens: + req_state.prev_num_draft_len = 0 + + self._apply_pp_sampled_tokens_from_scheduler_output(scheduler_output) + + return OmniGPUModelRunner._update_states(self, scheduler_output) @torch.inference_mode() def _dummy_run( @@ -130,14 +191,16 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): force_has_lora=num_active_loras > 0, force_num_active_loras=num_active_loras, ) - if self.use_cp: - self.pcp_manager.init_batch_info( + if self.use_dcp: + self.dcp_manager.init_batch_info( num_scheduled_tokens, num_reqs, + self.input_batch.num_computed_tokens_cpu, + self.input_batch.num_prompt_tokens, ) if self.speculative_config: - self.pcp_manager.query_lens_pcp_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) - self.pcp_manager.query_lens_pcp_full.copy_to_gpu() + self.dcp_manager.query_lens_full.cpu[:num_reqs] = torch.from_numpy(num_scheduled_tokens) + self.dcp_manager.query_lens_full.copy_to_gpu() if cudagraph_runtime_mode is None: cudagraph_runtime_mode = _cudagraph_mode else: @@ -151,66 +214,81 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # pad is needed if the pad of `num_tokens` is triggered inside CudagraphDispatcher num_tokens_across_dp[:] = num_tokens_padded num_scheduled_tokens = num_scheduled_tokens.repeat(num_reqs_padded) + + if self.dynamic_eplb: + self.update_eplb_heat_collection_status(num_tokens_padded) + # vllm-ascend does not support ubatch now ubatch_slices, ubatch_slices_padded = None, None attn_metadata: PerLayerAttnMetadata | None = None - # Build attention metadata for dummy_run - if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): - if create_mixed_batch: - raise NotImplementedError( - "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" - ) - self.attn_state = AscendAttentionState.DecodeOnly - if self.speculative_config and self.speculative_config.method == "mtp": - # `AscendAttentionState.SpecDecoding` is only designed for mla - if self.vllm_config.model_config.use_mla: - self.attn_state = AscendAttentionState.SpecDecoding + with self.synchronize_input_prep(): + # Build attention metadata for dummy_run + if self._should_build_dummy_attn_metadata(force_attention, is_profile, cudagraph_runtime_mode): + if create_mixed_batch: + raise NotImplementedError( + "create_mixed_batch is used for warmup deepgemm, vllm-ascend does not need it" + ) + self.attn_state = AscendAttentionState.DecodeOnly + if self.speculative_config and self.speculative_config.method == "mtp": + # `AscendAttentionState.SpecDecoding` is only designed for mla + if self.vllm_config.model_config.use_mla: + self.attn_state = AscendAttentionState.SpecDecoding + else: + self.attn_state = AscendAttentionState.ChunkedPrefill + # The reason why we use a fixed seq_len rather than max_query_len is that + # _npu_paged_attention_get_workspace only returns max workspace with specific + # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len + # in inference. This will be removed once npu_fused_infer_attention_score + # outperforms _npu_paged_attention on all cases. + if profile_seq_lens is not None: + seq_lens = profile_seq_lens else: - self.attn_state = AscendAttentionState.ChunkedPrefill - # The reason why we use a fixed seq_len rather than max_query_len is that - # _npu_paged_attention_get_workspace only returns max workspace with specific - # seq_lens. We use this seq_len only when capturing graph, and still use max_query_len - # in inference. This will be removed once npu_fused_infer_attention_score - # outperforms _npu_paged_attention on all cases. - if profile_seq_lens is not None: - seq_lens = profile_seq_lens - else: - seq_lens = ( - SEQ_LEN_WITH_MAX_PA_WORKSPACE - if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) - else max_query_len - ) # type: ignore[assignment] - self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens - self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) - self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + seq_lens = ( + SEQ_LEN_WITH_MAX_PA_WORKSPACE + if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config) + else max_query_len + ) # type: ignore[assignment] + self.optimistic_seq_lens_cpu[:num_reqs] = seq_lens + self.optimistic_seq_lens_cpu[num_reqs:].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) - cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) - self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.query_start_loc.copy_to_gpu() - if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens - self.gdn_query_start_loc.copy_to_gpu() + cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) + self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() + if self._has_gdn: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + self.gdn_query_start_loc.copy_to_gpu() - if not profile_cpp: - num_reqs_padded = self._pad_query_start_loc_for_fia( - self.query_start_loc, - num_tokens_padded, - num_reqs_padded, - num_reqs, - cudagraph_runtime_mode, - batch_desc.num_reqs, + if not profile_cpp: + num_reqs_padded = self._pad_query_start_loc_for_fia( + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_runtime_mode, + batch_desc.num_reqs, + ) + + # Dummy graph runs do not go through _prepare_inputs(), but GDN/Mamba + # metadata reads block_table[:num_reqs_padded] below. Sync padded + # rows as well so device-side metadata does not see stale block ids. + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, + for_cudagraph_capture=is_graph_capturing, + num_scheduled_tokens_np=num_scheduled_tokens, ) - - pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL - attn_metadata, _ = self._build_attention_metadata( - num_tokens=num_tokens_unpadded, - num_tokens_padded=num_tokens_padded, - num_reqs=num_reqs_padded, - max_query_len=max_query_len, - ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, - for_cudagraph_capture=is_graph_capturing, - num_scheduled_tokens_np=num_scheduled_tokens, - ) + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, @@ -224,15 +302,24 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens - if getattr(getattr(self, "model", None), "has_preprocess", False): - input_ids = self.input_ids.gpu[:num_tokens_padded] - inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] - elif self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds: + # ---------------------------------------Omni-new---------------------------------------------- + model_kwargs = self._init_model_kwargs() + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + elif getattr(getattr(self, "model", None), "has_preprocess", False): + input_ids = self.input_ids.gpu[:num_tokens_padded] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None + # ---------------------------------------Omni-new---------------------------------------------- if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] @@ -290,6 +377,9 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): aclgraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, model_instance=self.model, + has_sinks=self._has_sinks, + input_ids=input_ids, + eplb_heat_collection_status=self.eplb_heat_collection_status if self.dynamic_eplb else False, ): # ---------------------------------------Omni-new---------------------------------------------- if getattr(self.model, "talker", None) is not None and self.has_talker_mtp: @@ -309,6 +399,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, + **model_kwargs, ) # ---------------------------------------Omni-new---------------------------------------------- if self.use_aux_hidden_state_outputs: @@ -320,7 +411,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # ---------------------------------------Omni-new---------------------------------------------- dummy_compute_logits(hidden_states) - if self.drafter: + if self.drafter and not profile_cpp: self.drafter.dummy_run( num_tokens=num_tokens_padded, with_prefill=with_prefill, @@ -333,10 +424,10 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): is_profile=is_profile, ) if is_profile and self.dynamic_eplb: - target = self.model.language_model if hasattr(self.model, "language_model") else self.model - target.clear_all_moe_loads() - if self.dynamic_eplb: - self.eplb_updator.forward_end() + self.eplb_updator.adaptor.clear_all_moe_loads() + if not is_profile and self.dynamic_eplb: + self.eplb_updator.forward_end(self.eplb_heat_collection_status) + self._finalize_dump_data(dump=False) return hidden_states, hidden_states def _model_forward( @@ -375,11 +466,11 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): if self.enable_enpu: # The soft segmentation scenario requires event.record first, then event.wait. - self._update_full_graph_params_if_needed(forward_context, num_tokens_padded, positions) + self._update_full_graph_params_if_needed(forward_context, num_tokens_padded) model_output = run_model() else: model_output = run_model() - self._update_full_graph_params_if_needed(forward_context, num_tokens_padded, positions) + self._update_full_graph_params_if_needed(forward_context, num_tokens_padded) # Omni-specific: wrap output if needed if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):