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

[Refactor]: clean up diffusion scheduler interfaces (#5216)

Browse Source 
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Zeyu Huang | 黃澤宇Copilot Autofix powered by AICursor
 authored and GitHub committed 2026-07-22 14:05:49 +08:00
1 parent e2115fb642
commit e39f4dddfb
31 files changed
+319 -338
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/api/README.md
+1 -1
View File
@@ -131,7 +131,7 @@ Worker classes and model runners for distributed inference.
- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
- [vllm_omni.diffusion.worker.utils.DiffusionRequestState][]
- [vllm_omni.diffusion.worker.utils.StepRequestState][]
- [vllm_omni.diffusion.worker.utils.RunnerOutput][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
docs/design/feature/diffusion_continuous_batching.md
+6 -5
View File
@@ -71,7 +71,7 @@ The scheduler derives its batch capacity from `max_num_seqs` through
`max_num_running_reqs`.
Batch admission is gated by
[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
[`StepBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
the core correctness rule for batching: requests are only co-batched when they
share the same denoise tensor contract.
@@ -100,9 +100,10 @@ request-mode prompt semantics, see
## Runner
The runner keeps persistent per-request execution state in
[`DiffusionRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
while the scheduler owns a separate lightweight request state for queueing and
lifecycle tracking.
[`StepRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
while the scheduler owns a separate lightweight
[`SchedulerRequestState`](gh-file:vllm_omni/diffusion/sched/interface.py)
for queueing and lifecycle tracking.
For each step, the runner builds an
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
@@ -140,7 +141,7 @@ compatibility gating and runner-side `StepInputBatch` packing.
- Experimental feature; use `max_num_seqs=1` for the older conservative path.
- Only native pipelines that already support `step_execution=True`.
- Only homogeneous batches keyed by `SamplingParamsKey` are supported.
- Only homogeneous batches keyed by `StepBatchSamplingParamsKey` are supported.
- `cache_backend`, KV transfer, and other request-mode extras are not wired
into the batched step-wise path yet.
- Future work can relax the current same-shape restriction with richer
docs/design/feature/diffusion_request_level_batching.md
+2 -2
View File
@@ -82,7 +82,7 @@ The scheduler derives its capacity from `max_num_seqs` through
can decide whether admission wait is useful before scheduling a new wave.
Batch compatibility is controlled by
[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key
[`RequestBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key
contains shape-sensitive and guidance-sensitive fields, including output count
and LoRA identity. Requests with incompatible shapes, CFG settings, output
counts, LoRA adapters, or LoRA scales are kept in separate batches.
@@ -147,7 +147,7 @@ outputs do not fall back to pickle IPC for tensor payloads.
- Only pipelines that declare the request-batch contract use fused batch
execution.
- Batches are homogeneous under `SamplingParamsKey`; heterogeneous resolution or
- Batches are homogeneous under `RequestBatchSamplingParamsKey`; heterogeneous resolution or
incompatible guidance settings do not co-batch yet.
- FIFO scheduling can reduce batching opportunities when an incompatible
request is at the front of the queue.
docs/design/feature/diffusion_step_execution.md
+1 -1
View File
@@ -49,7 +49,7 @@ state object:
The state lives in
[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
as `DiffusionRequestState`. Store request-scoped tensors there, or use
as `StepRequestState`. Store request-scoped tensors there, or use
`state.extra` for model-specific fields that do not justify extending the core
dataclass.
docs/design/module/dit_module.md
+27 -30
View File
@@ -135,28 +135,7 @@ The scheduler is a **request-state scheduler**. It owns request lifecycle manage
### Key Components
#### 2.1 Scheduler Interface
```python
class SchedulerInterface(ABC):
def add_request(self, request: OmniDiffusionRequest) -> str: ...
def schedule(self) -> DiffusionSchedulerOutput: ...
def update_from_output(
self,
sched_output: DiffusionSchedulerOutput,
output: DiffusionOutput,
) -> set[str]: ...
```
**Responsibilities**:
- **Lifecycle contract**: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
- **Stable boundary**: `DiffusionSchedulerOutput` is the only scheduling result consumed by `DiffusionEngine`.
- **Pluggability**: Different scheduler policies can reuse the same engine integration path.
#### 2.2 Request State Model
#### 2.1 Request State Model
```python
class DiffusionRequestStatus(enum.IntEnum):
@@ -168,10 +147,12 @@ class DiffusionRequestStatus(enum.IntEnum):
FINISHED_ERROR = ...
@dataclass
class DiffusionRequestState:
class SchedulerRequestState:
request_id: str
req: OmniDiffusionRequest
sampling_params_key: StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey | None = None
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
error: str | None = None
```
**Design Features**:
@@ -180,18 +161,34 @@ class DiffusionRequestState:
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
#### 2.3 Shared Bookkeeping in `_BaseScheduler`
#### 2.2 Shared Bookkeeping in `BaseScheduler`
```python
class _BaseScheduler(SchedulerInterface):
class BaseScheduler:
def __init__(self) -> None:
self._request_states = {}
self._waiting = deque()
self._running = []
self._finished_req_ids = set()
self.max_num_running_reqs = 1
def add_request(self, request: OmniDiffusionRequest) -> str: ...
def schedule(self) -> DiffusionSchedulerOutput: ...
def update_from_output(
self,
sched_output: DiffusionSchedulerOutput,
output: BaseRunnerOutput,
) -> set[str]: ...
```
**Responsibilities**:
- **Lifecycle contract**: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
- **Stable boundary**: `DiffusionSchedulerOutput` is the only scheduling result consumed by `DiffusionEngine`.
- **Pluggability**: Different scheduler policies can reuse the same engine integration path.
**Design Features**:
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
@@ -199,12 +196,12 @@ class _BaseScheduler(SchedulerInterface):
- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
duplicated in each policy.
- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`. Request-mode diffusion can use that capacity for compatible independent requests when the configured pipeline declares `supports_request_batch = True`; step-wise diffusion uses the same scheduler capacity for compatible step batches.
- **Current constraint boundary**: `BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`. Request-mode diffusion can use that capacity for compatible independent requests when the configured pipeline declares `supports_request_batch = True`; step-wise diffusion uses the same scheduler capacity for compatible step batches.
#### 2.4 Current `RequestScheduler` Policy
#### 2.3 Current `RequestScheduler` Policy
```python
class RequestScheduler(_BaseScheduler):
class RequestScheduler(BaseScheduler):
def schedule(self) -> DiffusionSchedulerOutput:
# 1. keep existing RUNNING requests in the scheduling result
# 2. pull WAITING requests while capacity remains
@@ -215,11 +212,11 @@ class RequestScheduler(_BaseScheduler):
- **FIFO request scheduling**: Waiting requests are promoted in queue order.
- **Compatible request admission**: `RequestScheduler` admits waiting requests while capacity remains and the request's `SamplingParamsKey` is compatible with the active batch. Request-mode execution keeps each logical request independent, while batch-capable pipelines receive the scheduled requests as a runner-side `DiffusionRequestBatch`.
- **Compatible request admission**: `RequestScheduler` admits waiting requests while capacity remains and the request's `RequestBatchSamplingParamsKey` is compatible with the active batch. Request-mode execution keeps each logical request independent, while batch-capable pipelines receive the scheduled requests as a runner-side `DiffusionRequestBatch`.
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
#### 2.5 Engine-Driven Execution Loop
#### 2.4 Engine-Driven Execution Loop
```python
request_id = scheduler.add_request(request)
tests/ar_diffusion/test_engine_routing.py
+3 -3
View File
@@ -136,8 +136,8 @@ def test_runner_rejects_multi_sequence_config():
ARDiffusionModelRunner._preallocate_kv_cache(fake)
def test_execute_model_accepts_kv_prefetch_jobs_kwarg():
"""Regression: the worker passes ``kv_prefetch_jobs`` to every runner
def test_execute_model_accepts_kv_prefetch_job_kwarg():
"""Regression: the worker passes ``kv_prefetch_job`` to every runner
(added with the KV prefetch path, #4448); the AR-Diffusion override must
accept and forward it or DreamZero crashes on the warm-up dummy run."""
import inspect
@@ -147,4 +147,4 @@ def test_execute_model_accepts_kv_prefetch_jobs_kwarg():
for cls in (DiffusionModelRunner, ARDiffusionModelRunner):
params = inspect.signature(cls.execute_model).parameters
assert "kv_prefetch_jobs" in params, cls
assert "kv_prefetch_job" in params, cls
tests/ar_diffusion/test_pipeline_bridge.py
+1 -1
View File
@@ -203,7 +203,7 @@ def test_failed_forward_tears_down_session(monkeypatch):
runner.device = None
runner._perf_e2e_times = []
def boom(self, req, kv_prefetch_jobs=None):
def boom(self, req, kv_prefetch_job=None):
raise RuntimeError("layer 17 exploded")
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
+4 -4
View File
@@ -20,7 +20,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
)
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.diffusion.worker.utils import StepRequestState
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -38,8 +38,8 @@ def _pipeline():
return pipeline
def _state(request_id: str, step_index: int) -> DiffusionRequestState:
state = DiffusionRequestState(
def _state(request_id: str, step_index: int) -> StepRequestState:
state = StepRequestState(
request_id=request_id,
sampling=SimpleNamespace(),
prompt="prompt",
@@ -133,7 +133,7 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
pipeline.prepare_model_inputs = fake_prepare_model_inputs
state = DiffusionRequestState(
state = StepRequestState(
request_id="req-bot-task",
sampling=sampling,
prompt=prompt_item,
tests/diffusion/test_diffusion_engine_rpc_routing.py
+2 -2
View File
@@ -36,10 +36,10 @@ from vllm_omni.diffusion.sched.interface import RequestBatchSamplingParamsKey
from vllm_omni.diffusion.worker.utils import RunnerOutput
# Default values for every batch-key field, so SimpleNamespace-based
# sampling_params satisfy ``get_sampling_params_key``'s attribute lookups.
# sampling_params satisfy ``RequestScheduler._build_sampling_params_key``'s attribute lookups.
_SAMPLING_KEY_DEFAULTS = {f.name: f.default for f in _dc_fields(RequestBatchSamplingParamsKey)}
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu]
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
# ───────────────────────────────────────── helpers ─────────────────────────
tests/diffusion/test_diffusion_scheduler.py
+32 -30
View File
@@ -14,10 +14,10 @@ from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedErr
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import (
BaseScheduler,
DiffusionRequestStatus,
RequestScheduler,
Scheduler,
SchedulerInterface,
StepScheduler,
)
from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
@@ -85,7 +85,7 @@ def _cached_ids(sched_output) -> list[str]:
return list(sched_output.scheduled_cached_reqs.request_ids)
class _StubScheduler(SchedulerInterface):
class _StubScheduler:
def __init__(self, request: OmniDiffusionRequest, output) -> None:
self._request = request
self._output = output
@@ -154,8 +154,14 @@ class _StubScheduler(SchedulerInterface):
return None
class TestGetSamplingParamsKey:
"""Pure-function tests for the batch-compatibility key builder."""
class _ConcreteScheduler(BaseScheduler):
def update_from_output(self, sched_output, output) -> set[str]:
del sched_output, output
return set()
class TestGetStepBatchSamplingParamsKey:
"""Tests for the step-batch compatibility key builder on BaseScheduler."""
@staticmethod
def _make(lora_int_id: int | None = None, lora_scale: float = 1.0) -> OmniDiffusionRequest:
@@ -176,34 +182,32 @@ class TestGetSamplingParamsKey:
)
def test_distinguishes_lora_id(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
assert get_sampling_params_key(self._make(lora_int_id=1)) != get_sampling_params_key(self._make(lora_int_id=2))
scheduler = _ConcreteScheduler()
assert scheduler._build_sampling_params_key(self._make(lora_int_id=1)) != scheduler._build_sampling_params_key(
self._make(lora_int_id=2)
)
def test_distinguishes_lora_scale(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
assert get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) != get_sampling_params_key(
self._make(lora_int_id=1, lora_scale=1.0)
)
scheduler = _ConcreteScheduler()
assert scheduler._build_sampling_params_key(
self._make(lora_int_id=1, lora_scale=0.5)
) != scheduler._build_sampling_params_key(self._make(lora_int_id=1, lora_scale=1.0))
def test_treats_no_lora_as_distinct_bucket(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
assert get_sampling_params_key(self._make(lora_int_id=None)) != get_sampling_params_key(
self._make(lora_int_id=1)
)
scheduler = _ConcreteScheduler()
assert scheduler._build_sampling_params_key(
self._make(lora_int_id=None)
) != scheduler._build_sampling_params_key(self._make(lora_int_id=1))
def test_equal_for_same_lora_identity(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
a = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
b = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
scheduler = _ConcreteScheduler()
a = scheduler._build_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
b = scheduler._build_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
assert a == b
class TestGetRequestBatchSamplingParamsKey:
"""Pure-function tests for the request-batch compatibility key builder."""
"""Tests for the request-batch compatibility key builder on RequestScheduler."""
@staticmethod
def _make(
@@ -220,21 +224,19 @@ class TestGetRequestBatchSamplingParamsKey:
return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}")
def test_distinguishes_num_inference_steps(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_request_batch_sampling_params_key
assert get_request_batch_sampling_params_key(
scheduler = RequestScheduler()
assert scheduler._build_sampling_params_key(
self._make(num_inference_steps=2)
) != get_request_batch_sampling_params_key(self._make(num_inference_steps=4))
) != scheduler._build_sampling_params_key(self._make(num_inference_steps=4))
def test_ignores_seed_and_generator(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_request_batch_sampling_params_key
scheduler = RequestScheduler()
gen_a = torch.Generator(device="cpu").manual_seed(1)
gen_b = torch.Generator(device="cpu").manual_seed(2)
assert get_request_batch_sampling_params_key(
assert scheduler._build_sampling_params_key(
self._make(seed=1, generator=gen_a)
) == get_request_batch_sampling_params_key(self._make(seed=2, generator=gen_b))
) == scheduler._build_sampling_params_key(self._make(seed=2, generator=gen_b))
class TestRequestScheduler:
tests/diffusion/test_diffusion_step_pipeline.py
+3 -3
View File
@@ -43,7 +43,7 @@ from vllm_omni.diffusion.sched.interface import (
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput
from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
@@ -369,8 +369,8 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
)
def _make_input_batch_state(request_id: str, latent_value: float) -> DiffusionRequestState:
state = DiffusionRequestState(
def _make_input_batch_state(request_id: str, latent_value: float) -> StepRequestState:
state = StepRequestState(
request_id=request_id,
sampling=SimpleNamespace(),
prompt=None,
tests/diffusion/test_input_batch_txt_seq_lens.py
+3 -3
View File
@@ -2,14 +2,14 @@ import pytest
import torch
from vllm_omni.diffusion.worker.input_batch import _prepare_request_prompt_field
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.diffusion.worker.utils import StepRequestState
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.cpu]
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_prepare_request_prompt_field_refreshes_txt_seq_lens_after_padding():
state = DiffusionRequestState(
state = StepRequestState(
request_id="req-0",
sampling=OmniDiffusionSamplingParams(),
prompt_embeds=torch.randn(1, 10, 4),
tests/diffusion/test_worker_wrapper_base.py
+1 -1
View File
@@ -184,7 +184,7 @@ class TestWorkerWrapperBaseDelegation:
mock_reqs = [mocker.Mock()]
result = wrapper.execute_model(mock_reqs, mock_od_config)
wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_jobs=None)
wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_job=None)
assert result == mock_output
def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config):
vllm_omni/diffusion/diffusion_engine.py
+4 -6
View File
@@ -44,7 +44,7 @@ from vllm_omni.diffusion.registry import (
get_diffusion_pre_process_func,
)
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
from vllm_omni.diffusion.sched import BaseScheduler, RequestScheduler, StepScheduler
from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
from vllm_omni.errors import client_error_from_metadata, is_client_error_status
@@ -144,7 +144,7 @@ class DiffusionEngine:
def __init__(
self,
od_config: OmniDiffusionConfig,
scheduler: SchedulerInterface | None = None,
scheduler: BaseScheduler | None = None,
):
"""Initialize the diffusion engine.
@@ -167,9 +167,7 @@ class DiffusionEngine:
executor_class = DiffusionExecutor.get_class(od_config)
self.executor = executor_class(od_config)
self.scheduler: SchedulerInterface = scheduler or (
StepScheduler() if self.step_execution else RequestScheduler()
)
self.scheduler: BaseScheduler = scheduler or (StepScheduler() if self.step_execution else RequestScheduler())
self.scheduler.initialize(od_config)
self.supports_request_batch = False if self.step_execution else supports_request_batch(od_config)
self.main_loop: asyncio.AbstractEventLoop | None = None
@@ -618,7 +616,7 @@ class DiffusionEngine:
@staticmethod
def make_engine(
config: OmniDiffusionConfig,
scheduler: SchedulerInterface | None = None,
scheduler: BaseScheduler | None = None,
) -> DiffusionEngine:
"""Factory method to create the engine selected by ``config.engine_backend``.
vllm_omni/diffusion/executor/multiproc_executor.py
+1 -1
View File
@@ -329,7 +329,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
try:
result = self.collective_rpc(
"execute_model",
args=(req, self.od_config, scheduler_output.kv_prefetch_jobs),
args=(req, self.od_config, scheduler_output.kv_prefetch_job),
unique_reply_rank=0,
exec_all_ranks=True,
)
vllm_omni/diffusion/models/helios/pipeline_helios.py
+12 -12
View File
@@ -37,7 +37,7 @@ if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.diffusion.worker.utils import StepRequestState
logger = logging.getLogger(__name__)
@@ -279,9 +279,9 @@ class HeliosPipeline(
def prepare_encode(
self,
state: DiffusionRequestState,
state: StepRequestState,
**kwargs: Any,
) -> DiffusionRequestState:
) -> StepRequestState:
"""Initialize Helios request state for chunk-wise step execution."""
del kwargs
# Wrap the single request in a DiffusionRequestBatch so the batch
@@ -530,7 +530,7 @@ class HeliosPipeline(
self._prepare_next_chunk(state)
return state
def _prepare_next_chunk(self, state: DiffusionRequestState) -> None:
def _prepare_next_chunk(self, state: StepRequestState) -> None:
extra = state.extra
k = state.chunk_index
is_first_chunk = k == 0
@@ -604,7 +604,7 @@ class HeliosPipeline(
state.step_index = 0
self._num_timesteps = state.chunk_num_steps
def _prepare_stage2_chunk(self, state: DiffusionRequestState) -> None:
def _prepare_stage2_chunk(self, state: StepRequestState) -> None:
extra = state.extra
batch_size, num_channel, num_frames_lat, height, width = state.latents.shape
latents_flat = state.latents.permute(0, 2, 1, 3, 4).reshape(
@@ -634,7 +634,7 @@ class HeliosPipeline(
return num_steps * 2
return num_steps
def _set_stage2_timesteps(self, state: DiffusionRequestState) -> None:
def _set_stage2_timesteps(self, state: StepRequestState) -> None:
extra = state.extra
patch_size = self.transformer.config.patch_size
image_seq_len = (state.latents.shape[-1] * state.latents.shape[-2] * state.latents.shape[-3]) // (
@@ -655,7 +655,7 @@ class HeliosPipeline(
def denoise_step(
self,
input_batch: InputBatch,
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
**kwargs: Any,
) -> torch.Tensor | None:
del kwargs
@@ -668,7 +668,7 @@ class HeliosPipeline(
def _denoise_stage1_step(
self,
state: DiffusionRequestState,
state: StepRequestState,
latents: torch.Tensor,
timesteps: torch.Tensor,
) -> torch.Tensor:
@@ -727,7 +727,7 @@ class HeliosPipeline(
cfg_normalize=False,
)
def _denoise_stage2_step(self, state: DiffusionRequestState) -> torch.Tensor:
def _denoise_stage2_step(self, state: StepRequestState) -> torch.Tensor:
extra = state.extra
latents = state.latents
assert latents is not None
@@ -773,7 +773,7 @@ class HeliosPipeline(
def step_scheduler(
self,
state: DiffusionRequestState,
state: StepRequestState,
noise_pred: torch.Tensor,
**kwargs: Any,
) -> None:
@@ -800,7 +800,7 @@ class HeliosPipeline(
state.step_in_chunk += 1
state.step_index = state.step_in_chunk
def _step_scheduler_stage2(self, state: DiffusionRequestState, noise_pred: torch.Tensor) -> None:
def _step_scheduler_stage2(self, state: StepRequestState, noise_pred: torch.Tensor) -> None:
extra = state.extra
t = state.current_timestep
assert t is not None and state.latents is not None
@@ -866,7 +866,7 @@ class HeliosPipeline(
def post_decode(
self,
state: DiffusionRequestState,
state: StepRequestState,
**kwargs: Any,
) -> DiffusionOutput:
del kwargs
vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+20 -20
View File
@@ -54,7 +54,7 @@ from .system_prompt import get_system_prompt
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.diffusion.worker.utils import StepRequestState
logger = logging.getLogger(__name__)
@@ -486,7 +486,7 @@ class HunyuanImage3Pipeline(
self._pipeline = HunyuanImage3Text2ImagePipeline(model=self, scheduler=self.scheduler, vae=self.vae)
return self._pipeline
def _validate_step_request(self, state: "DiffusionRequestState") -> None:
def _validate_step_request(self, state: "StepRequestState") -> None:
prompt = state.prompt
sampling = state.sampling
if prompt is None:
@@ -584,7 +584,7 @@ class HunyuanImage3Pipeline(
def _extract_step_prompt_inputs(
self,
state: "DiffusionRequestState",
state: "StepRequestState",
) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]:
sampling = state.sampling
return self._extract_prompt_inputs(
@@ -610,7 +610,7 @@ class HunyuanImage3Pipeline(
def _restore_injected_ar_kv(
self,
states: list["DiffusionRequestState"],
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> None:
@@ -755,7 +755,7 @@ class HunyuanImage3Pipeline(
def _prompt_kv_prefix_lens(
self,
states: list["DiffusionRequestState"],
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> list[int] | None:
@@ -774,7 +774,7 @@ class HunyuanImage3Pipeline(
def _merge_step_model_inputs(
self,
states: list["DiffusionRequestState"],
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
first_step: bool,
@@ -828,7 +828,7 @@ class HunyuanImage3Pipeline(
def _restore_prompt_kv_cache(
self,
states: list["DiffusionRequestState"],
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> None:
@@ -855,7 +855,7 @@ class HunyuanImage3Pipeline(
def _capture_prompt_kv_cache(
self,
states: list["DiffusionRequestState"],
states: list["StepRequestState"],
row_state_indexes: list[int],
row_branches: list[int],
) -> None:
@@ -1838,9 +1838,9 @@ class HunyuanImage3Pipeline(
def prepare_encode(
self,
state: "DiffusionRequestState",
state: "StepRequestState",
**kwargs: Any,
) -> "DiffusionRequestState":
) -> "StepRequestState":
del kwargs
self._validate_step_request(state)
pipe = self.pipeline
@@ -1955,7 +1955,7 @@ class HunyuanImage3Pipeline(
}
return state
def _step_group_key(self, state: "DiffusionRequestState") -> tuple[Any, ...]:
def _step_group_key(self, state: "StepRequestState") -> tuple[Any, ...]:
if _STEP_CFG_FACTOR not in state.extra:
raise ValueError(f"Missing Hunyuan CFG factor for request {state.request_id}.")
if _STEP_MODEL_KWARGS not in state.extra:
@@ -1989,16 +1989,16 @@ class HunyuanImage3Pipeline(
def _split_step_groups(
self,
states: list["DiffusionRequestState"],
) -> list[list["DiffusionRequestState"]]:
grouped: dict[tuple[Any, ...], list[DiffusionRequestState]] = {}
states: list["StepRequestState"],
) -> list[list["StepRequestState"]]:
grouped: dict[tuple[Any, ...], list[StepRequestState]] = {}
for state in states:
grouped.setdefault(self._step_group_key(state), []).append(state)
return list(grouped.values())
@staticmethod
def _step_row_order(
states: list["DiffusionRequestState"],
states: list["StepRequestState"],
cfg_factor: int,
) -> tuple[list[int], list[int]]:
row_state_indexes = [state_idx for _ in range(cfg_factor) for state_idx in range(len(states))]
@@ -2006,7 +2006,7 @@ class HunyuanImage3Pipeline(
return row_state_indexes, row_branches
@staticmethod
def _validate_step_group_states(states: list["DiffusionRequestState"]) -> tuple[bool, int]:
def _validate_step_group_states(states: list["StepRequestState"]) -> tuple[bool, int]:
if not states:
raise ValueError("HunyuanImage3 denoise_step received an empty group.")
@@ -2043,7 +2043,7 @@ class HunyuanImage3Pipeline(
def _split_merged_kwargs_to_states(
self,
states: list["DiffusionRequestState"],
states: list["StepRequestState"],
merged_kwargs: dict[str, Any],
row_state_indexes: list[int],
row_branches: list[int],
@@ -2094,7 +2094,7 @@ class HunyuanImage3Pipeline(
state.extra[_STEP_MODEL_KWARGS] = next_kwargs
state.extra[_STEP_INPUT_IDS] = None
def _denoise_step_group(self, states: list["DiffusionRequestState"]) -> torch.Tensor:
def _denoise_step_group(self, states: list["StepRequestState"]) -> torch.Tensor:
first_step, cfg_factor = self._validate_step_group_states(states)
row_state_indexes, row_branches = self._step_row_order(states, cfg_factor)
latents = torch.cat([state.latents for state in states], dim=0)
@@ -2179,7 +2179,7 @@ class HunyuanImage3Pipeline(
def step_scheduler(
self,
state: "DiffusionRequestState",
state: "StepRequestState",
noise_pred: torch.Tensor,
**kwargs: Any,
) -> None:
@@ -2200,7 +2200,7 @@ class HunyuanImage3Pipeline(
def post_decode(
self,
state: "DiffusionRequestState",
state: "StepRequestState",
**kwargs: Any,
) -> DiffusionOutput:
output_type = kwargs.get("output_type", "pil")
vllm_omni/diffusion/models/interface.py
+4 -4
View File
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.diffusion.worker.utils import StepRequestState
@runtime_checkable
@@ -54,7 +54,7 @@ class SupportsStepExecution(Protocol):
supports_step_execution: ClassVar[bool] = True
def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState:
def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState:
"""Prepare request-level inputs and return initialized state."""
...
@@ -62,11 +62,11 @@ class SupportsStepExecution(Protocol):
"""Run one denoise forward on the runner-assembled batch."""
...
def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
"""Run one scheduler step."""
...
def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput:
def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput:
"""Decode output after denoise loop or at a partial chunk boundary."""
...
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+5 -5
View File
@@ -48,7 +48,7 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, spli
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.diffusion.worker.utils import StepRequestState
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
@@ -770,9 +770,9 @@ class QwenImagePipeline(
def prepare_encode(
self,
state: "DiffusionRequestState",
state: "StepRequestState",
**kwargs: Any,
) -> "DiffusionRequestState":
) -> "StepRequestState":
"""Populate *state* with encoded prompts, latents, timesteps, and CFG config."""
sampling = state.sampling
prompt, negative_prompt = self._extract_prompts([state.prompt] if state.prompt is not None else [])
@@ -959,7 +959,7 @@ class QwenImagePipeline(
def step_scheduler(
self,
state: "DiffusionRequestState",
state: "StepRequestState",
noise_pred: torch.Tensor,
**kwargs: Any,
) -> None:
@@ -980,7 +980,7 @@ class QwenImagePipeline(
def post_decode(
self,
state: "DiffusionRequestState",
state: "StepRequestState",
**kwargs: Any,
) -> DiffusionOutput:
"""Decode final latents from *state*."""
vllm_omni/diffusion/request.py
+2 -1
View File
@@ -4,6 +4,7 @@
import random
from dataclasses import dataclass
from typing import Any
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -28,7 +29,7 @@ class OmniDiffusionRequest:
prompt: OmniPromptType
sampling_params: OmniDiffusionSamplingParams
request_id: str
kv_sender_info: dict | None = None
kv_sender_info: dict[str, Any] | None = None
def __post_init__(self):
"""Initialize dependent fields after dataclass initialization."""
vllm_omni/diffusion/sched/__init__.py
+8 -3
View File
@@ -1,13 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler, SchedulerInterface
from vllm_omni.diffusion.sched.interface import (
CachedRequestData,
DiffusionRequestState,
DiffusionRequestStatus,
DiffusionSchedulerOutput,
KVPrefetchJob,
NewRequestData,
SchedulerInterface,
SchedulerRequestState,
StepBatchSamplingParamsKey,
)
from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
@@ -17,10 +19,13 @@ Scheduler = RequestScheduler
__all__ = [
"DiffusionRequestStatus",
"CachedRequestData",
"DiffusionRequestState",
"DiffusionSchedulerOutput",
"KVPrefetchJob",
"NewRequestData",
"SchedulerRequestState",
"BaseScheduler",
"SchedulerInterface",
"StepBatchSamplingParamsKey",
"RequestScheduler",
"StepScheduler",
"Scheduler",
vllm_omni/diffusion/sched/base_scheduler.py
+51 -38
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from collections import deque
from dataclasses import fields
@@ -12,54 +13,37 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import (
CachedRequestData,
DiffusionRequestState,
DiffusionRequestStatus,
DiffusionSchedulerOutput,
KVPrefetchJob,
NewRequestData,
RequestBatchSamplingParamsKey,
SamplingParamsKey,
SchedulerInterface,
SchedulerRequestState,
StepBatchSamplingParamsKey,
)
from vllm_omni.diffusion.worker.utils import RunnerOutput
logger = init_logger(__name__)
BatchSamplingParamsKey = StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey
# LoRA identity is derived from `sampling.lora_request`, not a same-named field
# on sampling params, so it must be resolved separately from the bulk lookup.
_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey)) - {"lora_int_id"}
_REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(f.name for f in fields(RequestBatchSamplingParamsKey)) - {
_STEP_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(field.name for field in fields(StepBatchSamplingParamsKey)) - {
"lora_int_id"
}
def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey:
"""Build a batch-compatibility key from the request's sampling params."""
sampling = request.sampling_params
lora_request = getattr(sampling, "lora_request", None)
return SamplingParamsKey(
lora_int_id=lora_request.lora_int_id if lora_request is not None else None,
**{name: getattr(sampling, name) for name in _SAMPLING_PARAMS_KEY_FIELD_NAMES},
)
def get_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
"""Build a request-batch compatibility key from the request's sampling params."""
sampling = request.sampling_params
lora_request = getattr(sampling, "lora_request", None)
key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
return RequestBatchSamplingParamsKey(**key_kwargs)
class _BaseScheduler(SchedulerInterface):
class BaseScheduler(ABC):
"""Shared queue/state bookkeeping for diffusion schedulers."""
def __init__(self) -> None:
self.od_config: OmniDiffusionConfig | None = None
self._request_states: dict[str, DiffusionRequestState] = {}
self._request_states: dict[str, SchedulerRequestState] = {}
self._step_id: int = 0
self._waiting: deque[str] = deque()
self._running: list[str] = []
self._running_sampling_params_key: SamplingParamsKey | RequestBatchSamplingParamsKey | None = None
self._running_sampling_params_key: BatchSamplingParamsKey | None = None
self._finished_req_ids: set[str] = set()
self.max_num_running_reqs: int = 1
self._prefetch_enabled: bool = False
@@ -129,13 +113,13 @@ class _BaseScheduler(SchedulerInterface):
# kv_sender_info (would target the wrong sender under multi-replica) or
# one already finished/aborted (would consume its sender buffer for
# nothing).
kv_prefetch_jobs: dict | None = None
kv_prefetch_job: KVPrefetchJob | None = None
if self._prefetch_enabled and self._waiting:
nxt = self._request_states.get(self._waiting[0])
if nxt is not None and not nxt.is_finished():
sender_info = getattr(nxt.req, "kv_sender_info", None)
if sender_info:
kv_prefetch_jobs = {
kv_prefetch_job = {
"request_id": nxt.request_id,
"kv_sender_info": sender_info,
}
@@ -147,7 +131,7 @@ class _BaseScheduler(SchedulerInterface):
finished_req_ids=set(self._finished_req_ids),
num_running_reqs=len(self._running),
num_waiting_reqs=len(self._waiting),
kv_prefetch_jobs=kv_prefetch_jobs,
kv_prefetch_job=kv_prefetch_job,
)
# update after schedule
@@ -155,6 +139,10 @@ class _BaseScheduler(SchedulerInterface):
self._finished_req_ids.clear()
return scheduler_output
@abstractmethod
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
pass
def has_requests(self) -> bool:
return bool(self._waiting or self._running)
@@ -164,10 +152,10 @@ class _BaseScheduler(SchedulerInterface):
def num_running_requests(self) -> int:
return len(self._running)
def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
def get_request_state(self, request_id: str) -> SchedulerRequestState | None:
return self._request_states.get(request_id)
def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
def pop_request_state(self, request_id: str) -> SchedulerRequestState | None:
self._pop_extra_request_state(request_id)
return self._request_states.pop(request_id, None)
@@ -262,21 +250,21 @@ class _BaseScheduler(SchedulerInterface):
def _pop_extra_request_state(self, request_id: str) -> None:
"""Remove subclass-owned per-request state before popping request state."""
def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> DiffusionRequestState:
return DiffusionRequestState(
def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> SchedulerRequestState:
return SchedulerRequestState(
request_id=request_id,
req=request,
sampling_params_key=self._build_sampling_params_key(request),
)
def _can_schedule_waiting(self, state: DiffusionRequestState) -> bool:
def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool:
if not self._running:
return True
current_key = self._current_sampling_params_key()
return current_key is not None and current_key == state.sampling_params_key
def _current_sampling_params_key(self) -> SamplingParamsKey | RequestBatchSamplingParamsKey | None:
def _current_sampling_params_key(self) -> BatchSamplingParamsKey | None:
if self._running_sampling_params_key is not None or not self._running:
return self._running_sampling_params_key
state = self._request_states.get(self._running[0])
@@ -285,5 +273,30 @@ class _BaseScheduler(SchedulerInterface):
def _build_sampling_params_key(
self, request: OmniDiffusionRequest
) -> SamplingParamsKey | RequestBatchSamplingParamsKey | None:
return get_sampling_params_key(request)
) -> StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey: # return type loosened for subclassing
"""Build a step-batch compatibility key from sampling parameters."""
sampling = request.sampling_params
# LoRA identity is optional on sampling params (and on test stubs).
lora_request = getattr(sampling, "lora_request", None)
return StepBatchSamplingParamsKey(
lora_int_id=lora_request.lora_int_id if lora_request is not None else None,
**{name: getattr(sampling, name) for name in _STEP_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES},
)
class SchedulerInterface(BaseScheduler):
"""Deprecated compatibility base for custom scheduler injection.
Prefer subclassing :class:`BaseScheduler` directly. Subclassing this name
still works but emits a :class:`DeprecationWarning`.
"""
def __init_subclass__(cls, **kwargs) -> None:
import warnings
warnings.warn(
"SchedulerInterface is deprecated; subclass BaseScheduler instead.",
DeprecationWarning,
stacklevel=2,
)
super().__init_subclass__(**kwargs)
vllm_omni/diffusion/sched/interface.py
+31 -78
View File
@@ -4,17 +4,12 @@
from __future__ import annotations
import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import cached_property
from typing import TYPE_CHECKING
from typing import Any, TypedDict
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.request import OmniDiffusionRequest
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.utils import RunnerOutput
class DiffusionRequestStatus(enum.IntEnum):
"""Request status tracked by diffusion scheduler."""
@@ -34,7 +29,7 @@ class DiffusionRequestStatus(enum.IntEnum):
@dataclass(frozen=True, eq=True)
class SamplingParamsKey:
class StepBatchSamplingParamsKey:
"""Denoise step level Batch-compatibility key derived from ``OmniDiffusionSamplingParams``.
Only requests with the same key can be batched together.
@@ -81,33 +76,33 @@ class RequestBatchSamplingParamsKey:
"""
# Spatial / temporal shape.
height: object = None
width: object = None
height: int | None = None
width: int | None = None
num_frames: int = 1
resolution: object = 640
fps: object = None
frame_rate: object = None
boundary_ratio: object = None
resolution: int = 640
fps: int | None = None
frame_rate: float | None = None
boundary_ratio: float | None = None
# CFG / guidance.
do_classifier_free_guidance: bool = False
guidance_scale: float = 0.0
guidance_scale_provided: bool = False
guidance_scale_2: object = None
guidance_scale_2: float | None = None
guidance_rescale: float = 0.0
true_cfg_scale: object = None
true_cfg_scale: float | None = None
cfg_normalize: bool = False
strength: object = None
strength: float | None = None
# Scheduling / output shape.
num_inference_steps: object = None
sigmas: object = None
max_sequence_length: object = None
num_inference_steps: int | None = None
sigmas: list[float] | None = None
max_sequence_length: int | None = None
num_outputs_per_prompt: int = 1
eta: float = 0.0
decode_timestep: object = None
decode_noise_scale: object = None
output_type: object = None
decode_timestep: float | list[float] | None = None
decode_noise_scale: float | list[float] | None = None
output_type: str | None = None
# Model-specific batch defaults used by request-mode pipelines.
layers: int = 4
@@ -119,12 +114,12 @@ class RequestBatchSamplingParamsKey:
@dataclass
class DiffusionRequestState:
class SchedulerRequestState:
"""Scheduler-owned state for one queued OmniDiffusionRequest."""
request_id: str
req: OmniDiffusionRequest
sampling_params_key: SamplingParamsKey | RequestBatchSamplingParamsKey | None = None
sampling_params_key: StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey | None = None
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
error: str | None = None
@@ -145,7 +140,7 @@ class NewRequestData:
req: OmniDiffusionRequest
@classmethod
def from_state(cls, state: DiffusionRequestState) -> NewRequestData:
def from_state(cls, state: SchedulerRequestState) -> NewRequestData:
return cls(request_id=state.request_id, req=state.req)
@@ -160,10 +155,20 @@ class CachedRequestData:
return cls(request_ids=[])
class KVPrefetchJob(TypedDict):
"""Descriptor for prefetching the next request's received KV cache."""
request_id: str
kv_sender_info: dict[str, Any]
@dataclass
class DiffusionSchedulerOutput:
"""Output of a single scheduling cycle."""
# Stable scheduler-cycle diagnostics for engines, injected schedulers, and
# instrumentation; they are intentionally retained even without a direct
# production consumer for every field.
step_id: int # global step index
scheduled_new_reqs: list[NewRequestData]
scheduled_cached_reqs: CachedRequestData
@@ -171,7 +176,7 @@ class DiffusionSchedulerOutput:
num_running_reqs: int
num_waiting_reqs: int
# next request to background-prefetch KV
kv_prefetch_jobs: dict | None = None
kv_prefetch_job: KVPrefetchJob | None = None
@cached_property
def scheduled_request_ids(self) -> list[str]:
@@ -190,55 +195,3 @@ class DiffusionSchedulerOutput:
@property
def is_empty(self) -> bool:
return self.num_scheduled_reqs == 0
class SchedulerInterface(ABC):
"""Abstract lifecycle contract for diffusion schedulers."""
@abstractmethod
def initialize(self, od_config: OmniDiffusionConfig) -> None:
"""Initialize or reset scheduler state."""
@abstractmethod
def add_request(self, request: OmniDiffusionRequest) -> str:
"""Add a request and return the scheduler-owned request id."""
@abstractmethod
def schedule(self) -> DiffusionSchedulerOutput:
"""Run one scheduling cycle."""
@abstractmethod
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
"""Update scheduler state from executor output."""
@abstractmethod
def get_request_state(self, request_id: str) -> DiffusionRequestState | None:
"""Return request state if present."""
@abstractmethod
def has_requests(self) -> bool:
"""Return whether the scheduler still owns runnable requests."""
@abstractmethod
def num_waiting_requests(self) -> int:
"""Return the number of requests waiting to be scheduled."""
@abstractmethod
def num_running_requests(self) -> int:
"""Return the number of requests currently running."""
@abstractmethod
def pop_request_state(self, request_id: str) -> DiffusionRequestState | None:
"""Remove and return request state if present."""
@abstractmethod
def preempt_request(self, request_id: str) -> bool:
"""Preempt a running request back to waiting."""
@abstractmethod
def finish_requests(self, request_ids: str | list[str], status: DiffusionRequestStatus) -> None:
"""Mark one or more requests finished."""
@abstractmethod
def close(self) -> None:
"""Release scheduler-owned state."""
vllm_omni/diffusion/sched/request_scheduler.py
+18 -10
View File
@@ -3,30 +3,38 @@
from __future__ import annotations
from dataclasses import fields
from typing import TYPE_CHECKING
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.base_scheduler import _BaseScheduler, get_request_batch_sampling_params_key
from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
from vllm_omni.diffusion.sched.interface import (
DiffusionRequestStatus,
DiffusionSchedulerOutput,
RequestBatchSamplingParamsKey,
)
if TYPE_CHECKING:
from vllm_omni.diffusion.worker.utils import RunnerOutput
# LoRA identity is derived from `sampling.lora_request`, not a same-named field
# on sampling params, so it must be resolved separately from the bulk lookup.
_REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(
field.name for field in fields(RequestBatchSamplingParamsKey)
) - {"lora_int_id"}
class RequestScheduler(_BaseScheduler):
class RequestScheduler(BaseScheduler):
"""Diffusion scheduler with vLLM-style waiting/running queues."""
def _build_sampling_params_key(self, request: OmniDiffusionRequest):
return get_request_batch_sampling_params_key(request)
def add_request(self, request: OmniDiffusionRequest) -> str:
return super().add_request(request)
def schedule(self) -> DiffusionSchedulerOutput:
return super().schedule()
def _build_sampling_params_key(self, request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
"""Build a request-batch compatibility key from sampling parameters."""
sampling = request.sampling_params
# LoRA identity is optional on sampling params (and on test stubs).
lora_request = getattr(sampling, "lora_request", None)
key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES}
key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None
return RequestBatchSamplingParamsKey(**key_kwargs)
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
scheduled_request_ids = sched_output.scheduled_request_ids
vllm_omni/diffusion/sched/step_scheduler.py
+3 -6
View File
@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any
from vllm.logger import init_logger
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.base_scheduler import _BaseScheduler
from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
from vllm_omni.diffusion.sched.interface import (
DiffusionRequestStatus,
DiffusionSchedulerOutput,
@@ -27,8 +27,8 @@ class _StepProgress:
total_steps: int
class StepScheduler(_BaseScheduler):
"""Placeholder scheduler that advances a request one denoise step per update."""
class StepScheduler(BaseScheduler):
"""Scheduler that advances each request by one denoise step per update."""
def __init__(self) -> None:
super().__init__()
@@ -62,9 +62,6 @@ class StepScheduler(_BaseScheduler):
)
return request_id
def schedule(self) -> DiffusionSchedulerOutput:
return super().schedule()
def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
scheduled_request_ids = sched_output.scheduled_request_ids
if not scheduled_request_ids:
vllm_omni/diffusion/worker/diffusion_model_runner.py
+21 -19
View File
@@ -36,13 +36,13 @@ from vllm_omni.diffusion.models.interface import supports_step_execution
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.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.diffusion.worker.utils import (
BatchRunnerOutput,
DiffusionRequestState,
RunnerOutput,
StepRequestState,
attach_stage_durations,
clear_pipeline_stage_durations,
consume_pipeline_stage_durations,
@@ -124,7 +124,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
self.prompt_embed_cache = None
# Cache for per-request stepwise state.
self.state_cache: dict[str, DiffusionRequestState] = {}
self.state_cache: dict[str, StepRequestState] = {}
# Initialize KV cache manager for connector management.
self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
@@ -336,7 +336,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
req: OmniDiffusionRequest,
*,
od_config: OmniDiffusionConfig,
kv_prefetch_jobs: dict | None = None,
kv_prefetch_job: KVPrefetchJob | None = None,
use_prefetch: bool = False,
) -> None:
# Receive AR KV. Single-request execution can use the prefetch path:
@@ -358,8 +358,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms)
# Kick off the next request's prefetch (+ H2D) to overlap this forward.
if use_prefetch and self._kv_prefetch_enabled and kv_prefetch_jobs is not None:
self.kv_transfer_manager.start_prefetch(kv_prefetch_jobs, self.target_device)
if use_prefetch and self._kv_prefetch_enabled and kv_prefetch_job is not None:
self.kv_transfer_manager.start_prefetch(kv_prefetch_job, self.target_device)
if req.sampling_params.generator is None and req.sampling_params.seed is not None:
if req.sampling_params.generator_device is not None:
@@ -385,7 +385,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
return
# Refresh cache context if needed. Batch admission groups requests by
# SamplingParamsKey, so the first request's num_inference_steps applies
# RequestBatchSamplingParamsKey, so the first request's num_inference_steps applies
# to the whole runner batch.
num_inference_steps = first_req.sampling_params.num_inference_steps
if num_inference_steps is None and od_config.cache_backend in (
@@ -432,7 +432,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
od_config: OmniDiffusionConfig,
allow_single_output: bool,
require_request_batch_support: bool,
kv_prefetch_jobs: dict | None = None,
kv_prefetch_job: KVPrefetchJob | None = None,
record_name: str,
) -> BatchRunnerOutput:
assert self.pipeline is not None, "Model not loaded. Call load_model() first."
@@ -454,7 +454,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
self._prepare_request_for_forward(
req,
od_config=od_config,
kv_prefetch_jobs=kv_prefetch_jobs,
kv_prefetch_job=kv_prefetch_job,
use_prefetch=allow_single_output,
)
@@ -497,7 +497,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
def _attach_stepwise_metrics(
self,
state: DiffusionRequestState,
state: StepRequestState,
output: DiffusionOutput,
*,
is_primary: bool,
@@ -508,7 +508,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
)
attach_stage_durations(state, output)
def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
def execute_model(
self,
req: OmniDiffusionRequest,
kv_prefetch_job: KVPrefetchJob | None = None,
) -> DiffusionOutput:
"""
Execute a forward pass for the given requests.
@@ -529,7 +533,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
od_config=self.od_config,
allow_single_output=True,
require_request_batch_support=False,
kv_prefetch_jobs=kv_prefetch_jobs,
kv_prefetch_job=kv_prefetch_job,
record_name="pipeline_forward",
)
output = runner_output.runner_outputs[0].result
@@ -564,14 +568,12 @@ 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[list[DiffusionRequestState], list[str]]:
def _update_states(self, scheduler_output: DiffusionSchedulerOutput) -> tuple[list[StepRequestState], list[str]]:
"""Step-before update: cleanup finished requests and get/create one running state."""
for request_id in scheduler_output.finished_req_ids:
self.state_cache.pop(request_id, None)
resolved: list[DiffusionRequestState] = []
resolved: list[StepRequestState] = []
new_request_ids: list[str] = []
try:
# process new requests
@@ -580,7 +582,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
new_request_ids.append(request_id)
if request_id in self.state_cache:
raise ValueError(f"Received duplicate new-request payload for cached request {request_id}.")
new_state = DiffusionRequestState(
new_state = StepRequestState(
request_id=request_id,
sampling=copy.deepcopy(sched_new_req.req.sampling_params),
prompt=sched_new_req.req.prompt,
@@ -609,7 +611,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
return resolved, new_request_ids
def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request_ids: list[str]) -> InputBatch:
def _prepare_batch_inputs(self, states: list[StepRequestState], new_request_ids: list[str]) -> InputBatch:
# process new reqs
for state in states:
if state.request_id in new_request_ids:
@@ -639,7 +641,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
def _update_states_after(
self,
states: list[DiffusionRequestState],
states: list[StepRequestState],
input_batch: InputBatch,
interrupted: bool = False,
):
vllm_omni/diffusion/worker/diffusion_worker.py
+8 -8
View File
@@ -46,7 +46,7 @@ from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusio
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput
from vllm_omni.engine.stage_init_utils import set_death_signal
@@ -416,7 +416,7 @@ class DiffusionWorker:
self,
req: OmniDiffusionRequest,
od_config: OmniDiffusionConfig,
kv_prefetch_jobs: dict | None = None,
kv_prefetch_job: KVPrefetchJob | None = None,
) -> DiffusionOutput:
"""Execute a forward pass by delegating to the model runner."""
assert self.model_runner is not None, "Model runner not initialized"
@@ -430,7 +430,7 @@ class DiffusionWorker:
profiler = self._get_profiler()
ctx = profiler.annotate_context_manager("diffusion_forward") if profiler else nullcontext()
with ctx:
output = self.model_runner.execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs)
output = self.model_runner.execute_model(req, kv_prefetch_job=kv_prefetch_job)
if profiler:
profiler.step()
return output
@@ -440,7 +440,7 @@ class DiffusionWorker:
) -> BatchRunnerOutput:
"""Batch forward: LoRA activate once, delegate to model runner."""
assert self.model_runner is not None, "Model runner not initialized"
# LoRA: same adapter/scale within batch guaranteed by SamplingParamsKey
# LoRA: same adapter/scale within batch guaranteed by RequestBatchSamplingParamsKey.
if self.lora_manager is not None and scheduler_output.scheduled_new_reqs:
sp = scheduler_output.scheduled_new_reqs[0].req.sampling_params
try:
@@ -475,7 +475,7 @@ class DiffusionWorker:
Newly scheduled requests register their (lora_request, lora_scale)
in ``_step_lora_state`` so cached requests can resolve to the same
adapter on later ticks. Finished requests are evicted. Batch
homogeneity is enforced by the scheduler via ``SamplingParamsKey``,
homogeneity is enforced by the scheduler via ``StepBatchSamplingParamsKey``,
so any scheduled request id resolves to the active LoRA identity.
"""
for request_id in scheduler_output.finished_req_ids:
@@ -1155,7 +1155,7 @@ class WorkerWrapperBase:
self,
reqs: list[OmniDiffusionRequest],
od_config: OmniDiffusionConfig,
kv_prefetch_jobs: dict | None = None,
kv_prefetch_job: KVPrefetchJob | None = None,
) -> DiffusionOutput:
"""
Execute a forward pass.
@@ -1163,12 +1163,12 @@ class WorkerWrapperBase:
Args:
reqs: List of diffusion requests
od_config: OmniDiffusionConfig configuration
kv_prefetch_jobs: Optional next-request KV prefetch descriptor.
kv_prefetch_job: Optional next-request KV prefetch descriptor.
Returns:
DiffusionOutput with generated results
"""
return self.worker.execute_model(reqs, od_config, kv_prefetch_jobs=kv_prefetch_jobs)
return self.worker.execute_model(reqs, od_config, kv_prefetch_job=kv_prefetch_job)
def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
"""Execute one diffusion step."""
vllm_omni/diffusion/worker/input_batch.py
+29 -29
View File
@@ -17,7 +17,7 @@ from dataclasses import dataclass, field
import numpy as np
import torch
from vllm_omni.diffusion.worker.utils import DiffusionRequestState
from vllm_omni.diffusion.worker.utils import StepRequestState
def _normalize_prompt_embeds(x: torch.Tensor) -> torch.Tensor:
@@ -59,9 +59,9 @@ def _pad_mask(x: torch.Tensor, target_seq_len: int) -> torch.Tensor:
def _select_states(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
idx_mapping: torch.Tensor | None,
) -> tuple[list[DiffusionRequestState], torch.Tensor, np.ndarray]:
) -> tuple[list[StepRequestState], torch.Tensor, np.ndarray]:
if not states:
raise ValueError("Cannot build InputBatch from empty states.")
@@ -73,7 +73,7 @@ def _select_states(
raise ValueError("idx_mapping must be a 1D tensor.")
idx_mapping = idx_mapping.to(dtype=torch.int32)
selected_states: list[DiffusionRequestState] = []
selected_states: list[StepRequestState] = []
for batch_idx, state_idx in enumerate(idx_mapping.tolist()):
if state_idx < 0 or state_idx >= len(states):
raise ValueError(f"idx_mapping[{batch_idx}]={state_idx} is out of range for states.")
@@ -81,12 +81,12 @@ def _select_states(
return selected_states, idx_mapping, idx_mapping.detach().cpu().numpy()
def _prepare_request_ids(states: Sequence[DiffusionRequestState]) -> list[str]:
def _prepare_request_ids(states: Sequence[StepRequestState]) -> list[str]:
return [state.request_id for state in states]
def _prepare_prompt_field_on_state(
state: DiffusionRequestState,
state: StepRequestState,
*,
embeds_attr: str,
mask_attr: str,
@@ -180,7 +180,7 @@ def _get_seq_lens_from_mask(mask: torch.Tensor) -> list[int]:
def _get_request_prompt_seq_lens(
state: DiffusionRequestState,
state: StepRequestState,
*,
embeds_attr: str,
mask_attr: str,
@@ -209,7 +209,7 @@ def _get_request_prompt_seq_lens(
def _prepare_request_prompt_field(
state: DiffusionRequestState,
state: StepRequestState,
*,
embeds_attr: str,
mask_attr: str,
@@ -266,7 +266,7 @@ def _prepare_request_prompt_field(
def _prepare_padded_prompt_fields(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
*,
embeds_attr: str,
mask_attr: str,
@@ -337,7 +337,7 @@ def _prepare_padded_prompt_fields(
def _prepare_latents(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor:
@@ -352,7 +352,7 @@ def _prepare_latents(
def _require_state_latents(
state: DiffusionRequestState,
state: StepRequestState,
*,
for_field: str,
) -> torch.Tensor:
@@ -382,7 +382,7 @@ def _expand_scalar_or_vector(
def _prepare_timesteps(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor:
@@ -410,9 +410,9 @@ def _prepare_timesteps(
def _prepare_cfg_scalars(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
) -> tuple[bool, float, bool]:
def _cfg_scalars(state: DiffusionRequestState) -> tuple[bool, float, bool]:
def _cfg_scalars(state: StepRequestState) -> tuple[bool, float, bool]:
true_cfg_scale = getattr(state.sampling, "true_cfg_scale", None) or 4.0
cfg_normalize = bool(getattr(state.sampling, "cfg_normalize", False))
return state.do_true_cfg, true_cfg_scale, cfg_normalize
@@ -425,7 +425,7 @@ def _prepare_cfg_scalars(
def _prepare_guidance(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor | None:
@@ -461,7 +461,7 @@ def _prepare_guidance(
def _prepare_image_latents(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
*,
out: torch.Tensor | None = None,
) -> torch.Tensor | None:
@@ -478,7 +478,7 @@ def _prepare_image_latents(
def _prepare_seq_lens(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
attr_name: str,
) -> list[int] | None:
values = [getattr(state, attr_name) for state in states]
@@ -489,7 +489,7 @@ def _prepare_seq_lens(
return [int(value[0]) for value in values if value is not None]
def _prepare_img_shapes(states: Sequence[DiffusionRequestState]) -> list | None:
def _prepare_img_shapes(states: Sequence[StepRequestState]) -> list | None:
values = [state.img_shapes for state in states]
if all(value is None for value in values):
return None
@@ -499,7 +499,7 @@ def _prepare_img_shapes(states: Sequence[DiffusionRequestState]) -> list | None:
def _prepare_prompt_embeds(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
*,
embeds_out: torch.Tensor | None = None,
mask_out: torch.Tensor | None = None,
@@ -516,7 +516,7 @@ def _prepare_prompt_embeds(
def _prepare_negative_prompt_embeds(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
*,
embeds_out: torch.Tensor | None = None,
mask_out: torch.Tensor | None = None,
@@ -544,7 +544,7 @@ def _same_composition(
def _scatter_batch_tensor_by_mapping(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
idx_mapping_np: np.ndarray,
*,
attr_name: str,
@@ -582,7 +582,7 @@ class InputBatch:
"""Ephemeral step-level batch view.
Static request-local tensors are normalized and padded onto
``DiffusionRequestState`` itself, making the request state the persistent
``StepRequestState`` itself, making the request state the persistent
source of truth. ``InputBatch`` only assembles a contiguous view for the
current step and refreshes dynamic fields in-place when composition is
unchanged.
@@ -613,7 +613,7 @@ class InputBatch:
img_shapes: list | None = None
txt_seq_lens: list[int] | None = None
negative_txt_seq_lens: list[int] | None = None
states: Sequence[DiffusionRequestState] = field(default_factory=tuple)
states: Sequence[StepRequestState] = field(default_factory=tuple)
def __post_init__(self) -> None:
if len(self.request_ids) != int(self.idx_mapping.numel()):
@@ -625,14 +625,14 @@ class InputBatch:
def _refresh_dynamic_fields(
self,
selected_states: Sequence[DiffusionRequestState],
selected_states: Sequence[StepRequestState],
) -> None:
self.latents = _prepare_latents(selected_states, out=self.latents)
self.timesteps = _prepare_timesteps(selected_states, out=self.timesteps)
def _refresh_static_fields(
self,
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
) -> None:
self.do_true_cfg, self.true_cfg_scale, self.cfg_normalize = _prepare_cfg_scalars(states)
self.guidance = _prepare_guidance(states, out=self.guidance)
@@ -656,7 +656,7 @@ class InputBatch:
def _repack_dynamic_fields(
self,
selected_states: Sequence[DiffusionRequestState],
selected_states: Sequence[StepRequestState],
) -> None:
# Same-composition cache hits reuse static prompt fields; request ids
# must keep the same encoded prompt metadata for the batch lifetime.
@@ -665,7 +665,7 @@ class InputBatch:
def _rebuild(
self,
selected_states: Sequence[DiffusionRequestState],
selected_states: Sequence[StepRequestState],
idx_mapping: torch.Tensor,
idx_mapping_np: np.ndarray,
request_ids: list[str],
@@ -685,7 +685,7 @@ class InputBatch:
@classmethod
def make_batch(
cls,
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
idx_mapping: torch.Tensor | None = None,
cached_batch: InputBatch | None = None,
) -> InputBatch:
@@ -737,7 +737,7 @@ class InputBatch:
def scatter_latents(
states: Sequence[DiffusionRequestState],
states: Sequence[StepRequestState],
input_batch: InputBatch,
) -> None:
"""Scatter the step-updated latents back into persistent request states.
vllm_omni/diffusion/worker/utils.py
+3 -3
View File
@@ -33,7 +33,7 @@ def consume_pipeline_stage_durations(pipeline: Any) -> dict[str, float]:
def merge_stage_durations(
state: DiffusionRequestState,
state: StepRequestState,
stage_durations: dict[str, float],
) -> None:
if not stage_durations:
@@ -43,7 +43,7 @@ def merge_stage_durations(
def attach_stage_durations(
state: DiffusionRequestState,
state: StepRequestState,
output: DiffusionOutput,
) -> None:
if state.stage_durations:
@@ -51,7 +51,7 @@ def attach_stage_durations(
@dataclass
class DiffusionRequestState:
class StepRequestState:
"""Per-request mutable state across all pipeline stages.
Owned by Runner and passed through all step-execution stages:
vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
+5 -6
View File
@@ -15,6 +15,7 @@ import torch
from vllm.logger import init_logger
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import KVPrefetchJob
from vllm_omni.platforms import current_omni_platform
from .factory import OmniConnectorFactory
@@ -1205,16 +1206,14 @@ class OmniKVTransferManager:
logger.exception("Failed to release KV pool buffer")
buffers.clear()
def start_prefetch(
self, kv_prefetch_jobs: dict[str, Any] | None, target_device: torch.device | None = None
) -> None:
def start_prefetch(self, kv_prefetch_job: KVPrefetchJob | None, target_device: torch.device | None = None) -> None:
"""Kick off a background KV load (non-blocking). No-op unless prefetch enabled."""
if not (self._async_prefetch and self.config.need_recv_cache) or not kv_prefetch_jobs:
if not (self._async_prefetch and self.config.need_recv_cache) or not kv_prefetch_job:
return
# Followers receive via collective distribute; bg pull would consume the owner's payload.
if self.topo_config.is_follower:
return
rid = kv_prefetch_jobs.get("request_id")
rid = kv_prefetch_job.get("request_id")
if not rid:
return
# Memory pressure → skip prefetch; sync receive handles it.
@@ -1228,7 +1227,7 @@ class OmniKVTransferManager:
self._discard_future(stale_rid)
if rid in self._prefetch_futures:
return
sender_info = kv_prefetch_jobs.get("kv_sender_info")
sender_info = kv_prefetch_job.get("kv_sender_info")
if not sender_info:
# No explicit endpoint → bg receive would target wrong sender under multi-replica.
logger.debug("Skip KV prefetch for %s: stub has no kv_sender_info", rid)
vllm_omni/experimental/ar_diffusion/runner.py
+8 -3
View File
@@ -25,6 +25,7 @@ from vllm.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import MAX_DREAMZERO_SESSIONS
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import KVPrefetchJob
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
@@ -234,10 +235,14 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
num_frame_per_block=num_frame_per_block,
)
def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
def execute_model(
self,
req: OmniDiffusionRequest,
kv_prefetch_job: KVPrefetchJob | None = None,
) -> DiffusionOutput:
# KV disabled -> base behavior, unchanged.
if self.kv_cache is None:
return super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs)
return super().execute_model(req, kv_prefetch_job=kv_prefetch_job)
kv = self.kv_cache
# DreamZero KV is session-scoped (the model state persists across a
@@ -273,7 +278,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
# sync below makes the stop reflect completed GPU work, not just dispatch.
_e2e_t0 = time.perf_counter()
try:
out = super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs)
out = super().execute_model(req, kv_prefetch_job=kv_prefetch_job)
except Exception:
# Transactional containment: a forward that died partway may have
# written some layers' K/V into allocated-but-uncommitted blocks