[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/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.