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

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>

docs/api/README.md

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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)