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

[NPU] Upgrade to v0.18.0

Browse Source 
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
gcanlin committed 2026-03-20 09:53:34 +00:00
1 parent 521f178619
commit 9538e3f30b
4 files changed
+115 -41
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/getting_started/installation/npu/npu.inc.md
+3 -3
View File
@@ -68,18 +68,18 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
```bash
# Pin vLLM version to 0.17.0
# Pin vLLM version to 0.18.0
cd /vllm-workspace/vllm
git pull origin main
git fetch origin --tags
git checkout v0.17.0
git checkout v0.18.0
VLLM_TARGET_DEVICE=empty pip install -v -e .
# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
cd /vllm-workspace/vllm-ascend
git pull origin main
git fetch origin --tags
git checkout v0.17.0
git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
pip install -v -e .
# Install vLLM-Omni from the latest main branch
vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+66 -22
View File
@@ -26,6 +26,7 @@ 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.ubatch_utils import maybe_create_ubatch_slices
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
@@ -237,6 +238,23 @@ class NPUARModelRunner(OmniNPUModelRunner):
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
# NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877,
# there should be a corresponding 'postprocess_mamba'. However, it is called inside
# '_update_states_after_model_execute', which is not overridden in vLLM-Ascend.
# We simply utilize the implementation in vLLM.
if self.cache_config.mamba_cache_mode == "align":
mamba_utils.preprocess_mamba(
scheduler_output,
self.kv_cache_config,
self.cache_config,
self.mamba_state_idx,
self.input_batch,
self.requests,
self.compilation_config.static_forward_context,
self.model.get_mamba_state_copy_func(),
self._get_mamba_copy_bufs(),
)
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
@@ -339,7 +357,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
skip_compiled=has_encoder_input,
),
self.maybe_get_kv_connector_output(
scheduler_output, clear_metadata=clear_kv_metadata
scheduler_output, defer_finalize=not clear_kv_metadata
) as kv_connector_output,
):
hidden_states = self._model_forward(
@@ -606,6 +624,34 @@ class NPUARModelRunner(OmniNPUModelRunner):
hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output
)
# Pre-copy multimodal tensors to CPU once (not per-request) to avoid
# redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU.
mm_cpu: dict[str, object] = {}
if isinstance(multimodal_outputs, dict) and multimodal_outputs:
for k, v in multimodal_outputs.items():
try:
if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
mm_cpu[k] = v.detach().to("cpu").contiguous()
elif isinstance(v, dict):
sub_dict: dict[str, torch.Tensor] = {}
for sk, sv in v.items():
if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
sub_dict[str(sk)] = sv.detach().to("cpu").contiguous()
if sub_dict:
mm_cpu[k] = sub_dict
elif isinstance(v, list):
if len(v) == 0:
continue
cpu_list = []
for elem in v:
if isinstance(elem, torch.Tensor):
cpu_list.append(elem.detach().to("cpu").contiguous())
else:
cpu_list.append(elem)
mm_cpu[k] = cpu_list
except Exception as e:
logger.error(f"Error in merge multimodal outputs: {e}")
pooler_output: list[dict[str, object]] = []
for rid in req_ids_output_copy:
idx = req_id_to_index_output_copy[rid]
@@ -614,28 +660,26 @@ class NPUARModelRunner(OmniNPUModelRunner):
end = start + sched
hidden_slice = hidden_states_cpu[start:end]
payload: dict[str, object] = {"hidden": hidden_slice}
if isinstance(multimodal_outputs, dict) and multimodal_outputs:
if mm_cpu:
mm_payload: dict[str, object] = {}
for k, v in multimodal_outputs.items():
try:
if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
mm_payload[k] = v.detach().to("cpu")[start:end].contiguous()
elif isinstance(v, dict):
sub_dict: dict[str, torch.Tensor] = {}
for sk, sv in v.items():
if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous()
if sub_dict:
mm_payload[k] = sub_dict
elif isinstance(v, list):
element = v[0]
if isinstance(element, torch.Tensor):
element = element.detach().to("cpu").contiguous()
mm_payload[k] = element
except Exception as e:
logger.error(f"Error in merge multimodal outputs: {e}")
if mm_payload:
payload.update(mm_payload)
for k, v in mm_cpu.items():
if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
mm_payload[k] = v[start:end].contiguous()
elif isinstance(v, dict):
mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()}
elif isinstance(v, list):
element = v[idx] if idx < len(v) else v[0]
# Clone tensors to avoid cross-request aliasing
if isinstance(element, torch.Tensor):
element = element.clone()
mm_payload[k] = element
elif isinstance(v, torch.Tensor):
# List-derived tensor payloads are request-invariant; clone to
# avoid accidental cross-request aliasing on downstream mutation.
mm_payload[k] = v.clone()
else:
mm_payload[k] = v
payload.update(mm_payload)
pooler_output.append(payload)
model_runner_output = OmniModelRunnerOutput(
vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+28 -6
View File
@@ -20,6 +20,7 @@ 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.ubatch_utils import maybe_create_ubatch_slices
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
@@ -195,6 +196,23 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
# NOTE(Angazenn): According to https://github.com/vllm-project/vllm/pull/30877,
# there should be a corresponding 'postprocess_mamba'. However, it is called inside
# '_update_states_after_model_execute', which is not overridden in vLLM-Ascend.
# We simply utilize the implementation in vLLM.
if self.cache_config.mamba_cache_mode == "align":
mamba_utils.preprocess_mamba(
scheduler_output,
self.kv_cache_config,
self.cache_config,
self.mamba_state_idx,
self.input_batch,
self.requests,
self.compilation_config.static_forward_context,
self.model.get_mamba_state_copy_func(),
self._get_mamba_copy_bufs(),
)
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
@@ -300,7 +318,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
skip_compiled=has_encoder_input,
),
self.maybe_get_kv_connector_output(
scheduler_output, clear_metadata=clear_kv_metadata
scheduler_output, defer_finalize=not clear_kv_metadata
) as kv_connector_output,
):
# -------------------------------------- Omni-new -------------------------------------------------
@@ -503,6 +521,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
remove_lora: bool = True,
is_graph_capturing: bool = False,
num_active_loras: int = 0,
profile_seq_lens: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
# only support eager mode and piecewise graph now
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
@@ -609,11 +628,14 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
# 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.
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]
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.seq_lens.np[:num_reqs_padded] = seq_lens
self.seq_lens.np[num_reqs_padded:] = 0
self.seq_lens.copy_to_gpu()
vllm_omni/platforms/npu/worker/npu_model_runner.py
+18 -10
View File
@@ -35,9 +35,11 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
enable_sp(self.vllm_config)
# TODO move this model specific logic to a separate class
# TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
self.has_talker_mtp = False
talker_mtp = getattr(self.model, "talker_mtp", None)
if talker_mtp is not None:
self.talker_mtp = talker_mtp # type: ignore[assignment]
self.has_talker_mtp = True
cudagraph_mode = self.compilation_config.cudagraph_mode
assert cudagraph_mode is not None
# Only wrap talker_mtp in CUDAGraphWrapper for Omni models that
@@ -73,6 +75,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
remove_lora: bool = True,
is_graph_capturing: bool = False,
num_active_loras: int = 0,
profile_seq_lens: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
# only support eager mode and piecewise graph now
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes()
@@ -179,11 +182,14 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
# 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.
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]
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.seq_lens.np[:num_reqs_padded] = seq_lens
self.seq_lens.np[num_reqs_padded:] = 0
self.seq_lens.copy_to_gpu()
@@ -283,7 +289,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
model_instance=self.model,
):
# ---------------------------------------Omni-new----------------------------------------------
if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
if getattr(self.model, "talker", None) is not None and self.has_talker_mtp:
num_tokens_padded_talker_mtp = num_tokens_padded
if num_tokens_padded_talker_mtp == self.max_num_tokens:
num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
@@ -362,7 +368,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
# Omni-specific: wrap output if needed
if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
model_output = self.model.make_omni_output(model_output, **model_kwargs_extra)
model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra)
# Omni-specific: cache model output for later sample_tokens
self._omni_last_model_output = model_output
@@ -415,12 +421,14 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
None, self.vllm_config, aclgraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
):
req_embeds, code_predictor_codes = self.talker_mtp(req_input_ids, req_embeds, last_talker_hidden, text_step)
# update the inputs_embeds and code_predictor_codes
code_predictor_codes_cpu = code_predictor_codes.detach().to("cpu").contiguous()
# code_predictor_codes stays on GPU here; _update_intermediate_buffer
# keeps it device-resident when the key is in gpu_resident_buffer_keys.
# D2H is deferred to sample_tokens where hidden_states.to("cpu") already
# syncs the stream, avoiding a per-step cudaStreamSynchronize.
out_key = getattr(self.model, "talker_mtp_output_key", "code_predictor_codes")
for idx, req_id in enumerate(decode_req_ids):
req_index = self.input_batch.req_ids.index(req_id)
start_offset = int(self.query_start_loc.cpu[req_index])
inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
update_dict = {out_key: code_predictor_codes_cpu[idx : idx + 1]}
update_dict = {out_key: code_predictor_codes[idx : idx + 1]}
self._merge_additional_information_update(req_id, update_dict)