381 lines
15 KiB
Python
381 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import gc
|
|
import logging
|
|
from typing import List, Optional, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from vllm.multimodal.inputs import MultiModalKwargs
|
|
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
from vllm.v1.worker.gpu_model_runner import (
|
|
EMPTY_MODEL_RUNNER_OUTPUT,
|
|
IntermediateTensors,
|
|
get_pp_group,
|
|
set_forward_context,
|
|
)
|
|
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
|
|
|
|
from vllm_omni.outputs import OmniModelRunnerOutput
|
|
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GPUDiffusionModelRunner(OmniGPUModelRunner):
|
|
"""Diffusion 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 diffusion process and returns tensors via `pooler_output`.
|
|
"""
|
|
|
|
@torch.inference_mode()
|
|
def execute_model(
|
|
self,
|
|
scheduler_output: "SchedulerOutput",
|
|
intermediate_tensors: Optional[IntermediateTensors] = None,
|
|
) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
|
|
self._update_states(scheduler_output)
|
|
|
|
if not scheduler_output.total_num_scheduled_tokens:
|
|
|
|
return EMPTY_MODEL_RUNNER_OUTPUT
|
|
|
|
# Prepare decoder inputs and attention metadata (for batch/order mapping)
|
|
(
|
|
attn_metadata,
|
|
attention_cuda_graphs,
|
|
logits_indices,
|
|
spec_decode_metadata,
|
|
num_scheduled_tokens_np,
|
|
spec_decode_common_attn_metadata,
|
|
) = self._prepare_inputs(scheduler_output)
|
|
|
|
# Input token count for this iteration (not used by diffusion, but
|
|
# retained to keep DP padding/ordering consistent)
|
|
num_input_tokens = scheduler_output.total_num_scheduled_tokens
|
|
num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
|
|
num_input_tokens += num_pad
|
|
|
|
# Multimodal conditioning (e.g., text/audio/video encoders)
|
|
if self.is_multimodal_model:
|
|
self._execute_mm_encoder(scheduler_output)
|
|
mm_embeds = self._gather_mm_embeddings(scheduler_output)
|
|
else:
|
|
mm_embeds = []
|
|
|
|
# Build inputs to mirror AR runner: input_ids/positions/embeds
|
|
if self.is_multimodal_model and get_pp_group().is_first_rank:
|
|
inputs_embeds_scheduled = self.model.get_input_embeddings(
|
|
input_ids=self.input_ids[: scheduler_output.total_num_scheduled_tokens],
|
|
multimodal_embeddings=mm_embeds or None,
|
|
)
|
|
self.inputs_embeds[: scheduler_output.total_num_scheduled_tokens].copy_(
|
|
inputs_embeds_scheduled
|
|
)
|
|
input_ids = self.input_ids[:num_input_tokens]
|
|
inputs_embeds = self.inputs_embeds[:num_input_tokens]
|
|
model_mm_kwargs = self._extract_mm_kwargs(scheduler_output)
|
|
else:
|
|
input_ids = self.input_ids[:num_input_tokens]
|
|
inputs_embeds = None
|
|
model_mm_kwargs = {}
|
|
|
|
# Positions (mrope or standard)
|
|
if self.uses_mrope:
|
|
positions = self.mrope_positions[:, :num_input_tokens]
|
|
else:
|
|
positions = self.positions[:num_input_tokens]
|
|
|
|
# Intermediate tensors sync for PP (if any)
|
|
if get_pp_group().is_first_rank:
|
|
intermediate_tensors = None
|
|
else:
|
|
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
|
|
num_input_tokens, intermediate_tensors, True
|
|
)
|
|
|
|
# Set forward context mainly for resource management and kv connector
|
|
skip_cuda_graphs = True # diffusion path does not rely on cuda graphs here
|
|
with set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_input_tokens,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
skip_cuda_graphs=skip_cuda_graphs,
|
|
), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output:
|
|
|
|
if not get_pp_group().is_last_rank:
|
|
# For non-last PP stages, pass through intermediate tensors.
|
|
assert intermediate_tensors is not None
|
|
intermediate_tensors.kv_connector_output = kv_connector_output
|
|
return intermediate_tensors
|
|
|
|
outputs = self._run_diffusion(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
multimodal_kwargs=model_mm_kwargs,
|
|
logits_indices=logits_indices,
|
|
)
|
|
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
|
|
|
|
# Ensure one tensor per request, map to CPU for output struct
|
|
pooler_output: List[Optional[torch.Tensor]] = []
|
|
if isinstance(multimodal_outputs, torch.Tensor):
|
|
# If model returned a single stacked tensor, split by requests
|
|
assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
|
|
for i in range(self.input_batch.num_reqs):
|
|
pooler_output.append(
|
|
multimodal_outputs[i].detach().to("cpu").contiguous()
|
|
)
|
|
elif isinstance(multimodal_outputs, list):
|
|
for out in multimodal_outputs:
|
|
pooler_output.append(
|
|
out.detach().to("cpu").contiguous() if out is not None else None
|
|
)
|
|
elif isinstance(multimodal_outputs, dict):
|
|
for out in multimodal_outputs.values():
|
|
pooler_output.append(
|
|
out.detach().to("cpu").contiguous() if out is not None else None
|
|
)
|
|
else:
|
|
raise RuntimeError("Unsupported diffusion output type")
|
|
|
|
self.eplb_step()
|
|
|
|
return OmniModelRunnerOutput(
|
|
req_ids=self.input_batch.req_ids,
|
|
req_id_to_index=self.input_batch.req_id_to_index,
|
|
sampled_token_ids=[],
|
|
spec_token_ids=None,
|
|
logprobs=None,
|
|
prompt_logprobs_dict={},
|
|
pooler_output=pooler_output,
|
|
kv_connector_output=kv_connector_output,
|
|
num_nans_in_logits={},
|
|
)
|
|
|
|
def _run_diffusion(
|
|
self,
|
|
*,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
intermediate_tensors: Optional[IntermediateTensors],
|
|
inputs_embeds: Optional[torch.Tensor],
|
|
multimodal_kwargs: dict,
|
|
logits_indices: torch.Tensor,
|
|
) -> Union[torch.Tensor, list[torch.Tensor]]:
|
|
"""Runs the diffusion process and returns per-request tensors.
|
|
|
|
Tries model interfaces in the following order for maximal compatibility:
|
|
1) model.sample(condition=..., **kwargs)
|
|
2) model.forward(condition=..., **kwargs)
|
|
3) model.diffuse(condition=..., **kwargs)
|
|
"""
|
|
# Keep inputs identical to AR runner
|
|
kwargs = dict(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
|
|
sampling_metadata=self.input_batch.sampling_metadata,
|
|
logits_index=logits_indices,
|
|
sampler=self.sampler,
|
|
)
|
|
|
|
# For Qwen 2.5 Omni's current implementation, we only support the forward method
|
|
if hasattr(self.model, "forward"):
|
|
return self.model.forward(**kwargs)
|
|
# TODO: add the diffuse method for other models
|
|
|
|
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_run(
|
|
self,
|
|
num_tokens: int,
|
|
capture_attn_cudagraph: bool = False,
|
|
skip_eplb: bool = False,
|
|
is_profile: bool = False,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
# Padding for DP
|
|
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
|
|
num_tokens += num_pad
|
|
|
|
# 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.scheduler_config.max_num_batched_tokens
|
|
max_num_reqs = self.scheduler_config.max_num_seqs
|
|
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)
|
|
|
|
attn_metadata: Optional[dict[str, dict]] = None
|
|
if capture_attn_cudagraph:
|
|
attn_metadata = {}
|
|
|
|
# Make sure max_model_len is used at the graph capture time.
|
|
self.seq_lens_np[:num_reqs] = self.max_model_len
|
|
self.seq_lens_np[num_reqs:] = 0
|
|
self.seq_lens[:num_reqs].copy_(
|
|
self.seq_lens_cpu[:num_reqs], non_blocking=True
|
|
)
|
|
|
|
for kv_cache_group_id, kv_cache_group_spec in enumerate(
|
|
self.kv_cache_config.kv_cache_groups
|
|
):
|
|
common_attn_metadata = CommonAttentionMetadata(
|
|
query_start_loc=self.query_start_loc[: num_reqs + 1],
|
|
query_start_loc_cpu=self.query_start_loc_cpu[: num_reqs + 1],
|
|
seq_lens=self.seq_lens[:num_reqs],
|
|
seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
|
|
num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501
|
|
:num_reqs
|
|
],
|
|
num_reqs=num_reqs,
|
|
num_actual_tokens=num_tokens,
|
|
max_query_len=num_tokens,
|
|
block_table_tensor=self.input_batch.block_table[
|
|
kv_cache_group_id
|
|
].get_device_tensor()[:num_reqs],
|
|
slot_mapping=self.input_batch.block_table[
|
|
kv_cache_group_id
|
|
].slot_mapping[:num_tokens],
|
|
causal=True,
|
|
)
|
|
|
|
for attn_group in self.attn_groups[kv_cache_group_id]:
|
|
attn_metadata_i = (
|
|
attn_group.metadata_builder.build_for_cudagraph_capture(
|
|
common_attn_metadata
|
|
)
|
|
)
|
|
for layer_name in kv_cache_group_spec.layer_names:
|
|
attn_metadata[layer_name] = attn_metadata_i
|
|
|
|
with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens):
|
|
if self.is_multimodal_model:
|
|
input_ids = None
|
|
inputs_embeds = self.inputs_embeds[:num_tokens]
|
|
model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
|
|
else:
|
|
input_ids = self.input_ids[:num_tokens]
|
|
inputs_embeds = None
|
|
model_mm_kwargs = {}
|
|
|
|
if self.uses_mrope:
|
|
positions = self.mrope_positions[:, :num_tokens]
|
|
else:
|
|
positions = self.positions[:num_tokens]
|
|
|
|
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_slice_intermediate_tensors(
|
|
num_tokens, None, False
|
|
)
|
|
|
|
# Diffusion path: avoid CUDA graphs; we only use context for resource wiring
|
|
with self.maybe_randomize_inputs(input_ids), set_forward_context(
|
|
attn_metadata,
|
|
self.vllm_config,
|
|
num_tokens=num_tokens,
|
|
num_tokens_across_dp=num_tokens_across_dp,
|
|
):
|
|
outputs = self.model(
|
|
input_ids=input_ids,
|
|
positions=positions,
|
|
intermediate_tensors=intermediate_tensors,
|
|
inputs_embeds=inputs_embeds,
|
|
**MultiModalKwargs.as_kwargs(
|
|
model_mm_kwargs,
|
|
device=self.device,
|
|
),
|
|
sampler=None,
|
|
)
|
|
|
|
if self.use_aux_hidden_state_outputs:
|
|
hidden_states, _ = outputs
|
|
else:
|
|
hidden_states = outputs
|
|
|
|
# Extract multimodal outputs if present; we ignore them here because
|
|
# dummy run returns tensors only. The actual diffusion runner returns
|
|
# multimodal outputs via pooler_output in execute_model.
|
|
text_hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
|
|
|
|
if not skip_eplb:
|
|
self.eplb_step(is_dummy=True, is_profile=is_profile)
|
|
|
|
# logit_indices = np.cumsum(num_scheduled_tokens) - 1 # unused variable
|
|
return text_hidden_states, None
|
|
|
|
@torch.inference_mode()
|
|
def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
|
|
logger.warning("Dummy sampler run is not implemented for diffusion model")
|
|
return None
|
|
|
|
def profile_run(self) -> None:
|
|
# Profile with multimodal encoder & encoder cache, similar to base but
|
|
# without any logits/sampler warming.
|
|
if self.is_multimodal_model:
|
|
mm_budget = self.mm_budget
|
|
assert mm_budget is not None
|
|
|
|
# TODO: handle encoder-decoder models once supported.
|
|
if (mm_budget.get_encoder_budget()) > 0: # encoder_budget unused
|
|
(
|
|
dummy_modality,
|
|
max_tokens,
|
|
) = mm_budget.get_modality_with_max_tokens()
|
|
(
|
|
max_mm_items_per_prompt,
|
|
max_mm_items_per_batch,
|
|
) = mm_budget.get_max_items(dummy_modality, max_tokens)
|
|
|
|
batched_dummy_mm_inputs = self._get_mm_dummy_batch(
|
|
dummy_modality,
|
|
max_mm_items_per_batch,
|
|
)
|
|
|
|
dummy_encoder_outputs = self.model.get_multimodal_embeddings(
|
|
**batched_dummy_mm_inputs
|
|
)
|
|
|
|
sanity_check_mm_encoder_outputs(
|
|
dummy_encoder_outputs,
|
|
expected_num_items=max_mm_items_per_batch,
|
|
)
|
|
|
|
self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
|
|
|
|
hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
|
|
if get_pp_group().is_last_rank:
|
|
pass # No sampler/pooler warmup for diffusion
|
|
self._sync_device()
|
|
del hidden_states
|
|
self.encoder_cache.clear()
|
|
gc.collect()
|