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

[Feat] DiffusionEngine Support async batch infer (#2729)

Browse Source 
Signed-off-by: Semmer <semmer@live.cn>
Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: jader <yjader@foxmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
Ting FUasukaqaq-sjaderHongsheng LiuGao Han
 authored and GitHub committed 2026-05-06 09:47:13 +12:00
1 parent 2b85474c23
commit e49fbd8a2d
32 files changed
+1951 -190
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/worker/diffusion_model_runner.py
+141 -71
View File
@@ -14,6 +14,7 @@ import copy
import time
from collections.abc import Iterable
from contextlib import nullcontext
from typing import Any
import torch
from torch.profiler import record_function
@@ -32,7 +33,8 @@ from vllm_omni.diffusion.offloader import get_offload_backend
from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput
from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, DiffusionRequestState, RunnerOutput
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
from vllm_omni.platforms import current_omni_platform
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
@@ -297,50 +299,107 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
"""Return whether current pipeline supports step execution."""
return self.pipeline is not None and supports_step_execution(self.pipeline)
def _update_states(self, scheduler_output: DiffusionSchedulerOutput) -> tuple[DiffusionRequestState, bool]:
def _update_states(
self, scheduler_output: DiffusionSchedulerOutput
) -> tuple[list[DiffusionRequestState], list[str]]:
"""Step-before update: cleanup finished requests and get/create one running state."""
for req_id in scheduler_output.finished_req_ids:
self.state_cache.pop(req_id, None)
if scheduler_output.num_scheduled_reqs != 1:
raise ValueError(
"Step mode currently supports batch_size=1, "
f"but got {scheduler_output.num_scheduled_reqs} scheduled requests."
)
resolved: list[DiffusionRequestState] = []
new_req_id: list[str] = []
try:
# process new requests
for sched_new_req in scheduler_output.scheduled_new_reqs:
# new_req_data = scheduler_output.scheduled_new_reqs[0]
req_id = sched_new_req.sched_req_id
req = sched_new_req.req
new_req_id.append(req_id)
if req_id in self.state_cache:
raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
request_ids = req.request_ids or [req_id]
if len(request_ids) != len(req.prompts):
raise ValueError(
f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
)
new_state = DiffusionRequestState(
req_id=req_id,
sampling=copy.deepcopy(req.sampling_params),
prompts=req.prompts,
)
self.state_cache[req_id] = new_state
resolved.append(new_state)
if scheduler_output.scheduled_new_reqs:
new_req_data = scheduler_output.scheduled_new_reqs[0]
req_id = new_req_data.sched_req_id
req = new_req_data.req
if req_id in self.state_cache:
raise ValueError(f"Received duplicate new-request payload for cached request {req_id}.")
else:
req_id = scheduler_output.scheduled_cached_reqs.sched_req_ids[0]
state = self.state_cache.get(req_id)
if state is None:
raise ValueError(f"Missing cached state for request {req_id}.")
return state, False
# process cached requests
for req_id in scheduler_output.scheduled_cached_reqs.sched_req_ids:
state = self.state_cache.get(req_id)
if state is None:
raise ValueError(f"Missing cached state for request {req_id}.")
resolved.append(state)
except Exception:
for req_id in new_req_id:
self.state_cache.pop(req_id, None)
raise
request_ids = req.request_ids or [req_id]
if len(request_ids) != len(req.prompts):
raise ValueError(
f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
)
return resolved, new_req_id
state = DiffusionRequestState(
req_id=req_id,
sampling=copy.deepcopy(req.sampling_params),
prompts=req.prompts,
def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request_ids: list[str]) -> InputBatch:
# process new reqs
for state in states:
if state.req_id in new_request_ids:
# set generator
if state.sampling.generator is None and state.sampling.seed is not None:
if state.sampling.generator_device is not None:
gen_device = state.sampling.generator_device
elif self.device.type == "cpu":
gen_device = "cpu"
else:
gen_device = self.device
state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
# encode
self.pipeline.prepare_encode(state)
input_batch = InputBatch.make_batch(
states,
cached_batch=getattr(self, "input_batch", None),
)
self.state_cache[req_id] = state
return state, True
self.input_batch = input_batch
return input_batch
def _update_states_after(self, state: DiffusionRequestState, finished: bool) -> None:
def _update_states_after(
self,
states: list[DiffusionRequestState],
input_batch: InputBatch,
interrupted: bool = False,
):
"""Step-after update: clear cached state for completed request."""
if finished:
self.state_cache.pop(state.req_id, None)
gathered_latents = torch.cat([state.latents for state in states], dim=0)
if (
input_batch.latents.size() == gathered_latents.size()
and input_batch.latents.dtype == gathered_latents.dtype
and input_batch.latents.device == gathered_latents.device
):
input_batch.latents.copy_(gathered_latents)
else:
input_batch.latents = gathered_latents.clone()
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
self.input_batch = input_batch
scatter_latents(states, input_batch)
for state in states:
if interrupted or state.denoise_completed:
self.state_cache.pop(state.req_id, None)
def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any:
model_state = getattr(self, "model_state", None)
if model_state is None:
return {}
prepare_attn = getattr(model_state, "prepare_attn", None)
if not callable(prepare_attn):
return {}
return prepare_attn(input_batch)
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput:
"""Execute one step for one scheduled request and return runner output."""
assert self.pipeline is not None, "Model not loaded. Call load_model() first."
if not self.supports_step_mode():
@@ -354,46 +413,57 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
use_hsdp = self.od_config.parallel_config.use_hsdp
grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
with grad_context:
state, is_new_request = self._update_states(scheduler_output)
states, new_request_ids = self._update_states(scheduler_output)
input_batch = self._prepare_batch_inputs(states, new_request_ids)
attn_metadata = self._prepare_attn_metadata(input_batch)
if is_new_request:
# TODO: support kv manager recv
# TODO: support cache backend
if state.sampling.generator is None and state.sampling.seed is not None:
if state.sampling.generator_device is not None:
gen_device = state.sampling.generator_device
elif self.device.type == "cpu":
gen_device = "cpu"
else:
gen_device = self.device
state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
with set_forward_context(
vllm_config=self.vllm_config,
omni_diffusion_config=self.od_config,
attn_metadata=attn_metadata,
):
noise_pred = self.pipeline.denoise_step(input_batch)
with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
# step0/new request: encode
if is_new_request:
self.pipeline.prepare_encode(state)
runner_output_list = []
pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
if noise_pred is None and pipeline_interrupted:
for state in states:
runner_output_list.append(
RunnerOutput(
req_id=state.req_id,
step_index=state.step_index,
finished=True,
result=DiffusionOutput(error="stepwise denoise interrupted"),
)
)
noise_pred = self.pipeline.denoise_step(state)
finished = False
# In CFG parallel mode, only rank 0 gets the actual noise_pred; non-rank-0 workers receive None.
# A true interrupt (all ranks return None) is detected by checking self.pipeline.interrupt.
if noise_pred is None and getattr(self.pipeline, "interrupt", False):
finished = True
result = DiffusionOutput(error="stepwise denoise interrupted")
else:
self.pipeline.step_scheduler(state, noise_pred)
finished = state.denoise_completed
if finished:
result = self.pipeline.post_decode(state)
else:
result = None
offset = 0
for req in states:
row_num = req.latents.shape[0]
self.pipeline.step_scheduler(
req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
)
offset = offset + row_num
if req.denoise_completed:
result = self.pipeline.post_decode(req)
else:
result = None
runner_output_list.append(
RunnerOutput(
req_id=req.req_id,
step_index=req.step_index,
finished=req.denoise_completed,
result=result,
)
)
self._update_states_after(state, finished)
if noise_pred is not None and offset != noise_pred.shape[0]:
raise ValueError(
f"Stepwise noise_pred consumed {offset} rows, "
f"but batched noise_pred has {noise_pred.shape[0]} rows."
)
return RunnerOutput(
req_id=state.req_id,
step_index=state.step_index,
finished=finished,
result=result,
)
self._update_states_after(states, input_batch, pipeline_interrupted)
return BatchRunnerOutput.from_list(runner_output_list)