[Feature][LingBot World]stepwise execution (#6844)

Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Signed-off-by: liuqihao <liuqihao970610@gmail.com>

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """CPU contracts for capability-driven AR-Diffusion sessions."""
 
 from __future__ import annotations
@@ -10,6 +11,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.experimental.ar_diffusion.capability import (
     ARDiffusionCrossAttentionKVSpec,
@@ -27,6 +29,22 @@ NEG = "negative"
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def platform_synchronize_calls(monkeypatch) -> list[None]:
+    """Stub the platform barrier the runner takes at every chunk boundary.
+
+    These tests run on a CPU host, where the resolved platform is
+    ``UnspecifiedOmniPlatform`` and ``synchronize()`` raises. Tests that need
+    a failing barrier override this with their own fake.
+    """
+    calls: list[None] = []
+    monkeypatch.setattr(
+        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
+        SimpleNamespace(synchronize=lambda: calls.append(None)),
+    )
+    return calls
+
+
 def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
     """Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV."""
     return ARDiffusionKVCacheSpec(
@@ -114,6 +132,43 @@ class BatchCapablePipeline(CapablePipeline):
     supports_request_batch = True
 
 
+class StepCapablePipeline(CapablePipeline):
+    supports_step_execution = True
+
+    def prepare_encode(self, state, **kwargs):
+        del kwargs
+        return state
+
+    def denoise_step(self, input_batch, *, states=None, **kwargs):
+        del input_batch, states, kwargs
+        return None
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del state, noise_pred, kwargs
+
+    def post_decode(self, state, **kwargs):
+        del state, kwargs
+        from vllm_omni.diffusion.data import DiffusionOutput
+
+        return DiffusionOutput()
+
+
+def scheduler_output(
+    request_ids: tuple[str, ...] = ("req-1",),
+    *,
+    finished: tuple[str, ...] = (),
+) -> DiffusionSchedulerOutput:
+    """Build the real scheduler payload; stepwise requests arrive as cached rows."""
+    return DiffusionSchedulerOutput(
+        step_id=0,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData(request_ids=list(request_ids)),
+        finished_req_ids=set(finished),
+        num_running_reqs=len(request_ids),
+        num_waiting_reqs=0,
+    )
+
+
 def make_runner(
     pipeline: object,
     *,
@@ -140,6 +195,7 @@ def make_runner(
     runner._sessions = OrderedDict()
     runner._session_capacity = 0
     runner._perf_e2e_times = []
+    runner._stepwise_chunk_started = {}
     runner._preallocate_kv_cache(available_bytes=available_bytes)
     return runner
 
@@ -395,12 +451,16 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
         ctx.ensure_video_slots(torch.device("cpu"))
         return object()
 
-    def synchronize_boom(device):
+    def synchronize_boom():
         raise RuntimeError("asynchronous kernel failed")
 
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
-    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
-    runner.device = torch.device("cuda")
+    # The runner defers the device barrier to the platform, so the failure is
+    # injected there rather than on a device-specific torch entry point.
+    monkeypatch.setattr(
+        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
+        SimpleNamespace(synchronize=synchronize_boom),
+    )
     request = SimpleNamespace(
         request_id="broken-request",
         sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
@@ -424,6 +484,11 @@ def test_ar_runner_rejects_step_and_request_batch_modes():
         make_runner(BatchCapablePipeline(lingbot_like_spec()))
 
 
+def test_ar_runner_allows_step_execution_when_pipeline_implements_the_contract():
+    runner = make_runner(StepCapablePipeline(lingbot_like_spec()), step_execution=True)
+    assert runner.kv_cache is not None
+
+
 def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
     runner = object.__new__(ARDiffusionModelRunner)
     with pytest.raises(RuntimeError, match="request-batch execution"):
@@ -432,6 +497,144 @@ def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
         runner.execute_stepwise(None)
 
 
+def test_execute_stepwise_binds_request_id_session_and_unbinds_after_call(monkeypatch):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    bound_during = []
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        bound_during.append(pipeline.bound_state is not None)
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False)])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    output = runner.execute_stepwise(scheduler_output())
+
+    assert bound_during == [True]
+    assert pipeline.bound_state is None
+    assert pipeline.binds == ["req-1"]
+    assert "req-1" in runner._sessions
+    assert output.get_request_output("req-1") is not None
+    assert not pipeline.closes
+
+
+def test_execute_stepwise_closes_session_when_request_finishes(monkeypatch):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=True)])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    runner.execute_stepwise(scheduler_output())
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["req-1"]
+    assert "req-1" not in runner._sessions
+
+
+def test_execute_stepwise_exception_releases_kv_and_does_not_resume_pages(monkeypatch):
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    first = runner._get_or_create_session("req-1")
+
+    def boom(self, scheduler_output):
+        del self, scheduler_output
+        raise RuntimeError("step failed")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", boom)
+    with pytest.raises(RuntimeError, match="step failed"):
+        runner.execute_stepwise(scheduler_output())
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["req-1"]
+    assert "req-1" not in runner._sessions
+    second = runner._get_or_create_session("req-1")
+    assert second is not first
+
+
+def test_execute_stepwise_closes_session_for_scheduler_aborted_request(monkeypatch):
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    runner._get_or_create_session("req-1")
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    # An aborted request is never scheduled again and never reports finished,
+    # so the retire path is the only thing that can free its KV.
+    runner.execute_stepwise(scheduler_output(request_ids=(), finished=("req-1",)))
+
+    assert "req-1" not in runner._sessions
+    assert pipeline.closes == ["req-1"]
+
+
+def test_execute_stepwise_times_once_per_chunk_not_once_per_step(monkeypatch, platform_synchronize_calls):
+    from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    results = [None, None, None, DiffusionOutput()]
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=results.pop(0))])
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    for _ in range(4):
+        runner.execute_stepwise(scheduler_output())
+
+    # Four denoise steps, one emitted AR block: one timing sample, matching the
+    # one-sample-per-block meaning request mode already has.
+    assert len(runner._perf_e2e_times) == 1
+    assert len(platform_synchronize_calls) == 1
+    assert "req-1" not in runner._stepwise_chunk_started
+
+
+def test_execute_stepwise_synchronize_exception_fails_closed(monkeypatch):
+    """The chunk barrier is where an asynchronous device error surfaces. It must
+    clean up like an in-step failure instead of leaving the session, the cached
+    step state and the timing entry to the scheduler's next cycle."""
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    pipeline = StepCapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline, step_execution=True)
+    runner.state_cache = {"req-1": object()}
+
+    def fake_execute(self, scheduler_output):
+        del self, scheduler_output
+        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=object())])
+
+    def synchronize_boom():
+        raise RuntimeError("asynchronous kernel failed")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
+    monkeypatch.setattr(
+        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
+        SimpleNamespace(synchronize=synchronize_boom),
+    )
+
+    with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
+        runner.execute_stepwise(scheduler_output())
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["req-1"]
+    assert "req-1" not in runner._sessions
+    assert "req-1" not in runner.state_cache
+    assert "req-1" not in runner._stepwise_chunk_started
+    assert not runner._perf_e2e_times
+
+
 def test_model_specific_warmup_provider_is_consumed(monkeypatch):
     requests = [object(), object()]
     pipeline = WarmupPipeline(lingbot_like_spec(), requests)

tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -13,6 +14,7 @@ from vllm_omni.diffusion.models.lingbot_world.actions import (
     LingBotCameraControlReducer,
     integrate_lingbot_camera_actions,
     parse_lingbot_camera_action_frames,
+    parse_lingbot_camera_action_script,
 )
 from vllm_omni.experimental.ar_diffusion.session import (
     ARDiffusionSession,
@@ -292,3 +294,8 @@ def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
         assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
 
     asyncio.run(exercise())
+
+
+def test_parse_camera_action_script_rejects_wrong_chunk_width() -> None:
+    with pytest.raises(ValueError, match="exactly 3 per-latent-frame"):
+        parse_lingbot_camera_action_script([[["w"], ["w"]]], frames_per_chunk=3)

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -15,8 +15,10 @@ from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
 from PIL import Image
 from torch import nn
 
+import vllm_omni.diffusion.models.lingbot_world.dmd_block as lingbot_dmd_block
 import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline
 from tests.diffusion.models.wan2_2.conftest import noop_progress_bar
+from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
 from vllm_omni.diffusion.models.lingbot_world.actions import (
     integrate_lingbot_camera_actions,
 )
@@ -24,6 +26,8 @@ from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as
 from vllm_omni.diffusion.models.lingbot_world.camera import (
     build_plucker_embedding as _real_build_plucker_embedding,
 )
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.utils import StepRequestState
 from vllm_omni.experimental.ar_diffusion.tick_protocol import (
     ARDiffusionControlInput,
     ARDiffusionTickRequest,
@@ -253,6 +257,8 @@ class _SamplingParams:
         if extra_args is not None:
             self.extra_args.update(extra_args)
         self.latents = None
+        self.guidance_scale = None
+        self.guidance_scale_2 = None
 
 
 class _RequestBatch:
@@ -345,7 +351,6 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
         "FlowUniPCMultistepScheduler": _FakeScheduler,
         "CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory,
         "get_local_device": lambda: torch.device("cpu"),
-        "set_forward_context_denoise_step_idx": lambda index: None,
         "prefetch_subfolders": prefetch_subfolders,
         "from_pretrained_with_prefetch": from_pretrained_with_prefetch,
         "load_transformer_config": load_transformer_config,
@@ -358,6 +363,9 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
     }
     for name, value in replacements.items():
         monkeypatch.setattr(lingbot_pipeline, name, value)
+    # The DMD block math lives in its own module; stub the symbols it imports.
+    monkeypatch.setattr(lingbot_dmd_block, "set_forward_context_denoise_step_idx", lambda index: None)
+    monkeypatch.setattr(lingbot_dmd_block, "randn_tensor", _diffusers_randn_tensor)
     monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False)
 
 
@@ -774,6 +782,7 @@ def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() ->
         return torch.zeros(shape, device=device, dtype=dtype)
 
     module.randn_tensor = randn
+    lingbot_dmd_block.randn_tensor = randn
     result = pipeline(_request())
 
     assert requested_noise_dtypes == [torch.float32] * 4
@@ -1139,7 +1148,7 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N
     module = _load_pipeline_module()
     transformer = _RecordingTransformer()
     pipeline = _pipeline(module, transformer=transformer)
-    module.randn_tensor = lambda shape, **kwargs: torch.full(
+    module.randn_tensor = lingbot_dmd_block.randn_tensor = lambda shape, **kwargs: torch.full(
         shape,
         -99.0,
         device=kwargs["device"],
@@ -1745,6 +1754,288 @@ def test_request_cache_becomes_unreachable_after_transformer_error() -> None:
     assert not hasattr(pipeline, "transformer_cache")
 
 
+class _FakeARState:
+    def __init__(self, session_id: str = "req-1") -> None:
+        self.session_id = session_id
+        self.commits: list[str] = []
+
+    def get_kv_caches(self, branch, *, seq_len, commit_current):
+        del branch, seq_len, commit_current
+        return [SimpleNamespace()]
+
+    def commit_paged_context(self, branch):
+        self.commits.append(branch)
+
+    def clear_cross_attention(self):
+        return None
+
+    def is_cross_attention_populated(self, branch, name):
+        del branch, name
+        return True
+
+    def get_cross_attention_kv(self, branch, name):
+        del branch, name
+        zeros = torch.zeros(1, 1, 2, 4)
+        return [{"k": zeros, "v": zeros} for _ in range(2)]
+
+
+def _empty_action_script(num_chunks: int):
+    return tuple(((), (), ()) for _ in range(num_chunks))
+
+
+def _stepwise_state(
+    *,
+    request_id: str = "req-1",
+    num_frames: int = 21,
+    script=None,
+    seed: int = 17,
+):
+    num_chunks = ((num_frames - 1) // 4 + 1) // 3
+    sampling = _SamplingParams(num_frames=num_frames, seed=seed)
+    sampling.extra_args["_lingbot_camera_trajectory"] = None
+    sampling.extra_args["_lingbot_camera_actions"] = None
+    sampling.extra_args["_lingbot_camera_action_script"] = (
+        script if script is not None else _empty_action_script(num_chunks)
+    )
+    return StepRequestState(
+        request_id=request_id,
+        sampling=sampling,
+        prompt=_prompt(),
+    )
+
+
+def _run_stepwise(pipeline, state):
+    outputs = []
+    pipeline.prepare_encode(state)
+    while not state.request_denoise_completed:
+        noise = pipeline.denoise_step(None, states=[state])
+        pipeline.step_scheduler(state, noise)
+        if state.chunk_denoise_completed:
+            outputs.append(pipeline.post_decode(state))
+    return outputs
+
+
+def test_pipeline_declares_step_execution_support() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+
+    assert pipeline.supports_step_execution is True
+    assert isinstance(pipeline, SupportsStepExecution)
+    assert supports_step_execution(pipeline) is True
+
+
+def test_preprocess_materializes_camera_action_script_without_action_path() -> None:
+    module = _load_pipeline_module()
+    sampling = _SamplingParams(include_action=False)
+    sampling.extra_args["camera_action_script"] = [[["w"], ["w"], ["w"]], [["a"], [], []]]
+    request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=sampling, request_id="req-1")
+
+    result = module.get_lingbot_world_pre_process_func(_od_config())(request)
+
+    assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None
+    assert result.sampling_params.extra_args["_lingbot_camera_actions"] is None
+    assert result.sampling_params.extra_args["_lingbot_camera_action_script"] == (
+        (("w",), ("w",), ("w",)),
+        (("a",), (), ()),
+    )
+
+
+@pytest.mark.parametrize(
+    "extra_args",
+    [
+        pytest.param({}, id="request_mode"),
+        pytest.param(_tick_extra_args(chunk_index=0), id="tick_mode"),
+    ],
+)
+def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None:
+    """A script only steers step execution, so request mode must not drop it silently."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    sampling = _SamplingParams(
+        include_action=False,
+        extra_args={
+            **extra_args,
+            "_lingbot_camera_trajectory": None,
+            "_lingbot_camera_action_script": _empty_action_script(1),
+        },
+    )
+
+    with pytest.raises(ValueError, match="camera_action_script is read only by LingBot step execution"):
+        pipeline(_request(sampling=sampling))
+
+
+def test_stepwise_progress_metadata_and_commit_trace() -> None:
+    module = _load_pipeline_module()
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert pipeline._ar_diffusion_kv_state is None
+    assert len(outputs) == 2
+    assert [output.chunk_index for output in outputs] == [0, 1]
+    assert all(output.total_chunks == 2 for output in outputs)
+    assert outputs[-1].finished is True
+    assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2
+    assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
+    assert fake.commits == ["main", "main"]
+    for chunk_index, output in enumerate(outputs):
+        metadata = output.output["metadata"]["ar_diffusion"]
+        assert metadata == {
+            "session_id": "req-1",
+            "request_id": "req-1",
+            "chunk_index": chunk_index,
+            "applied_event_ids": [],
+        }
+        assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2)
+
+
+def test_stepwise_matches_tick_transformer_trace_and_latents() -> None:
+    module = _load_pipeline_module()
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",)))
+    fake = _FakeARState("world-1")
+
+    with pipeline.bind_ar_diffusion_state("world-1", fake):
+        tick_latents = []
+        for chunk_index, frames in enumerate(script):
+            tick = ARDiffusionTickRequest(
+                session_id="world-1",
+                request_id="legacy-lingbot-request",
+                chunk_index=chunk_index,
+                controls=(
+                    ARDiffusionControlInput(
+                        track="camera",
+                        schema="lingbot.camera_actions.v1",
+                        data={"mode": "frames", "frames": [list(frame) for frame in frames]},
+                    ),
+                ),
+            )
+            sampling = _SamplingParams(extra_args=tick.to_extra_args(), seed=17)
+            sampling.extra_args["_lingbot_camera_trajectory"] = None
+            sampling.extra_args["_lingbot_camera_actions"] = frames
+            result = pipeline(_request(sampling=sampling))
+            tick_latents.append(result.output["payload"]["latents"].clone())
+    tick_calls = list(transformer.calls)
+    transformer.calls.clear()
+    pipeline.close_ar_diffusion_session("world-1")
+
+    stepwise_state = _stepwise_state(request_id="req-1", num_frames=21, script=script, seed=17)
+    stepwise_fake = _FakeARState(stepwise_state.request_id)
+    with pipeline.bind_ar_diffusion_state(stepwise_state.request_id, stepwise_fake):
+        stepwise_outputs = _run_stepwise(pipeline, stepwise_state)
+
+    assert [call["update_cache"] for call in transformer.calls] == [call["update_cache"] for call in tick_calls]
+    assert [call["start_frame"] for call in transformer.calls] == [call["start_frame"] for call in tick_calls]
+    for tick_latent, output in zip(tick_latents, stepwise_outputs, strict=True):
+        torch.testing.assert_close(output.output["payload"]["latents"], tick_latent)
+
+
+def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None:
+    """A trajectory that slows down between blocks must condition every stepwise
+    chunk exactly as request mode does. Request mode embeds the whole trajectory
+    once; embedding per chunk would re-normalize each block's framewise
+    translations by its own largest step and present a slow block as full-speed
+    motion."""
+    from vllm_omni.diffusion.models.lingbot_world import camera as camera_module
+
+    module = _load_pipeline_module()
+    # The fixture stubs the ray embedding with a pose-blind ramp; parity needs
+    # the real geometry so a scale mismatch actually shows up.
+    monkeypatch.setattr(module, "build_plucker_embedding", camera_module.build_plucker_embedding)
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+
+    # Six latent frames -> two 3-frame blocks. The first block moves at full
+    # speed and everything after it at a tenth of that, so a per-chunk
+    # normalization would rescale the second block by 10x.
+    steps = torch.tensor([1.0] * 3 + [0.1] * 18)
+    poses = torch.eye(4).repeat(21, 1, 1)
+    poses[:, 2, 3] = torch.cumsum(steps, dim=0) - steps[0]
+    trajectory = _CameraTrajectory(
+        poses=poses,
+        intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(21, 1),
+    )
+
+    sampling = _SamplingParams(num_frames=21)
+    sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
+    pipeline(_request(sampling=sampling))
+    request_cameras = [call["camera_hidden_states"] for call in transformer.calls]
+    transformer.calls.clear()
+
+    state = _stepwise_state(num_frames=21)
+    state.sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
+    state.sampling.extra_args.pop("_lingbot_camera_action_script")
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        _run_stepwise(pipeline, state)
+    stepwise_cameras = [call["camera_hidden_states"] for call in transformer.calls]
+
+    # Two blocks x (four probes + one commit), in the same order on both paths.
+    assert len(request_cameras) == len(stepwise_cameras) == 10
+    for stepwise, request in zip(stepwise_cameras, request_cameras, strict=True):
+        torch.testing.assert_close(stepwise, request)
+
+
+def test_stepwise_emits_latents_without_touching_the_vae() -> None:
+    """Latent mode must stay latent: streaming decode is opt-in per request."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert [set(output.output["payload"]) for output in outputs] == [{"latents"}, {"latents"}]
+    assert pipeline.vae.decode_inputs == []
+
+
+def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
+    """A non-latent request must stream pixels under the primary "video" key."""
+    import diffusers.video_processor as video_processor_module
+
+    module = _load_pipeline_module()
+    processed = []
+
+    class VideoProcessor:
+        def __init__(self, *, vae_scale_factor):
+            assert vae_scale_factor == 8
+
+        def postprocess_video(self, video, *, output_type):
+            processed.append((tuple(video.shape), output_type))
+            return f"frames-{len(processed)}"
+
+    monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    # One decode per AR block, each seeing only that block's three latent frames.
+    assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
+    assert [output_type for _, output_type in processed] == ["np", "np"]
+    assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
+    # Identity metadata survives the switch to pixel output.
+    assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
+
+
 def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
     module = _load_pipeline_module()
     resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
@@ -1755,7 +2046,7 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
     assert resolved is module.LingBotWorldCausalDMDPipeline
     assert cache_acceleration_disabled
     assert preprocess_name == "get_lingbot_world_pre_process_func"
-    request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams())
+    request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=_SamplingParams(), request_id="req-1")
     assert preprocess(request) is request
     assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory)
     assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline"