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

[Bugfix][NPU] Port OmniConnectorModelRunnerMixin to NPU runners to fix async_chunk=false full-payload hang (#5234) (#5290)

Browse Source 
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
This commit is contained in:
Weiming Liao authored and GitHub committed 2026-07-22 13:53:04 +00:00
1 parent 23d0ee5c0c
commit 8219c01b57
2 files changed
+85 -11
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+42 -5
View File
@@ -46,6 +46,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTran
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
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
@@ -100,7 +101,7 @@ class ExecuteModelState(NamedTuple):
batch_desc: BatchDescriptor
multimodal_outputs: Any # Omni-Specific
class NPUARModelRunner(OmniNPUModelRunner):
class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
"""Autoregressive NPU model runner that returns hidden states per request."""
def __init__(self, *args, **kwargs):
@@ -112,6 +113,22 @@ class NPUARModelRunner(OmniNPUModelRunner):
# 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_CONNECTOR_INIT_ARCHS = {
"Qwen3OmniMoeForConditionalGeneration",
"Qwen2_5OmniForConditionalGeneration",
"CovoAudioForConditionalGeneration",
"MiMoAudioModel",
"Qwen3TTSTalkerForConditionalGeneration",
"Qwen3TTSCode2Wav",
"CosyVoice3Model",
"DyninOmniForConditionalGeneration",
"IndexTTS2TalkerForConditionalGeneration",
}
if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
self.init_omni_connectors(
model_config=self.model_config,
kv_transfer_manager=self.kv_transfer_manager,
)
self._downstream_payload_cache: dict[str, bool] = {}
def _make_buffer(self, *size, dtype, numpy=True):
@@ -380,6 +397,15 @@ class NPUARModelRunner(OmniNPUModelRunner):
request_id_resolver=self._resolve_global_request_id,
)
# -------------------------------------- Omni-new -------------------------------------------------
if hasattr(self, "_omni_connector"):
for request in getattr(scheduler_output, "pending_input_registrations", []):
self.register_chunk_recv(request)
self.recv_full_payload_inputs(scheduler_output)
if self._pending_full_payload_send:
flush_ids = set(getattr(scheduler_output, "finished_req_ids", set()))
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)
# 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.
@@ -435,7 +461,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
if kv_ids:
output = copy(output)
output.kv_extracted_req_ids = kv_ids
return output
return self.attach_omni_connector_output(output)
if not num_scheduled_tokens:
if (
@@ -462,7 +488,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
output = copy(output)
output.kv_extracted_req_ids = kv_ids
return output
return self.attach_omni_connector_output(output)
if self.cache_config.kv_sharing_fast_prefill:
assert not self.num_prompt_logprobs, (
"--kv-sharing-fast-prefill produces incorrect "
@@ -879,11 +905,11 @@ class NPUARModelRunner(OmniNPUModelRunner):
# In case of PP with kv transfer, we need to pass through the
# kv_connector_output
if kv_connector_output.is_empty():
return EMPTY_MODEL_RUNNER_OUTPUT
return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
output.kv_connector_output = kv_connector_output
return output
return self.attach_omni_connector_output(output)
# Unpack ephemeral state.
(
@@ -1166,6 +1192,15 @@ class NPUARModelRunner(OmniNPUModelRunner):
# inter_stage_outputs (the NPU runner has no separate full-payload
# accumulate). #4527's (None, pooler_output) starved it. (PR #4792)
pooler_inter, pooler_client = pooler_output, pooler_output
# [Omni] Full-payload send-side accumulation. Mirrors gpu_ar_model_runner.py.
if pooler_inter and self._should_accumulate_full_payload_output():
with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"):
for i, rid in enumerate(req_ids_output_copy):
req_state = self.requests.get(rid)
if req_state is not None and pooler_inter[i]:
self.accumulate_full_payload_output(rid, pooler_inter[i], req_state)
inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
multimodal_outputs = self._build_multimodal_outputs(pooler_client)
model_runner_output = OmniModelRunnerOutput(
@@ -1183,6 +1218,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
)
model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
model_runner_output.routed_experts = routed_experts_lists
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"):
vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+43 -6
View File
@@ -36,14 +36,30 @@ 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
from vllm_omni.utils.mm_outputs import partition_payload_list
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
class NPUGenerationModelRunner(OmniNPUModelRunner):
class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin):
"""Generation model runner for vLLM-omni on NPU (non-autoregressive)."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._async_chunk = getattr(self.model_config, "async_chunk", False)
_OMNI_CONNECTOR_INIT_ARCHS = {
"Qwen3OmniMoeForConditionalGeneration",
"Qwen2_5OmniForConditionalGeneration",
"CovoAudioForConditionalGeneration",
"MiMoAudioModel",
"Qwen3TTSTalkerForConditionalGeneration",
"Qwen3TTSCode2Wav",
"CosyVoice3Model",
"DyninOmniForConditionalGeneration",
"IndexTTS2S2MelDecoder",
}
if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
self.init_omni_connectors(
model_config=self.model_config,
)
def _update_request_states(self, scheduler_output: SchedulerOutput):
# remove requests
@@ -113,6 +129,16 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
if kv_connector_metadata is not None:
get_kv_transfer_group().handle_preemptions(kv_connector_metadata)
if hasattr(self, "_omni_connector"):
for request in getattr(scheduler_output, "pending_input_registrations", []):
self.register_chunk_recv(request)
self.recv_full_payload_inputs(scheduler_output)
if self._pending_full_payload_send:
flush_ids = set(getattr(scheduler_output, "finished_req_ids", set()))
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)
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
with record_function_or_nullcontext("prepare input"):
# -------------------------------------- Omni-new -------------------------------------------------
@@ -129,7 +155,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
encoder_cache=self.encoder_cache,
) as ec_connector_output:
self._execute_mm_encoder(scheduler_output)
return make_empty_encoder_model_runner_output(scheduler_output)
return self.attach_omni_connector_output(
make_empty_encoder_model_runner_output(scheduler_output)
)
if not num_scheduled_tokens:
if (
@@ -145,8 +173,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
self._dummy_run(1)
if not has_kv_transfer_group():
# Return empty ModelRunnerOutput if no work to do.
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
return self.attach_omni_connector_output(
self.kv_connector_no_forward(scheduler_output, self.vllm_config)
)
if self.cache_config.kv_sharing_fast_prefill:
assert not self.num_prompt_logprobs, (
"--kv-sharing-fast-prefill produces incorrect "
@@ -416,11 +446,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
# In case of PP with kv transfer, we need to pass through the
# kv_connector_output
if kv_connector_output.is_empty():
return EMPTY_MODEL_RUNNER_OUTPUT
return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
output.kv_connector_output = kv_connector_output
return output
return self.attach_omni_connector_output(output)
# Unpack ephemeral state.
(
@@ -489,6 +519,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
# [Omni] Copy req_id mappings to avoid async scheduling mutation.
req_ids_output_copy = self.input_batch.req_ids.copy()
req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
# [Omni] Full-payload send-side accumulation. Mirrors gpu_generation_model_runner.py.
if inter_stage_outputs and self._should_accumulate_full_payload_output():
for i, rid in enumerate(req_ids_output_copy):
req_state = self.requests.get(rid)
if req_state is not None and inter_stage_outputs[i]:
self.accumulate_full_payload_output(rid, inter_stage_outputs[i], req_state)
routed_experts_lists = None
if self.vllm_config.model_config.enable_return_routed_experts and hasattr(
self.input_batch, "num_tokens_no_spec"
@@ -509,6 +545,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
)
output.routed_experts = routed_experts_lists
output.omni_connector_output = self.get_omni_connector_output()
# -------------------------------------- Omni-new -------------------------------------------------
if self.speculative_config is not None: