Merge branch 'main' into dev/rebase-v0.18.0

docs/.nav.yml

@@ -60,6 +60,7 @@ nav:
         - FP8: user_guide/diffusion/quantization/fp8.md
         - Int8: user_guide/diffusion/quantization/int8.md
         - GGUF: user_guide/diffusion/quantization/gguf.md
+      - Step Execution: user_guide/diffusion/step_execution.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
       - LoRA: user_guide/diffusion/lora.md
@@ -91,6 +92,7 @@ nav:
       - design/feature/cache_dit.md
       - design/feature/teacache.md
       - design/feature/async_chunk_design.md
+      - design/feature/diffusion_step_execution.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/contributing/model/adding_diffusion_model.md

@@ -739,6 +739,22 @@ See detailed guide: [How to add Sequence Parallel support](../../design/feature/
 omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
 ```
 
+### Step Execution
+
+See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
+
+Use this only when your pipeline can be split into stable request-scoped and
+step-scoped phases. The reference implementation is
+`QwenImagePipeline`, which maps its request-level `forward()` into:
+
+1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup.
+2. `denoise_step()` for one transformer/noise prediction.
+3. `step_scheduler()` for one scheduler update and `step_index` advance.
+4. `post_decode()` for the final VAE decode.
+
+Do not enable `step_execution=True` until those four methods are implemented
+and validated against the request-level path.
+
 ### Cache Acceleration
 
 #### TeaCache

docs/design/feature/diffusion_step_execution.md

@@ -0,0 +1,121 @@
+# Diffusion Step Execution
+
+This guide documents vLLM-Omni's stepwise diffusion contract for model authors
+and contributors implementing `step_execution=True` support for a diffusion
+pipeline.
+
+For end-user enablement, supported models, and current limitations, see
+[Step Execution](../../user_guide/diffusion/step_execution.md).
+
+## Current Support Scope
+
+`step_execution` is **not** a generic diffusion toggle. It only works for
+pipelines that implement the segmented stateful contract in
+[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
+
+Current in-tree support:
+
+| Pipeline | Example models | Step execution |
+|----------|----------------|----------------|
+| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
+| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
+
+Current engine/runtime limitations:
+
+- `StepScheduler` only schedules `batch_size=1`.
+- `cache_backend` is not supported in step mode.
+- Request-mode extras such as KV transfer are not wired into step mode yet.
+- Unsupported pipelines now fail early during model loading instead of failing on the first request.
+
+## Execution Contract
+
+Step mode is driven by four pipeline methods plus the shared mutable request
+state object:
+
+- `prepare_encode(state)`: one-time request preparation.
+- `denoise_step(state)`: compute the noise prediction for the current step.
+- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
+- `post_decode(state)`: decode the final output after denoising is complete.
+
+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
+`state.extra` for model-specific fields that do not justify extending the core
+dataclass.
+
+The worker-side step loop lives in
+[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
+
+1. `prepare_encode()` runs once for a new request.
+2. `denoise_step()` runs every scheduler tick.
+3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
+4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
+
+## Recommended Split
+
+When converting an existing request-level `forward()` pipeline, keep the split
+strict and mechanical:
+
+| Request-level phase | Stepwise method | What belongs there |
+|---------------------|-----------------|--------------------|
+| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
+| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
+| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
+| VAE decode / postprocess | `post_decode()` | Final decode only |
+
+Keep the stepwise path reusing the same helpers as the request-level path
+whenever possible. Reimplementing the denoise loop from scratch is the easiest
+way to introduce behavioral drift.
+
+## Qwen-Image Reference
+
+[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
+is the reference implementation and is split correctly for the current
+contract:
+
+- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
+  latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
+  with `forward()`.
+- `prepare_encode()` deep-copies `self.scheduler` **after**
+  `prepare_timesteps()` so request-specific scheduler state is isolated.
+- `denoise_step()` reuses `_build_denoise_kwargs()` plus
+  `predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
+  non-CFG behavior stay identical to the request-level path.
+- `step_scheduler()` only calls
+  `scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
+  and increments `state.step_index`.
+- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
+  the normal `forward()` path.
+
+That decomposition is the target pattern for future models.
+
+## Rules For New Pipelines
+
+- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
+  `state.scheduler` during `prepare_encode()`.
+- Do not mutate `state.step_index` inside `denoise_step()`. Only
+  `step_scheduler()` should advance the step.
+- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
+- If the request-level pipeline has condition latents, masks, or edit-specific
+  tensors, store them in `state` or `state.extra`, not in global pipeline
+  attributes.
+- Preserve CFG behavior by sharing the same helper path used by `forward()`.
+- Keep `post_decode()` equivalent to the tail of `forward()`.
+
+## Validation Checklist
+
+Before marking a pipeline as `supports_step_execution = True`, verify:
+
+- Stepwise output matches request-level output for the same seed and sampling params.
+- Per-request scheduler state is isolated across concurrent requests.
+- Abort during denoise does not leak cached state.
+- `step_index` reported by `RunnerOutput` matches the scheduler progress.
+- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
+
+## Related Files
+
+- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
+- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
+- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
+- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
+- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)

docs/user_guide/diffusion/step_execution.md

@@ -0,0 +1,61 @@
+# Step Execution
+
+Step execution is an opt-in diffusion execution mode enabled with
+`step_execution=True` when constructing `Omni`.
+
+It is not a generic diffusion toggle for every pipeline. Only pipelines that
+implement the stepwise contract support it today.
+
+## Quick Start
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    step_execution=True,
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
+```
+
+## Supported Pipelines
+
+| Pipeline | Example models | Step execution |
+|----------|----------------|----------------|
+| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
+| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
+
+## Current Limitations
+
+- `step_execution` currently supports `batch_size=1` only.
+- `cache_backend` is not supported together with step execution.
+- Unsupported pipelines fail early during model loading.
+- Request-mode extras such as KV transfer are not wired into step mode yet.
+
+## When To Use It
+
+Use step execution only when you specifically need the pipeline to run through
+its stepwise request state machine. For normal diffusion inference, leave it
+disabled unless your workflow depends on this mode.
+
+If you are looking for general diffusion speedups, see
+[Diffusion Acceleration Overview](../diffusion_acceleration.md).
+
+## Troubleshooting
+
+If model loading fails with a message mentioning `prepare_encode()`,
+`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected
+pipeline does not support step execution.
+
+## For Model Authors
+
+If you want to add step execution support to a new diffusion pipeline, see the
+implementation guide:
+[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).

tests/diffusion/test_diffusion_step_pipeline.py

@@ -0,0 +1,507 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for step-level diffusion runner and worker execution."""
+
+import os
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+from tests.utils import hardware_test
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    get_sp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.ipc import (
+    pack_diffusion_output_shm,
+    unpack_diffusion_output_shm,
+)
+from vllm_omni.diffusion.sched.interface import (
+    DiffusionRequestState as SchedulerRequestState,
+)
+from vllm_omni.diffusion.sched.interface import (
+    DiffusionSchedulerOutput,
+)
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+# ---------------------------------------------------------------------------
+# Helpers & fixtures
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def _noop_forward_context(*args, **kwargs):
+    del args, kwargs
+    yield
+
+
+def _update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for key, value in envs_dict.items():
+        os.environ[key] = value
+
+
+class _StepPipeline:
+    """Minimal pipeline stub that supports step-wise execution."""
+
+    supports_step_execution = True
+
+    def __init__(self):
+        self.prepare_calls = 0
+        self.denoise_calls = 0
+        self.scheduler_calls = 0
+        self.decode_calls = 0
+
+    def prepare_encode(self, state, **kwargs):
+        del kwargs
+        self.prepare_calls += 1
+        state.timesteps = [torch.tensor(10), torch.tensor(5)]
+        state.latents = torch.tensor([0.0])
+        return state
+
+    def denoise_step(self, state, **kwargs):
+        del state, kwargs
+        self.denoise_calls += 1
+        return torch.tensor([1.0])
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del noise_pred, kwargs
+        self.scheduler_calls += 1
+        state.step_index += 1
+
+    def post_decode(self, state, **kwargs):
+        del kwargs
+        self.decode_calls += 1
+        return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
+
+
+class _IdentityNoiseTransformer(torch.nn.Module):
+    def forward(self, x: torch.Tensor, **kwargs):
+        del kwargs
+        return (x,)
+
+
+class _AdditiveScheduler:
+    def step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, return_dict: bool = False):
+        del t, return_dict
+        return (latents + noise_pred,)
+
+
+class _DistributedStepPipeline(CFGParallelMixin):
+    supports_step_execution = True
+
+    def __init__(self, mode: str, device: torch.device):
+        self.mode = mode
+        self.device = device
+        self._interrupt = False
+        self.scheduler = _AdditiveScheduler()
+        self.transformer = _IdentityNoiseTransformer()
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def prepare_encode(self, state, **kwargs):
+        del kwargs
+        state.timesteps = [torch.tensor(1.0, device=self.device)]
+        state.latents = torch.ones((1, 1), device=self.device)
+        state.step_index = 0
+        state.scheduler = self.scheduler
+        state.do_true_cfg = self.mode == "cfg"
+        return state
+
+    def denoise_step(self, state, **kwargs):
+        del kwargs
+        if self.mode == "ulysses":
+            sp_group = get_sp_group().ulysses_group
+            seq_world_size = torch.distributed.get_world_size(sp_group)
+            input_tensor = torch.randn(1, 2, 2 * seq_world_size, 2, device=self.device)
+            original = input_tensor.clone()
+            intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, False)
+            output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, False)
+            torch.testing.assert_close(output, original, rtol=1e-5, atol=1e-5)
+            return torch.ones_like(state.latents)
+
+        if self.mode == "ring":
+            ring_group = get_sp_group().ring_group
+            rank = torch.distributed.get_rank(ring_group)
+            world_size = torch.distributed.get_world_size(ring_group)
+            comm = RingComm(ring_group)
+            input_tensor = torch.full((1, 2, 2), float(rank + 1), device=self.device)
+            recv_tensor = comm.send_recv(input_tensor)
+            comm.commit()
+            comm.wait()
+            expected = torch.full_like(recv_tensor, float(((rank - 1) % world_size) + 1))
+            torch.testing.assert_close(recv_tensor, expected, rtol=1e-5, atol=1e-5)
+            return torch.ones_like(state.latents)
+
+        positive_kwargs = {"x": state.latents + 1}
+        negative_kwargs = {"x": state.latents - 1}
+        return self.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=1.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del kwargs
+        if self.mode == "cfg":
+            state.latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred,
+                state.current_timestep,
+                state.latents,
+                do_true_cfg=True,
+                per_request_scheduler=state.scheduler,
+            )
+        else:
+            state.latents = state.latents + noise_pred
+        state.step_index += 1
+
+    def post_decode(self, state, **kwargs):
+        del kwargs
+        return DiffusionOutput(output=state.latents.detach().cpu())
+
+
+def _make_step_request(num_inference_steps: int = 2):
+    return SimpleNamespace(
+        prompts=["a prompt"],
+        request_ids=["req-1"],
+        sampling_params=SimpleNamespace(
+            generator=None,
+            seed=None,
+            generator_device=None,
+            num_inference_steps=num_inference_steps,
+        ),
+    )
+
+
+def _make_runner():
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.od_config = SimpleNamespace(
+        cache_backend=None,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    runner.device = torch.device("cpu")
+    runner.pipeline = _StepPipeline()
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.state_cache = {}
+    runner.kv_transfer_manager = SimpleNamespace()
+    return runner
+
+
+def _make_distributed_runner(mode: str, device: torch.device):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.od_config = SimpleNamespace(
+        cache_backend=None,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    runner.device = device
+    runner.pipeline = _DistributedStepPipeline(mode=mode, device=device)
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.state_cache = {}
+    runner.kv_transfer_manager = SimpleNamespace()
+    return runner
+
+
+def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
+    return DiffusionSchedulerOutput(
+        step_id=step_id,
+        req_states=[SchedulerRequestState(sched_req_id=sched_req_id, req=req)],
+        finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
+        num_running_reqs=1,
+        num_waiting_reqs=0,
+    )
+
+
+def _expected_output_for_mode(mode: str) -> torch.Tensor:
+    if mode == "cfg":
+        return torch.tensor([[3.0]])
+    return torch.tensor([[2.0]])
+
+
+def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master_port: str):
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    _update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": master_port,
+        }
+    )
+    model_runner_module.set_forward_context = _noop_forward_context
+
+    try:
+        init_distributed_environment()
+        if mode == "ulysses":
+            initialize_model_parallel(ulysses_degree=world_size)
+        elif mode == "ring":
+            initialize_model_parallel(ring_degree=world_size)
+        elif mode == "cfg":
+            initialize_model_parallel(cfg_parallel_size=world_size)
+        else:
+            raise ValueError(f"Unsupported distributed test mode: {mode}")
+
+        runner = _make_distributed_runner(mode, device)
+        output = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
+        )
+
+        assert output.finished is True
+        assert output.result is not None
+        torch.testing.assert_close(output.result.output, _expected_output_for_mode(mode), rtol=1e-5, atol=1e-5)
+        assert "req-1" not in runner.state_cache
+    finally:
+        destroy_distributed_env()
+
+
+# ---------------------------------------------------------------------------
+# Runner / Worker
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.cpu
+class TestRunner:
+    """DiffusionModelRunner.execute_stepwise"""
+
+    def test_completes_request_and_clears_state(self, monkeypatch):
+        runner = _make_runner()
+        req = _make_step_request()
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+        first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        assert first.req_id == "req-1"
+        assert first.step_index == 1
+        assert first.finished is False
+        assert first.result is None
+        assert "req-1" in runner.state_cache
+
+        second = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=1))
+        assert second.req_id == "req-1"
+        assert second.step_index == 2
+        assert second.finished is True
+        assert second.result is not None
+        assert second.result.error is None
+        assert torch.equal(second.result.output, torch.tensor([2.0]))
+        assert "req-1" not in runner.state_cache
+
+        assert runner.pipeline.prepare_calls == 1
+        assert runner.pipeline.denoise_calls == 2
+        assert runner.pipeline.scheduler_calls == 2
+        assert runner.pipeline.decode_calls == 1
+
+    def test_load_model_rejects_unsupported_step_execution(self, monkeypatch):
+        class _RequestOnlyPipeline:
+            pass
+
+        class _FakeLoader:
+            def __init__(self, *args, **kwargs):
+                del args, kwargs
+
+            def load_model(self, **kwargs):
+                del kwargs
+                return _RequestOnlyPipeline()
+
+        class _FakeProfiler:
+            consumed_memory = 0
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                del exc_type, exc, tb
+                return False
+
+        runner = object.__new__(DiffusionModelRunner)
+        runner.vllm_config = object()
+        runner.od_config = SimpleNamespace(
+            enable_cpu_offload=False,
+            enable_layerwise_offload=False,
+            enforce_eager=True,
+            cache_backend=None,
+            cache_config=None,
+            step_execution=True,
+            model_class_name="RequestOnlyPipeline",
+            parallel_config=SimpleNamespace(use_hsdp=False),
+        )
+        runner.device = torch.device("cpu")
+        runner.pipeline = None
+        runner.cache_backend = None
+        runner.offload_backend = None
+        runner.state_cache = {}
+        runner.kv_transfer_manager = SimpleNamespace()
+
+        monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader)
+        monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler)
+        monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
+        monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
+
+        with pytest.raises(ValueError, match="RequestOnlyPipeline"):
+            DiffusionModelRunner.load_model(runner)
+
+
+@pytest.mark.cpu
+class TestWorker:
+    """DiffusionWorker.execute_stepwise"""
+
+    def test_delegates_to_model_runner(self):
+        worker = object.__new__(DiffusionWorker)
+        expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
+        scheduler_output = SimpleNamespace(
+            req_states=[
+                SimpleNamespace(
+                    req=SimpleNamespace(
+                        sampling_params=SimpleNamespace(lora_request=None),
+                    )
+                )
+            ]
+        )
+        worker.lora_manager = None
+        worker.model_runner = SimpleNamespace(
+            execute_stepwise=lambda arg: expected if arg is scheduler_output else None
+        )
+
+        output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert output is expected
+
+    def test_clears_active_lora_before_stepwise_execution(self):
+        worker = object.__new__(DiffusionWorker)
+        scheduler_output = SimpleNamespace(
+            req_states=[
+                SimpleNamespace(
+                    req=SimpleNamespace(
+                        sampling_params=SimpleNamespace(lora_request=None),
+                    )
+                )
+            ]
+        )
+        calls: list[object | None] = []
+
+        class _FakeLoRAManager:
+            def set_active_adapter(self, adapter):
+                calls.append(adapter)
+
+        worker.lora_manager = _FakeLoRAManager()
+        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
+
+        DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert calls == [None]
+
+    def test_rejects_lora_requests_in_step_mode(self):
+        worker = object.__new__(DiffusionWorker)
+        scheduler_output = SimpleNamespace(
+            req_states=[
+                SimpleNamespace(
+                    req=SimpleNamespace(
+                        sampling_params=SimpleNamespace(lora_request=object()),
+                    )
+                )
+            ]
+        )
+        worker.lora_manager = None
+        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
+
+        with pytest.raises(ValueError, match="does not support LoRA"):
+            DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+
+@pytest.mark.cpu
+class TestIPC:
+    def test_pack_unpack_runner_output_shm(self):
+        tensor = torch.zeros(300_000, dtype=torch.float32)
+        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+
+        packed = pack_diffusion_output_shm(output)
+        assert isinstance(packed.result.output, dict)
+        assert packed.result.output["__tensor_shm__"] is True
+
+        unpacked = unpack_diffusion_output_shm(packed)
+        assert isinstance(unpacked.result.output, torch.Tensor)
+        torch.testing.assert_close(unpacked.result.output, tensor)
+
+
+@pytest.mark.cpu
+class TestSupportedPipelines:
+    """Step-execution protocol checks for supported pipelines."""
+
+    def test_qwen_image_supports_step_execution(self):
+        from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
+        from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline
+
+        # Avoid loading model weights; protocol membership depends on the class contract.
+        pipeline = object.__new__(QwenImagePipeline)
+
+        assert pipeline.supports_step_execution is True
+        assert supports_step_execution(pipeline) is True
+        assert isinstance(pipeline, SupportsStepExecution) is True
+
+
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_execute_stepwise_with_ulysses_parallel():
+    world_size = 2
+    if current_omni_platform.get_device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} devices")
+
+    torch.multiprocessing.spawn(
+        _distributed_step_worker,
+        args=(world_size, "ulysses", "29540"),
+        nprocs=world_size,
+    )
+
+
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_execute_stepwise_with_ring_parallel():
+    world_size = 2
+    if current_omni_platform.get_device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} devices")
+
+    torch.multiprocessing.spawn(
+        _distributed_step_worker,
+        args=(world_size, "ring", "29541"),
+        nprocs=world_size,
+    )
+
+
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_execute_stepwise_with_cfg_parallel():
+    world_size = 2
+    if current_omni_platform.get_device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} devices")
+
+    torch.multiprocessing.spawn(
+        _distributed_step_worker,
+        args=(world_size, "cfg", "29542"),
+        nprocs=world_size,
+    )

vllm_omni/diffusion/data.py

@@ -464,6 +464,9 @@ class OmniDiffusionConfig:
     # Diffusion pipeline Profiling config
     enable_diffusion_pipeline_profiler: bool = False
 
+    # Step mode settings
+    step_execution: bool = False
+
     @property
     def is_moe(self) -> bool:
         num_experts = self.tf_model_config.get("num_experts", None)

vllm_omni/diffusion/distributed/cfg_parallel.py

@@ -182,7 +182,13 @@ class CFGParallelMixin(metaclass=ABCMeta):
         """
         raise NotImplementedError("Subclasses must implement diffuse")
 
-    def scheduler_step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor) -> torch.Tensor:
+    def scheduler_step(
+        self,
+        noise_pred: torch.Tensor,
+        t: torch.Tensor,
+        latents: torch.Tensor,
+        per_request_scheduler: Any | None = None,
+    ) -> torch.Tensor:
         """
         Step the scheduler.
 
@@ -190,14 +196,30 @@ class CFGParallelMixin(metaclass=ABCMeta):
             noise_pred: Predicted noise
             t: Current timestep
             latents: Current latents
+            per_request_scheduler: Optional request-scoped scheduler that
+                overrides ``self.scheduler`` for this call. This is
+                primarily used by step-wise execution, where each request
+                may keep scheduler state in its own runner-managed state
+                object. Request-level execution should usually leave this
+                as ``None`` and continue using ``self.scheduler``.
 
         Returns:
             Updated latents after scheduler step
         """
-        return self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        sched = per_request_scheduler if per_request_scheduler is not None else getattr(self, "scheduler", None)
+        if sched is None:
+            raise ValueError("No scheduler is available. Set self.scheduler or pass per_request_scheduler.")
+        if not callable(getattr(sched, "step", None)):
+            raise TypeError("per_request_scheduler must provide a callable step(...) method.")
+        return sched.step(noise_pred, t, latents, return_dict=False)[0]
 
     def scheduler_step_maybe_with_cfg(
-        self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, do_true_cfg: bool
+        self,
+        noise_pred: torch.Tensor,
+        t: torch.Tensor,
+        latents: torch.Tensor,
+        do_true_cfg: bool,
+        per_request_scheduler: Any | None = None,
     ) -> torch.Tensor:
         """
         Step the scheduler with (maybe) automatic CFG parallel synchronization.
@@ -210,6 +232,11 @@ class CFGParallelMixin(metaclass=ABCMeta):
             t: Current timestep
             latents: Current latents
             do_true_cfg: Whether CFG is enabled
+            per_request_scheduler: Optional request-scoped scheduler that
+                overrides ``self.scheduler`` for this call. This is mainly
+                needed by step-wise execution, where scheduler state may be
+                stored per request. Request-level execution should normally
+                leave this as ``None``.
 
         Returns:
             Updated latents (synchronized across all CFG ranks)
@@ -223,13 +250,23 @@ class CFGParallelMixin(metaclass=ABCMeta):
 
             # Only rank 0 computes the scheduler step
             if cfg_rank == 0:
-                latents = self.scheduler_step(noise_pred, t, latents)
+                latents = self.scheduler_step(
+                    noise_pred,
+                    t,
+                    latents,
+                    per_request_scheduler=per_request_scheduler,
+                )
 
             # Broadcast the updated latents to all ranks
             latents = latents.contiguous()
             cfg_group.broadcast(latents, src=0)
         else:
             # No CFG parallel: directly compute scheduler step
-            latents = self.scheduler_step(noise_pred, t, latents)
+            latents = self.scheduler_step(
+                noise_pred,
+                t,
+                latents,
+                per_request_scheduler=per_request_scheduler,
+            )
 
         return latents

vllm_omni/diffusion/ipc.py

@@ -65,12 +65,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
     return tensor
 
 
-def pack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
-    """Replace large tensors in *output* with shared-memory handles.
-
-    The DiffusionOutput is modified **in-place** so that the (now lightweight)
-    object can be serialised cheaply through a MessageQueue.
-    """
+def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
     if output.output is not None and isinstance(output.output, torch.Tensor):
         if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD:
             output.output = _tensor_to_shm(output.output)
@@ -80,10 +75,35 @@ def pack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
     return output
 
 
-def unpack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
-    """Reconstruct tensors from shared-memory handles produced by ``pack_diffusion_output_shm``."""
+def pack_diffusion_output_shm(output: object) -> object:
+    """Replace large tensors in diffusion worker outputs with SHM handles.
+
+    Supports either a bare ``DiffusionOutput`` or a wrapper object carrying one
+    in ``.result`` (for example ``RunnerOutput``).
+    """
+    if isinstance(output, DiffusionOutput):
+        return _pack_diffusion_fields(output)
+
+    result = getattr(output, "result", None)
+    if isinstance(result, DiffusionOutput):
+        output.result = _pack_diffusion_fields(result)
+    return output
+
+
+def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
     if isinstance(output.output, dict) and output.output.get("__tensor_shm__"):
         output.output = _tensor_from_shm(output.output)
     if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
         output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
     return output
+
+
+def unpack_diffusion_output_shm(output: object) -> object:
+    """Reconstruct tensors from SHM handles in diffusion worker outputs."""
+    if isinstance(output, DiffusionOutput):
+        return _unpack_diffusion_fields(output)
+
+    result = getattr(output, "result", None)
+    if isinstance(result, DiffusionOutput):
+        output.result = _unpack_diffusion_fields(result)
+    return output

vllm_omni/diffusion/models/interface.py

@@ -1,11 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
 from typing import (
+    TYPE_CHECKING,
+    Any,
     ClassVar,
     Protocol,
     runtime_checkable,
 )
 
+if TYPE_CHECKING:
+    import torch
+
+    from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+
 
 @runtime_checkable
 class SupportImageInput(Protocol):
@@ -21,3 +31,34 @@ class SupportAudioInput(Protocol):
 @runtime_checkable
 class SupportAudioOutput(Protocol):
     support_audio_output: ClassVar[bool] = True
+
+
+@runtime_checkable
+class SupportsStepExecution(Protocol):
+    """State-driven step-level execution protocol for diffusion pipelines.
+
+    Pipelines should split request-level ``forward()`` into:
+    ``prepare_encode()`` (one-time request setup), ``denoise_step()``
+    (one denoise forward), ``step_scheduler()`` (one scheduler update),
+    and ``post_decode()`` (final decode).
+    """
+
+    supports_step_execution: ClassVar[bool] = True
+
+    def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState:
+        """Prepare request-level inputs and return initialized state."""
+
+    def denoise_step(self, state: DiffusionRequestState, **kwargs: Any) -> torch.Tensor | None:
+        """Run one denoise step."""
+
+    def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
+        """Run one scheduler step."""
+
+    def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput:
+        """Decode output after denoise loop."""
+
+
+def supports_step_execution(pipeline: object) -> bool:
+    """Return whether `pipeline` implements :class:`SupportsStepExecution`."""
+
+    return isinstance(pipeline, SupportsStepExecution)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -1,13 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import copy
 import inspect
 import json
 import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 import torch
@@ -35,6 +36,10 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -238,6 +243,8 @@ def apply_rotary_emb_qwen(
 
 
 class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin):
+    supports_step_execution: ClassVar[bool] = True
+
     def __init__(
         self,
         *,
@@ -539,6 +546,372 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
     def interrupt(self):
         return self._interrupt
 
+    def _extract_prompts(self, prompts):
+        """Extract prompt and negative_prompt from OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    def _prepare_generation_context(
+        self,
+        *,
+        prompt,
+        negative_prompt,
+        height,
+        width,
+        num_inference_steps,
+        sigmas,
+        guidance_scale,
+        num_images_per_prompt,
+        generator,
+        true_cfg_scale,
+        max_sequence_length,
+        prompt_embeds=None,
+        prompt_embeds_mask=None,
+        negative_prompt_embeds=None,
+        negative_prompt_embeds_mask=None,
+        latents=None,
+        attention_kwargs=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        """Shared preparation logic for forward() and prepare_encode().
+
+        Validates inputs, encodes prompts, prepares latents, computes timesteps,
+        and returns all intermediate values as a dict.
+        """
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_embeds_mask,
+            negative_prompt_embeds_mask,
+            callback_on_step_end_tensor_inputs,
+            max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs or {}
+        self._current_timestep = None
+        self._interrupt = False
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        elif prompt_embeds is not None:
+            batch_size = prompt_embeds.shape[0]
+        else:
+            batch_size = 1
+
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
+
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_embeds_mask = None
+
+        num_channels_latents = self.transformer.in_channels // 4
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+
+        img_shapes = [[(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]] * batch_size
+
+        timesteps, num_inference_steps = self.prepare_timesteps(
+            num_inference_steps,
+            sigmas,
+            latents.shape[1],
+        )
+        self._num_timesteps = len(timesteps)
+
+        if self.transformer.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        negative_txt_seq_lens = (
+            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+        )
+
+        return {
+            "prompt_embeds": prompt_embeds,
+            "prompt_embeds_mask": prompt_embeds_mask,
+            "negative_prompt_embeds": negative_prompt_embeds,
+            "negative_prompt_embeds_mask": negative_prompt_embeds_mask,
+            "latents": latents,
+            "img_shapes": img_shapes,
+            "timesteps": timesteps,
+            "do_true_cfg": do_true_cfg,
+            "guidance": guidance,
+            "txt_seq_lens": txt_seq_lens,
+            "negative_txt_seq_lens": negative_txt_seq_lens,
+        }
+
+    def prepare_encode(
+        self,
+        state: "DiffusionRequestState",
+        **kwargs: Any,
+    ) -> "DiffusionRequestState":
+        """Populate *state* with encoded prompts, latents, timesteps, and CFG config."""
+        sampling = state.sampling
+        prompt, negative_prompt = self._extract_prompts(state.prompts or [])
+
+        ctx = self._prepare_generation_context(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=sampling.height or self.default_sample_size * self.vae_scale_factor,
+            width=sampling.width or self.default_sample_size * self.vae_scale_factor,
+            num_inference_steps=sampling.num_inference_steps or 50,
+            sigmas=sampling.sigmas,
+            guidance_scale=sampling.guidance_scale if sampling.guidance_scale_provided else 1.0,
+            num_images_per_prompt=sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1,
+            generator=sampling.generator,
+            true_cfg_scale=sampling.true_cfg_scale or 4.0,
+            max_sequence_length=sampling.max_sequence_length or 512,
+            attention_kwargs=kwargs.get("attention_kwargs"),
+        )
+
+        # prepare_timesteps() has already materialized request-specific timestep
+        # state on self.scheduler, so deepcopy preserves dynamic-shifting state
+        # without replaying set_timesteps() on the per-request scheduler.
+        # Per-request scheduler (must not share state with self.scheduler)
+        req_scheduler = copy.deepcopy(self.scheduler)
+        req_scheduler.set_begin_index(0)
+
+        # Populate state from generation context
+        state.prompt_embeds = ctx["prompt_embeds"]
+        state.prompt_embeds_mask = ctx["prompt_embeds_mask"]
+        state.negative_prompt_embeds = ctx["negative_prompt_embeds"]
+        state.negative_prompt_embeds_mask = ctx["negative_prompt_embeds_mask"]
+        state.latents = ctx["latents"]
+        state.timesteps = ctx["timesteps"]
+        state.step_index = 0
+        state.scheduler = req_scheduler
+        state.do_true_cfg = ctx["do_true_cfg"]
+        state.guidance = ctx["guidance"]
+        state.img_shapes = ctx["img_shapes"]
+        state.txt_seq_lens = ctx["txt_seq_lens"]
+        state.negative_txt_seq_lens = ctx["negative_txt_seq_lens"]
+        # QwenImage always normalizes CFG output (matching forward())
+        state.sampling.cfg_normalize = True
+
+        return state
+
+    def _build_denoise_kwargs(
+        self,
+        latents: torch.Tensor,
+        timestep: torch.Tensor,
+        guidance: torch.Tensor | None,
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        img_shapes: list,
+        txt_seq_lens: list[int] | None,
+        do_true_cfg: bool,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds_mask: torch.Tensor | None,
+        negative_txt_seq_lens: list[int] | None,
+        image_latents: torch.Tensor | None = None,
+        extra_transformer_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any] | None, int | None]:
+        """Build positive/negative kwargs and output_slice for one denoise step.
+
+        Returns:
+            (positive_kwargs, negative_kwargs, output_slice)
+        """
+        extra_transformer_kwargs = extra_transformer_kwargs or {}
+
+        # Broadcast timestep to match batch size
+        t_for_model = timestep.expand(latents.shape[0]).to(
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+
+        # Concatenate image latents if available (editing pipelines)
+        latent_model_input = latents
+        if image_latents is not None:
+            latent_model_input = torch.cat([latents, image_latents], dim=1)
+
+        positive_kwargs = {
+            "hidden_states": latent_model_input,
+            "timestep": t_for_model / 1000,
+            "guidance": guidance,
+            "encoder_hidden_states_mask": prompt_embeds_mask,
+            "encoder_hidden_states": prompt_embeds,
+            "img_shapes": img_shapes,
+            "txt_seq_lens": txt_seq_lens,
+            **extra_transformer_kwargs,
+        }
+        if do_true_cfg:
+            negative_kwargs = {
+                "hidden_states": latent_model_input,
+                "timestep": t_for_model / 1000,
+                "guidance": guidance,
+                "encoder_hidden_states_mask": negative_prompt_embeds_mask,
+                "encoder_hidden_states": negative_prompt_embeds,
+                "img_shapes": img_shapes,
+                "txt_seq_lens": negative_txt_seq_lens,
+                **extra_transformer_kwargs,
+            }
+        else:
+            negative_kwargs = None
+
+        output_slice = latents.size(1) if image_latents is not None else None
+        return positive_kwargs, negative_kwargs, output_slice
+
+    def _decode_latents(
+        self,
+        latents: torch.Tensor,
+        height: int,
+        width: int,
+        output_type: str = "pil",
+    ) -> DiffusionOutput:
+        """Unpack, normalize, and VAE-decode latents into a DiffusionOutput."""
+        if output_type == "latent":
+            return DiffusionOutput(
+                output=latents,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            )
+
+        latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents = latents.to(self.vae.dtype)
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+            latents.device, latents.dtype
+        )
+        latents = latents / latents_std + latents_mean
+        image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+        return DiffusionOutput(
+            output=image,
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
+    def denoise_step(
+        self,
+        state: "DiffusionRequestState",
+        **kwargs: Any,
+    ) -> torch.Tensor | None:
+        """One denoise step: read from *state*, delegate to CFGParallelMixin.
+
+        Reuses ``predict_noise_maybe_with_cfg`` so that CFG-parallel,
+        sequential-CFG, and no-CFG paths are handled identically to
+        ``diffuse()``.
+        """
+        if self.interrupt:
+            return None
+
+        t = state.current_timestep
+        self._current_timestep = t
+        self.transformer.do_true_cfg = state.do_true_cfg
+
+        # Normalize timestep to [batch_size] tensor
+        if not torch.is_tensor(t):
+            t = torch.tensor(t, device=state.latents.device, dtype=state.latents.dtype)
+
+        positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs(
+            latents=state.latents,
+            timestep=t,
+            guidance=state.guidance,
+            prompt_embeds=state.prompt_embeds,
+            prompt_embeds_mask=state.prompt_embeds_mask,
+            img_shapes=state.img_shapes,
+            txt_seq_lens=state.txt_seq_lens,
+            do_true_cfg=state.do_true_cfg,
+            negative_prompt_embeds=state.negative_prompt_embeds,
+            negative_prompt_embeds_mask=state.negative_prompt_embeds_mask,
+            negative_txt_seq_lens=state.negative_txt_seq_lens,
+            image_latents=state.sampling.image_latent,
+            extra_transformer_kwargs={
+                "attention_kwargs": self.attention_kwargs,
+                "return_dict": False,
+            },
+        )
+
+        true_cfg_scale = state.sampling.true_cfg_scale or 4.0
+        cfg_normalize = state.sampling.cfg_normalize
+
+        return self.predict_noise_maybe_with_cfg(
+            state.do_true_cfg,
+            true_cfg_scale,
+            positive_kwargs,
+            negative_kwargs,
+            cfg_normalize,
+            output_slice,
+        )
+
+    def step_scheduler(
+        self,
+        state: "DiffusionRequestState",
+        noise_pred: torch.Tensor,
+        **kwargs: Any,
+    ) -> None:
+        """One scheduler step: update ``state.latents`` and advance ``step_index``."""
+        if self.interrupt:
+            return
+
+        t = state.current_timestep
+        state.latents = self.scheduler_step_maybe_with_cfg(
+            noise_pred,
+            t,
+            state.latents,
+            state.do_true_cfg,
+            per_request_scheduler=state.scheduler,
+        )
+
+        state.step_index += 1
+
+    def post_decode(
+        self,
+        state: "DiffusionRequestState",
+        **kwargs: Any,
+    ) -> DiffusionOutput:
+        """Decode final latents from *state*."""
+        self._current_timestep = None
+
+        height = state.sampling.height or self.default_sample_size * self.vae_scale_factor
+        width = state.sampling.width or self.default_sample_size * self.vae_scale_factor
+        output_type = kwargs.get("output_type", "pil")
+
+        return self._decode_latents(state.latents, height, width, output_type)
+
     def forward(
         self,
         req: OmniDiffusionRequest,
@@ -562,13 +935,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         callback_on_step_end_tensor_inputs: list[str] = ["latents"],
         max_sequence_length: int = 512,
     ) -> DiffusionOutput:
-        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
-        # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
-        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
-            negative_prompt = None
-        elif req.prompts:
-            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt
 
         height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
@@ -584,113 +952,40 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
             if req.sampling_params.num_outputs_per_prompt > 0
             else num_images_per_prompt
         )
-        # 1. check inputs
-        # 2. encode prompts
-        # 3. prepare latents and timesteps
-        # 4. diffusion process
-        # 5. decode latents
-        # 6. post-process outputs
-        self.check_inputs(
-            prompt,
-            height,
-            width,
-            negative_prompt,
-            prompt_embeds,
-            negative_prompt_embeds,
-            prompt_embeds_mask,
-            negative_prompt_embeds_mask,
-            callback_on_step_end_tensor_inputs,
-            max_sequence_length,
-        )
 
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._current_timestep = None
-        self._interrupt = False
-
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
-
-        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
-        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
-
-        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+        ctx = self._prepare_generation_context(
             prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            sigmas=sigmas,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+            generator=generator,
+            true_cfg_scale=true_cfg_scale,
+            max_sequence_length=max_sequence_length,
             prompt_embeds=prompt_embeds,
             prompt_embeds_mask=prompt_embeds_mask,
-            num_images_per_prompt=num_images_per_prompt,
-            max_sequence_length=max_sequence_length,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            latents=latents,
+            attention_kwargs=attention_kwargs,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
         )
-        if do_true_cfg:
-            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
-                prompt=negative_prompt,
-                prompt_embeds=negative_prompt_embeds,
-                prompt_embeds_mask=negative_prompt_embeds_mask,
-                num_images_per_prompt=num_images_per_prompt,
-                max_sequence_length=max_sequence_length,
-            )
-
-        num_channels_latents = self.transformer.in_channels // 4
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            self.device,
-            generator,
-            latents,
-        )
-        img_shapes = [
-            [
-                (
-                    1,
-                    height // self.vae_scale_factor // 2,
-                    width // self.vae_scale_factor // 2,
-                )
-            ]
-        ] * batch_size
-
-        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
-        # num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-        self._num_timesteps = len(timesteps)
-
-        # handle guidance
-        if self.transformer.guidance_embeds:
-            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
-            guidance = guidance.expand(latents.shape[0])
-        else:
-            guidance = None
-
-        if self.attention_kwargs is None:
-            self._attention_kwargs = {}
-
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
-        # print inputp params
 
         latents = self.diffuse(
-            prompt_embeds,
-            prompt_embeds_mask,
-            negative_prompt_embeds,
-            negative_prompt_embeds_mask,
-            latents,
-            img_shapes,
-            txt_seq_lens,
-            negative_txt_seq_lens,
-            timesteps,
-            do_true_cfg,
-            guidance,
+            ctx["prompt_embeds"],
+            ctx["prompt_embeds_mask"],
+            ctx["negative_prompt_embeds"],
+            ctx["negative_prompt_embeds_mask"],
+            ctx["latents"],
+            ctx["img_shapes"],
+            ctx["txt_seq_lens"],
+            ctx["negative_txt_seq_lens"],
+            ctx["timesteps"],
+            ctx["do_true_cfg"],
+            ctx["guidance"],
             true_cfg_scale,
             image_latents=None,
             cfg_normalize=True,
@@ -701,26 +996,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         )
 
         self._current_timestep = None
-        if output_type == "latent":
-            image = latents
-        else:
-            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
-            latents = latents.to(self.vae.dtype)
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-                latents.device, latents.dtype
-            )
-            latents = latents / latents_std + latents_mean
-            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
-            # processed_image = self.image_processor.postprocess(image, output_type=output_type)
-
-        return DiffusionOutput(
-            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
-        )
+        return self._decode_latents(latents, height, width, output_type)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)

vllm_omni/diffusion/sched/interface.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass
+
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+class DiffusionRequestStatus(enum.IntEnum):
+    """Request status tracked by diffusion scheduler."""
+
+    WAITING = enum.auto()
+    RUNNING = enum.auto()
+    PREEMPTED = enum.auto()
+
+    # if any status is after or equal to FINISHED_COMPLETED, it is considered finished
+    FINISHED_COMPLETED = enum.auto()
+    FINISHED_ABORTED = enum.auto()
+    FINISHED_ERROR = enum.auto()
+
+    @staticmethod
+    def is_finished(status: DiffusionRequestStatus) -> bool:
+        return status >= DiffusionRequestStatus.FINISHED_COMPLETED
+
+
+@dataclass
+class DiffusionRequestState:
+    """Scheduler-owned state for one queued OmniDiffusionRequest."""
+
+    sched_req_id: str
+    req: OmniDiffusionRequest
+    status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
+    error: str | None = None
+
+    def is_finished(self) -> bool:
+        return DiffusionRequestStatus.is_finished(self.status)
+
+
+@dataclass
+class DiffusionSchedulerOutput:
+    """Output of a single scheduling cycle.
+
+    Kept intentionally small so step-execution components can share a stable
+    transport shape while scheduler policy continues to evolve.
+    """
+
+    step_id: int
+    req_states: list[DiffusionRequestState]
+    finished_req_ids: set[str]
+    num_running_reqs: int
+    num_waiting_reqs: int

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -10,6 +10,7 @@ model-related operations.
 
 from __future__ import annotations
 
+import copy
 import time
 from collections.abc import Iterable
 from contextlib import nullcontext
@@ -26,9 +27,12 @@ from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+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.worker.utils import DiffusionRequestState, RunnerOutput
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.platforms import current_omni_platform
 
@@ -65,6 +69,9 @@ class DiffusionModelRunner:
         self.cache_backend = None
         self.offload_backend = None
 
+        # Cache for per-request stepwise state.
+        self.state_cache: dict[str, DiffusionRequestState] = {}
+
         # Initialize KV cache manager for connector management
         self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config)
 
@@ -138,6 +145,13 @@ class DiffusionModelRunner:
         )
         logger.info("Model runner: Model loaded successfully.")
 
+        if getattr(self.od_config, "step_execution", False) and not self.supports_step_mode():
+            raise ValueError(
+                "step_execution=True requires a pipeline implementing "
+                "prepare_encode(), denoise_step(), step_scheduler(), and post_decode(); "
+                f"{self.od_config.model_class_name} does not support that contract."
+            )
+
         # Apply CPU offloading
         self.offload_backend = get_offload_backend(self.od_config, device=self.device)
         if self.offload_backend is not None:
@@ -239,3 +253,106 @@ class DiffusionModelRunner:
                 cache_summary(self.pipeline, details=True)
 
             return output
+
+    # ------------------------------------------------------------------
+    # Step-wise execution
+    # ------------------------------------------------------------------
+
+    def supports_step_mode(self) -> bool:
+        """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) -> DiffusionRequestState:
+        """Step-before update: cleanup finished requests and get/create one running state."""
+        for req_id in scheduler_output.finished_req_ids:
+            self.state_cache.pop(req_id, None)
+
+        req_states = scheduler_output.req_states
+        if len(req_states) != 1:
+            raise ValueError(f"Step mode currently supports batch_size=1, but got {len(req_states)} req_states.")
+
+        # TODO: remove req state from SchedulerOutput
+        # Stepwise mode currently trusts runner-owned cached state more than
+        # re-validating scheduler-provided request content on every step.
+        sched_req_state = req_states[0]
+        req_id = sched_req_state.sched_req_id
+        if req_id in self.state_cache:
+            return self.state_cache[req_id]
+
+        req = sched_req_state.req
+        request_ids = req.request_ids or [req_id]
+        if len(request_ids) != len(req.prompts):
+            raise ValueError(
+                f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})"
+            )
+
+        state = DiffusionRequestState(
+            req_id=req_id,
+            sampling=copy.deepcopy(req.sampling_params),
+            prompts=req.prompts,
+        )
+        self.state_cache[req_id] = state
+        return state
+
+    def _update_states_after(self, state: DiffusionRequestState, finished: bool) -> None:
+        """Step-after update: clear cached state for completed request."""
+        if finished:
+            self.state_cache.pop(state.req_id, None)
+
+    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        """Execute one step for one scheduled request and return runner output."""
+        assert self.pipeline is not None, "Model not loaded. Call load_model() first."
+        if not self.supports_step_mode():
+            raise ValueError("Current pipeline does not support step execution.")
+        # Stepwise mode only supports the basic state-driven denoise path for now.
+        # Request-mode extras such as cache backends, KV transfer, editing inputs,
+        # and similar features are not supported here yet.
+        if self.od_config.cache_backend not in (None, "none"):
+            raise ValueError("Step mode does not support cache_backend yet.")
+
+        use_hsdp = self.od_config.parallel_config.use_hsdp
+        grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
+        with grad_context:
+            state = self._update_states(scheduler_output)
+
+            if state.new_request:
+                # TODO: support kv manager recv
+                # TODO: support cache backend
+                if state.sampling.generator is None and state.sampling.seed is not None:
+                    if state.sampling.generator_device is not None:
+                        gen_device = state.sampling.generator_device
+                    elif self.device.type == "cpu":
+                        gen_device = "cpu"
+                    else:
+                        gen_device = self.device
+                    state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
+
+            with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
+                # step0/new request: encode
+                if state.new_request:
+                    self.pipeline.prepare_encode(state)
+
+                noise_pred = self.pipeline.denoise_step(state)
+                finished = False
+
+                # In CFG parallel mode, only rank 0 gets the actual noise_pred; non-rank-0 workers receive None.
+                # A true interrupt (all ranks return None) is detected by checking self.pipeline.interrupt.
+                if noise_pred is None and getattr(self.pipeline, "interrupt", False):
+                    finished = True
+                    result = DiffusionOutput(error="stepwise denoise interrupted")
+                else:
+                    self.pipeline.step_scheduler(state, noise_pred)
+                    finished = state.denoise_completed
+                    if finished:
+                        result = self.pipeline.post_decode(state)
+                    else:
+                        result = None
+
+                self._update_states_after(state, finished)
+
+                return RunnerOutput(
+                    req_id=state.req_id,
+                    step_index=state.step_index,
+                    finished=finished,
+                    result=result,
+                )

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -38,7 +38,9 @@ from vllm_omni.diffusion.ipc import pack_diffusion_output_shm
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.profiler import CurrentProfiler
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.utils import RunnerOutput
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
@@ -194,6 +196,19 @@ class DiffusionWorker:
                 logger.warning("LoRA activation skipped: %s", exc)
         return self.model_runner.execute_model(req)
 
+    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        """Execute one diffusion step by delegating to the model runner."""
+        assert self.model_runner is not None, "Model runner not initialized"
+        if self.lora_manager is not None:
+            # Step mode does not support LoRA yet. Clear any previously active
+            # adapter first so worker-local LoRA state cannot leak in.
+            self.lora_manager.set_active_adapter(None)
+
+        if any(req_state.req.sampling_params.lora_request is not None for req_state in scheduler_output.req_states):
+            raise ValueError("Step mode does not support LoRA yet.")
+
+        return self.model_runner.execute_stepwise(scheduler_output)
+
     def load_weights(self, weights) -> set[str]:
         """Load weights by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"
@@ -376,7 +391,7 @@ class WorkerProc:
         )
         return wrapper
 
-    def return_result(self, output: DiffusionOutput):
+    def return_result(self, output: object):
         """Reply to client, only on rank 0."""
         if self.result_mq is not None:
             try:
@@ -623,6 +638,10 @@ class WorkerWrapperBase:
         """
         return self.worker.execute_model(reqs, od_config)
 
+    def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
+        """Execute one diffusion step."""
+        return self.worker.execute_stepwise(scheduler_output)
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """
         Load model weights.

vllm_omni/diffusion/worker/utils.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-request mutable state for step-wise diffusion execution."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+
+
+@dataclass
+class DiffusionRequestState:
+    """Per-request mutable state across all pipeline stages.
+
+    Owned by Runner and passed through all step-execution stages:
+    ``prepare_encode()`` initializes/updates fields, ``denoise_step()`` and
+    ``step_scheduler()`` mutate per-step fields, and ``post_decode()``
+    consumes final latents. This state object is also the cache unit for
+    future continuous batching.
+
+    This dataclass keeps only the minimal cross-model state required by the
+    step-execution contract. Pipeline-specific state should be stored in
+    ``extra`` and promoted here only when it becomes shared across models.
+
+    Examples:
+    - Wan-style pipelines may keep ``condition``, ``first_frame_mask``, or
+      ``image_embeds`` in ``extra``.
+    - Bagel-style pipelines may keep ``gen_context``,
+      ``cfg_text_context``, ``cfg_img_context``, or ``image_shape`` in
+      ``extra``.
+    """
+
+    # ── Identity / request-level inputs ──
+    req_id: str
+    sampling: OmniDiffusionSamplingParams
+    prompts: list[OmniPromptType] | None = None
+
+    # ── Encoded prompts (set once by prepare_encode) ──
+    prompt_embeds: torch.Tensor | None = None
+    prompt_embeds_mask: torch.Tensor | None = None
+    negative_prompt_embeds: torch.Tensor | None = None
+    negative_prompt_embeds_mask: torch.Tensor | None = None
+
+    # ── Latent state (mutated every step by step_scheduler) ──
+    latents: torch.Tensor | None = None
+
+    # ── Timestep schedule (set once by prepare_encode) ──
+    timesteps: torch.Tensor | list[torch.Tensor] | None = None
+    step_index: int = 0
+
+    # ── Per-request scheduler instance (set once by prepare_encode) ──
+    scheduler: Any | None = None
+
+    # ── CFG config (set once by prepare_encode) ──
+    do_true_cfg: bool = False
+    guidance: torch.Tensor | None = None
+
+    # ── Spatial / sequence metadata (set once by prepare_encode) ──
+    img_shapes: list | None = None
+    txt_seq_lens: list[int] | None = None
+    negative_txt_seq_lens: list[int] | None = None
+
+    # Pipeline-specific extras. Keep model-private fields here unless they
+    # become part of the shared step-execution contract.
+    # For example: Wan condition tensors / masks, or Bagel KV contexts.
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    # ── Properties ──
+
+    @property
+    def current_timestep(self) -> torch.Tensor | None:
+        if self.timesteps is None:
+            return None
+        if self.step_index >= self.total_steps:
+            return None
+        if isinstance(self.timesteps, torch.Tensor):
+            if self.timesteps.ndim == 0:
+                return self.timesteps
+            return self.timesteps[self.step_index]
+        return self.timesteps[self.step_index]
+
+    @property
+    def total_steps(self) -> int:
+        if self.timesteps is None:
+            return 0
+        if isinstance(self.timesteps, torch.Tensor):
+            if self.timesteps.ndim == 0:
+                return 1
+            return int(self.timesteps.shape[0])
+        return len(self.timesteps)
+
+    @property
+    def denoise_completed(self) -> bool:
+        total_steps = self.total_steps
+        if total_steps == 0:
+            return False
+        return self.step_index >= total_steps
+
+    @property
+    def new_request(self) -> bool:
+        # TODO: this is only an approximation for current stepwise mode.
+        # A real "new request" signal should eventually come from scheduler/runner state transitions.
+        return self.step_index == 0 or self.timesteps is None
+
+
+@dataclass
+class RunnerOutput:
+    """Output of a single denoising step for a request.
+
+    NOTE: `latents` may be None when returned through IPC to avoid
+    serialization overhead. The actual latents are kept in Worker's
+    _request_state_cache.
+    """
+
+    req_id: str
+    step_index: int | None = None
+    finished: bool = False
+    result: DiffusionOutput | None = None

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -159,6 +159,12 @@ class Qwen3OmniMoeForConditionalGeneration(
             # suppress tokens by setting their probability to ~1e-9 (finite very small)
             self.suppressed_tokens = self._get_talker_suppressed_tokens()
             self.requires_raw_input_tokens = True
+            # Keys that should stay on GPU in model_intermediate_buffer to avoid CPU↔GPU round-trips
+            self.gpu_resident_buffer_keys: set[str] = {
+                "last_talker_hidden",
+                "trailing_text_hidden",
+                "tts_pad_embed_projected",
+            }
 
         elif self.model_stage == "code2wav":
             self.enable_update_additional_information = True
@@ -223,14 +229,16 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     # ==================== Forward Pass ====================
     def _get_talker_suppressed_tokens(self):
-        return [
-            i
-            for i in range(
-                self.config.talker_config.text_config.vocab_size - 1024,
-                self.config.talker_config.text_config.vocab_size,
-            )
-            if i != self.config.talker_config.codec_eos_token_id
-        ]
+        """Return a boolean mask on GPU for suppressed token positions."""
+        vocab_size = self.config.talker_config.text_config.vocab_size
+        mask = torch.zeros(vocab_size, dtype=torch.bool)
+        start = vocab_size - 1024
+        eos_id = self.config.talker_config.codec_eos_token_id
+        for i in range(start, vocab_size):
+            if i != eos_id:
+                mask[i] = True
+        # Will be moved to the correct device on first use
+        return mask
 
     def get_mrope_input_positions(
         self,
@@ -578,7 +586,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         Postprocess the talker hidden states.
         """
         update_dict = {}
-        update_dict["last_talker_hidden"] = hidden_states[-1, :].detach().to("cpu").contiguous()
+        update_dict["last_talker_hidden"] = hidden_states[-1, :].detach()
         return update_dict
 
     def talker_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
@@ -632,9 +640,11 @@ class Qwen3OmniMoeForConditionalGeneration(
         if inputs_embeds.shape[-1] == 2048:
             inputs_embeds = self.text_projection(inputs_embeds)
         code_predictor_codes, summed_embeddings = self.talker.code_predictor_forward(
-            input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden
+            input_ids, inputs_embeds, last_talker_hidden=last_talker_hidden
         )
-        inputs_embeds = summed_embeddings.clone()
+        # summed_embeddings is [B, seq_len, H] (3D) while text_step is [B, H] (2D).
+        # Flatten to 2D first to avoid wrong broadcasting: [B,1,H]+[B,H] → [B,B,H]
+        inputs_embeds = summed_embeddings.reshape(-1, self.talker_config.text_config.hidden_size)
         inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size)
         return inputs_embeds, code_predictor_codes.squeeze(-1)
 
@@ -758,7 +768,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     # compatible with old shape [1,S,D]
                     rem_tail = trailing_text_hidden.squeeze(0)
                 if rem_tail.shape[0] > 0:
-                    update_dict["trailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous()
+                    update_dict["trailing_text_hidden"] = rem_tail.detach()
             # Also persist projected tts_pad for decode fallback if needed
             if isinstance(tts_pad_thinker, torch.Tensor):
                 pad_in = tts_pad_thinker
@@ -767,7 +777,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 if pad_in.ndim == 1:
                     pad_in = pad_in.view(1, 1, -1)
                 pad_proj = self.talker.text_projection(pad_in.to(self._module_device(self.talker)))
-                update_dict["tts_pad_embed_projected"] = pad_proj.detach().to("cpu").contiguous()
+                update_dict["tts_pad_embed_projected"] = pad_proj.detach()
         except Exception:
             pass
         self._talker_cache_thinker_decode_embeds(info_dict, update_dict)
@@ -926,7 +936,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:
                     use_vec = q_tail[0:1, :]
                     new_q_tail = (
-                        q_tail[1:, :].detach().to("cpu").contiguous()
+                        q_tail[1:, :].detach()
                         if q_tail.shape[0] > 1
                         else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
                     )
@@ -1122,16 +1132,10 @@ class Qwen3OmniMoeForConditionalGeneration(
         # implemented by assigning their logits to log(1e-9).
 
         if getattr(self, "model_stage", None) == "talker" and isinstance(logits, torch.Tensor):
-            try:
-                logits_cpu = logits.cpu()
-                logits_cpu[:, self.suppressed_tokens] = -1e9
-                logits = logits_cpu.to(logits.device)
-            except Exception as e:
-                print(f"Error in logits suppression: {e}")
-                print(f"logits.shape: {logits.shape}")
-                print(f"self.suppressed_tokens: {self.suppressed_tokens}")
-                raise e
-            logits[:, self.suppressed_tokens] = -1e9
+            # Move mask to device once (lazy), then reuse every step
+            if self.suppressed_tokens.device != logits.device:
+                self.suppressed_tokens = self.suppressed_tokens.to(logits.device)
+            logits.masked_fill_(self.suppressed_tokens.unsqueeze(0), -1e9)
         return logits
 
     def sample(

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py

@@ -1,27 +1,16 @@
-"""Qwen3-Omni Code Predictor with MTP (Multi-Token Prediction) support.
+"""Qwen3-Omni Code Predictor -- optimized re-prefill, no KV cache.
 
-This module implements the code predictor component for Qwen3-Omni talker models.
-
-The code predictor generates residual RVQ (Residual Vector Quantization) codes
-autoregressively, predicting layers 1 to N based on layer-0 codes from the talker.
+* SDPA attention (F.scaled_dot_product_attention) with native GQA support
+* Per-call embedding buffer to avoid cross-request aliasing
+* Pre-allocated position_ids (read-only, safe to persist)
+* torch.compile on inner transformer by default
+* Inline sampling (top-k + top-p) -- no custom op overhead
 """
 
-from collections import namedtuple
-from typing import Any
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import Cache, PretrainedConfig
-from transformers.generation.logits_process import (
-    LogitsProcessorList,
-    TopKLogitsWarper,
-    TopPLogitsWarper,
-)
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
-from vllm.compilation.decorators import support_torch_compile
-from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
-from vllm.forward_context import get_forward_context
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -32,47 +21,53 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from vllm.utils.torch_utils import direct_register_custom_op
+
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
-# ============================================================================
-# Code Predictor Attention Layer
-# ============================================================================
+
+# ===================================================================
+#  Standalone Attention (SDPA, no KV cache, no HF backend fallback)
+# ===================================================================
 
 
 class Qwen3OmniCodePredictorAttention(nn.Module):
-    """Multi-head self-attention for code predictor with vLLM optimization."""
+    """Multi-head self-attention for code predictor.
+
+    Uses ``F.scaled_dot_product_attention`` directly.  No KV cache -- the code
+    predictor always re-prefills the full (short) sequence each AR step.
+
+    Input : [B, seq_len, hidden_size]
+    Output: [B, seq_len, hidden_size]
+    """
 
     def __init__(
         self,
         config,
-        layer_idx: int,
-        vllm_config: VllmConfig = None,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
         super().__init__()
-
-        self.num_heads = config.code_predictor_config.num_attention_heads
-        self.num_key_value_heads = config.code_predictor_config.num_key_value_heads
+        cp_cfg = config.code_predictor_config
+        self.num_heads = cp_cfg.num_attention_heads
+        self.num_kv_heads = cp_cfg.num_key_value_heads
         self.head_dim = getattr(
-            config.code_predictor_config,
+            cp_cfg,
             "head_dim",
-            config.code_predictor_config.hidden_size // config.code_predictor_config.num_attention_heads,
+            cp_cfg.hidden_size // cp_cfg.num_attention_heads,
         )
-        self.hidden_size = config.code_predictor_config.hidden_size
-
-        if self.num_heads % self.num_key_value_heads != 0:
-            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
-
-        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.hidden_size = cp_cfg.hidden_size
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self._use_gqa = self.num_kv_heads != self.num_heads
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size=self.hidden_size,
             head_size=self.head_dim,
             total_num_heads=self.num_heads,
-            total_num_kv_heads=self.num_key_value_heads,
+            total_num_kv_heads=self.num_kv_heads,
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
@@ -88,120 +83,55 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
         )
         self.rotary_emb = get_rope(
             self.head_dim,
-            max_position=config.code_predictor_config.max_position_embeddings,
+            max_position=cp_cfg.max_position_embeddings,
             rope_parameters=None,
             dual_chunk_attention_config=None,
         )
-
-        self.q_size = self.num_heads * self.head_dim
-        self.kv_size = self.num_key_value_heads * self.head_dim
-
-        # Query/Key normalization
-        self.q_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps)
-        self.k_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps)
-        self.is_causal = True
-        self.config = config
-
-        self.attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
-        cudagraph_mode = get_current_vllm_config().compilation_config.cudagraph_mode
-        if "flash_attention_2" in ALL_ATTENTION_FUNCTIONS and cudagraph_mode.has_full_cudagraphs():
-            logger.warning(
-                f"CUDAGraphMode.{cudagraph_mode.name} is currently not supported "
-                f"with flash attention for Qwen3-Omni talker MTP."
-                f"removing flash attention from attention_backends"
-            )
-            self.attention_backends.remove("flash_attention_2")
+        self.q_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
 
     def forward(
         self,
         hidden_states: torch.Tensor,
-        past_key_values: Cache | None = None,
-        cache_position: torch.LongTensor | None = None,
-        use_cache: bool = False,
-        position_ids: torch.LongTensor | None = None,
+        position_ids: torch.Tensor,
     ) -> torch.Tensor:
         bsz, seq_len, _ = hidden_states.shape
 
-        qkv, _ = self.qkv_proj(hidden_states)
+        # Flatten to 2-D so vLLM rotary_emb gets [num_tokens, size]
+        qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1))
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
 
-        # Reshape for attention
-        q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim)
-        k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
-        v = v.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
-
-        # Apply normalization
-        q = self.q_norm(q).contiguous()
-        k = self.k_norm(k).contiguous()
-        q = q.reshape(-1, self.q_size)
-        k = k.reshape(-1, self.kv_size)
-
-        # Apply RoPE
+        # QK-norm -> RoPE (both 2-D)
+        q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
+        k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
         q, k = self.rotary_emb(position_ids, q, k)
 
-        # Reshape for attention
-        q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim)
-        k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+        # [B, heads, seq, head_dim]
+        q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
 
-        v_heads = v.transpose(1, 2).contiguous()
-        q_heads = q.transpose(1, 2).contiguous()
-        k_heads = k.transpose(1, 2).contiguous()
+        attn_out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            scale=self.scaling,
+            is_causal=True,
+            enable_gqa=self._use_gqa,
+        )
 
-        if past_key_values is not None:
-            sin, cos = self.rotary_emb.get_cos_sin(seq_len)
-            # sin and cos are specific to RoPE models; cache_position needed for the static cache
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs)
-
-        # Try attention backends in order of preference, with runtime error handling
-        # This handles cases where the backend is registered but not actually available
-        attn_output = None
-        last_error = None
-
-        for backend_name in self.attention_backends:
-            if backend_name not in ALL_ATTENTION_FUNCTIONS:
-                continue
-
-            try:
-                attention_interface = ALL_ATTENTION_FUNCTIONS[backend_name]
-                attn_output, _ = attention_interface(
-                    self,
-                    q_heads,
-                    k_heads,
-                    v_heads,
-                    None,
-                    dropout=0.0 if not self.training else getattr(self, "attention_dropout", 0.0),
-                    scaling=self.head_dim**-0.5,
-                    sliding_window=None,
-                    use_cache=use_cache,
-                    position_ids=position_ids[:seq_len].unsqueeze(0),
-                    output_hidden_states=True,
-                    output_attentions=False,
-                )
-                break
-            except (ValueError, ImportError, RuntimeError, AttributeError) as e:
-                # Store error and try next backend
-                last_error = e
-                continue
-
-        if attn_output is None:
-            raise RuntimeError(
-                f"All attention backends failed. Last error: {last_error}. "
-                "Please install flash-attn, or ensure PyTorch's scaled_dot_product_attention is available."
-            )
-        attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous()
-
-        attn_output, _ = self.o_proj(attn_output)
-        return attn_output
+        attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1)
+        output, _ = self.o_proj(attn_out)
+        return output.view(bsz, seq_len, -1)
 
 
-# ============================================================================
-# Code Predictor MLP Layer
-# ============================================================================
+# ===================================================================
+#  MLP
+# ===================================================================
 
 
 class Qwen3OmniCodePredictorMLP(nn.Module):
-    """Feed-forward network for code predictor with fused gate/up projection."""
+    """SiLU-gated MLP for code predictor."""
 
     def __init__(
         self,
@@ -221,7 +151,6 @@ class Qwen3OmniCodePredictorMLP(nn.Module):
             prefix=f"{prefix}.gate_up_proj",
             disable_tp=True,
         )
-
         self.down_proj = RowParallelLinear(
             input_size=intermediate_size,
             output_size=hidden_size,
@@ -238,35 +167,23 @@ class Qwen3OmniCodePredictorMLP(nn.Module):
         return down
 
 
-# ============================================================================
-# MTP Layer (Multi-Token Prediction Layer)
-# ============================================================================
+# ===================================================================
+#  Decoder Layer
+# ===================================================================
 
 
-class Qwen3OmniCodePredictorMTPLayer(nn.Module):
-    """MTP layer for speculative decoding - predicts next residual code layer."""
+class Qwen3OmniCodePredictorDecoderLayer(nn.Module):
+    """Transformer decoder layer (SDPA, no KV cache)."""
 
     def __init__(
         self,
-        config: PretrainedConfig,
-        prefix: str,
-        model_config: ModelConfig,
-        layer_idx: int,
-        cache_config: CacheConfig | None = None,
+        config,
         quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
-        self.layer_idx = layer_idx
-        self.config = config
-
         self.self_attn = Qwen3OmniCodePredictorAttention(
             config,
-            layer_idx,
-            vllm_config=type(
-                "VllmConfig",
-                (),
-                {"cache_config": cache_config, "quant_config": quant_config, "model_config": model_config},
-            )(),
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
         )
@@ -275,260 +192,258 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.mlp",
         )
-        self.input_layernorm = RMSNorm(
-            config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps
-        )
-        self.post_attention_layernorm = RMSNorm(
-            config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps
-        )
+        cp_cfg = config.code_predictor_config
+        self.input_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
 
-    def mtp_block(
+    def forward(
         self,
         hidden_states: torch.Tensor,
-        past_key_values: Cache | None = None,
-        cache_position: torch.LongTensor | None = None,
-        use_cache: bool = False,
-        position_ids: torch.LongTensor | None = None,
+        position_ids: torch.Tensor,
     ) -> torch.Tensor:
-        # Self-attention with residual
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(hidden_states, past_key_values, cache_position, use_cache, position_ids)
+        hidden_states = self.self_attn(hidden_states, position_ids)
         hidden_states = residual + hidden_states
 
-        # MLP with residual
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
-
         return hidden_states
 
 
-class Qwen3OmniCodePredictorBaseModel(nn.Module):
-    """
-    Base model for code predictor - matches HF Qwen3OmniMoeTalkerCodePredictorModel structure.
+# ===================================================================
+#  Base Transformer Model (re-prefill, no KV cache)
+# ===================================================================
 
-    This is a simple transformer that processes inputs_embeds and outputs hidden states.
+
+class Qwen3OmniCodePredictorBaseModel(nn.Module):
+    """Inner transformer for code predictor.
+
+    Signature: ``forward(inputs_embeds, position_ids) -> hidden_states``
+    -- plain Tensor in, plain Tensor out (no namedtuple).
     """
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         config = vllm_config.model_config.hf_config.code_predictor_config
-
         self.config = config
-        self.vocab_size = config.vocab_size
-        self.num_code_groups = config.num_code_groups
 
-        # Codec embeddings (for layers 1-num_code_groups-1)
         self.codec_embedding = nn.ModuleList(
-            [
-                VocabParallelEmbedding(
-                    config.vocab_size,
-                    config.hidden_size,
-                )
-                for _ in range(config.num_code_groups - 1)
-            ]
+            [VocabParallelEmbedding(config.vocab_size, config.hidden_size) for _ in range(config.num_code_groups - 1)]
         )
 
-        # Decoder layers
         self.layers = nn.ModuleList(
             [
-                Qwen3OmniCodePredictorMTPLayer(
+                Qwen3OmniCodePredictorDecoderLayer(
                     vllm_config.model_config.hf_config,
-                    f"{prefix}.layers.{idx}",
-                    model_config=vllm_config.model_config,
-                    layer_idx=idx,
-                    cache_config=vllm_config.cache_config,
                     quant_config=vllm_config.quant_config,
+                    prefix=f"{prefix}.layers.{idx}",
                 )
                 for idx in range(config.num_hidden_layers)
             ]
         )
-
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
         self,
         inputs_embeds: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: Any | None = None,
-        use_cache: bool | None = False,
-        cache_position: torch.LongTensor | None = None,
-        **kwargs: Any,
-    ) -> Any:
-        """
-        Forward pass matching HF structure.
-
-        Args:
-            inputs_embeds: [batch, seq_len, hidden_size]
-            position_ids: Optional position IDs tensor
-            past_key_values: Optional cached key-value pairs
-            use_cache: Whether to use cache
-            cache_position: Optional cache position tensor
-            **kwargs: Additional keyword arguments
-
-        Returns:
-            Named tuple with .last_hidden_state and .past_key_values attributes
-        """
-        batch_size, seq_len, _ = inputs_embeds.shape
-        # Forward through decoder layers
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
         hidden_states = inputs_embeds
-
         for layer in self.layers:
-            hidden_states = layer.mtp_block(hidden_states, past_key_values, cache_position, use_cache, position_ids)
-
-        # Final norm
+            hidden_states = layer(hidden_states, position_ids)
         hidden_states = self.norm(hidden_states)
-
-        # Return in HF-compatible format
-        Output = namedtuple("Output", ["last_hidden_state", "past_key_values"])
-        return Output(last_hidden_state=hidden_states, past_key_values=None)  # [batch, num_code_groups-1, hidden_size]
-
-    def get_input_embeddings(self):
-        """Return codec embeddings for HF compatibility."""
-        return self.codec_embedding
+        return hidden_states
 
 
-def code_predictor_sample(
-    logits: torch.Tensor,
-    layer_name: str,
-) -> torch.Tensor:
-    forward_context = get_forward_context()
-    self = forward_context.no_compile_layers[layer_name]
-    logits = self.logits_processors(None, logits[:, -1])
-    probs = F.softmax(logits, dim=-1)
-    code = torch.multinomial(probs.squeeze(1), num_samples=1)  # [batch, 1]
-    return code
+# ===================================================================
+#  Code Predictor Wrapper (optimized re-prefill, persistent buffers)
+# ===================================================================
 
 
-def code_predictor_sample_fake(
-    logits: torch.Tensor,
-    layer_name: str,
-) -> torch.Tensor:
-    return torch.empty((logits.shape[0], 1), dtype=torch.int64, device=logits.device)
-
-
-direct_register_custom_op(
-    op_name="qwen3_omni_code_predictor_sample",
-    op_func=code_predictor_sample,
-    fake_impl=code_predictor_sample_fake,
-)
-
-
-@support_torch_compile
 class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
-    """
-    Code predictor wrapper matching HF structure.
+    """Optimized code predictor -- re-prefill approach, no KV cache.
 
-    Structure:
-    - self.model: Qwen3OmniCodePredictorBaseModel (transformer)
-    - self.lm_head: ModuleList of output heads
+    Each AR step forwards the full growing sequence (len 2 -> num_code_groups+1)
+    through the transformer.  The extra O(T^2) FLOPs are negligible for
+    short sequences, and this avoids all KV-cache management overhead.
+
+    Optimizations:
+      1. Per-call embedding buffer -- avoids cross-request aliasing.
+      2. Pre-allocated position_ids -- no torch.arange per step.
+      3. Cached module references -- bypass ModuleList indexing.
+      4. torch.compile on inner transformer.
+      5. Inline sampling (top-k + top-p) -- no custom op overhead.
     """
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
-        talker_code_predictor_config = vllm_config.model_config.hf_config
+        config = vllm_config.model_config.hf_config
+        self.config = config
         self.quant_config = vllm_config.quant_config
         self.prefix = prefix
 
-        self.config = talker_code_predictor_config
-        self.vocab_size = self.config.code_predictor_config.vocab_size
-        self.num_code_groups = self.config.code_predictor_config.num_code_groups
+        self.num_code_groups = config.code_predictor_config.num_code_groups
+        self._hidden_size = config.code_predictor_config.hidden_size
 
-        # Base transformer model (matches HF structure)
-        self.model = Qwen3OmniCodePredictorBaseModel(vllm_config=vllm_config, prefix=prefix)
+        self.model = Qwen3OmniCodePredictorBaseModel(
+            vllm_config=vllm_config,
+            prefix=prefix,
+        )
 
-        # Output heads for each residual layer (1-num_layers-1)
+        # One lm_head per residual layer (layers 1 .. G-1)
         self.lm_head = nn.ModuleList(
             [
                 nn.Linear(
-                    self.config.code_predictor_config.hidden_size,
-                    self.config.code_predictor_config.vocab_size,
+                    config.code_predictor_config.hidden_size,
+                    config.code_predictor_config.vocab_size,
                     bias=False,
                 )
                 for _ in range(self.num_code_groups - 1)
             ]
         )
-        self.logits_processors = LogitsProcessorList(
-            [
-                TopKLogitsWarper(top_k=50),
-                TopPLogitsWarper(top_p=0.8),
-            ]
-        )
 
-        compilation_config = get_current_vllm_config().compilation_config
-        if prefix in compilation_config.static_forward_context:
-            raise ValueError(f"Duplicate layer name: {prefix}")
-        compilation_config.static_forward_context[prefix] = self
-        self.layer_name = prefix
+        self.set_sampling_params()
 
+        # Lazily initialised position ids (read-only, safe to persist)
+        self._pos_ids: torch.Tensor | None = None
+
+        # Cached plain-list refs (set once)
+        self._lm_heads: list | None = None
+        self._codec_embeds: list | None = None
+
+        # Model forward (optionally compiled)
+        self._model_fwd: object | None = None
+
+    def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8):
+        """Configure sampling parameters to maintain consistency with previous implementation."""
+        self._top_k = top_k
+        self._top_p = top_p
+        logger.debug(f"Sampling parameters updated: top_k={top_k}, top_p={top_p}s")
+
+    # ------------------------------------------------------------------
+    #  Lazy-init helpers
+    # ------------------------------------------------------------------
+
+    def _ensure_pos_ids(self, device: torch.device) -> None:
+        if self._pos_ids is not None and self._pos_ids.device == device:
+            return
+        max_seq = self.num_code_groups + 1
+        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device)
+
+    def _ensure_cached_refs(self) -> None:
+        if self._lm_heads is not None:
+            return
+        self._lm_heads = list(self.lm_head)
+        self._codec_embeds = list(self.model.codec_embedding)
+
+    def _ensure_model_fwd(self) -> None:
+        if self._model_fwd is not None:
+            return
+        if current_omni_platform.supports_torch_inductor():
+            # torch.compile fuses the 5-layer transformer's small kernels,
+            # reducing BF16 intermediate round-trips and improving precision.
+            # mode="default" avoids Inductor's own CUDA graph capture so it
+            # doesn't conflict with the outer CUDAGraphWrapper.
+            self._model_fwd = torch.compile(
+                self.model.forward,
+                mode="default",
+                dynamic=True,
+            )
+            logger.info("code_predictor: torch.compile enabled (mode=default)")
+        else:
+            self._model_fwd = self.model.forward
+            logger.info("code_predictor: using eager mode (no torch.compile)")
+
+    # ------------------------------------------------------------------
+    #  Forward -- re-prefill + inline sampling
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
     def forward(
         self,
         layer0_code: torch.Tensor,
         layer0_embed: torch.Tensor,
         last_talker_hidden: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass for code predictor.
+        """Predict residual codebooks 1..G-1 autoregressively via re-prefill.
 
         Args:
-            layer0_code:
-                Code index for code-group (layer) 0.
-                Shape: [batch_size, 1], dtype typically int64.
-
-            last_talker_hidden:
-
-                Shape: [batch_size, hidden_size].
+            layer0_code:        [bsz, 1]  int64
+            layer0_embed:       [bsz, 1, hidden_size]
+            last_talker_hidden: [bsz, 1, hidden_size]
 
         Returns:
-            pos_all_layers:
-                Predicted codes for all code groups, including `layer0_code`.
-                Shape: [batch_size, num_code_groups, 1].
-
-            current_input:
-                The final input embedding sequence after appending embeddings of all
-                predicted codes (one token per predicted layer).
-                Shape: [batch_size, num_code_groups + 2, hidden_size].
+            all_codes: [bsz, num_code_groups, 1]
+            proj_buf:  [bsz, num_code_groups + 1, hidden_size]
+                pos 0   = last_talker_hidden (NOT a codec embed)
+                pos 1   = layer0_embed
+                pos 2.. = codec_embedding[i](predicted_code_i)
         """
-        pos_codes = [layer0_code]  # Start with layer 0: [batch, 1]
-        try:
-            current_input = torch.cat([last_talker_hidden, layer0_embed], dim=1)  # [batch, 2, hidden_size]
-        except Exception as e:
-            print(f"Error in current_input: {e}")
-            print(f"last_talker_hidden shape: {last_talker_hidden.shape}")
-            print(f"prev_embed shape: {layer0_embed.shape}")
-            raise e
-        batch_size = current_input.shape[0]
+        bsz = int(layer0_code.shape[0])
+        device = layer0_code.device
+        dtype = last_talker_hidden.dtype
+        num_groups = self.num_code_groups
 
-        # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
-        for layer_idx in range(self.num_code_groups - 1):
-            seq_len = layer_idx + 2
-            # Compute position_ids dynamically to avoid torch.compile specializing batch_size
-            position_ids = torch.arange(seq_len, device=current_input.device, dtype=torch.int64).repeat(batch_size)
-            # Forward through code_predictor model
-            outputs = self.model(
-                inputs_embeds=current_input,
-                attention_mask=None,
-                position_ids=position_ids,
-                past_key_values=None,
-                use_cache=False,
-                cache_position=None,
-            )
-            hidden_state = outputs.last_hidden_state  # [batch, 2, hidden_size]
+        # Lazy init (read-only caches only)
+        self._ensure_pos_ids(device)
+        self._ensure_model_fwd()
+        self._ensure_cached_refs()
 
-            # Use the corresponding lm_head for this layer
-            logits = self.lm_head[layer_idx](hidden_state[:, -1:, :])
-            code = torch.ops.vllm.qwen3_omni_code_predictor_sample(logits, self.layer_name)
-            pos_codes.append(code)
-            # Update prev_embed for next layer (if not last layer)
-            # layer_idx=0 predicts layer 1, embed with codec_embedding[1]
-            new_embed = self.model.codec_embedding[layer_idx](code)  # [batch, 1, hidden_size]
-            current_input = torch.cat([current_input, new_embed], dim=1)  # [batch, 3~n, hidden_size]
-        pos_all_layers = torch.stack(pos_codes, dim=1)  # [batch, num_code_groups, 1]
-        return pos_all_layers, current_input
+        # Allocate proj_buf locally each call to avoid cross-call aliasing
+        max_seq = num_groups + 1
+        proj_buf = torch.zeros(bsz, max_seq, self._hidden_size, dtype=dtype, device=device)
+        pos_ids = self._pos_ids
+        model_fwd = self._model_fwd
+        lm_heads = self._lm_heads
+        codec_embeds = self._codec_embeds
+
+        # Output codes
+        all_codes = torch.empty(bsz, num_groups, 1, dtype=torch.int64, device=device)
+        all_codes[:, 0] = layer0_code
+
+        # Fill buffer positions 0 & 1
+        proj_buf[:bsz, 0:1, :] = last_talker_hidden
+        proj_buf[:bsz, 1:2, :] = layer0_embed
+
+        # Autoregressive loop: predict layers 1..G-1
+        for step in range(1, num_groups):
+            seq_len = step + 1
+            projected = proj_buf[:bsz, :seq_len, :]
+            step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz)
+            assert step_pos_ids.shape[0] == bsz * seq_len
+
+            hidden_out = model_fwd(projected, step_pos_ids)
+
+            # Inline sampling: top-k -> top-p -> softmax -> multinomial
+            logits = lm_heads[step - 1](hidden_out[:, -1, :])  # [bsz, vocab]
+            if self._top_k > 0:
+                topk_vals, _ = logits.topk(self._top_k, dim=-1)
+                logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf"))
+            if self._top_p < 1.0:
+                sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
+                cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
+                # Remove tokens with cumulative probability above top_p
+                remove_mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= self._top_p
+                sorted_logits[remove_mask] = float("-inf")
+                logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
+            probs = F.softmax(logits, dim=-1)
+            code = torch.multinomial(probs, num_samples=1)  # [bsz, 1]
+
+            all_codes[:, step] = code
+
+            # Embed predicted code -> next buffer position
+            new_embed = codec_embeds[step - 1](code)  # [batch, 1, hidden_size]
+            proj_buf[:bsz, step + 1 : step + 2, :] = new_embed
+
+        return all_codes, proj_buf[:bsz]
+
+    # ------------------------------------------------------------------
+    #  Weight loading
+    # ------------------------------------------------------------------
 
     def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights with mapping for fused QKV and gate_up projections.

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -129,111 +129,70 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         self.code_predictor = Qwen3OmniMoeTalkerCodePredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor")
         )
-        max_batch_size = max(
-            vllm_config.scheduler_config.max_num_seqs, vllm_config.compilation_config.max_cudagraph_capture_size
-        )
-        self.layer0_embed_buffer = torch.zeros(
-            (max_batch_size, 1, self.config.text_config.hidden_size),
-            dtype=vllm_config.model_config.dtype,
-        )
 
     def code_predictor_forward(
         self,
         input_ids: torch.Tensor,
         inputs_embeds: torch.Tensor | None = None,
         *,
-        temperature: float = 1.0,
-        top_k: int = 50,  # Match transformers default
-        top_p: float = 0.8,  # Match transformers default
-        generation_steps: int | None = None,
         last_talker_hidden: torch.Tensor | None = None,
         **_: object,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Generate full RVQ codec codes for the provided sequence.
+        """Generate full RVQ codes + summed embeddings (single-loop, no KV cache).
 
-        The code predictor consumes the layer-0 codec codes produced by the talker
-        alongside the talker's hidden states, and autoregressively predicts the remaining
-        residual layers (to num_codec_groups).
+        The code predictor uses re-prefill: each AR step re-forwards the full
+        (short) sequence through the transformer. The returned ``proj_buf``
+        already contains all codec embeddings at positions 1..G,
+        so summed_embeddings = proj_buf[:, 1:, :].sum(dim=1)  — no second
+        loop or re-embedding needed.
 
         Returns:
-            tuple containing:
-                - residual_codes: A tensor of shape [batch, num_code_groups, seq_len] containing
-                  the complete set of codec codes
-                - summed_embeddings: A tensor of shape [batch, seq_len, hidden_size]
-                  Sum of all layer embeddings at each position (like Transformers)
+            result_codes:      [batch, num_code_groups, seq_len]
+            summed_embeddings: [batch, seq_len, hidden_size]
         """
         if input_ids is None:
-            raise ValueError("`input_ids` containing layer-0 codec codes must be provided.")
+            raise ValueError("`input_ids` (layer-0 codes) must be provided.")
         if inputs_embeds is None:
-            raise ValueError("`inputs_embeds` containing talker hidden states must be provided.")
+            raise ValueError("`inputs_embeds` (talker hidden states) must be provided.")
 
         if inputs_embeds.ndim == 2:
             inputs_embeds = inputs_embeds.unsqueeze(0)
         if input_ids.ndim == 1:
             input_ids = input_ids.unsqueeze(0)
 
-        # Ensure the tensors are contiguous for the autoregressive sampling loop
-        inputs_embeds = inputs_embeds.contiguous()
-        input_ids = input_ids.contiguous()
-
-        # Generate full codec codes using MTP
-        # This will be the parallel prediction implementation
         batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+        embed_fn = self.language_model.model.codec_embedding
+        hidden_size = self.config.code_predictor_config.hidden_size
 
-        # For now, use sequential generation (TODO: implement parallel)
-        # Result will be [batch, num_code_groups, seq_len]
-        # - all_codes_per_position will collect [batch, num_code_groups, 1] for each position
-        all_codes_per_position = []
-        middle_hidden_states = []  # Collect hidden states for each position
-
-        # Generate residual layers for each position
-        for pos in range(seq_len):
-            layer0_code = input_ids[:, pos : pos + 1]  # [batch, 1]
-
-            # Initial input: [last_talker_hidden, layer0_embed]
-            layer0_embed = self.embed_input_ids(layer0_code)
-            self.layer0_embed_buffer[:batch_size].copy_(layer0_embed)
-            pos_all_layers, current_input = self.code_predictor(
-                layer0_code, self.layer0_embed_buffer[:batch_size], last_talker_hidden
-            )
-
-            # Stack all layers for this position: [batch, num_code_groups, 1]
-            all_codes_per_position.append(pos_all_layers)
-            middle_hidden_states.append(current_input[:, 2:-1, :])
-
-        # Concatenate across positions: [batch, num_code_groups, seq_len]
-        result_codes = torch.cat(all_codes_per_position, dim=2)
-
-        # Build summed embeddings for each position (like Transformers)
-        # This combines layer-0 embed, mid layers hidden states, and last layer embed
-        all_summed_embeddings = []
+        result_codes = torch.empty(
+            batch_size,
+            self.num_code_groups,
+            seq_len,
+            dtype=torch.int64,
+            device=device,
+        )
+        summed_embeddings = torch.empty(
+            batch_size,
+            seq_len,
+            hidden_size,
+            dtype=inputs_embeds.dtype,
+            device=device,
+        )
 
         for pos in range(seq_len):
-            # Layer 0 embedding
-            layer0_code = result_codes[:, 0, pos : pos + 1]  # [batch, 1]
-            layer0_embed = self.embed_input_ids(layer0_code)  # [batch, 1, hidden_size]
+            layer0_code = input_ids[:, pos : pos + 1]
+            layer0_embed = embed_fn(layer0_code)
 
-            # mid layers hidden states (from CodePredictor)
-            mid_residual_hiddens = middle_hidden_states[pos]  # [batch, num_code_groups-2, hidden_size]
-            mid_list = list(mid_residual_hiddens.split(1, dim=1))
-
-            # last layer embedding
-            last_layer_code = result_codes[:, -1, pos : pos + 1]  # [batch, 1]
-            last_residual_hidden = self.code_predictor.model.codec_embedding[-1](last_layer_code)
-
-            # Concatenate all layers: [batch, num_code_groups, hidden_size]
-            pos_codec_hiddens = torch.cat(
-                [layer0_embed] + mid_list + [last_residual_hidden],
-                dim=1,
+            pos_all_layers, proj_buf = self.code_predictor(
+                layer0_code,
+                layer0_embed,
+                last_talker_hidden,
             )
 
-            # Sum across layers: [batch, 1, hidden_size] (like Transformers)
-            pos_summed = pos_codec_hiddens.sum(dim=1, keepdim=True)
-            all_summed_embeddings.append(pos_summed)
-
-        # Concatenate across positions: [batch, seq_len, hidden_size]
-        summed_embeddings = torch.cat(all_summed_embeddings, dim=1).squeeze(1)
+            result_codes[:, :, pos : pos + 1] = pos_all_layers
+            # proj_buf layout: [0]=talker_hidden, [1..G]=codec embeds
+            summed_embeddings[:, pos, :] = proj_buf[:, 1:, :].sum(dim=1)
 
         return result_codes, summed_embeddings
 

vllm_omni/worker/gpu_model_runner.py

@@ -1377,10 +1377,17 @@ class OmniGPUModelRunner(GPUModelRunner):
         req_state = self.requests.get(req_id)
         if req_state is None:
             return
+        # Check if the model declares keys that should stay on GPU
+        gpu_keys: set[str] = set()
+        if hasattr(self, "model") and hasattr(self.model, "gpu_resident_buffer_keys"):
+            gpu_keys = self.model.gpu_resident_buffer_keys
         existing = self.model_intermediate_buffer.setdefault(req_id, {})
         for k, v in upd.items():
             if isinstance(v, torch.Tensor):
-                existing[k] = v.detach().to("cpu").contiguous()
+                if k in gpu_keys:
+                    existing[k] = v.detach().clone()
+                else:
+                    existing[k] = v.detach().to("cpu").contiguous()
             elif isinstance(v, list):
                 existing[k] = [
                     (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v