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

Dependency refractored to vLLM v0.11.0

Browse Source 
Signed-off-by: Gaohan123 <gaohan19@huawei.com>
This commit is contained in:
Gaohan123 committed 2025-11-04 22:06:41 +08:00
1 parent d4a364f052
commit 6e1478222e
38 files changed
+1500 -1719
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/worker/gpu_diffusion_model_runner.py
+459 -142
View File
@@ -2,21 +2,33 @@ from __future__ import annotations
import gc
import logging
from typing import List, Optional, Union
from typing import Any, List, Optional, Union
import numpy as np
import torch
from vllm.config import CUDAGraphMode
from vllm.forward_context import BatchDescriptor
from vllm.multimodal.inputs import MultiModalKwargs
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
from vllm.utils import cdiv
from vllm.v1.attention.backends.utils import (
CommonAttentionMetadata,
split_attn_metadata,
)
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm.v1.utils import record_function_or_nullcontext
from vllm.v1.worker.gpu_model_runner import (
EMPTY_MODEL_RUNNER_OUTPUT,
AsyncGPUModelRunnerOutput,
IntermediateTensors,
PerLayerAttnMetadata,
get_pp_group,
set_forward_context,
)
from vllm.v1.worker.ubatch_splitting import ubatch_split
from vllm.v1.worker.ubatch_utils import UBatchSlices
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
@@ -31,65 +43,97 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
- Executes diffusion process and returns tensors via `pooler_output`.
"""
@torch.inference_mode()
def execute_model(
def _preprocess(
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)
ubatch_slices: Optional[UBatchSlices] = None,
num_tokens_after_padding: Optional[torch.Tensor] = None,
) -> tuple[
int,
int,
Optional[torch.Tensor],
Optional[torch.Tensor],
Optional[torch.Tensor],
torch.Tensor,
Optional[IntermediateTensors],
dict[str, Any],
]:
# 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_pad, num_tokens_after_padding = 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:
# _prepare_inputs may reorder the batch, so we must gather multi
# modal outputs after that to ensure the correct order
if (
self.supports_mm_inputs
and get_pp_group().is_first_rank
and not self.model_config.is_encoder_decoder
):
# Run the multimodal encoder if any.
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:
# NOTE(woosuk): To unify token ids and soft tokens (vision
# embeddings), we always use embeddings (rather than token ids)
# as input to the multimodal model, even when the input is text.
inputs_embeds_scheduled = self.model.get_input_embeddings(
input_ids=self.input_ids[: scheduler_output.total_num_scheduled_tokens],
input_ids=self.input_ids.gpu[:num_input_tokens],
multimodal_embeddings=mm_embeds or None,
)
self.inputs_embeds[: scheduler_output.total_num_scheduled_tokens].copy_(
inputs_embeds_scheduled
# TODO(woosuk): Avoid the copy. Optimize.
self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
input_ids = self.input_ids.gpu[:num_input_tokens]
inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
model_kwargs = {
**self._init_model_kwargs(num_input_tokens),
**self._extract_mm_kwargs(scheduler_output),
}
elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
# Get the input embeddings for the tokens that are not input embeds,
# then put them into the appropriate positions.
# TODO(qthequartermasterman): Since even when prompt embeds are
# enabled, (a) not all requests will use prompt embeds, and (b)
# after the initial prompt is processed, the rest of the generated
# tokens will be token ids, it is not desirable to have the
# embedding layer outside of the CUDA graph all the time. The v0
# engine avoids this by "double compiling" the CUDA graph, once
# with input_ids and again with inputs_embeds, for all num_tokens.
# If a batch only has token ids, then including the embedding layer
# in the CUDA graph will be more performant (like in the else case
# below).
token_ids_idx = (
self.is_token_ids.gpu[:num_input_tokens]
.nonzero(as_tuple=False)
.squeeze(1)
)
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)
# Some tokens ids may need to become embeds
if token_ids_idx.numel() > 0:
token_ids = self.input_ids.gpu[token_ids_idx]
tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
model_kwargs = self._init_model_kwargs(num_input_tokens)
input_ids = self.input_ids.gpu[:num_input_tokens]
else:
input_ids = self.input_ids[:num_input_tokens]
# For text-only models, we use token ids as input.
# While it is possible to use embeddings as input just like the
# multimodal models, it is not desirable for performance since
# then the embedding layer is not included in the CUDA graph.
input_ids = self.input_ids.gpu[:num_input_tokens]
inputs_embeds = None
model_mm_kwargs = {}
# Positions (mrope or standard)
model_kwargs = self._init_model_kwargs(num_input_tokens)
if self.uses_mrope:
positions = self.mrope_positions[:, :num_input_tokens]
positions = self.mrope_positions.gpu[:, :num_input_tokens]
else:
positions = self.positions[:num_input_tokens]
positions = self.positions.gpu[:num_input_tokens]
# Intermediate tensors sync for PP (if any)
if get_pp_group().is_first_rank:
intermediate_tensors = None
else:
@@ -97,32 +141,94 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
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 (
self.model_config.is_encoder_decoder
and scheduler_output.scheduled_encoder_inputs
):
encoder_inputs = self._extract_encoder_inputs(scheduler_output)
model_kwargs.update(encoder_inputs)
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
return (
num_input_tokens,
num_input_tokens,
num_tokens_after_padding,
input_ids,
inputs_embeds,
positions,
intermediate_tensors,
model_kwargs,
)
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
with record_function_or_nullcontext("Preprocess"):
with self.synchronize_input_prep():
# Update persistent batch states.
self._update_states(scheduler_output)
if not scheduler_output.total_num_scheduled_tokens:
return EMPTY_MODEL_RUNNER_OUTPUT
# Prepare the decoder inputs.
(
attn_metadata,
logits_indices,
spec_decode_metadata,
num_scheduled_tokens_np,
spec_decode_common_attn_metadata,
max_query_len,
ubatch_slices,
num_tokens_after_padding,
) = self._prepare_inputs(scheduler_output)
# Input token count for this iteration (not used by diffusion, but
# retained to keep DP padding/ordering consistent)
(
num_scheduled_tokens,
num_input_tokens,
num_tokens_across_dp,
input_ids,
inputs_embeds,
positions,
intermediate_tensors,
model_kwargs,
) = self._preprocess(
scheduler_output,
intermediate_tensors,
ubatch_slices,
num_tokens_after_padding,
)
cudagraph_runtime_mode = CUDAGraphMode.NONE
# Run the model.
# Use persistent buffers for CUDA graphs.
with (
set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
cudagraph_runtime_mode=cudagraph_runtime_mode,
batch_descriptor=None,
ubatch_slices=ubatch_slices,
),
record_function_or_nullcontext("Forward"),
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
):
outputs = self._run_diffusion(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
multimodal_kwargs=model_mm_kwargs,
multimodal_kwargs=model_kwargs,
logits_indices=logits_indices,
)
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
_, 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):
@@ -145,13 +251,10 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
else:
raise RuntimeError("Unsupported diffusion output type")
self.eplb_step()
return OmniModelRunnerOutput(
output = 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,
@@ -159,6 +262,16 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
num_nans_in_logits={},
)
if not self.use_async_scheduling:
return output
return AsyncGPUModelRunnerOutput(
model_runner_output=output,
sampled_token_ids=[],
invalid_req_indices=[],
async_output_copy_stream=self.async_output_copy_stream,
)
def _run_diffusion(
self,
*,
@@ -202,84 +315,217 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
def _dummy_run(
self,
num_tokens: int,
capture_attn_cudagraph: bool = False,
cudagraph_runtime_mode: Optional[CUDAGraphMode] = 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,
) -> tuple[torch.Tensor, torch.Tensor]:
# Padding for DP
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
num_tokens += num_pad
"""
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
"""
assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
CUDAGraphMode.NONE,
CUDAGraphMode.PIECEWISE,
CUDAGraphMode.FULL,
}
# 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.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
if create_mixed_batch:
assert not uniform_decode
# Create mixed batch:
# first half decode tokens, second half one prefill
num_decode_tokens = 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 = 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)
attn_metadata: Optional[dict[str, dict]] = None
if capture_attn_cudagraph:
attn_metadata = {}
total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
# 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
ubatch_slices = None
num_tokens_after_padding = None
# We currently only microbatch if the number of tokens is
# over a certain threshold.
if self.parallel_config.enable_dbo and allow_microbatching:
ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split(
num_scheduled_tokens,
total_num_scheduled_tokens,
total_num_scheduled_tokens,
uniform_decode=uniform_decode,
vllm_config=self.vllm_config,
)
# Currently when DBO is enabled `ubatch_split` returns
# the num_tokens_after_padding for a single ubatch, but we have 2
# TODO(sage,lucas): this is cruft that should be addressed in the
# padding refactor.
if ubatch_num_tokens_after_padding is not None:
num_tokens_after_padding = ubatch_num_tokens_after_padding * 2
# If we failed to microbatch, currently need to resynchronize
# TODO(lucas,sage): we should be able to avoid this second sync by
# refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
# a single `coordinate_batch_across_dp` function.
if num_tokens_after_padding is None:
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
num_tokens_after_padding = num_tokens + num_pad
else:
num_tokens_across_dp = num_tokens_after_padding
num_tokens_after_padding = int(num_tokens_after_padding[0].item())
attn_metadata: Optional[PerLayerAttnMetadata] = None
# 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:
attn_metadata = {}
if ubatch_slices is not None:
attn_metadata = [dict() for _ in range(len(ubatch_slices))]
if create_mixed_batch:
# In the mixed batch mode (used for FI warmup), we use
# shorter sequence lengths to run faster.
# TODO(luka) better system for describing dummy batches
seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
else:
seq_lens = max_query_len
self.seq_lens.np[:num_reqs] = seq_lens
self.seq_lens.np[num_reqs:] = 0
self.seq_lens.copy_to_gpu()
cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
self.query_start_loc.copy_to_gpu()
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],
query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
seq_lens=self.seq_lens.gpu[: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,
max_query_len=max_query_len,
max_seq_len=self.max_model_len,
block_table_tensor=self.input_batch.block_table[
kv_cache_group_id
].get_device_tensor()[:num_reqs],
].get_device_tensor(num_reqs),
slot_mapping=self.input_batch.block_table[
kv_cache_group_id
].slot_mapping[:num_tokens],
].slot_mapping.gpu[: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(
if ubatch_slices is not None:
common_attn_metadata_list = split_attn_metadata(
ubatch_slices, common_attn_metadata
)
for ubid, common_attn_metadata in enumerate(
common_attn_metadata_list
):
assert common_attn_metadata.max_query_len == 1
attn_metadata_i = attn_group.get_metadata_builder(
ubatch_id=ubid
).build_for_cudagraph_capture(common_attn_metadata)
for layer_name in attn_group.layer_names:
assert type(attn_metadata) is list
attn_metadata[ubid][layer_name] = attn_metadata_i
else:
assert type(attn_metadata) is dict
attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture( # noqa: E501
common_attn_metadata
)
)
for layer_name in kv_cache_group_spec.layer_names:
attn_metadata[layer_name] = attn_metadata_i
for layer_name in attn_group.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:
with self.maybe_dummy_run_with_lora(
self.lora_config, num_scheduled_tokens, remove_lora
):
model_kwargs = self._init_model_kwargs(num_tokens)
if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
input_ids = None
inputs_embeds = self.inputs_embeds[:num_tokens]
model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
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]
model_kwargs = self._init_model_kwargs(num_tokens)
else:
input_ids = self.input_ids[:num_tokens]
input_ids = self.input_ids.gpu[:num_tokens]
inputs_embeds = None
model_mm_kwargs = {}
if self.uses_mrope:
positions = self.mrope_positions[:, :num_tokens]
positions = self.mrope_positions.gpu[:, :num_tokens]
else:
positions = self.positions[:num_tokens]
positions = self.positions.gpu[:num_tokens]
if get_pp_group().is_first_rank:
intermediate_tensors = None
@@ -297,23 +543,53 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
num_tokens, None, False
)
# Diffusion path: avoid CUDA graphs; we only use context for resource wiring
# filter out the valid batch descriptor
_cg_mode, batch_descriptor = (
self.cudagraph_dispatcher.dispatch(
BatchDescriptor(
num_tokens=num_tokens_after_padding,
uniform_decode=uniform_decode,
)
)
if not is_profile
else (CUDAGraphMode.NONE, None)
)
if cudagraph_runtime_mode is not None:
# we allow forcing NONE when the dispatcher disagrees to support
# warm ups for cudagraph capture
assert (
cudagraph_runtime_mode == CUDAGraphMode.NONE
or cudagraph_runtime_mode == _cg_mode
), (
f"Cudagraph runtime mode mismatch at dummy_run. "
f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
)
else:
cudagraph_runtime_mode = _cg_mode
if ubatch_slices 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_after_padding = ubatch_slices[0].num_tokens
if num_tokens_across_dp is not None:
num_tokens_across_dp[:] = num_tokens_after_padding
with self.maybe_randomize_inputs(input_ids), set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_tokens,
num_tokens=num_tokens_after_padding,
num_tokens_across_dp=num_tokens_across_dp,
cudagraph_runtime_mode=cudagraph_runtime_mode,
batch_descriptor=batch_descriptor,
ubatch_slices=ubatch_slices,
):
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,
**model_kwargs,
)
if self.use_aux_hidden_state_outputs:
@@ -321,16 +597,23 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
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 self.speculative_config and self.speculative_config.use_eagle():
assert isinstance(self.drafter, EagleProposer)
self.drafter.dummy_run(num_tokens)
if not skip_eplb:
self.eplb_step(is_dummy=True, is_profile=is_profile)
# 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)
# logit_indices = np.cumsum(num_scheduled_tokens) - 1 # unused variable
return text_hidden_states, None
# logit_indices = np.cumsum(num_scheduled_tokens) - 1
hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
return hidden_states, None
@torch.inference_mode()
def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
@@ -338,42 +621,76 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
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,
# Profile with multimodal encoder & encoder cache.
if self.supports_mm_inputs:
if self.model_config.multimodal_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
dummy_encoder_outputs = self.model.get_multimodal_embeddings(
**batched_dummy_mm_inputs
)
if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
# 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.max_items_per_batch_by_modality[
dummy_modality
]
sanity_check_mm_encoder_outputs(
dummy_encoder_outputs,
expected_num_items=max_mm_items_per_batch,
)
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,
)
self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
# 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.get_multimodal_embeddings(
**batched_dummy_mm_inputs
)
sanity_check_mm_encoder_outputs(
dummy_encoder_outputs,
expected_num_items=max_mm_items_per_batch,
)
# NOTE: This happens when encoder cache needs to store
# the embeddings that encoder outputs are scattered onto.
# In this case we create dummy embeddings of size
# (encode_budget, hidden_size) and scatter encoder
# output into it.
encoder_output_shape = dummy_encoder_outputs[0].shape
if encoder_output_shape[0] < encoder_budget:
expanded_outputs = []
for output in dummy_encoder_outputs:
expanded = output.new_zeros(
(encoder_budget, encoder_output_shape[-1])
)
num_tokens = output.shape[0]
expanded[:num_tokens].copy_(output)
expanded_outputs.append(expanded)
dummy_encoder_outputs = expanded_outputs
# Cache the dummy encoder outputs.
self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
# Add `is_profile` here to pre-allocate communication buffers
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
pass
self._sync_device()
del hidden_states
self.encoder_cache.clear()