968 lines
46 KiB
Python
968 lines
46 KiB
Python
"""Code2Wav GPU Model Runner for vLLM-Omni.
|
|
|
|
Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav.
|
|
This is a non-autoregressive model that doesn't require sampling or logits computation.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import logging
|
|
from collections.abc import Mapping
|
|
from dataclasses import replace
|
|
|
|
import numpy as np
|
|
import torch
|
|
from vllm.config import CUDAGraphMode
|
|
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
|
|
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
|
|
from vllm.distributed.parallel_state import get_pp_group
|
|
from vllm.forward_context import set_forward_context
|
|
from vllm.utils.math_utils import cdiv
|
|
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
|
|
from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
|
|
from vllm.v1.spec_decode.dflash import DFlashProposer
|
|
from vllm.v1.spec_decode.draft_model import DraftModelProposer
|
|
from vllm.v1.spec_decode.eagle import EagleProposer
|
|
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
|
|
from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
|
|
from vllm.v1.utils import record_function_or_nullcontext
|
|
from vllm.v1.worker.gpu_model_runner import (
|
|
EMPTY_MODEL_RUNNER_OUTPUT,
|
|
AsyncGPUModelRunnerOutput,
|
|
IntermediateTensors,
|
|
PerLayerAttnMetadata,
|
|
)
|
|
from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
|
|
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
|
|
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.utils.mm_outputs import partition_payload_list
|
|
from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values
|
|
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
|
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
|
|
"""Generation model runner for vLLM-Omni (non-autoregressive).
|
|
|
|
- Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
|
|
- Does not compute logits or perform token sampling.
|
|
- Executes generation process and returns tensors via `pooler_output`.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._async_chunk = getattr(self.model_config, "async_chunk", False)
|
|
# Mirrors the init allowlist in gpu_ar_model_runner.py.
|
|
_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(
|
|
vllm_config=self.vllm_config,
|
|
model_config=self.model_config,
|
|
)
|
|
|
|
def _update_request_states(self, scheduler_output: SchedulerOutput):
|
|
# remove requests
|
|
# Some stateful vocoder model may need to clean the state
|
|
# to avoid the leak of slots when the requests have been aborted.
|
|
if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
|
|
self.model.on_requests_finished(scheduler_output.finished_req_ids)
|
|
for req_id in scheduler_output.finished_req_ids:
|
|
self.input_batch.remove_request(req_id)
|
|
scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
|
|
cached_req_ids = self.input_batch.req_id_to_index.keys()
|
|
resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
|
|
unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids)
|
|
for req_id in unscheduled_req_ids:
|
|
self.input_batch.remove_request(req_id)
|
|
cached_reqs = scheduler_output.scheduled_cached_reqs
|
|
req_states = []
|
|
for req_id in cached_reqs.req_ids:
|
|
req_state = self.requests.get(req_id)
|
|
assert req_state is not None
|
|
req_state.prompt_token_ids = cached_reqs.prompt_token_ids.get(req_id)
|
|
req_states.append(req_state)
|
|
# Remove the request from the current input batch only if it is still present.
|
|
if req_id in self.input_batch.req_id_to_index:
|
|
self.input_batch.remove_request(req_id)
|
|
for req_state in req_states:
|
|
# update the request state in self.input_batch
|
|
self.input_batch.add_request(req_state)
|
|
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
|
|
if self.uses_mrope:
|
|
self._init_mrope_positions(req_state)
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
intermediate_tensors: IntermediateTensors | None = None,
|
|
) -> OmniModelRunnerOutput | IntermediateTensors:
|
|
if self.execute_model_state is not None:
|
|
raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
|
|
|
|
if self.routed_experts_initialized:
|
|
self.routed_experts_capturer.clear_buffer()
|
|
|
|
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)
|
|
|
|
if self.routed_experts_initialized:
|
|
capturer = self.routed_experts_capturer
|
|
if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
|
|
capturer.finalize_pending_copy()
|
|
|
|
# If ngram_gpu is used, we need to copy the scheduler_output to avoid
|
|
# the modification has influence on the scheduler_output in engine core process.
|
|
# The replace is much faster than deepcopy.
|
|
if self.speculative_config is not None and self.speculative_config.use_ngram_gpu():
|
|
num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy()
|
|
spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy()
|
|
scheduler_output = replace(
|
|
scheduler_output,
|
|
num_scheduled_tokens=num_scheduled_tokens_copy,
|
|
scheduled_spec_decode_tokens=spec_decode_tokens_copy,
|
|
)
|
|
|
|
if has_kv_transfer_group():
|
|
kv_connector_metadata = scheduler_output.kv_connector_metadata
|
|
if kv_connector_metadata is not None:
|
|
get_kv_transfer_group().handle_preemptions(kv_connector_metadata)
|
|
|
|
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
|
|
with (
|
|
record_function_or_nullcontext("gpu_model_runner: preprocess"),
|
|
self.synchronize_input_prep(),
|
|
):
|
|
if self.model_config.async_chunk and num_scheduled_tokens:
|
|
self._update_request_states(scheduler_output)
|
|
deferred_state_corrections_fn = self._update_states(scheduler_output)
|
|
if not scheduler_output.total_num_scheduled_tokens:
|
|
return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
|
|
|
|
if has_ec_transfer() and not get_ec_transfer().is_consumer:
|
|
with self.maybe_get_ec_connector_output(
|
|
scheduler_output,
|
|
encoder_cache=self.encoder_cache,
|
|
) as ec_connector_output:
|
|
self._execute_mm_encoder(scheduler_output)
|
|
return self.attach_omni_connector_output(make_empty_encoder_model_runner_output(scheduler_output))
|
|
|
|
if not num_scheduled_tokens:
|
|
if (
|
|
self.parallel_config.distributed_executor_backend == "external_launcher"
|
|
and self.parallel_config.data_parallel_size > 1
|
|
):
|
|
# this is a corner case when both external launcher
|
|
# and DP are enabled, num_scheduled_tokens could be
|
|
# 0, and has_unfinished_requests in the outer loop
|
|
# returns True. before returning early here we call
|
|
# dummy run to ensure coordinate_batch_across_dp
|
|
# is called into to avoid out of sync issues.
|
|
self._dummy_run(1)
|
|
if not has_kv_transfer_group():
|
|
# Return empty ModelRunnerOutput if no work to do.
|
|
return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
|
|
|
|
result = self.kv_connector_no_forward(scheduler_output, self.vllm_config)
|
|
return self.attach_omni_connector_output(result)
|
|
|
|
if self.cache_config.kv_sharing_fast_prefill:
|
|
assert not self.num_prompt_logprobs, (
|
|
"--kv-sharing-fast-prefill produces incorrect "
|
|
"logprobs for prompt tokens, tokens, please disable "
|
|
"it when the requests need prompt logprobs"
|
|
)
|
|
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]
|
|
num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
|
|
max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
|
|
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
|
|
|
|
logits_indices, spec_decode_metadata = self._prepare_inputs(
|
|
scheduler_output,
|
|
num_scheduled_tokens_np,
|
|
)
|
|
|
|
cascade_attn_prefix_lens = None
|
|
# Disable cascade attention when using microbatching (DBO)
|
|
if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
|
|
# Pre-compute cascade attention prefix lengths
|
|
cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
|
|
num_scheduled_tokens_np,
|
|
self.input_batch.num_computed_tokens_cpu[:num_reqs],
|
|
scheduler_output.num_common_prefix_blocks,
|
|
)
|
|
|
|
(
|
|
cudagraph_mode,
|
|
batch_desc,
|
|
should_ubatch,
|
|
num_tokens_across_dp,
|
|
cudagraph_stats,
|
|
) = self._determine_batch_execution_and_padding(
|
|
num_tokens=num_tokens_unpadded,
|
|
num_reqs=num_reqs,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
max_num_scheduled_tokens=max_num_scheduled_tokens,
|
|
use_cascade_attn=cascade_attn_prefix_lens is not None,
|
|
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,
|
|
)
|
|
|
|
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
|
|
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
|
|
should_ubatch,
|
|
num_scheduled_tokens_np,
|
|
num_tokens_padded,
|
|
num_reqs_padded,
|
|
self.parallel_config.num_ubatches,
|
|
)
|
|
|
|
logger.debug(
|
|
"ubatch_slices: %s, ubatch_slices_padded: %s",
|
|
ubatch_slices,
|
|
ubatch_slices_padded,
|
|
)
|
|
|
|
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
|
|
|
|
use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
|
|
ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
|
|
|
|
# OMNI: True if any attention backend handles KV cache update separately
|
|
# from forward() (i.e., forward_includes_kv_cache_update=False). When true,
|
|
# slot_mappings must use padded dimensions to match the key/value tensors.
|
|
from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec
|
|
|
|
has_separate_kv_update = not all(
|
|
all(g.backend.forward_includes_kv_cache_update for g in self.attn_groups[id])
|
|
for id, spec in enumerate(self.kv_cache_config.kv_cache_groups)
|
|
if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec)
|
|
)
|
|
|
|
slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
|
|
num_tokens_padded=num_tokens_padded if pad_attn or has_separate_kv_update else num_tokens_unpadded,
|
|
num_reqs_padded=(num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs),
|
|
num_tokens_unpadded=num_tokens_unpadded,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
)
|
|
|
|
attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
|
|
num_tokens=num_tokens_unpadded,
|
|
num_tokens_padded=num_tokens_padded if pad_attn else None,
|
|
num_reqs=num_reqs,
|
|
num_reqs_padded=num_reqs_padded if pad_attn else None,
|
|
max_query_len=max_num_scheduled_tokens,
|
|
ubatch_slices=ubatch_slices_attn,
|
|
logits_indices=logits_indices,
|
|
use_spec_decode=use_spec_decode,
|
|
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
|
|
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
|
|
slot_mappings=slot_mappings_by_group,
|
|
)
|
|
self._maybe_attach_attention_metadata_extensions(
|
|
attn_metadata=attn_metadata,
|
|
num_reqs=num_reqs,
|
|
num_reqs_padded=num_reqs_padded,
|
|
max_query_len=max_num_scheduled_tokens,
|
|
pad_attn=pad_attn,
|
|
num_scheduled_tokens_np=num_scheduled_tokens_np,
|
|
)
|
|
|
|
(
|
|
input_ids,
|
|
inputs_embeds,
|
|
positions,
|
|
intermediate_tensors,
|
|
model_kwargs,
|
|
ec_connector_output,
|
|
) = self._preprocess(
|
|
scheduler_output,
|
|
num_tokens_padded,
|
|
intermediate_tensors,
|
|
)
|
|
# [Omni] Pass token counts per request for code2wav output slicing
|
|
model_kwargs["seq_token_counts"] = tokens
|
|
|
|
# Set cudagraph mode to none if calc_kv_scales is true.
|
|
# KV scales calculation involves dynamic operations that are incompatible
|
|
# with CUDA graph capture.
|
|
if self.calculate_kv_scales:
|
|
cudagraph_mode = CUDAGraphMode.NONE
|
|
# Mark KV scales as calculated after the first forward pass
|
|
self.calculate_kv_scales = False
|
|
|
|
# Run the model.
|
|
# Use persistent buffers for CUDA graphs.
|
|
# When spec decode is enabled, defer connector finalization
|
|
# (wait_for_save + clear metadata) until after draft model runs.
|
|
defer_kv_connector_finalize = self.speculative_config is not None
|
|
with (
|
|
set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_tokens_padded,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
cudagraph_runtime_mode=cudagraph_mode,
|
|
batch_descriptor=batch_desc,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
slot_mapping=slot_mappings, # OMNI: required for KV cache operations
|
|
),
|
|
record_function_or_nullcontext("Forward"),
|
|
self.maybe_get_kv_connector_output(
|
|
scheduler_output,
|
|
defer_finalize=defer_kv_connector_finalize,
|
|
) as kv_connector_output,
|
|
):
|
|
outputs = self._run_generation_model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
model_kwargs=model_kwargs,
|
|
logits_indices=logits_indices,
|
|
)
|
|
|
|
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
|
|
self.execute_model_state = ExecuteModelState(
|
|
scheduler_output,
|
|
None,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
ec_connector_output,
|
|
cudagraph_stats,
|
|
multimodal_outputs,
|
|
slot_mappings, # OMNI: pass slot_mappings for upstream v1 API compatibility
|
|
)
|
|
self.kv_connector_output = kv_connector_output
|
|
|
|
if deferred_state_corrections_fn:
|
|
deferred_state_corrections_fn()
|
|
|
|
if self.routed_experts_initialized and hasattr(self, "_positions_cpu"):
|
|
self._omni_routed_experts_d2h(scheduler_output)
|
|
|
|
return None
|
|
|
|
@torch.inference_mode()
|
|
def sample_tokens(
|
|
self,
|
|
grammar_output: GrammarOutput | None = None,
|
|
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
|
|
# NOTE: Even though the model is non-autoregressive, we still need
|
|
# this function to match the interface of the engine core.
|
|
# In this case, this function
|
|
kv_connector_output = self.kv_connector_output
|
|
self.kv_connector_output = None
|
|
|
|
if self.execute_model_state is None:
|
|
# receive sampled token ids from the last PP rank.
|
|
if self.use_async_scheduling and not get_pp_group().is_last_rank:
|
|
self._pp_receive_prev_sampled_token_ids_to_input_batch()
|
|
# In case of PP with kv transfer, we need to pass through the
|
|
# kv_connector_output
|
|
return self.attach_omni_connector_output(
|
|
OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
|
|
)
|
|
|
|
# Unpack ephemeral state.
|
|
(
|
|
scheduler_output,
|
|
logits,
|
|
spec_decode_metadata,
|
|
spec_decode_common_attn_metadata,
|
|
hidden_states,
|
|
_hidden_states_cpu,
|
|
sample_hidden_states,
|
|
aux_hidden_states,
|
|
ec_connector_output,
|
|
cudagraph_stats,
|
|
multimodal_outputs_raw,
|
|
slot_mappings, # OMNI: unpack slot_mappings for upstream v1 API compatibility
|
|
) = self.execute_model_state
|
|
self.execute_model_state = None
|
|
|
|
# Finalize KV connector (wait_for_save + clear metadata) after
|
|
# draft model runs. Deferred from target model forward.
|
|
if self.speculative_config is not None:
|
|
self.finalize_kv_connector()
|
|
|
|
# Build per-request multimodal_outputs list (dedicated channel).
|
|
# pooler_output is no longer used for multimodal data.
|
|
per_req_payloads: list[dict[str, object]] = []
|
|
if isinstance(multimodal_outputs_raw, torch.Tensor):
|
|
assert multimodal_outputs_raw.shape[0] == 1, (
|
|
"model should return a single tensor, to return multiple tensors, use a dict"
|
|
)
|
|
assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs
|
|
for i in range(self.input_batch.num_reqs):
|
|
per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()})
|
|
elif isinstance(multimodal_outputs_raw, list):
|
|
assert len(multimodal_outputs_raw) == 1, (
|
|
"model should return a single list, to return multiple lists, use a dict"
|
|
)
|
|
for out in multimodal_outputs_raw:
|
|
per_req_payloads.append(
|
|
{"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
|
|
)
|
|
elif isinstance(multimodal_outputs_raw, Mapping):
|
|
num_reqs = self.input_batch.num_reqs
|
|
for i in range(num_reqs):
|
|
mm_payload = {}
|
|
for key, out in multimodal_outputs_raw.items():
|
|
if isinstance(out, list):
|
|
if len(out) != num_reqs:
|
|
raise ValueError(
|
|
f"Multimodal output list for key '{key}' has length {len(out)} "
|
|
f"but expected {num_reqs} (one entry per request)."
|
|
)
|
|
mm_payload[key] = out[i].detach().to("cpu").contiguous()
|
|
elif isinstance(out, torch.Tensor):
|
|
mm_payload[key] = out.detach().to("cpu").contiguous()
|
|
else:
|
|
logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
|
|
per_req_payloads.append(_ensure_tensor_values(mm_payload))
|
|
else:
|
|
raise RuntimeError("Unsupported diffusion output type")
|
|
|
|
if self._async_chunk:
|
|
inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
|
|
else:
|
|
# See gpu_ar_model_runner: non-async-chunk ships the full payload to the next
|
|
# stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792)
|
|
inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads
|
|
|
|
# [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()
|
|
routed_experts_lists = None
|
|
if self.routed_experts_initialized:
|
|
routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
|
|
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)
|
|
|
|
output = OmniModelRunnerOutput(
|
|
req_ids=req_ids_output_copy,
|
|
req_id_to_index=req_id_to_index_output_copy,
|
|
sampled_token_ids=[],
|
|
logprobs=None,
|
|
prompt_logprobs_dict={},
|
|
pooler_output=None,
|
|
multimodal_outputs=multimodal_outputs,
|
|
inter_stage_outputs=inter_stage_outputs,
|
|
kv_connector_output=kv_connector_output,
|
|
num_nans_in_logits={},
|
|
cudagraph_stats=cudagraph_stats,
|
|
ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
|
|
)
|
|
output.omni_connector_output = self.get_omni_connector_output()
|
|
output.routed_experts = routed_experts_lists
|
|
|
|
if not self.use_async_scheduling:
|
|
return output
|
|
|
|
return AsyncGPUModelRunnerOutput(
|
|
model_runner_output=output,
|
|
sampled_token_ids=torch.tensor([], device=self.device),
|
|
invalid_req_indices=[],
|
|
async_output_copy_stream=self.async_output_copy_stream,
|
|
vocab_size=self.input_batch.vocab_size,
|
|
logprobs_tensors=None,
|
|
)
|
|
|
|
def _run_generation_model(
|
|
self,
|
|
*,
|
|
input_ids: torch.Tensor | None,
|
|
positions: torch.Tensor | None,
|
|
intermediate_tensors: IntermediateTensors | None,
|
|
inputs_embeds: torch.Tensor | None,
|
|
model_kwargs: dict,
|
|
logits_indices: torch.Tensor,
|
|
) -> torch.Tensor | list[torch.Tensor]:
|
|
"""Run generation from codec codes to waveforms.
|
|
|
|
Args:
|
|
scheduler_output: Contains codec codes in input_ids or additional info
|
|
intermediate_tensors: PP intermediate tensors if applicable
|
|
|
|
Returns:
|
|
Audio waveforms: [batch, 1, waveform_len] or list of tensors
|
|
"""
|
|
# Keep inputs identical to AR runner
|
|
kwargs = dict(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**model_kwargs,
|
|
sampling_metadata=self.input_batch.sampling_metadata,
|
|
logits_index=logits_indices,
|
|
sampler=self.sampler,
|
|
)
|
|
|
|
if hasattr(self.model, "forward"):
|
|
return self._model_forward(**kwargs)
|
|
|
|
raise RuntimeError(
|
|
"The loaded model does not expose diffusion interfaces 'sample', "
|
|
"'forward', or 'diffuse'. Please implement one of them or adapt the runner."
|
|
)
|
|
|
|
@torch.inference_mode()
|
|
def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
logger.warning("Dummy sampler run is not implemented for generation model")
|
|
return torch.tensor([])
|
|
|
|
@torch.inference_mode()
|
|
def _dummy_run(
|
|
self,
|
|
num_tokens: int,
|
|
cudagraph_runtime_mode: CUDAGraphMode | None = None,
|
|
force_attention: bool = False,
|
|
uniform_decode: bool = False,
|
|
allow_microbatching: bool = True,
|
|
skip_eplb: bool = False,
|
|
is_profile: bool = False,
|
|
create_mixed_batch: bool = False,
|
|
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]:
|
|
"""
|
|
Run a dummy forward pass to warm up/profile run or capture the
|
|
CUDA graph for the model.
|
|
|
|
Args:
|
|
num_tokens: Number of tokens to run the dummy forward pass.
|
|
cudagraph_runtime_mode: used to control the behavior.
|
|
- if not set will determine the cudagraph mode based on using
|
|
the self.cudagraph_dispatcher.
|
|
- CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
|
|
- CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
|
|
- CUDAGraphMode.FULL: Full cudagraph, attention metadata is
|
|
needed.
|
|
force_attention: If True, always create attention metadata. Used to
|
|
warm up attention backend when mode is NONE.
|
|
uniform_decode: If True, the batch is a uniform decode batch.
|
|
skip_eplb: If True, skip EPLB state update.
|
|
is_profile: If True, this is a profile run.
|
|
create_mixed_batch: If True, create a mixed batch with both decode
|
|
(1 token) and prefill (multiple tokens) requests.
|
|
remove_lora: If False, dummy LoRAs are not destroyed after the run
|
|
num_active_loras: Number of distinct active LoRAs to capture for.
|
|
LoRA is activated when num_active_loras > 0.
|
|
profile_seq_lens: If provided, use this value for seq_lens instead
|
|
of max_query_len. Used to profile attention workspace that
|
|
scales with context length.
|
|
"""
|
|
mm_config = self.vllm_config.model_config.multimodal_config
|
|
if mm_config and mm_config.mm_encoder_only:
|
|
# The current dummy run only covers LM execution, so we can skip it.
|
|
# mm encoder dummy run may need to add in the future.
|
|
return torch.tensor([]), torch.tensor([])
|
|
|
|
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode()
|
|
|
|
# If cudagraph_mode.decode_mode() == FULL and
|
|
# cudagraph_mode.separate_routine(). This means that we are using
|
|
# different graphs and/or modes for mixed prefill-decode batches vs.
|
|
# uniform decode batches. A uniform decode batch means that all
|
|
# requests have identical query length, except a potential virtual
|
|
# request (shorter) in the batch account for padding.
|
|
# Uniform decode batch could either be common pure decode, where
|
|
# max_query_len == 1, or speculative decode, where
|
|
# max_query_len == 1 + num_spec_decode_tokens.
|
|
|
|
# When setting max_query_len = 1, we switch to and capture the optimized
|
|
# routine of FA2 for pure decode, i.e., Flashdecode + an optimization
|
|
# for GQA/MQA.
|
|
max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
|
|
|
|
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
|
|
# for dummy run with LoRA so that the num_reqs collectively
|
|
# has num_tokens in total.
|
|
assert num_tokens <= self.max_num_tokens
|
|
max_num_reqs = self.scheduler_config.max_num_seqs
|
|
if create_mixed_batch:
|
|
assert not uniform_decode
|
|
# Create mixed batch:
|
|
# first half decode tokens, second half one prefill
|
|
num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2)
|
|
num_prefill_tokens = num_tokens - num_decode_tokens
|
|
num_reqs = num_decode_tokens + 1
|
|
|
|
# Create decode requests (1 token each) followed by prefill request
|
|
num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens]
|
|
# Note: Overriding max_query_len to be the prefill tokens
|
|
max_query_len = num_prefill_tokens
|
|
elif uniform_decode:
|
|
assert not create_mixed_batch
|
|
num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len))
|
|
num_scheduled_tokens_list = [max_query_len] * num_reqs
|
|
if num_tokens % max_query_len != 0:
|
|
num_scheduled_tokens_list[-1] = num_tokens % max_query_len
|
|
else:
|
|
num_reqs = min(num_tokens, max_num_reqs)
|
|
min_tokens_per_req = num_tokens // num_reqs
|
|
num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
|
|
num_scheduled_tokens_list[-1] += num_tokens % num_reqs
|
|
|
|
assert sum(num_scheduled_tokens_list) == num_tokens
|
|
assert len(num_scheduled_tokens_list) == num_reqs
|
|
num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
|
|
num_tokens_unpadded = int(num_scheduled_tokens.sum())
|
|
|
|
num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
|
|
|
|
_cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = (
|
|
self._determine_batch_execution_and_padding(
|
|
num_tokens=num_tokens_unpadded,
|
|
num_reqs=num_reqs,
|
|
num_scheduled_tokens_np=num_scheduled_tokens,
|
|
max_num_scheduled_tokens=max_query_len,
|
|
use_cascade_attn=False,
|
|
allow_microbatching=allow_microbatching,
|
|
force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE),
|
|
# `force_uniform_decode` is used for cudagraph capture; because for
|
|
# capturing mixed prefill-decode batches, we sometimes use
|
|
# num_tokens == num_reqs which looks like a uniform decode batch to the
|
|
# dispatcher; but we actually want to capture a piecewise cudagraph
|
|
force_uniform_decode=uniform_decode,
|
|
# `force_has_lora` is used for cudagraph capture; because LoRA is
|
|
# activated later in the context manager, but we need to know the
|
|
# LoRA state when determining the batch descriptor for capture
|
|
force_has_lora=num_active_loras > 0,
|
|
# `force_num_active_loras` is used for cudagraph capture; because we
|
|
# need to capture graphs for specific num_active_loras counts
|
|
force_num_active_loras=num_active_loras,
|
|
)
|
|
)
|
|
|
|
if cudagraph_runtime_mode is None:
|
|
cudagraph_runtime_mode = _cudagraph_mode
|
|
else:
|
|
assert cudagraph_runtime_mode == _cudagraph_mode, (
|
|
f"Cudagraph runtime mode mismatch in dummy_run. "
|
|
f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}."
|
|
)
|
|
|
|
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
|
|
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
|
|
should_ubatch,
|
|
num_scheduled_tokens,
|
|
num_tokens_padded,
|
|
num_reqs_padded,
|
|
self.vllm_config.parallel_config.num_ubatches,
|
|
)
|
|
logger.debug(
|
|
"ubatch_slices: %s, ubatch_slices_padded: %s",
|
|
ubatch_slices,
|
|
ubatch_slices_padded,
|
|
)
|
|
|
|
attn_metadata: PerLayerAttnMetadata | None = None
|
|
|
|
# OMNI: Get slot mappings before building attention metadata
|
|
slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
|
|
num_tokens_padded=num_tokens_padded,
|
|
num_reqs_padded=num_reqs_padded,
|
|
num_tokens_unpadded=num_tokens_unpadded,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
)
|
|
|
|
if slot_mappings_by_group is not None:
|
|
for sm in slot_mappings_by_group.values():
|
|
sm.fill_(-1)
|
|
|
|
with self.synchronize_input_prep():
|
|
# If force_attention is True, we always capture attention.
|
|
# Otherwise, it only happens for cudagraph_runtime_mode=FULL.
|
|
if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
|
|
if profile_seq_lens is not None:
|
|
seq_lens = profile_seq_lens # type: ignore[assignment]
|
|
elif create_mixed_batch:
|
|
seq_lens = torch.tensor( # type: ignore[assignment]
|
|
[1] * num_decode_tokens + [num_prefill_tokens + 1],
|
|
dtype=torch.int,
|
|
)
|
|
else:
|
|
seq_lens = 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 + 1] = cum_num_tokens
|
|
self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1])
|
|
self.query_start_loc.copy_to_gpu()
|
|
|
|
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 if pad_attn else None,
|
|
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,
|
|
slot_mappings=slot_mappings_by_group,
|
|
use_spec_decode=self.speculative_config is not None,
|
|
)
|
|
self._maybe_attach_attention_metadata_extensions(
|
|
attn_metadata=attn_metadata,
|
|
num_reqs=num_reqs_padded,
|
|
num_reqs_padded=num_reqs_padded,
|
|
max_query_len=max_query_len,
|
|
pad_attn=True,
|
|
for_cudagraph_capture=is_graph_capturing,
|
|
)
|
|
|
|
with self.maybe_dummy_run_with_lora(
|
|
self.lora_config,
|
|
num_scheduled_tokens,
|
|
num_sampled_tokens,
|
|
remove_lora,
|
|
num_active_loras,
|
|
):
|
|
# Make sure padding doesn't exceed max_num_tokens
|
|
assert num_tokens_padded <= self.max_num_tokens
|
|
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]
|
|
model_kwargs = self._init_model_kwargs()
|
|
else:
|
|
input_ids = self.input_ids.gpu[:num_tokens_padded]
|
|
inputs_embeds = None
|
|
|
|
# Some generation-stage models (e.g. MammothModa2DiTPipeline) require
|
|
# model-specific runtime information (such as image size and conditioning
|
|
# embeddings) even during the dummy profiling run that vLLM uses to
|
|
# estimate KV-cache capacity. get_dummy_runtime_additional_information
|
|
# provides placeholder values of the correct shape so that the profiling
|
|
# run does not raise an error due to missing inputs.
|
|
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
|
|
# Generation-stage code2wav models split flat input_ids using this
|
|
# shared contract. Dummy/profile runs use padded synthetic inputs,
|
|
# so expose one synthetic segment covering the actual tensor length.
|
|
model_kwargs["seq_token_counts"] = [int(num_tokens_padded)]
|
|
|
|
if self.uses_mrope:
|
|
positions = self.mrope_positions.gpu[:, :num_tokens_padded]
|
|
elif self.uses_xdrope_dim > 0:
|
|
positions = self.xdrope_positions.gpu[:, :num_tokens_padded]
|
|
else:
|
|
positions = self.positions[:num_tokens_padded]
|
|
|
|
if get_pp_group().is_first_rank:
|
|
intermediate_tensors = None
|
|
else:
|
|
if self.intermediate_tensors is None:
|
|
self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
|
|
batch_size=self.max_num_tokens,
|
|
dtype=self.model_config.dtype,
|
|
device=self.device,
|
|
)
|
|
|
|
intermediate_tensors = self.sync_and_gather_intermediate_tensors(num_tokens_padded, None, False)
|
|
|
|
if ubatch_slices_padded is not None:
|
|
# Adjust values to reflect a single ubatch.
|
|
# TODO(sage,lucas): this is cruft that should be addressed in
|
|
# the padding refactor.
|
|
num_tokens_padded = ubatch_slices_padded[0].num_tokens
|
|
if num_tokens_across_dp is not None:
|
|
num_tokens_across_dp[:] = num_tokens_padded
|
|
|
|
with (
|
|
self.maybe_randomize_inputs(input_ids, inputs_embeds),
|
|
set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_tokens_padded,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
cudagraph_runtime_mode=cudagraph_runtime_mode,
|
|
batch_descriptor=batch_desc,
|
|
ubatch_slices=ubatch_slices_padded,
|
|
slot_mapping=slot_mappings, # OMNI: required for KV cache operations
|
|
),
|
|
):
|
|
outputs = self.model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**model_kwargs,
|
|
)
|
|
|
|
if self.use_aux_hidden_state_outputs:
|
|
hidden_states, _ = outputs
|
|
else:
|
|
hidden_states = outputs
|
|
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
|
|
if self.speculative_config and (
|
|
self.speculative_config.use_eagle()
|
|
or self.speculative_config.uses_draft_model()
|
|
or self.speculative_config.uses_extract_hidden_states()
|
|
):
|
|
assert isinstance(
|
|
self.drafter,
|
|
EagleProposer | DFlashProposer | DraftModelProposer | ExtractHiddenStatesProposer | Gemma4Proposer,
|
|
)
|
|
assert self.speculative_config is not None
|
|
# Eagle currently only supports PIECEWISE cudagraphs.
|
|
# Therefore only use cudagraphs if the main model uses PIECEWISE
|
|
# NOTE(lucas): this is a hack, need to clean up.
|
|
use_cudagraphs = (
|
|
(is_graph_capturing and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE)
|
|
or (not is_graph_capturing and cudagraph_runtime_mode != CUDAGraphMode.NONE)
|
|
) and not self.speculative_config.enforce_eager
|
|
|
|
# Note(gnovack) - We need to disable cudagraphs for one of the two
|
|
# lora cases when cudagraph_specialize_lora is enabled. This is a
|
|
# short term mitigation for issue mentioned in
|
|
# https://github.com/vllm-project/vllm/issues/28334
|
|
if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0:
|
|
use_cudagraphs = False
|
|
|
|
self.drafter.dummy_run(
|
|
num_tokens,
|
|
use_cudagraphs=use_cudagraphs,
|
|
is_graph_capturing=is_graph_capturing,
|
|
slot_mappings=slot_mappings, # OMNI: pass slot_mappings (upstream v1 API)
|
|
)
|
|
|
|
# We register layerwise NVTX hooks here after the first dynamo tracing is
|
|
# done to avoid nvtx operations in hook functions being traced by
|
|
# torch dynamo and causing graph breaks.
|
|
# Note that for DYNAMO_ONCE and VLLM_COMPILE mode,
|
|
# compiled model's dynamo tracing is only done once and the compiled model's
|
|
# __call__ function is replaced by calling the compiled function.
|
|
# So it's safe to register hooks here. Hooks will be registered to
|
|
# both compiled and uncompiled models but they will never
|
|
# be called on the compiled model execution path.
|
|
self._register_layerwise_nvtx_hooks()
|
|
|
|
# This is necessary to avoid blocking DP.
|
|
# For dummy runs, we typically skip EPLB since we don't have any real
|
|
# requests to process.
|
|
# However, in DP settings, there may be cases when some DP ranks do
|
|
# not have any requests to process, so they're executing dummy batches.
|
|
# In such cases, we still have to trigger EPLB to make sure
|
|
# ranks execute the rearrangement in synchronization.
|
|
if not skip_eplb:
|
|
self.eplb_step(is_dummy=True, is_profile=is_profile)
|
|
|
|
return hidden_states, None
|
|
|
|
def profile_run(self) -> None:
|
|
# Profile with multimodal encoder & encoder cache.
|
|
if self.supports_mm_inputs:
|
|
mm_config = self.model_config.multimodal_config
|
|
if mm_config is not None and mm_config.skip_mm_profiling:
|
|
logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
|
|
else:
|
|
mm_budget = self.mm_budget
|
|
assert mm_budget is not None
|
|
|
|
if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
|
|
if not mm_budget.mm_max_toks_per_item:
|
|
# All modality limits are 0 — embedding-only mode.
|
|
# Budget is non-zero for embedding storage, but
|
|
# there's no encoder to profile.
|
|
logger.info(
|
|
"Skipping encoder profiling for embedding-only "
|
|
"mode (all modality limits=0 with "
|
|
"enable_mm_embeds=True).",
|
|
)
|
|
else:
|
|
# NOTE: Currently model is profiled with a single non-text
|
|
# modality with the max possible input tokens even when
|
|
# it supports multiple.
|
|
dummy_modality = mm_budget.get_modality_with_max_tokens()
|
|
max_mm_items_per_batch = mm_budget.mm_max_items_per_batch[dummy_modality]
|
|
|
|
logger.info(
|
|
"Encoder cache will be initialized with a budget of "
|
|
"%s tokens, and profiled with %s %s items of the "
|
|
"maximum feature size.",
|
|
encoder_budget,
|
|
max_mm_items_per_batch,
|
|
dummy_modality,
|
|
)
|
|
|
|
# Create dummy batch of multimodal inputs.
|
|
batched_dummy_mm_inputs = self._get_mm_dummy_batch(
|
|
dummy_modality,
|
|
max_mm_items_per_batch,
|
|
)
|
|
|
|
# Run multimodal encoder.
|
|
dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs)
|
|
|
|
sanity_check_mm_encoder_outputs(
|
|
dummy_encoder_outputs,
|
|
expected_num_items=max_mm_items_per_batch,
|
|
)
|
|
for i, output in enumerate(dummy_encoder_outputs):
|
|
self.encoder_cache[f"tmp_{i}"] = output
|
|
|
|
# Add `is_profile` here to pre-allocate communication buffers
|
|
hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
|
|
output = None
|
|
self._sync_device()
|
|
del hidden_states
|
|
self.encoder_cache.clear()
|
|
gc.collect()
|