anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[NPU] upgrade to v0.26.0 (#5490)

Browse Source 
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
This commit is contained in:
Weiming Liao authored and GitHub committed 2026-07-29 11:21:04 +08:00
1 parent ab0da662c0
commit 39663d5287
10 files changed
+482 -298
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docker/Dockerfile.npu
+1 -1
View File
@@ -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
docker/Dockerfile.npu.a3
+1 -1
View File
@@ -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
docker/Dockerfile.npu.ci
+1 -1
View File
@@ -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
docker/Dockerfile.npu.ci.a3
+1 -1
View File
@@ -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
docs/getting_started/installation/npu/npu.inc.md
+6 -6
View File
@@ -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
vllm_omni/deploy/minicpmo_4_5_batching.yaml
+1
View File
@@ -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
vllm_omni/platforms/npu/platform.py
+2
View File
@@ -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:
vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+155 -99
View File
@@ -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(
vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+155 -121
View File
@@ -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
vllm_omni/platforms/npu/worker/npu_model_runner.py
+159 -68
View File
@@ -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"):