From e39f4dddfb2319c1c19790ec4361a78cfb0534bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?=
 <11222265+fhfuih@users.noreply.github.com>
Date: Wed, 22 Jul 2026 14:05:49 +0800
Subject: [PATCH] [Refactor]: clean up diffusion scheduler interfaces (#5216)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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                            |   2 +-
 .../feature/diffusion_continuous_batching.md  |  11 +-
 .../diffusion_request_level_batching.md       |   4 +-
 .../feature/diffusion_step_execution.md       |   2 +-
 docs/design/module/dit_module.md              |  57 +++++----
 tests/ar_diffusion/test_engine_routing.py     |   6 +-
 tests/ar_diffusion/test_pipeline_bridge.py    |   2 +-
 .../test_hunyuan_image3_step_execution.py     |   8 +-
 .../test_diffusion_engine_rpc_routing.py      |   4 +-
 tests/diffusion/test_diffusion_scheduler.py   |  62 +++++-----
 .../diffusion/test_diffusion_step_pipeline.py |   6 +-
 .../test_input_batch_txt_seq_lens.py          |   6 +-
 tests/diffusion/test_worker_wrapper_base.py   |   2 +-
 vllm_omni/diffusion/diffusion_engine.py       |  10 +-
 .../diffusion/executor/multiproc_executor.py  |   2 +-
 .../models/helios/pipeline_helios.py          |  24 ++--
 .../hunyuan_image3/pipeline_hunyuan_image3.py |  40 +++----
 vllm_omni/diffusion/models/interface.py       |   8 +-
 .../models/qwen_image/pipeline_qwen_image.py  |  10 +-
 vllm_omni/diffusion/request.py                |   3 +-
 vllm_omni/diffusion/sched/__init__.py         |  11 +-
 vllm_omni/diffusion/sched/base_scheduler.py   |  89 ++++++++------
 vllm_omni/diffusion/sched/interface.py        | 109 +++++-------------
 .../diffusion/sched/request_scheduler.py      |  28 +++--
 vllm_omni/diffusion/sched/step_scheduler.py   |   9 +-
 .../worker/diffusion_model_runner.py          |  40 ++++---
 .../diffusion/worker/diffusion_worker.py      |  16 +--
 vllm_omni/diffusion/worker/input_batch.py     |  58 +++++-----
 vllm_omni/diffusion/worker/utils.py           |   6 +-
 .../omni_connectors/kv_transfer_manager.py    |  11 +-
 vllm_omni/experimental/ar_diffusion/runner.py |  11 +-
 31 files changed, 319 insertions(+), 338 deletions(-)

diff --git a/docs/api/README.md b/docs/api/README.md
index cfea6d175..09f6a6e24 100644
--- a/docs/api/README.md
+++ b/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][]
diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md
index f90969369..223b88518 100644
--- a/docs/design/feature/diffusion_continuous_batching.md
+++ b/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
diff --git a/docs/design/feature/diffusion_request_level_batching.md b/docs/design/feature/diffusion_request_level_batching.md
index 57287ada4..8c1cdbd33 100644
--- a/docs/design/feature/diffusion_request_level_batching.md
+++ b/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.
diff --git a/docs/design/feature/diffusion_step_execution.md b/docs/design/feature/diffusion_step_execution.md
index 797a1cb48..01936e223 100644
--- a/docs/design/feature/diffusion_step_execution.md
+++ b/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.
 
diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md
index 02d50181b..97d1db499 100644
--- a/docs/design/module/dit_module.md
+++ b/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)
diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py
index fb55e01e9..5972ffff4 100644
--- a/tests/ar_diffusion/test_engine_routing.py
+++ b/tests/ar_diffusion/test_engine_routing.py
@@ -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
diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py
index 28c3c959d..2c4284531 100644
--- a/tests/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/ar_diffusion/test_pipeline_bridge.py
@@ -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)
diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
index b4f1e7493..d9d310ee4 100644
--- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
+++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
@@ -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,
diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py
index 4b7d0efbb..8c4f08152 100644
--- a/tests/diffusion/test_diffusion_engine_rpc_routing.py
+++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py
@@ -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 ─────────────────────────
diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py
index 7d6d2d11b..68b237c67 100644
--- a/tests/diffusion/test_diffusion_scheduler.py
+++ b/tests/diffusion/test_diffusion_scheduler.py
@@ -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:
diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py
index 7e18996ed..0964833d7 100644
--- a/tests/diffusion/test_diffusion_step_pipeline.py
+++ b/tests/diffusion/test_diffusion_step_pipeline.py
@@ -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,
diff --git a/tests/diffusion/test_input_batch_txt_seq_lens.py b/tests/diffusion/test_input_batch_txt_seq_lens.py
index 315f0f703..88fb60a5e 100644
--- a/tests/diffusion/test_input_batch_txt_seq_lens.py
+++ b/tests/diffusion/test_input_batch_txt_seq_lens.py
@@ -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),
diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py
index ba22743fc..e9c4d613f 100644
--- a/tests/diffusion/test_worker_wrapper_base.py
+++ b/tests/diffusion/test_worker_wrapper_base.py
@@ -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):
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index d6e7b35de..e8ab578db 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -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``.
 
diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py
index d0ed76bb9..f90624ed7 100644
--- a/vllm_omni/diffusion/executor/multiproc_executor.py
+++ b/vllm_omni/diffusion/executor/multiproc_executor.py
@@ -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,
                 )
diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py
index 12ecda368..655ef8821 100644
--- a/vllm_omni/diffusion/models/helios/pipeline_helios.py
+++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py
@@ -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
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
index c0476b879..095a420c8 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
@@ -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")
diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py
index 657a45bf8..bd26c0395 100644
--- a/vllm_omni/diffusion/models/interface.py
+++ b/vllm_omni/diffusion/models/interface.py
@@ -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."""
         ...
 
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
index 474ea946c..0aeb8b85c 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
@@ -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*."""
diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py
index 3ad3f1541..0edb30d23 100644
--- a/vllm_omni/diffusion/request.py
+++ b/vllm_omni/diffusion/request.py
@@ -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."""
diff --git a/vllm_omni/diffusion/sched/__init__.py b/vllm_omni/diffusion/sched/__init__.py
index e02637338..29b9cc667 100644
--- a/vllm_omni/diffusion/sched/__init__.py
+++ b/vllm_omni/diffusion/sched/__init__.py
@@ -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",
diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py
index 9e4b6af5b..5ea91b167 100644
--- a/vllm_omni/diffusion/sched/base_scheduler.py
+++ b/vllm_omni/diffusion/sched/base_scheduler.py
@@ -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)
diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py
index 695641b5d..25ec722d5 100644
--- a/vllm_omni/diffusion/sched/interface.py
+++ b/vllm_omni/diffusion/sched/interface.py
@@ -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."""
diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py
index 9a02eb8f3..9e7442d3a 100644
--- a/vllm_omni/diffusion/sched/request_scheduler.py
+++ b/vllm_omni/diffusion/sched/request_scheduler.py
@@ -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
diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py
index 6cb6a9a28..a5a7a04ae 100644
--- a/vllm_omni/diffusion/sched/step_scheduler.py
+++ b/vllm_omni/diffusion/sched/step_scheduler.py
@@ -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:
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index cb16bdeac..4ac9afe7a 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -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,
     ):
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 006bbae48..1500ccd3f 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -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."""
diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py
index 511f2dd69..c0d10198a 100644
--- a/vllm_omni/diffusion/worker/input_batch.py
+++ b/vllm_omni/diffusion/worker/input_batch.py
@@ -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.
diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py
index ba6165a9d..8d85fae70 100644
--- a/vllm_omni/diffusion/worker/utils.py
+++ b/vllm_omni/diffusion/worker/utils.py
@@ -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:
diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
index 51e604336..b75a7089a 100644
--- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
+++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
@@ -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)
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 862b3d9e8..8c968f8ce 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -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