[Core] Add realtime AR-Diffusion tick sessions for LingBot World 2.0 (#5491)

Signed-off-by: zhangbin.327 <zhangbin.327@bytedance.com>
Signed-off-by: Jack47 <scsvip@gmail.com>
Co-authored-by: zhangbin.327 <zhangbin.327@bytedance.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -18,6 +18,7 @@ from vllm_omni.experimental.ar_diffusion.capability import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
 from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
 
 BLOCK = 16
 POS = "positive"
@@ -56,6 +57,21 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
     )
 
 
+def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
+    return ARDiffusionKVCacheSpec(
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=1,
+        tokens_per_frame=1,
+        frames_per_block=3,
+        window_frames=3,
+        sink_frames=3,
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=capacity,
+        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
+    )
+
+
 class CapablePipeline:
     def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
         self.spec = spec
@@ -103,6 +119,7 @@ def make_runner(
     *,
     available_bytes: int = 1 << 28,
     step_execution: bool = False,
+    gpu_memory_fraction: float = 0.1,
 ) -> ARDiffusionModelRunner:
     runner = object.__new__(ARDiffusionModelRunner)
     runner.od_config = SimpleNamespace(
@@ -113,7 +130,10 @@ def make_runner(
     )
     runner.device = torch.device("cpu")
     runner.pipeline = pipeline
-    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
+    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
+        enable=True,
+        gpu_memory_fraction=gpu_memory_fraction,
+    )
     runner.kv_cache = None
     runner._ar_diffusion_capability = None
     runner._ar_diffusion_kv_cache_spec = None
@@ -132,6 +152,16 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
     return state
 
 
+def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
+    state = runner._get_or_create_session(session_id)
+    assert runner.kv_cache is not None
+    seq_len = BLOCK * runner.kv_cache.frames_per_block
+    ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    state.commit_paged_context(kv_branch)
+    return state
+
+
 def test_ar_runner_rejects_pipeline_without_capability():
     runner = object.__new__(ARDiffusionModelRunner)
     runner.od_config = SimpleNamespace(max_num_seqs=1)
@@ -140,6 +170,43 @@ def test_ar_runner_rejects_pipeline_without_capability():
         runner._preallocate_kv_cache(available_bytes=1 << 20)
 
 
+def test_runner_uses_typed_tick_as_authoritative_session_contract():
+    tick = ARDiffusionTickRequest(
+        session_id="world-7",
+        request_id="request-3",
+        chunk_index=3,
+        reset=True,
+    )
+    req = SimpleNamespace(
+        request_id="request-3",
+        sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
+    )
+
+    session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
+
+    assert session_id == "world-7"
+    assert extra_args == tick.to_extra_args()
+    assert parsed == tick
+
+
+def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
+    tick = ARDiffusionTickRequest(
+        session_id="world-7",
+        request_id="client-request-3",
+        chunk_index=3,
+    )
+    req = SimpleNamespace(
+        request_id="engine-request-uuid",
+        sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
+    )
+
+    session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
+
+    assert session_id == "world-7"
+    assert req.request_id == "engine-request-uuid"
+    assert parsed.request_id == "client-request-3"
+
+
 def test_lingbot_like_single_branch_session_reuse_reset_and_close():
     pipeline = CapablePipeline(lingbot_like_spec())
     runner = make_runner(pipeline)
@@ -170,6 +237,31 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
     assert pipeline.closes == ["s1"]
 
 
+def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
+    runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
+    kv = runner.kv_cache
+    assert kv is not None
+
+    session_a = commit_one_block(runner, "world-a", "main")
+    a0_blocks = kv.window_block_ids(session_a.adapter("main"))
+    session_b = commit_one_block(runner, "world-b", "main")
+    b0_blocks = kv.window_block_ids(session_b.adapter("main"))
+    session_a_again = commit_one_block(runner, "world-a", "main")
+    a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
+
+    assert session_a_again is session_a
+    assert session_a.adapter("main").request_id == "ar::world-a::main"
+    assert session_b.adapter("main").request_id == "ar::world-b::main"
+    assert session_a.adapter("main").completed_chunks == 6
+    assert session_b.adapter("main").completed_chunks == 3
+    assert len(a0_blocks) == 3
+    assert len(b0_blocks) == 3
+    assert len(a1_blocks) == 6
+    assert set(a0_blocks) <= set(a1_blocks)
+    assert set(a1_blocks).isdisjoint(b0_blocks)
+    assert tuple(runner._sessions) == ("world-b", "world-a")
+
+
 def test_lingbot_like_sink_survives_sliding_window_eviction():
     runner = make_runner(CapablePipeline(lingbot_like_spec()))
     kv = runner.kv_cache
@@ -198,12 +290,12 @@ def test_dreamzero_like_two_branches_are_independent():
 
 
 def test_lru_eviction_releases_blocks_and_notifies_pipeline():
-    pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
+    pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
     runner = make_runner(pipeline)
     kv = runner.kv_cache
     assert kv is not None
-    assert runner._session_capacity == 1
-    assert kv.session_capacity == 1
+    assert runner._session_capacity == 2
+    assert kv.session_capacity == 2
     free_total = kv.manager.block_pool.get_num_free_blocks()
     old = commit_one_frame(runner, "old", "main")
     k = torch.randn(1, 8, 4, 64)
@@ -211,13 +303,56 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
     assert kv.manager.block_pool.get_num_free_blocks() < free_total
 
     runner._get_or_create_session("new")
+    assert tuple(runner._sessions) == ("old", "new")
+    runner._get_or_create_session("newest")
 
-    assert tuple(runner._sessions) == ("new",)
+    assert tuple(runner._sessions) == ("new", "newest")
     assert pipeline.closes == ["old"]
     assert "old" not in kv._cross_sessions
     assert kv.manager.block_pool.get_num_free_blocks() == free_total
 
 
+def test_budget_reduced_capacity_drives_runner_lru():
+    pipeline = CapablePipeline(tiny_spec(capacity=2))
+    # One tiny LingBot-like session requires 128 bytes; two require 192.
+    runner = make_runner(
+        pipeline,
+        available_bytes=128,
+        gpu_memory_fraction=1.0,
+    )
+    kv = runner.kv_cache
+    assert kv is not None
+    assert kv.requested_session_capacity == 2
+    assert kv.session_capacity == 1
+    assert runner._session_capacity == 1
+
+    runner._get_or_create_session("old")
+    runner._get_or_create_session("new")
+
+    assert tuple(runner._sessions) == ("new",)
+    assert pipeline.closes == ["old"]
+
+
+def test_dreamzero_like_requested_capacity_is_capped_by_budget():
+    spec = dreamzero_like_spec(capacity=64)
+    # Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
+    # managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
+    # cross-attention=1 page/session, for 44 pages total.
+    page_bytes = 65_536
+    pipeline = CapablePipeline(spec)
+    runner = make_runner(
+        pipeline,
+        available_bytes=44 * page_bytes,
+        gpu_memory_fraction=1.0,
+    )
+    kv = runner.kv_cache
+    assert kv is not None
+    assert kv.requested_session_capacity == 64
+    assert kv.session_capacity == 2
+    assert runner._session_capacity == 2
+    assert kv.cross_attention_reserved_bytes == 2 * page_bytes
+
+
 def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
     pipeline = CapablePipeline(lingbot_like_spec())
     runner = make_runner(pipeline)
@@ -232,7 +367,10 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
         raise RuntimeError("layer exploded")
 
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
-    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+    request = SimpleNamespace(
+        request_id="broken-request",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+    )
 
     with pytest.raises(RuntimeError, match="layer exploded"):
         runner.execute_model(request)
@@ -263,7 +401,10 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
     monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
     runner.device = torch.device("cuda")
-    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+    request = SimpleNamespace(
+        request_id="broken-request",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+    )
 
     with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
         runner.execute_model(request)

tests/diffusion/ar_diffusion/test_consumer.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
+from vllm_omni.experimental.ar_diffusion.session import (
+    ARDiffusionSession,
+    ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionChunkMetadata,
+    ARDiffusionControlInput,
+    ARDiffusionTickRequest,
+)
+from vllm_omni.inputs.data import (
+    OmniDiffusionSamplingParams,
+    OmniPromptType,
+    OmniSamplingParams,
+)
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeStagePool:
+    def __init__(self, num_replicas: int = 1) -> None:
+        self.num_replicas = num_replicas
+
+
+class FakeEngine:
+    def __init__(self, num_replicas: int = 1) -> None:
+        self.stage_pools = [FakeStagePool(num_replicas)]
+
+
+class FakeGenerateClient:
+    def __init__(
+        self,
+        *,
+        include_metadata: bool = True,
+        applied_event_ids: list[int] | None = None,
+        output_error: str | None = None,
+        num_replicas: int = 1,
+    ) -> None:
+        self.engine = FakeEngine(num_replicas)
+        self.include_metadata = include_metadata
+        self.applied_event_ids = applied_event_ids
+        self.output_error = output_error
+        self.calls: list[dict[str, Any]] = []
+
+    async def generate(
+        self,
+        prompt: OmniPromptType,
+        *,
+        request_id: str,
+        sampling_params_list: Sequence[OmniSamplingParams],
+        output_modalities: list[str] | None = None,
+    ) -> AsyncIterator[OmniRequestOutput]:
+        self.calls.append(
+            {
+                "prompt": prompt,
+                "request_id": request_id,
+                "sampling_params_list": sampling_params_list,
+                "output_modalities": output_modalities,
+            }
+        )
+        params = sampling_params_list[0]
+        assert isinstance(params, OmniDiffusionSamplingParams)
+        tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
+        assert tick is not None
+        chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
+        if self.applied_event_ids is not None:
+            chunk_metadata["applied_event_ids"] = self.applied_event_ids
+        metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
+        yield OmniRequestOutput(
+            request_id=request_id,
+            stage_id=0,
+            final_output_type="video",
+            images=["chunk"],  # type: ignore[list-item]
+            _multimodal_output=metadata,
+            finished=True,
+            error=self.output_error,
+        )
+
+
+class FakeLifecycle:
+    async def reset_session(self, session_id: str) -> None:
+        pass
+
+    async def close_session(self, session_id: str) -> None:
+        pass
+
+
+def make_tick() -> ARDiffusionTickRequest:
+    return ARDiffusionTickRequest(
+        session_id="world-1",
+        request_id="world-1-chunk-3",
+        chunk_index=3,
+        applied_event_ids=(7, 8),
+        prompt="turn left",
+        controls=(
+            ARDiffusionControlInput(
+                track="camera",
+                schema="example.camera.v1",
+                data={"samples": [{"x": 1.0}]},
+            ),
+        ),
+    )
+
+
+def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
+    client = FakeGenerateClient(num_replicas=2)
+
+    with pytest.raises(
+        ValueError,
+        match=r"require exactly one replica.*num_replicas=2",
+    ):
+        ARDiffusionOmniTickConsumer(
+            client,
+            prompt_provider=lambda tick: tick.prompt or "",
+            sampling_params_list=[OmniDiffusionSamplingParams()],
+            diffusion_stage_id=0,
+        )
+
+    assert client.calls == []
+
+
+@pytest.mark.asyncio
+async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
+    client = FakeGenerateClient()
+    template = OmniDiffusionSamplingParams(
+        output_type="latent",
+        extra_args={"deployment_option": "kept"},
+    )
+    consumer = ARDiffusionOmniTickConsumer(
+        client,
+        prompt_provider=lambda tick: {
+            "prompt": tick.prompt,
+            "multi_modal_data": {"image": "initial-image"},
+        },
+        sampling_params_list=[template],
+        diffusion_stage_id=0,
+        output_modalities=["video"],
+    )
+    tick = make_tick()
+
+    output = await consumer.execute_tick(tick)
+
+    assert output.images == ["chunk"]
+    assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
+    assert len(client.calls) == 1
+    call = client.calls[0]
+    assert call["request_id"] == tick.request_id
+    assert call["output_modalities"] == ["video"]
+    assert call["prompt"] == {
+        "prompt": "turn left",
+        "multi_modal_data": {"image": "initial-image"},
+    }
+    submitted_params = call["sampling_params_list"][0]
+    assert submitted_params.extra_args["deployment_option"] == "kept"
+    assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
+    assert "ar_diffusion_tick" not in template.extra_args
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_missing_chunk_metadata() -> None:
+    consumer = ARDiffusionOmniTickConsumer(
+        FakeGenerateClient(include_metadata=False),
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+
+    with pytest.raises(ValueError, match="multimodal_output.metadata"):
+        await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
+    consumer = ARDiffusionOmniTickConsumer(
+        FakeGenerateClient(applied_event_ids=[8, 7]),
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+
+    with pytest.raises(ValueError, match="unique and strictly increasing"):
+        await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
+    consumer = ARDiffusionOmniTickConsumer(
+        FakeGenerateClient(include_metadata=False, output_error="model failed"),
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+
+    with pytest.raises(RuntimeError, match="model failed"):
+        await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
+    client = FakeGenerateClient()
+    consumer = ARDiffusionOmniTickConsumer(
+        client,
+        prompt_provider=lambda tick: tick.prompt or "",
+        sampling_params_list=[OmniDiffusionSamplingParams()],
+        diffusion_stage_id=0,
+    )
+    session = ARDiffusionSession(
+        "world-1",
+        tick_consumer=consumer,
+        lifecycle=FakeLifecycle(),
+        max_pending_events=4,
+        request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
+    )
+    await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
+
+    output = await session.next_chunk()
+
+    assert output.request_id == "world-1-chunk-0"
+    assert session.chunk_index == 1
+    assert session.pending_event_count == 0
+    assert client.calls[0]["request_id"] == "world-1-chunk-0"

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -30,6 +30,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 BLOCK = 16
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
     return ChunkWindowSpec(
@@ -95,7 +97,6 @@ def test_sliding_replace_snaps_to_chunk_boundary():
 def test_sink_chunks_protected():
     # sink = 1 chunk (16 tokens) is never skipped.
     skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
-    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
     assert skip == 32
 
 
@@ -201,12 +202,149 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
         kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
         session_capacity=1,
         cross_attention_lengths={"text": L},
-        device=torch.device("cpu"),
     )
     cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
-    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    page_bytes = kv.spec.page_size_bytes * 2
+    scratch_bytes = kv.scratch_num_blocks * page_bytes
+    expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
     assert kv.num_blocks == expected
-    assert expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested
+    assert kv.cross_attention_reserved_bytes == cross_bytes
+    assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
+    assert expected > 1 * (2 + 1) + 2  # above the local-slot minimum, so the cross deduction is what's tested
+
+
+def _make_tiny_capacity_kv(
+    *,
+    requested_capacity: int,
+    available_bytes: int,
+    gpu_memory_fraction: float = 1.0,
+    model_owned_state_bytes_per_session: int = 0,
+) -> ARDiffusionKVCache:
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=1,
+            window_chunks=3,
+            sink_chunks=3,
+            gpu_memory_fraction=gpu_memory_fraction,
+        ),
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        block_size=1,
+        max_model_len=128,
+        available_bytes=available_bytes,
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=requested_capacity,
+        cross_attention_lengths={"text": 2},
+        frames_per_block=3,
+        model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
+        device=torch.device("cpu"),
+    )
+
+
+def test_capacity_two_retains_both_windows_and_allocates_next_block():
+    # LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
+    # 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
+    # reservations: 192 bytes exactly.
+    kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
+    assert kv.requested_session_capacity == 2
+    assert kv.session_capacity == 2
+    assert kv.managed_num_blocks == 17
+
+    adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
+    for adapter in adapters:
+        for _ in range(6):  # three sink + three window blocks
+            kv.allocate_token_slots(adapter, 1)
+            kv.commit_chunk(adapter)
+        assert len(kv.window_block_ids(adapter)) == 6
+
+    # Only one request is in flight, so its three-frame block is counted once.
+    kv.allocate_token_slots(adapters[0], 3)
+
+
+def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
+    kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
+
+    assert kv.requested_session_capacity == 64
+    assert kv.session_capacity == 2
+    assert kv.cross_attention_bytes_per_session == 16
+    assert kv.cross_attention_reserved_bytes == 32
+    assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
+
+
+def test_capacity_rejects_budget_that_cannot_fit_one_session():
+    # One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
+    # 14 * 8 + 16 = 128 bytes.
+    with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
+        _make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
+
+
+def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
+    kv = _make_tiny_capacity_kv(
+        requested_capacity=2,
+        available_bytes=1280,
+        gpu_memory_fraction=0.05,
+    )
+
+    assert kv.configured_memory_budget_bytes == 64
+    assert kv.memory_budget_bytes == 128
+    assert kv.session_capacity == 1
+
+
+def test_model_owned_state_reduces_effective_session_capacity():
+    kv = _make_tiny_capacity_kv(
+        requested_capacity=2,
+        available_bytes=200,
+        model_owned_state_bytes_per_session=16,
+    )
+
+    assert kv.session_capacity == 1
+    assert kv.model_owned_state_reserved_bytes == 16
+    assert (
+        kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
+        <= kv.memory_budget_bytes
+    )
+
+
+def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=1560,
+            window_chunks=9,
+            sink_chunks=9,
+            gpu_memory_fraction=gpu_memory_fraction,
+        ),
+        num_layers=40,
+        num_kv_heads=40,
+        head_size=128,
+        dtype=torch.bfloat16,
+        block_size=1560,
+        max_model_len=1 << 20,
+        available_bytes=100 * (1 << 30),
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=2,
+        cross_attention_lengths={"text": 512},
+        frames_per_block=3,
+        model_owned_state_bytes_per_session=7_488_000,
+    )
+
+
+def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
+    kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
+
+    assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
+    assert kv.memory_budget_bytes == 33_653_670_400
+    assert kv.session_capacity == 1
+
+
+def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
+    kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
+
+    assert kv.memory_budget_bytes == 60 * (1 << 30)
+    assert kv.session_capacity == 2
 
 
 def _make_kv(
@@ -221,15 +359,24 @@ def _make_kv(
         if local_branches == 1
         else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
     )
+    page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
+    declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
+    scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
+    managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
     return ARDiffusionKVCache(
-        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        ARDiffusionKVConfig(
+            enable=True,
+            chunk_size=BLOCK,
+            window_chunks=window_chunks,
+            gpu_memory_fraction=1.0,
+        ),
         num_layers=1,
         num_kv_heads=4,
         head_size=64,
         dtype=torch.float32,
         block_size=BLOCK,
         max_model_len=4096,
-        available_bytes=1 << 16,  # tiny -> the floor binds
+        available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
         device=torch.device("cpu"),
         kv_branches=kv_branches,
         session_capacity=1,

tests/diffusion/ar_diffusion/test_kv_orchestration.py

@@ -54,13 +54,15 @@ def test_full_request_lifecycle_and_eviction():
         used = {int(s) // kv.block_size for s in slots}
         assert kv.null_block_id not in used
         assert used <= set(block_table)
-        # Resident window stays bounded by window + the in-flight chunk.
+        # Allocation may transiently add one in-flight chunk.
         assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
-        free_after.append(kv.manager.block_pool.get_num_free_blocks())
         kv.commit_chunk(adapter)
+        # A successful commit immediately prunes physical ownership.
+        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
+        free_after.append(kv.manager.block_pool.get_num_free_blocks())
 
     # Pool memory plateaus once the window is full (eviction recycles blocks).
-    assert free_after[-1] == free_after[kv.spec.window_chunks]
+    assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
 
     kv.end_request(adapter)
     assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -332,7 +332,7 @@ def test_block_table_padded_to_fixed_width():
 
 def test_prepare_is_idempotent_and_layers_share_metadata():
     device = torch.device("cpu")
-    _, st = make_state(num_layers=2)
+    kv, st = make_state(num_layers=2)
     contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
     fctx = contexts[0].forward_ctx
     fctx.prepare(device=device, action_len=0, query_len=BLOCK)
@@ -349,6 +349,10 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
     assert i0.block_table is i1.block_table
     assert i0.seq_lens is i1.seq_lens
     assert i0.video_slots is i1.video_slots
+    assert i0.key_pool is kv._k_pools[0]
+    assert i0.value_pool is kv._v_pools[0]
+    assert i1.key_pool is kv._k_pools[1]
+    assert i1.value_pool is kv._v_pools[1]
 
 
 def test_layer_inputs_before_prepare_raises():

tests/diffusion/ar_diffusion/test_pipeline_bridge.py

@@ -174,6 +174,19 @@ def test_window_reset_can_keep_named_cross_attention_cache():
     assert not st.is_cross_attention_populated(POS, "text")
 
 
+def test_clear_cross_attention_preserves_self_attention_history():
+    kv, st = make_state(num_layers=1, cross_attn_length=8)
+    _prepare_and_commit(st, POS, 2)
+    k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
+    st.populate_cross_attention(POS, "text", [(k, k)])
+    resident_before = kv.window_block_ids(st.adapter(POS))
+
+    st.clear_cross_attention()
+
+    assert kv.window_block_ids(st.adapter(POS)) == resident_before
+    assert not st.is_cross_attention_populated(POS, "text")
+
+
 def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
     _, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
     assert st.kv_cache.num_local_kv_branches == 1

tests/diffusion/ar_diffusion/test_session.py

@@ -0,0 +1,505 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.session import (
+    ARDiffusionChunkMetadataError,
+    ARDiffusionEventAcceptanceStatus,
+    ARDiffusionEventQueueFullError,
+    ARDiffusionPreparedControls,
+    ARDiffusionSession,
+    ARDiffusionSessionError,
+    ARDiffusionSessionEvent,
+    ARDiffusionSessionManager,
+    ARDiffusionSessionStateError,
+    ARDiffusionSessionStatus,
+    ARDiffusionStaleEventError,
+    ARDiffusionWorkerLifecycle,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionChunkMetadata,
+    ARDiffusionControlInput,
+    ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeTickConsumer:
+    def __init__(self) -> None:
+        self.ticks: list[ARDiffusionTickRequest] = []
+        self.failures_remaining = 0
+        self.started: asyncio.Event | None = None
+        self.release: asyncio.Event | None = None
+        self.metadata_override: ARDiffusionChunkMetadata | None = None
+
+    async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+        self.ticks.append(tick)
+        if self.started is not None:
+            self.started.set()
+        if self.release is not None:
+            await self.release.wait()
+        if self.failures_remaining:
+            self.failures_remaining -= 1
+            raise RuntimeError("fake tick failed")
+        return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
+
+    def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
+        return output
+
+
+class FakeLifecycle:
+    def __init__(self) -> None:
+        self.resets: list[str] = []
+        self.closes: list[str] = []
+        self.close_failures_remaining = 0
+
+    async def reset_session(self, session_id: str) -> None:
+        self.resets.append(session_id)
+
+    async def close_session(self, session_id: str) -> None:
+        self.closes.append(session_id)
+        if self.close_failures_remaining:
+            self.close_failures_remaining -= 1
+            raise RuntimeError("fake close failed")
+
+
+class FailingCommitReducer:
+    def __init__(self) -> None:
+        self.commit_calls = 0
+        self.reset_calls = 0
+
+    def prepare(
+        self,
+        *,
+        current_controls: dict[str, ARDiffusionControlInput],
+        events: tuple[ARDiffusionSessionEvent, ...],
+        chunk_index: int,
+    ) -> ARDiffusionPreparedControls:
+        del current_controls, events, chunk_index
+        return ARDiffusionPreparedControls(controls=())
+
+    def commit(self, prepared: ARDiffusionPreparedControls) -> None:
+        del prepared
+        self.commit_calls += 1
+        raise RuntimeError("fake reducer commit failed")
+
+    def reset(self) -> None:
+        self.reset_calls += 1
+
+
+def make_session(
+    *,
+    max_pending_events: int = 8,
+    consumer: FakeTickConsumer | None = None,
+    lifecycle: FakeLifecycle | None = None,
+) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
+    resolved_consumer = consumer or FakeTickConsumer()
+    resolved_lifecycle = lifecycle or FakeLifecycle()
+    session = ARDiffusionSession(
+        "world-1",
+        tick_consumer=resolved_consumer,
+        lifecycle=resolved_lifecycle,
+        max_pending_events=max_pending_events,
+        request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+    )
+    return session, resolved_consumer, resolved_lifecycle
+
+
+@pytest.mark.asyncio
+async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
+    session, consumer, _ = make_session()
+    event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
+    event_one = ARDiffusionSessionEvent(
+        event_id=1,
+        controls=(
+            ARDiffusionControlInput(
+                track="camera",
+                schema="example.camera.v1",
+                data={"samples": [{"x": 1.0}]},
+            ),
+        ),
+    )
+
+    accepted_two = await session.accept_event(event_two)
+    accepted_one = await session.accept_event(event_one)
+    duplicate = await session.accept_event(event_two)
+
+    assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+    assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+    assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+    assert consumer.ticks == []
+
+    metadata = await session.next_chunk()
+
+    assert metadata.applied_event_ids == (1, 2)
+    assert consumer.ticks[0] == ARDiffusionTickRequest(
+        session_id="world-1",
+        request_id="world-1-request-0",
+        chunk_index=0,
+        applied_event_ids=(1, 2),
+        prompt="turn",
+        controls=(
+            ARDiffusionControlInput(
+                track="camera",
+                schema="example.camera.v1",
+                data={"samples": [{"x": 1.0}]},
+            ),
+        ),
+    )
+    assert session.pending_event_count == 0
+
+    applied_duplicate = await session.accept_event(event_two)
+    assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+
+    next_metadata = await session.next_chunk()
+    assert next_metadata.chunk_index == 1
+    assert next_metadata.applied_event_ids == ()
+    assert consumer.ticks[1].prompt == "turn"
+    assert consumer.ticks[1].controls == consumer.ticks[0].controls
+
+
+@pytest.mark.asyncio
+async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
+    session, consumer, _ = make_session(max_pending_events=1)
+    payload = {"samples": [{"x": 1.0}]}
+    composite = ARDiffusionSessionEvent(
+        event_id=5,
+        prompt="new prompt",
+        controls=(
+            ARDiffusionControlInput(
+                track="camera",
+                schema="example.camera.v1",
+                data=payload,
+            ),
+            ARDiffusionControlInput(
+                track="action",
+                schema="example.action.v1",
+                data={"keys": ["left"]},
+            ),
+        ),
+    )
+    await session.accept_event(composite)
+    payload["samples"][0]["x"] = 99.0
+
+    duplicate = await session.accept_event(composite)
+    assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+    with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
+        await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
+
+    await session.next_chunk()
+    assert consumer.ticks[0].applied_event_ids == (5,)
+    assert consumer.ticks[0].prompt == "new prompt"
+    assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
+    assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
+    assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
+    assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
+
+
+def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
+    with pytest.raises(ValueError, match="at most once"):
+        ARDiffusionSessionEvent(
+            event_id=1,
+            controls=(
+                ARDiffusionControlInput(track="camera", schema="one", data={}),
+                ARDiffusionControlInput(track="camera", schema="two", data={}),
+            ),
+        )
+
+    with pytest.raises(ValueError, match="transport-safe"):
+        ARDiffusionControlInput(
+            track="camera",
+            schema="example.camera.v1",
+            data={"object": object()},
+        )
+
+
+@pytest.mark.asyncio
+async def test_failed_tick_closes_worker_and_requires_reset() -> None:
+    consumer = FakeTickConsumer()
+    consumer.failures_remaining = 1
+    lifecycle = FakeLifecycle()
+    session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
+    await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
+
+    with pytest.raises(RuntimeError, match="fake tick failed"):
+        await session.next_chunk()
+
+    assert session.status is ARDiffusionSessionStatus.FAILED
+    assert session.chunk_index == 0
+    assert session.pending_event_count == 1
+    assert lifecycle.closes == ["world-1"]
+    with pytest.raises(ARDiffusionSessionStateError, match="failed"):
+        await session.next_chunk()
+
+    await session.reset()
+    assert session.status is ARDiffusionSessionStatus.ACTIVE
+    assert session.pending_event_count == 0
+    assert lifecycle.resets == ["world-1"]
+    await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
+    metadata = await session.next_chunk()
+    assert metadata.chunk_index == 0
+    assert metadata.applied_event_ids == (4,)
+    assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
+
+
+@pytest.mark.asyncio
+async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
+    consumer = FakeTickConsumer()
+    consumer.started = asyncio.Event()
+    consumer.release = asyncio.Event()
+    session, _, _ = make_session(consumer=consumer)
+    await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
+
+    first_chunk = asyncio.create_task(session.next_chunk())
+    await consumer.started.wait()
+    accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
+    assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+    assert consumer.ticks[0].prompt == "first"
+    assert consumer.ticks[0].applied_event_ids == (10,)
+
+    with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
+        await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
+
+    consumer.release.set()
+    first_metadata = await first_chunk
+    second_metadata = await session.next_chunk()
+    assert first_metadata.applied_event_ids == (10,)
+    assert second_metadata.applied_event_ids == (11,)
+    assert consumer.ticks[1].prompt == "second"
+
+
+@pytest.mark.asyncio
+async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
+    consumer = FakeTickConsumer()
+    consumer.metadata_override = ARDiffusionChunkMetadata(
+        session_id="another-world",
+        request_id="wrong",
+        chunk_index=0,
+    )
+    lifecycle = FakeLifecycle()
+    session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
+    await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
+
+    with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
+        await session.next_chunk()
+
+    assert session.chunk_index == 0
+    assert session.pending_event_count == 1
+    assert session.status is ARDiffusionSessionStatus.FAILED
+    assert lifecycle.closes == ["world-1"]
+
+
+@pytest.mark.asyncio
+async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
+    consumer = FakeTickConsumer()
+    lifecycle = FakeLifecycle()
+    reducer = FailingCommitReducer()
+    session = ARDiffusionSession(
+        "world-1",
+        tick_consumer=consumer,
+        lifecycle=lifecycle,
+        max_pending_events=4,
+        request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+        control_reducer=reducer,
+    )
+    await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
+
+    with pytest.raises(RuntimeError, match="fake reducer commit failed"):
+        await session.next_chunk()
+
+    assert reducer.commit_calls == 1
+    assert session.status is ARDiffusionSessionStatus.FAILED
+    assert session.chunk_index == 0
+    assert session.pending_event_count == 1
+    assert lifecycle.closes == ["world-1"]
+    with pytest.raises(ARDiffusionSessionStateError, match="failed"):
+        await session.next_chunk()
+
+    await session.reset()
+    assert reducer.reset_calls == 1
+    assert session.status is ARDiffusionSessionStatus.ACTIVE
+    assert session.chunk_index == 0
+    assert session.pending_event_count == 0
+
+
+@pytest.mark.asyncio
+async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
+    consumer = FakeTickConsumer()
+    lifecycle = FakeLifecycle()
+    manager = ARDiffusionSessionManager(
+        tick_consumer=consumer,
+        lifecycle=lifecycle,
+        max_pending_events=4,
+        request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
+    )
+    session = await manager.create_session("world-reset")
+    await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
+    await session.next_chunk()
+    await manager.reset_session("world-reset")
+
+    assert lifecycle.resets == ["world-reset"]
+    assert session.status is ARDiffusionSessionStatus.ACTIVE
+    assert session.chunk_index == 0
+    duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
+    assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+
+    reset_tick = await session.next_chunk()
+    assert reset_tick.chunk_index == 0
+    assert consumer.ticks[-1].prompt is None
+    assert consumer.ticks[-1].controls == ()
+
+    await manager.close_session("world-reset")
+    assert lifecycle.closes == ["world-reset"]
+    assert session.status is ARDiffusionSessionStatus.CLOSED
+    with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
+        await manager.get_session("world-reset")
+
+    disconnected = await manager.create_session("world-disconnect")
+    await manager.disconnect("world-disconnect")
+    assert disconnected.status is ARDiffusionSessionStatus.CLOSED
+    assert lifecycle.closes[-1] == "world-disconnect"
+
+
+@pytest.mark.asyncio
+async def test_manager_keeps_interleaved_sessions_independent() -> None:
+    consumer = FakeTickConsumer()
+    manager = ARDiffusionSessionManager(
+        tick_consumer=consumer,
+        lifecycle=FakeLifecycle(),
+        max_pending_events=4,
+        request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+    )
+    session_a = await manager.create_session("world-a")
+    session_b = await manager.create_session("world-b")
+
+    await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
+    await session_a.next_chunk()
+    await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
+    await session_b.next_chunk()
+    await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
+    await session_a.next_chunk()
+
+    assert [
+        (
+            tick.session_id,
+            tick.request_id,
+            tick.chunk_index,
+            tick.applied_event_ids,
+        )
+        for tick in consumer.ticks
+    ] == [
+        ("world-a", "world-a-request-0", 0, (10,)),
+        ("world-b", "world-b-request-0", 0, (1,)),
+        ("world-a", "world-a-request-1", 1, (11,)),
+    ]
+    assert session_a.chunk_index == 2
+    assert session_b.chunk_index == 1
+    with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
+        await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
+    assert (
+        await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
+    ).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+
+
+@pytest.mark.asyncio
+async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
+    consumer = FakeTickConsumer()
+    lifecycle = FakeLifecycle()
+    lifecycle.close_failures_remaining = 1
+    manager = ARDiffusionSessionManager(
+        tick_consumer=consumer,
+        lifecycle=lifecycle,
+        max_pending_events=4,
+    )
+    session = await manager.create_session("world-close-retry")
+
+    with pytest.raises(RuntimeError, match="fake close failed"):
+        await manager.close_session("world-close-retry")
+
+    assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
+    assert await manager.get_session("world-close-retry") is session
+    with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
+        await manager.create_session("world-close-retry")
+    with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
+        await manager.reset_session("world-close-retry")
+
+    await manager.close_session("world-close-retry")
+
+    assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
+    assert session.status is ARDiffusionSessionStatus.CLOSED
+    with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
+        await manager.get_session("world-close-retry")
+    recreated = await manager.create_session("world-close-retry")
+    assert recreated is not session
+
+
+@pytest.mark.asyncio
+async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
+    lifecycle = FakeLifecycle()
+    lifecycle.close_failures_remaining = 1
+    manager = ARDiffusionSessionManager(
+        tick_consumer=FakeTickConsumer(),
+        lifecycle=lifecycle,
+        max_pending_events=4,
+    )
+    session = await manager.create_session("world-disconnect-retry")
+
+    with pytest.raises(RuntimeError, match="fake close failed"):
+        await manager.disconnect("world-disconnect-retry")
+
+    assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
+    assert await manager.get_session("world-disconnect-retry") is session
+
+    await manager.close_session("world-disconnect-retry")
+    assert session.status is ARDiffusionSessionStatus.CLOSED
+
+
+class FakeRPCClient:
+    def __init__(self, result: list[Any]) -> None:
+        self.result = result
+        self.calls: list[dict[str, Any]] = []
+
+    async def collective_rpc(self, **kwargs: Any) -> list[Any]:
+        self.calls.append(kwargs)
+        return self.result
+
+
+@pytest.mark.asyncio
+async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
+    client = FakeRPCClient([[True]])
+    lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
+
+    await lifecycle.reset_session("world-1")
+    await lifecycle.close_session("world-1")
+
+    assert client.calls == [
+        {
+            "method": "reset_ar_diffusion_session",
+            "timeout": 3.0,
+            "args": ("world-1",),
+            "stage_ids": [2],
+        },
+        {
+            "method": "close_ar_diffusion_session",
+            "timeout": 3.0,
+            "args": ("world-1",),
+            "stage_ids": [2],
+        },
+    ]
+
+
+@pytest.mark.asyncio
+async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
+    lifecycle = ARDiffusionWorkerLifecycle(
+        FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
+        stage_ids=[1],
+    )
+
+    with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
+        await lifecycle.close_session("world-1")

tests/diffusion/ar_diffusion/test_tick_protocol.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionChunkMetadata,
+    ARDiffusionControlInput,
+    ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
+    tick = ARDiffusionTickRequest(
+        session_id="world-1",
+        request_id="chunk-request-7",
+        chunk_index=7,
+        applied_event_ids=(10, 11),
+        prompt="turn left",
+        controls=(
+            ARDiffusionControlInput(
+                track="camera",
+                schema="lingbot.camera_trajectory.v1",
+                data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
+            ),
+        ),
+    )
+
+    restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
+
+    assert restored == tick
+    assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
+        "session_id": "world-1",
+        "request_id": "chunk-request-7",
+        "chunk_index": 7,
+        "applied_event_ids": [10, 11],
+    }
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        ("session_id", "", "session_id"),
+        ("request_id", "", "request_id"),
+        ("chunk_index", -1, "chunk_index"),
+        ("applied_event_ids", (2, 1), "strictly increasing"),
+        ("applied_event_ids", (1, 1), "strictly increasing"),
+    ],
+)
+def test_tick_rejects_invalid_identity(
+    field: str,
+    value: object,
+    message: str,
+) -> None:
+    values = {
+        "session_id": "world-1",
+        "request_id": "request-1",
+        "chunk_index": 0,
+        "applied_event_ids": (),
+    }
+    values[field] = value
+
+    with pytest.raises(ValueError, match=message):
+        ARDiffusionTickRequest(**values)
+
+
+def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
+    payload = {"samples": [{"keys": ["w"]}]}
+    event_ids = [1, 2]
+    controls = [
+        ARDiffusionControlInput(
+            track="camera",
+            schema="example.camera.v1",
+            data=payload,
+        )
+    ]
+    tick = ARDiffusionTickRequest(
+        session_id="world-1",
+        request_id="request-a",
+        chunk_index=0,
+        applied_event_ids=event_ids,  # type: ignore[arg-type]
+        controls=controls,  # type: ignore[arg-type]
+    )
+    serialized = tick.to_extra_args()
+
+    payload["samples"][0]["keys"].append("j")
+    event_ids.append(3)
+    controls.clear()
+    serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
+
+    assert tick.applied_event_ids == (1, 2)
+    assert isinstance(tick.controls, tuple)
+    assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
+        {
+            "track": "camera",
+            "schema": "example.camera.v1",
+            "data": {"samples": [{"keys": ["w"]}]},
+        }
+    ]
+    with pytest.raises(TypeError):
+        tick.controls[0].data["samples"] = ()  # type: ignore[index]
+    with pytest.raises(AttributeError):
+        tick.controls[0].data["samples"][0]["keys"].append("j")
+
+
+def test_legacy_request_has_no_typed_tick() -> None:
+    assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None

tests/diffusion/ar_diffusion/test_worker_lifecycle.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeARDiffusionRunner:
+    def __init__(self) -> None:
+        self.resets: list[str] = []
+        self.closes: list[str] = []
+
+    def reset_session(self, session_id: str) -> None:
+        self.resets.append(session_id)
+
+    def close_session(self, session_id: str) -> None:
+        self.closes.append(session_id)
+
+
+class FakeRegularRunner:
+    pass
+
+
+def make_worker(runner: object) -> DiffusionWorker:
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = runner
+    return worker
+
+
+def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
+    runner = FakeARDiffusionRunner()
+    worker = make_worker(runner)
+
+    assert worker.reset_ar_diffusion_session("world-1") is True
+    assert worker.close_ar_diffusion_session("world-1") is True
+    assert runner.resets == ["world-1"]
+    assert runner.closes == ["world-1"]
+
+
+def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
+    worker = make_worker(FakeRegularRunner())
+
+    assert worker.reset_ar_diffusion_session("world-1") is False
+    assert worker.close_ar_diffusion_session("world-1") is False
+
+
+def test_worker_rejects_empty_session_id() -> None:
+    worker = make_worker(FakeARDiffusionRunner())
+
+    with pytest.raises(ValueError, match="session_id"):
+        worker.close_ar_diffusion_session("")

tests/diffusion/models/dreamzero/test_pipeline_state.py

@@ -6,7 +6,10 @@ from types import SimpleNamespace
 
 import pytest
 
-from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
+    DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
+    DreamZeroPipeline,
+)
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,6 +35,37 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None:
     assert session_b.call_count == 3
 
 
+def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None:
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline.transformer = SimpleNamespace(
+        frame_seqlen=16,
+        blocks=[
+            SimpleNamespace(
+                self_attn=SimpleNamespace(
+                    max_attention_size=96,
+                    tp_num_heads=4,
+                )
+            )
+        ],
+        text_len=8,
+        model_type="t2v",
+        num_layers=2,
+        num_heads=4,
+        dim=256,
+        num_frame_per_block=4,
+        num_action_per_block=2,
+        num_state_per_block=1,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size",
+        lambda: 1,
+    )
+
+    spec = pipeline.ar_diffusion_kv_cache_spec()
+
+    assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION
+
+
 def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
     pipeline = _empty_pipeline()
 

tests/diffusion/models/lingbot_world/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture

@@ -0,0 +1,13 @@
+{
+  "_class_name": "LingBotWorldCausalDMDPipeline",
+  "_diffusers_version": "0.35.0.dev0",
+  "boundary_ratio": 0.9,
+  "image_encoder": [null, null],
+  "image_processor": [null, null],
+  "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+  "text_encoder": ["transformers", "UMT5EncoderModel"],
+  "tokenizer": ["transformers", "T5TokenizerFast"],
+  "transformer": ["diffusers", "CausalLingBotWorldTransformer3DModel"],
+  "transformer_2": [null, null],
+  "vae": ["diffusers", "AutoencoderKLWan"]
+}

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture

@@ -0,0 +1,42 @@
+{
+  "sources": {
+    "config": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/config.json",
+    "shard_1_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00001-of-00008.safetensors",
+    "shard_8_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00008-of-00008.safetensors"
+  },
+  "dtype": "F32",
+  "shard_representatives": {
+    "diffusion_pytorch_model-00001-of-00008.safetensors": "blocks.0.cam_injector_layer1.bias",
+    "diffusion_pytorch_model-00002-of-00008.safetensors": "blocks.4.cam_injector_layer1.bias",
+    "diffusion_pytorch_model-00003-of-00008.safetensors": "blocks.10.cam_injector_layer1.bias",
+    "diffusion_pytorch_model-00004-of-00008.safetensors": "blocks.14.cam_scale_layer.bias",
+    "diffusion_pytorch_model-00005-of-00008.safetensors": "blocks.20.cam_injector_layer1.bias",
+    "diffusion_pytorch_model-00006-of-00008.safetensors": "blocks.25.cam_injector_layer1.bias",
+    "diffusion_pytorch_model-00007-of-00008.safetensors": "blocks.30.cam_injector_layer1.bias",
+    "diffusion_pytorch_model-00008-of-00008.safetensors": "blocks.35.cam_injector_layer1.bias"
+  },
+  "parameter_shapes": {
+    "blocks.0.cam_injector_layer1.bias": [5120],
+    "blocks.4.cam_injector_layer1.bias": [5120],
+    "blocks.10.cam_injector_layer1.bias": [5120],
+    "blocks.14.cam_scale_layer.bias": [5120],
+    "blocks.20.cam_injector_layer1.bias": [5120],
+    "blocks.25.cam_injector_layer1.bias": [5120],
+    "blocks.30.cam_injector_layer1.bias": [5120],
+    "blocks.35.cam_injector_layer1.bias": [5120],
+    "patch_embedding.weight": [5120, 36, 1, 2, 2],
+    "patch_embedding_wancamctrl.weight": [5120, 1536],
+    "c2ws_hidden_states_layer1.weight": [5120, 5120],
+    "c2ws_hidden_states_layer2.weight": [5120, 5120],
+    "blocks.0.modulation": [1, 6, 5120],
+    "blocks.0.self_attn.q.weight": [5120, 5120],
+    "blocks.0.ffn.0.weight": [13824, 5120],
+    "blocks.0.ffn.2.weight": [5120, 13824],
+    "blocks.0.cam_injector_layer1.weight": [5120, 5120],
+    "blocks.0.cam_injector_layer2.weight": [5120, 5120],
+    "blocks.0.cam_scale_layer.weight": [5120, 5120],
+    "blocks.0.cam_shift_layer.weight": [5120, 5120],
+    "head.modulation": [1, 2, 5120],
+    "head.head.weight": [64, 5120]
+  }
+}

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture

@@ -0,0 +1,30 @@
+{
+  "_class_name": "UniPCMultistepScheduler",
+  "_diffusers_version": "0.35.0.dev0",
+  "beta_end": 0.02,
+  "beta_schedule": "linear",
+  "beta_start": 0.0001,
+  "disable_corrector": [],
+  "dynamic_thresholding_ratio": 0.995,
+  "final_sigmas_type": "zero",
+  "flow_shift": 5.0,
+  "lower_order_final": true,
+  "num_train_timesteps": 1000,
+  "predict_x0": true,
+  "prediction_type": "flow_prediction",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "solver_order": 2,
+  "solver_p": null,
+  "solver_type": "bh2",
+  "steps_offset": 0,
+  "thresholding": false,
+  "time_shift_type": "exponential",
+  "timestep_spacing": "linspace",
+  "trained_betas": null,
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_flow_sigmas": true,
+  "use_karras_sigmas": false
+}

tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture

@@ -0,0 +1,112 @@
+{
+  "source": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model.safetensors.index.json",
+  "index_parameter_count": 1421,
+  "canonical_name_set_sha256": "5d31b7f64fd5e5e271b789b2dbe50fb7bbf146bd8b750646bdc9abb4e5b8139d",
+  "canonical_name_to_shard_sha256": "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992",
+  "canonical_serialization": "sorted lines formatted as parameter_name=shard_filename with LF separators and no trailing LF",
+  "shards": {
+    "diffusion_pytorch_model-00001-of-00008.safetensors": {
+      "parameter_count": 180,
+      "parameter_names_sha256": "81a3d4778de604e75609aae7dac3d8dcbc081811e7748267fe953bbeb375e79d"
+    },
+    "diffusion_pytorch_model-00002-of-00008.safetensors": {
+      "parameter_count": 180,
+      "parameter_names_sha256": "4e36e2596a6468f29179fa5795396f2e9a7f7d8debc912b34a6f6c2c7c7c610d"
+    },
+    "diffusion_pytorch_model-00003-of-00008.safetensors": {
+      "parameter_count": 180,
+      "parameter_names_sha256": "6f74a7b2cf7444c2daa2b3fc32bd3787589b35543c230c5c13b89f0f0f2a0998"
+    },
+    "diffusion_pytorch_model-00004-of-00008.safetensors": {
+      "parameter_count": 180,
+      "parameter_names_sha256": "ec5c967ed5271ba7834337c4ee93a73f8ed0eead4492d4920624251c65c058ac"
+    },
+    "diffusion_pytorch_model-00005-of-00008.safetensors": {
+      "parameter_count": 180,
+      "parameter_names_sha256": "ad1fdab5830dd9ce43ea007f112b4be223327b7c8197c766d9cf9219e74d619b"
+    },
+    "diffusion_pytorch_model-00006-of-00008.safetensors": {
+      "parameter_count": 180,
+      "parameter_names_sha256": "fd58f0d86be42d8afcf55d3e07d8a22f7f6ef5064da931a600c890b71b6dfbae"
+    },
+    "diffusion_pytorch_model-00007-of-00008.safetensors": {
+      "parameter_count": 180,
+      "parameter_names_sha256": "305a1714b5f4fb54d5877c593b451273ab1490d90fd0896a299d96e140295047"
+    },
+    "diffusion_pytorch_model-00008-of-00008.safetensors": {
+      "parameter_count": 161,
+      "parameter_names_sha256": "4c18d94d96168682eb212f93e3cdf7313928776d8d9e6e4de0255fe899decb25"
+    }
+  },
+  "top_level_modules": [
+    "blocks",
+    "c2ws_hidden_states_layer1",
+    "c2ws_hidden_states_layer2",
+    "head",
+    "patch_embedding",
+    "patch_embedding_wancamctrl",
+    "text_embedding",
+    "time_embedding",
+    "time_projection"
+  ],
+  "non_block_parameters": [
+    "c2ws_hidden_states_layer1.bias",
+    "c2ws_hidden_states_layer1.weight",
+    "c2ws_hidden_states_layer2.bias",
+    "c2ws_hidden_states_layer2.weight",
+    "head.head.bias",
+    "head.head.weight",
+    "head.modulation",
+    "patch_embedding.bias",
+    "patch_embedding.weight",
+    "patch_embedding_wancamctrl.bias",
+    "patch_embedding_wancamctrl.weight",
+    "text_embedding.0.bias",
+    "text_embedding.0.weight",
+    "text_embedding.2.bias",
+    "text_embedding.2.weight",
+    "time_embedding.0.bias",
+    "time_embedding.0.weight",
+    "time_embedding.2.bias",
+    "time_embedding.2.weight",
+    "time_projection.1.bias",
+    "time_projection.1.weight"
+  ],
+  "block_parameter_suffixes": [
+    "cam_injector_layer1.bias",
+    "cam_injector_layer1.weight",
+    "cam_injector_layer2.bias",
+    "cam_injector_layer2.weight",
+    "cam_scale_layer.bias",
+    "cam_scale_layer.weight",
+    "cam_shift_layer.bias",
+    "cam_shift_layer.weight",
+    "cross_attn.k.bias",
+    "cross_attn.k.weight",
+    "cross_attn.norm_k.weight",
+    "cross_attn.norm_q.weight",
+    "cross_attn.o.bias",
+    "cross_attn.o.weight",
+    "cross_attn.q.bias",
+    "cross_attn.q.weight",
+    "cross_attn.v.bias",
+    "cross_attn.v.weight",
+    "ffn.0.bias",
+    "ffn.0.weight",
+    "ffn.2.bias",
+    "ffn.2.weight",
+    "modulation",
+    "norm3.bias",
+    "norm3.weight",
+    "self_attn.k.bias",
+    "self_attn.k.weight",
+    "self_attn.norm_k.weight",
+    "self_attn.norm_q.weight",
+    "self_attn.o.bias",
+    "self_attn.o.weight",
+    "self_attn.q.bias",
+    "self_attn.q.weight",
+    "self_attn.v.bias",
+    "self_attn.v.weight"
+  ]
+}

tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.lingbot_world.actions import (
+    LINGBOT_CAMERA_ACTION_SCHEMA,
+    LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
+    LingBotCameraControlReducer,
+    integrate_lingbot_camera_actions,
+    parse_lingbot_camera_action_frames,
+)
+from vllm_omni.experimental.ar_diffusion.session import (
+    ARDiffusionSession,
+    ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionChunkMetadata,
+    ARDiffusionControlInput,
+    ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _action_control(mode: str, **data) -> ARDiffusionControlInput:
+    return ARDiffusionControlInput(
+        track="camera",
+        schema=LINGBOT_CAMERA_ACTION_SCHEMA,
+        data={"mode": mode, **data},
+    )
+
+
+def _prepared_frames(prepared) -> tuple[tuple[str, ...], ...]:
+    assert len(prepared.controls) == 1
+    return prepared.controls[0].data["frames"]
+
+
+def test_state_controls_hold_release_and_preserve_short_pulse() -> None:
+    reducer = LingBotCameraControlReducer()
+    pressed = reducer.prepare(
+        current_controls={},
+        events=(
+            ARDiffusionSessionEvent(
+                event_id=1,
+                controls=(
+                    _action_control(
+                        "state",
+                        transitions=[{"actions": ["W"], "client_ts_ms": 10}],
+                    ),
+                ),
+            ),
+        ),
+        chunk_index=0,
+    )
+    assert _prepared_frames(pressed) == (("w",), ("w",), ("w",))
+    assert pressed.controls[0].to_dict()["data"]["frames"] == [["w"], ["w"], ["w"]]
+    reducer.commit(pressed)
+
+    held = reducer.prepare(
+        current_controls={"camera": pressed.controls[0]},
+        events=(),
+        chunk_index=1,
+    )
+    assert _prepared_frames(held) == (("w",), ("w",), ("w",))
+    reducer.commit(held)
+
+    pulse = reducer.prepare(
+        current_controls={"camera": held.controls[0]},
+        events=(
+            ARDiffusionSessionEvent(
+                event_id=2,
+                controls=(
+                    _action_control(
+                        "state",
+                        transitions=[
+                            {"actions": ["j"], "client_ts_ms": 20},
+                            {"actions": [], "client_ts_ms": 21},
+                        ],
+                    ),
+                ),
+            ),
+        ),
+        chunk_index=2,
+    )
+    assert _prepared_frames(pulse) == (("j",), (), ())
+    reducer.commit(pulse)
+
+    released = reducer.prepare(
+        current_controls={"camera": pulse.controls[0]},
+        events=(),
+        chunk_index=3,
+    )
+    assert _prepared_frames(released) == ((), (), ())
+
+
+def test_script_is_transactional_consumed_once_and_neutral_padded() -> None:
+    reducer = LingBotCameraControlReducer()
+    event = ARDiffusionSessionEvent(
+        event_id=1,
+        controls=(
+            _action_control(
+                "script",
+                frames=[["w"], ["d"], ["l"], ["i"]],
+            ),
+        ),
+    )
+    first_attempt = reducer.prepare(
+        current_controls={},
+        events=(event,),
+        chunk_index=0,
+    )
+    assert _prepared_frames(first_attempt) == (("w",), ("d",), ("l",))
+
+    # An uncommitted failed attempt must not consume the reducer's script.
+    second_attempt = reducer.prepare(
+        current_controls={},
+        events=(event,),
+        chunk_index=0,
+    )
+    assert _prepared_frames(second_attempt) == (("w",), ("d",), ("l",))
+    reducer.commit(second_attempt)
+
+    tail = reducer.prepare(
+        current_controls={"camera": second_attempt.controls[0]},
+        events=(),
+        chunk_index=1,
+    )
+    assert _prepared_frames(tail) == (("i",), (), ())
+    reducer.commit(tail)
+    exhausted = reducer.prepare(
+        current_controls={"camera": tail.controls[0]},
+        events=(),
+        chunk_index=2,
+    )
+    assert _prepared_frames(exhausted) == ((), (), ())
+
+
+def test_trajectory_control_remains_an_expert_override() -> None:
+    reducer = LingBotCameraControlReducer()
+    trajectory = ARDiffusionControlInput(
+        track="camera",
+        schema=LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
+        data={"poses": [[[1.0]]], "intrinsics": [[1.0]]},
+    )
+    prepared = reducer.prepare(
+        current_controls={},
+        events=(
+            ARDiffusionSessionEvent(
+                event_id=1,
+                controls=(trajectory,),
+            ),
+        ),
+        chunk_index=0,
+    )
+    assert prepared.controls == (trajectory,)
+
+
+def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
+    trajectory, pitch = integrate_lingbot_camera_actions(
+        [["w"], ["d"], ["l"]],
+        width=832,
+        height=480,
+    )
+
+    assert pitch == 0.0
+    torch.testing.assert_close(
+        trajectory.poses[0, :3, 3],
+        torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64),
+        atol=1e-6,
+        rtol=0,
+    )
+    torch.testing.assert_close(
+        trajectory.poses[1, :3, 3],
+        torch.tensor([0.05, 0.0, 0.05], dtype=torch.float64),
+        atol=1e-6,
+        rtol=0,
+    )
+    expected_yaw = torch.tensor(
+        [
+            [0.9945219, 0.0, 0.10452846],
+            [0.0, 1.0, 0.0],
+            [-0.10452846, 0.0, 0.9945219],
+        ],
+        dtype=torch.float64,
+    )
+    torch.testing.assert_close(
+        trajectory.poses[2, :3, :3],
+        expected_yaw,
+        atol=1e-6,
+        rtol=0,
+    )
+    torch.testing.assert_close(
+        trajectory.intrinsics[0],
+        torch.tensor([500.0, 500.0, 416.0, 240.0]),
+    )
+
+    resized, _ = integrate_lingbot_camera_actions(
+        [[]],
+        width=416,
+        height=240,
+    )
+    scaled = resized.intrinsics[0] * torch.tensor([416 / 832, 240 / 480, 416 / 832, 240 / 480])
+    torch.testing.assert_close(
+        scaled,
+        torch.tensor([500.0, 500.0, 208.0, 120.0]),
+    )
+
+
+@pytest.mark.parametrize(
+    ("data", "message"),
+    [
+        ({"mode": "frames", "frames": [["x"], [], []]}, "W/A/S/D/I/J/K/L"),
+        ({"mode": "frames", "frames": [["w"]]}, "exactly one action list"),
+        ({"mode": "state", "frames": [[], [], []]}, "mode='frames'"),
+    ],
+)
+def test_model_tick_action_payload_validation(data, message: str) -> None:
+    with pytest.raises(ValueError, match=message):
+        parse_lingbot_camera_action_frames(data, expected_frames=3)
+
+
+class _TickConsumer:
+    def __init__(self) -> None:
+        self.ticks: list[ARDiffusionTickRequest] = []
+
+    async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+        self.ticks.append(tick)
+        return ARDiffusionChunkMetadata.from_tick(tick)
+
+    def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
+        return output
+
+
+class _Lifecycle:
+    async def reset_session(self, session_id: str) -> None:
+        del session_id
+
+    async def close_session(self, session_id: str) -> None:
+        del session_id
+
+
+def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
+    async def exercise() -> None:
+        consumer = _TickConsumer()
+        reducer = LingBotCameraControlReducer()
+        session = ARDiffusionSession(
+            "world-actions",
+            tick_consumer=consumer,
+            lifecycle=_Lifecycle(),
+            max_pending_events=8,
+            request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
+            control_reducer=reducer,
+        )
+        await session.accept_event(
+            ARDiffusionSessionEvent(
+                event_id=1,
+                prompt="turn toward the window",
+                controls=(
+                    _action_control(
+                        "state",
+                        transitions=[{"actions": ["d"]}],
+                    ),
+                ),
+            )
+        )
+        await session.next_chunk()
+
+        tick = consumer.ticks[0]
+        assert tick.prompt == "turn toward the window"
+        assert tick.applied_event_ids == (1,)
+        assert tick.controls[0].data["frames"] == (("d",), ("d",), ("d",))
+
+        await session.reset()
+        await session.accept_event(
+            ARDiffusionSessionEvent(
+                event_id=2,
+                controls=(
+                    _action_control(
+                        "state",
+                        transitions=[{"actions": []}],
+                    ),
+                ),
+            )
+        )
+        await session.next_chunk()
+        assert consumer.ticks[1].chunk_index == 0
+        assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
+
+    asyncio.run(exercise())

tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py

@@ -0,0 +1,543 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import sys
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/transformer.py"
+_MODULE_NAME = "_lingbot_world_attention_under_test"
+_MISSING_MODULE = object()
+_STUBBED_MODULE_NAMES = (
+    "vllm",
+    "vllm.distributed",
+    "vllm.model_executor",
+    "vllm.model_executor.layers",
+    "vllm.model_executor.layers.conv",
+    "vllm.model_executor.layers.linear",
+    "vllm.model_executor.model_loader",
+    "vllm.model_executor.model_loader.weight_utils",
+    "vllm.model_executor.utils",
+    "vllm_omni",
+    "vllm_omni.diffusion",
+    "vllm_omni.diffusion.attention",
+    "vllm_omni.diffusion.attention.layer",
+    "vllm_omni.diffusion.layers",
+    "vllm_omni.diffusion.layers.norm",
+    "vllm_omni.diffusion.layers.rope",
+    "vllm_omni.experimental",
+    "vllm_omni.experimental.ar_diffusion",
+    "vllm_omni.experimental.ar_diffusion.kv_cache",
+    "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention",
+)
+
+
+@pytest.fixture(autouse=True)
+def _inference_context():
+    with torch.inference_mode():
+        yield
+
+
+def _install_vllm_stubs() -> None:
+    for name in _STUBBED_MODULE_NAMES:
+        sys.modules[name] = ModuleType(name)
+
+    distributed = sys.modules["vllm.distributed"]
+    distributed.get_tensor_model_parallel_rank = lambda: 0
+    distributed.get_tensor_model_parallel_world_size = lambda: 1
+    distributed.tensor_model_parallel_all_reduce = lambda value: value
+
+    def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
+        for name, value in attrs.items():
+            setattr(weight, name, value)
+
+    sys.modules["vllm.model_executor.utils"].set_weight_attrs = set_weight_attrs
+    sys.modules["vllm.model_executor.model_loader.weight_utils"].default_weight_loader = lambda param, loaded_weight: (
+        param.data.copy_(loaded_weight)
+    )
+
+    class _Linear(nn.Module):
+        def __init__(
+            self,
+            input_size: int,
+            output_size: int,
+            *,
+            bias: bool = True,
+            return_bias: bool = False,
+            **kwargs,
+        ) -> None:
+            super().__init__()
+            del kwargs
+            self.return_bias = return_bias
+            self.weight = nn.Parameter(torch.empty(output_size, input_size))
+            self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
+            nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+            if self.bias is not None:
+                nn.init.zeros_(self.bias)
+            self.calls = 0
+
+        def forward(self, value: torch.Tensor):
+            self.calls += 1
+            output = F.linear(value, self.weight, self.bias)
+            return (output, self.bias) if self.return_bias else output
+
+    linear = sys.modules["vllm.model_executor.layers.linear"]
+    linear.ColumnParallelLinear = _Linear
+    linear.RowParallelLinear = _Linear
+
+    class _QKVLinear(nn.Module):
+        def __init__(
+            self,
+            hidden_size: int,
+            head_size: int,
+            total_num_heads: int,
+            total_num_kv_heads: int | None = None,
+            *,
+            bias: bool = True,
+            **kwargs,
+        ) -> None:
+            super().__init__()
+            del kwargs
+            self.num_heads = total_num_heads
+            self.num_kv_heads = total_num_kv_heads or total_num_heads
+            output_size = (self.num_heads + 2 * self.num_kv_heads) * head_size
+            self.weight = nn.Parameter(torch.empty(output_size, hidden_size))
+            self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
+            nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+            if self.bias is not None:
+                nn.init.zeros_(self.bias)
+            self.calls = 0
+
+            def load_shard(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
+                shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
+                shard_size = param.shape[0] // 3
+                param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
+
+            self.weight.weight_loader = load_shard
+            if self.bias is not None:
+                self.bias.weight_loader = load_shard
+
+        def forward(self, value: torch.Tensor):
+            self.calls += 1
+            return F.linear(value, self.weight, self.bias), None
+
+    linear.QKVParallelLinear = _QKVLinear
+
+    class _Conv3dLayer(nn.Conv3d):
+        def __init__(self, in_channels: int, out_channels: int, kernel_size, stride, **kwargs) -> None:
+            super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs)
+
+    sys.modules["vllm.model_executor.layers.conv"].Conv3dLayer = _Conv3dLayer
+
+    class _LayerNorm(nn.LayerNorm):
+        def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True) -> None:
+            super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
+
+    sys.modules["vllm_omni.diffusion.layers.norm"].LayerNorm = _LayerNorm
+
+    class _Attention(nn.Module):
+        def __init__(self, *args, softmax_scale: float, **kwargs) -> None:
+            super().__init__()
+            del args
+            self.softmax_scale = softmax_scale
+            self.causal = kwargs["causal"]
+            self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+        def forward(
+            self,
+            query: torch.Tensor,
+            key: torch.Tensor,
+            value: torch.Tensor,
+            attn_metadata=None,
+        ) -> torch.Tensor:
+            del attn_metadata
+            self.calls.append((query.detach().clone(), key.detach().clone(), value.detach().clone()))
+            scores = torch.einsum("bqhd,bkhd->bhqk", query, key) * self.softmax_scale
+            weights = scores.softmax(dim=-1)
+            return torch.einsum("bhqk,bkhd->bqhd", weights, value)
+
+    sys.modules["vllm_omni.diffusion.attention.layer"].Attention = _Attention
+
+    class _RotaryEmbeddingWan(nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__()
+            del args, kwargs
+            self.calls: list[tuple[torch.Tensor, torch.Tensor]] = []
+
+        def forward(self, value: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+            self.calls.append((cos.detach().clone(), sin.detach().clone()))
+            return value
+
+    sys.modules["vllm_omni.diffusion.layers.rope"].RotaryEmbeddingWan = _RotaryEmbeddingWan
+
+    paged_attention = sys.modules["vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"]
+
+    class _PagedLayerContext:
+        pass
+
+    class _PagedLayerInputs:
+        pass
+
+    paged_attention.ARDiffusionPagedLayerContext = _PagedLayerContext
+    paged_attention.ARDiffusionPagedLayerInputs = _PagedLayerInputs
+    paged_attention.ar_diffusion_paged_attention = lambda *args, **kwargs: (_ for _ in ()).throw(
+        AssertionError("CUDA paged attention was not expected in this CPU test")
+    )
+    paged_attention.paged_write_attn = lambda *args, **kwargs: (_ for _ in ()).throw(
+        AssertionError("paged path was not expected in this direct-cache test")
+    )
+
+
+def _load_module():
+    assert _MODULE_PATH.exists(), "LingBot attention module has not been implemented"
+    module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
+    previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
+    try:
+        _install_vllm_stubs()
+        spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+        assert spec is not None and spec.loader is not None
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[spec.name] = module
+        spec.loader.exec_module(module)
+        return module
+    finally:
+        for name, original in previous.items():
+            if original is _MISSING_MODULE:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = original
+
+
+def test_load_module_restores_process_module_state() -> None:
+    module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
+    previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
+
+    _load_module()
+
+    for name, original in previous.items():
+        assert sys.modules.get(name, _MISSING_MODULE) is original
+
+
+def test_allocate_lingbot_cache_creates_request_local_layer_storage() -> None:
+    module = _load_module()
+
+    cache = module.allocate_lingbot_cache(
+        batch_size=2,
+        num_layers=3,
+        max_tokens=7,
+        num_local_heads=2,
+        head_dim=4,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    assert len(cache.self_attention) == 3
+    assert len(cache.cross_attention) == 3
+    assert cache.cross_attention == [None, None, None]
+    assert cache.self_attention[0].key.shape == (2, 7, 2, 4)
+    assert cache.self_attention[0].value.shape == (2, 7, 2, 4)
+    assert cache.self_attention[0].key.dtype == torch.float32
+    assert cache.self_attention[0].end == 0
+    assert cache.self_attention[0].key.data_ptr() != cache.self_attention[1].key.data_ptr()
+
+
+def _allocate_single_layer(module, *, max_tokens: int):
+    return module.allocate_lingbot_cache(
+        batch_size=1,
+        num_layers=1,
+        max_tokens=max_tokens,
+        num_local_heads=1,
+        head_dim=2,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+
+def _set_identity_attention(attention: nn.Module) -> None:
+    qkv = getattr(attention, "qkv", None)
+    if qkv is not None:
+        projections = (qkv, attention.o)
+    else:
+        projections = (attention.q, attention.k, attention.v, attention.o)
+    for projection in projections:
+        with torch.no_grad():
+            identity = torch.eye(2)
+            projection.weight.copy_(torch.cat((identity, identity, identity)) if projection is qkv else identity)
+            if projection.bias is not None:
+                projection.bias.zero_()
+    attention.norm_q = nn.Identity()
+    attention.norm_k = nn.Identity()
+
+
+def _tokens(*values: float) -> torch.Tensor:
+    return torch.tensor([[[value, 0.0] for value in values]])
+
+
+def _record_inputs(module: nn.Module) -> list[tuple]:
+    records: list[tuple] = []
+    module.register_forward_pre_hook(lambda _module, args: records.append(args))
+    return records
+
+
+def _record_outputs(module: nn.Module) -> list[torch.Tensor]:
+    records: list[torch.Tensor] = []
+
+    def record(_module, _args, output) -> None:
+        records.append(output[0] if isinstance(output, tuple) else output)
+
+    module.register_forward_hook(record)
+    return records
+
+
+def _cache_snapshot(cache) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]:
+    return (
+        cache.key.detach().view(torch.uint8).clone(),
+        cache.value.detach().view(torch.uint8).clone(),
+        (cache.end, cache.absolute_end, cache.last_start, cache.sink_end),
+    )
+
+
+def _assert_cache_unchanged(cache, snapshot) -> None:
+    key, value, metadata = snapshot
+    assert torch.equal(cache.key.detach().view(torch.uint8), key)
+    assert torch.equal(cache.value.detach().view(torch.uint8), value)
+    assert (cache.end, cache.absolute_end, cache.last_start, cache.sink_end) == metadata
+
+
+def test_self_attention_repeated_offset_overwrites_then_later_offset_appends() -> None:
+    module = _load_module()
+    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+    _set_identity_attention(attention)
+    cache = _allocate_single_layer(module, max_tokens=6).self_attention[0]
+
+    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+    attention(_tokens(10, 20), cache=cache, current_start=0, sink_tokens=0)
+
+    assert cache.end == 2
+    assert cache.absolute_end == 2
+    torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0]))
+
+    attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
+
+    assert cache.end == 4
+    assert cache.absolute_end == 4
+    torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0, 3.0, 4.0]))
+
+
+def test_self_attention_uses_one_fused_qkv_projection() -> None:
+    module = _load_module()
+    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+    cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
+
+    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+
+    assert attention.qkv.calls == 1
+
+
+def test_self_attention_routes_paged_inputs_through_engine_owned_cache() -> None:
+    module = _load_module()
+    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+    _set_identity_attention(attention)
+    paged_inputs = module.ARDiffusionPagedLayerInputs()
+    calls = []
+
+    def paged_write(inputs, query, key, value, k_act, v_act, scale):
+        calls.append((inputs, query.clone(), key.clone(), value.clone(), k_act, v_act, scale))
+        return query
+
+    module.paged_write_attn = paged_write
+    output = attention(
+        _tokens(1, 2),
+        cache=paged_inputs,
+        current_start=0,
+        sink_tokens=0,
+        update_cache=False,
+    )
+
+    assert output.shape == (1, 2, 2)
+    assert len(calls) == 1
+    assert calls[0][0] is paged_inputs
+    assert calls[0][4:6] == (None, None)
+    assert attention.attn.calls == []
+    assert not any(hasattr(attention, name) for name in ("q", "k", "v"))
+
+
+def test_self_attention_is_chunk_causal_without_masking_inside_current_chunk() -> None:
+    module = _load_module()
+    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+    _set_identity_attention(attention)
+    cache = _allocate_single_layer(module, max_tokens=8).self_attention[0]
+    attention_inputs = _record_inputs(attention.attn)
+
+    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+    first_keys = attention_inputs[-1][1].detach().clone()
+    attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
+    second_keys = attention_inputs[-1][1].detach().clone()
+
+    assert attention.attn.causal is False
+    torch.testing.assert_close(first_keys[0, :, 0, 0], torch.tensor([1.0, 2.0]))
+    torch.testing.assert_close(second_keys[0, :, 0, 0], torch.tensor([1.0, 2.0, 3.0, 4.0]))
+
+
+def test_self_attention_retains_sink_and_latest_local_history_after_eviction() -> None:
+    module = _load_module()
+    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+    _set_identity_attention(attention)
+    cache = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+    attention_inputs = _record_inputs(attention.attn)
+
+    attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=1)
+    attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=1)
+    attention(_tokens(5, 6), cache=cache, current_start=4, sink_tokens=1)
+
+    visible_keys = attention_inputs[-1][1].detach().clone()
+    assert cache.end == 4
+    assert cache.absolute_end == 6
+    torch.testing.assert_close(visible_keys[0, :, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
+    torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
+
+
+def test_cross_attention_projects_encoder_kv_once_per_request() -> None:
+    module = _load_module()
+    attention = module.LingBotCrossAttention(dim=2, num_heads=1)
+    _set_identity_attention(attention)
+    key_outputs = _record_outputs(attention.k)
+    value_outputs = _record_outputs(attention.v)
+    attention_inputs = _record_inputs(attention.attn)
+
+    output, cache = attention(_tokens(1), _tokens(2, 3), cache=None)
+    second_output, reused_cache = attention(_tokens(4), _tokens(20, 30), cache=cache)
+
+    assert output.shape == second_output.shape == (1, 1, 2)
+    assert reused_cache is cache
+    assert len(key_outputs) == 1
+    assert len(value_outputs) == 1
+    assert cache.end == 2
+    torch.testing.assert_close(cache.key[0, :, 0, 0], torch.tensor([2.0, 3.0]))
+    torch.testing.assert_close(attention_inputs[-1][1], cache.key)
+
+
+def test_self_attention_cache_is_isolated_between_requests() -> None:
+    module = _load_module()
+    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+    _set_identity_attention(attention)
+    first = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+    second = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+
+    attention(_tokens(1, 2), cache=first, current_start=0, sink_tokens=0)
+
+    assert first.end == 2
+    assert second.end == 0
+    assert torch.count_nonzero(second.key) == 0
+    attention(_tokens(8), cache=second, current_start=0, sink_tokens=0)
+    torch.testing.assert_close(first.key[0, : first.end, 0, 0], torch.tensor([1.0, 2.0]))
+    torch.testing.assert_close(second.key[0, : second.end, 0, 0], torch.tensor([8.0]))
+
+
+def test_cross_attention_keeps_checkpoint_qkvo_and_norm_parameter_names() -> None:
+    module = _load_module()
+    attention = module.LingBotCrossAttention(dim=4, num_heads=2)
+
+    names = set(attention.state_dict())
+
+    assert {
+        "q.weight",
+        "k.weight",
+        "v.weight",
+        "o.weight",
+        "norm_q.weight",
+        "norm_k.weight",
+    } <= names
+    assert callable(attention.norm_q.weight.weight_loader)
+
+
+def test_tp_world_size_one_attention_output_shapes() -> None:
+    module = _load_module()
+    self_attention = module.LingBotSelfAttention(dim=4, num_heads=2)
+    cross_attention = module.LingBotCrossAttention(dim=4, num_heads=2)
+    cache = module.allocate_lingbot_cache(
+        batch_size=2,
+        num_layers=1,
+        max_tokens=5,
+        num_local_heads=2,
+        head_dim=2,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+    hidden_states = torch.randn(2, 3, 4)
+    encoder_hidden_states = torch.randn(2, 5, 4)
+
+    self_output = self_attention(hidden_states, cache=cache.self_attention[0], current_start=0, sink_tokens=0)
+    cross_output, cache.cross_attention[0] = cross_attention(
+        hidden_states,
+        encoder_hidden_states,
+        cache=cache.cross_attention[0],
+    )
+
+    assert self_output.shape == hidden_states.shape
+    assert cross_output.shape == hidden_states.shape
+    assert cache.cross_attention[0] is not None
+
+
+def test_self_attention_applies_rotary_embedding_to_current_query_and_key() -> None:
+    module = _load_module()
+    attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+    cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
+    cos = torch.ones(2, 1)
+    sin = torch.zeros(2, 1)
+    rotary_inputs = _record_inputs(attention.rotary_embedding)
+
+    attention(_tokens(1, 2), cache=cache, current_start=0, rotary_emb=(cos, sin), sink_tokens=0)
+
+    assert len(rotary_inputs) == 2
+    torch.testing.assert_close(rotary_inputs[0][1], cos)
+    torch.testing.assert_close(rotary_inputs[1][2], sin)
+
+
+def test_attention_block_owns_checkpoint_parent_namespaces() -> None:
+    module = _load_module()
+    block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+
+    names = set(block.state_dict())
+    expected = {
+        *(f"self_attn.{child}.weight" for child in ("qkv", "o", "norm_q", "norm_k")),
+        *(f"cross_attn.{child}.weight" for child in ("q", "k", "v", "o", "norm_q", "norm_k")),
+    }
+
+    assert expected <= names
+
+
+def test_cross_attention_caches_are_isolated_between_requests() -> None:
+    module = _load_module()
+    attention = module.LingBotCrossAttention(dim=2, num_heads=1)
+
+    _, first = attention(_tokens(1), _tokens(2, 3), cache=None)
+    _, second = attention(_tokens(1), _tokens(8, 9), cache=None)
+
+    assert first.key.untyped_storage().data_ptr() != second.key.untyped_storage().data_ptr()
+    assert first.value.untyped_storage().data_ptr() != second.value.untyped_storage().data_ptr()
+    second_key = second.key.clone()
+    first.key.zero_()
+    torch.testing.assert_close(second.key, second_key)
+
+
+def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _load_module()
+    monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
+    norm = module._LingBotRMSNorm(2, eps=1e-6)
+
+    norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
+
+    torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))

tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py

@@ -0,0 +1,525 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/camera.py"
+_SPEC = importlib.util.spec_from_file_location("_lingbot_world_camera_under_test", _MODULE_PATH)
+assert _SPEC is not None and _SPEC.loader is not None
+_CAMERA_MODULE = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = _CAMERA_MODULE
+_SPEC.loader.exec_module(_CAMERA_MODULE)
+
+CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
+build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding
+interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory
+load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory
+resolve_trusted_action_directory = _CAMERA_MODULE.resolve_trusted_action_directory
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _identity_poses(num_frames: int) -> np.ndarray:
+    return np.repeat(np.eye(4, dtype=np.float64)[None], num_frames, axis=0)
+
+
+def _write_trajectory(
+    directory: Path,
+    poses: np.ndarray,
+    intrinsics: np.ndarray,
+) -> None:
+    np.save(directory / "poses.npy", poses)
+    np.save(directory / "intrinsics.npy", intrinsics)
+
+
+def _trusted(directory: Path):
+    return resolve_trusted_action_directory(".", directory)
+
+
+def _trajectory(
+    poses: torch.Tensor,
+    intrinsics: torch.Tensor | None = None,
+) -> CameraTrajectory:
+    if intrinsics is None:
+        intrinsics = torch.tensor([[832.0 / 3.0, 160.0, 416.0, 240.0]]).repeat(poses.shape[0], 1)
+    return CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+
+@pytest.mark.parametrize("present_file", [None, "poses.npy"])
+def test_load_camera_trajectory_rejects_missing_files(tmp_path: Path, present_file: str | None) -> None:
+    if present_file == "poses.npy":
+        np.save(tmp_path / present_file, _identity_poses(1))
+
+    missing_file = "poses.npy" if present_file is None else "intrinsics.npy"
+    with pytest.raises(FileNotFoundError, match=missing_file):
+        load_camera_trajectory(_trusted(tmp_path))
+
+
+@pytest.mark.parametrize(
+    ("poses", "intrinsics", "message"),
+    [
+        (np.zeros((2, 3, 4)), np.zeros((2, 4)), "poses.npy"),
+        (_identity_poses(2), np.zeros((2, 3)), "intrinsics.npy"),
+        (_identity_poses(0), np.zeros((0, 4)), "1 and 4096 source frames"),
+        (_identity_poses(2), np.zeros((1, 4)), "same number of frames"),
+    ],
+)
+def test_load_camera_trajectory_rejects_invalid_shapes(
+    tmp_path: Path,
+    poses: np.ndarray,
+    intrinsics: np.ndarray,
+    message: str,
+) -> None:
+    _write_trajectory(tmp_path, poses, intrinsics)
+
+    with pytest.raises(ValueError, match=message):
+        load_camera_trajectory(_trusted(tmp_path))
+
+
+@pytest.mark.parametrize("array_name", ["poses", "intrinsics"])
+def test_load_camera_trajectory_rejects_non_finite_values(tmp_path: Path, array_name: str) -> None:
+    poses = _identity_poses(1)
+    intrinsics = np.ones((1, 4), dtype=np.float64)
+    if array_name == "poses":
+        poses[0, 0, 0] = np.nan
+    else:
+        intrinsics[0, 0] = np.inf
+    _write_trajectory(tmp_path, poses, intrinsics)
+
+    with pytest.raises(ValueError, match="finite"):
+        load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_preserves_raw_c2w_poses(tmp_path: Path) -> None:
+    first_c2w = np.array(
+        [
+            [0.0, -1.0, 0.0, 2.0],
+            [1.0, 0.0, 0.0, 3.0],
+            [0.0, 0.0, 1.0, 4.0],
+            [0.0, 0.0, 0.0, 1.0],
+        ]
+    )
+    second_c2w = np.array(
+        [
+            [0.0, 0.0, 1.0, 5.0],
+            [0.0, 1.0, 0.0, 4.0],
+            [-1.0, 0.0, 0.0, 6.0],
+            [0.0, 0.0, 0.0, 1.0],
+        ]
+    )
+    poses = np.stack((first_c2w, second_c2w))
+    intrinsics = np.array([[100.0, 110.0, 120.0, 130.0]] * 2)
+    _write_trajectory(tmp_path, poses, intrinsics)
+
+    trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+    assert trajectory.poses.dtype == torch.float32
+    assert trajectory.intrinsics.dtype == torch.float32
+    torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
+    torch.testing.assert_close(trajectory.intrinsics, torch.from_numpy(intrinsics).float())
+
+
+def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_path: Path) -> None:
+    poses = _identity_poses(1).astype(object)
+    intrinsics = np.ones((1, 4), dtype=np.float32)
+    _write_trajectory(tmp_path, poses, intrinsics)
+
+    with pytest.raises(ValueError, match="poses.npy.*real numeric"):
+        load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None:
+    _write_trajectory(
+        tmp_path,
+        _identity_poses(269),
+        np.arange(269 * 4, dtype=np.float32).reshape(269, 4),
+    )
+
+    trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+    assert trajectory.poses.shape == (117, 4, 4)
+    assert trajectory.intrinsics.shape == (117, 4)
+    np.testing.assert_array_equal(
+        trajectory.intrinsics.numpy(),
+        np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117],
+    )
+
+
+def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None:
+    _write_trajectory(
+        tmp_path,
+        _identity_poses(4097),
+        np.ones((4097, 4), dtype=np.float32),
+    )
+
+    with pytest.raises(ValueError, match="poses.npy.*bounded camera-array limit|4096 source frames"):
+        load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_truncated_payload_before_loading(tmp_path: Path) -> None:
+    _write_trajectory(tmp_path, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+    poses_path = tmp_path / "poses.npy"
+    poses_path.write_bytes(poses_path.read_bytes()[:-1])
+
+    with pytest.raises(ValueError, match="poses.npy byte size"):
+        load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_oversized_npy_header(tmp_path: Path) -> None:
+    header_size = 20_000
+    (tmp_path / "poses.npy").write_bytes(
+        np.lib.format.magic(2, 0) + header_size.to_bytes(4, "little") + b" " * (header_size - 1) + b"\n"
+    )
+    np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
+
+    with pytest.raises(ValueError, match="poses.npy.*oversized NPY header"):
+        load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_symlink_swap_after_resolution(tmp_path: Path) -> None:
+    trusted_root = tmp_path / "trusted"
+    action_directory = trusted_root / "action"
+    outside = tmp_path / "outside"
+    action_directory.mkdir(parents=True)
+    outside.mkdir()
+    _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+    _write_trajectory(outside, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+    resolved = resolve_trusted_action_directory("action", trusted_root)
+    action_directory.rename(trusted_root / "original-action")
+    action_directory.symlink_to(outside, target_is_directory=True)
+
+    with pytest.raises(ValueError, match="securely open.*action directory"):
+        load_camera_trajectory(resolved)
+
+
+def test_load_camera_trajectory_rejects_symlinked_camera_file(tmp_path: Path) -> None:
+    trusted_root = tmp_path / "trusted"
+    action_directory = trusted_root / "action"
+    outside = tmp_path / "outside"
+    action_directory.mkdir(parents=True)
+    outside.mkdir()
+    _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+    np.save(outside / "poses.npy", _identity_poses(1))
+    resolved = resolve_trusted_action_directory("action", trusted_root)
+    (action_directory / "poses.npy").unlink()
+    (action_directory / "poses.npy").symlink_to(outside / "poses.npy")
+
+    with pytest.raises(ValueError, match="poses.npy.*opened securely"):
+        load_camera_trajectory(resolved)
+
+
+@pytest.mark.skipif(not hasattr(__import__("os"), "mkfifo"), reason="requires POSIX FIFO support")
+def test_load_camera_trajectory_rejects_fifo_without_blocking(tmp_path: Path) -> None:
+    import os
+
+    os.mkfifo(tmp_path / "poses.npy")
+    np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
+    script = """
+import importlib.util
+import sys
+
+spec = importlib.util.spec_from_file_location("_lingbot_camera_fifo_test", sys.argv[1])
+module = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)
+action = module.resolve_trusted_action_directory(".", sys.argv[2])
+try:
+    module.load_camera_trajectory(action)
+except ValueError as error:
+    print(error)
+else:
+    raise AssertionError("FIFO camera input was accepted")
+"""
+
+    completed = subprocess.run(
+        [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
+        capture_output=True,
+        text=True,
+        timeout=3,
+        check=False,
+    )
+
+    assert completed.returncode == 0, completed.stderr
+    assert "poses.npy must be a regular NPY file" in completed.stdout
+
+
+def test_load_camera_trajectory_decodes_the_validated_file_snapshot(tmp_path: Path, monkeypatch) -> None:
+    original_poses = _identity_poses(1)
+    replacement_poses = original_poses.copy()
+    replacement_poses[0, 0, 3] = 123.0
+    _write_trajectory(tmp_path, original_poses, np.ones((1, 4), dtype=np.float32))
+    poses_path = tmp_path / "poses.npy"
+    original_load = _CAMERA_MODULE.np.load
+    mutation_count = 0
+
+    def mutate_source_then_decode(source, *args, **kwargs):
+        nonlocal mutation_count
+        if mutation_count == 0:
+            mutation_count += 1
+            np.save(poses_path, replacement_poses)
+        return original_load(source, *args, **kwargs)
+
+    monkeypatch.setattr(_CAMERA_MODULE.np, "load", mutate_source_then_decode)
+
+    trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+    assert mutation_count == 1
+    torch.testing.assert_close(trajectory.poses, torch.from_numpy(original_poses).float())
+
+
+def test_load_camera_trajectory_ignores_non_camera_event_files(tmp_path: Path) -> None:
+    poses = _identity_poses(1)
+    intrinsics = np.ones((1, 4), dtype=np.float32)
+    _write_trajectory(tmp_path, poses, intrinsics)
+    np.save(tmp_path / "wasd_events.npy", np.array([{"key": "w"}], dtype=object))
+
+    trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+    torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
+
+
+def test_interpolate_camera_trajectory_uses_linear_and_spherical_interpolation() -> None:
+    poses = torch.eye(4).repeat(2, 1, 1)
+    poses[1, :3, :3] = torch.tensor(
+        [
+            [0.0, -1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    poses[1, :3, 3] = torch.tensor([2.0, 4.0, 6.0])
+    intrinsics = torch.tensor([[10.0, 20.0, 30.0, 40.0], [20.0, 40.0, 60.0, 80.0]])
+    trajectory = CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+    result = interpolate_camera_trajectory(trajectory, 3)
+
+    torch.testing.assert_close(result.poses[0], poses[0])
+    torch.testing.assert_close(result.poses[-1], poses[-1])
+    torch.testing.assert_close(result.poses[1, :3, 3], torch.tensor([1.0, 2.0, 3.0]))
+    expected_mid_rotation = torch.tensor(
+        [
+            [math.sqrt(0.5), -math.sqrt(0.5), 0.0],
+            [math.sqrt(0.5), math.sqrt(0.5), 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    torch.testing.assert_close(result.poses[1, :3, :3], expected_mid_rotation)
+    torch.testing.assert_close(result.intrinsics[1], torch.tensor([15.0, 30.0, 45.0, 60.0]))
+
+
+def test_build_plucker_embedding_has_forward_center_ray() -> None:
+    trajectory = _trajectory(torch.eye(4).unsqueeze(0))
+
+    result = build_plucker_embedding(
+        trajectory,
+        height=3,
+        width=3,
+        target_height=3,
+        target_width=3,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    torch.testing.assert_close(result[0, :3, 1, 1], torch.zeros(3))
+    torch.testing.assert_close(result[0, 3:, 1, 1], torch.tensor([0.0, 0.0, 1.0]))
+
+
+def test_build_plucker_embedding_scales_reference_intrinsics() -> None:
+    trajectory = _trajectory(
+        torch.eye(4).unsqueeze(0),
+        torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
+    )
+
+    result = build_plucker_embedding(
+        trajectory,
+        height=4,
+        width=8,
+        target_height=4,
+        target_width=8,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    expected_corner_direction = torch.tensor([-0.875, -0.75, 1.0])
+    expected_corner_direction /= torch.linalg.vector_norm(expected_corner_direction)
+    torch.testing.assert_close(result[0, 3:, 0, 0], expected_corner_direction)
+
+
+def test_build_plucker_embedding_samples_target_pixel_centers() -> None:
+    trajectory = _trajectory(
+        torch.eye(4).unsqueeze(0),
+        torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
+    )
+
+    result = build_plucker_embedding(
+        trajectory,
+        height=6,
+        width=10,
+        target_height=2,
+        target_width=5,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    # Target centers map through the source scale: (0.5 * 2, 0.5 * 3).
+    expected_direction = torch.tensor([-0.8, -0.5, 1.0])
+    expected_direction /= torch.linalg.vector_norm(expected_direction)
+    torch.testing.assert_close(result[0, 3:, 0, 0], expected_direction)
+
+
+def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_translations() -> None:
+    poses = torch.eye(4).repeat(3, 1, 1)
+    poses[0, :3, 3] = torch.tensor([5.0, 7.0, 11.0])
+    poses[1, :3, :3] = torch.tensor(
+        [
+            [0.0, -1.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    poses[1, :3, 3] = torch.tensor([5.0, 9.0, 11.0])
+    poses[2, :3, :3] = torch.tensor(
+        [
+            [-1.0, 0.0, 0.0],
+            [0.0, -1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    poses[2, :3, 3] = torch.tensor([3.0, 9.0, 11.0])
+    trajectory = _trajectory(
+        poses,
+        torch.tensor([[416.0, 480.0, 208.0, 240.0]]).repeat(3, 1),
+    )
+
+    result = build_plucker_embedding(
+        trajectory,
+        height=1,
+        width=2,
+        target_height=1,
+        target_width=2,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    expected_direction = torch.tensor([0.0, 1.0, 1.0]) / math.sqrt(2.0)
+    torch.testing.assert_close(result[0, :3, 0, 1], torch.zeros(3))
+    torch.testing.assert_close(result[1, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
+    torch.testing.assert_close(result[2, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
+    torch.testing.assert_close(result[1, 3:, 0, 1], expected_direction)
+    torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
+
+
+def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
+    poses = torch.eye(4).repeat(3, 1, 1)
+    poses[:, :3, 3] = torch.tensor(
+        [
+            [10.0, -2.0, 5.0],
+            [13.0, 2.0, 5.0],
+            [19.0, 10.0, 5.0],
+        ]
+    )
+
+    result = build_plucker_embedding(
+        _trajectory(poses),
+        height=1,
+        width=1,
+        target_height=1,
+        target_width=1,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
+    torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
+
+
+def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:
+    poses = torch.eye(4).repeat(2, 1, 1)
+    poses[0, :3, 3] = torch.tensor([20.0, 4.0, -2.0])
+    poses[1, :3, 3] = torch.tensor([26.0, 12.0, -2.0])
+    raw = _trajectory(poses)
+
+    interpolated = interpolate_camera_trajectory(raw, 3)
+    result = build_plucker_embedding(
+        interpolated,
+        height=1,
+        width=1,
+        target_height=1,
+        target_width=1,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    expected_delta = torch.tensor([0.6, 0.8, 0.0])
+    torch.testing.assert_close(result[0, :3, 0, 0], torch.zeros(3))
+    torch.testing.assert_close(result[1, :3, 0, 0], expected_delta)
+    torch.testing.assert_close(result[2, :3, 0, 0], expected_delta)
+
+
+def test_build_plucker_embedding_uses_half_pixel_centers_and_origin_direction_channel_order() -> None:
+    poses = torch.eye(4).repeat(2, 1, 1)
+    poses[1, :3, 3] = torch.tensor([2.0, 0.0, 0.0])
+    trajectory = _trajectory(
+        poses,
+        torch.tensor([[832.0, 480.0, 0.0, 0.0]]).repeat(2, 1),
+    )
+
+    result = build_plucker_embedding(
+        trajectory,
+        height=1,
+        width=2,
+        target_height=1,
+        target_width=2,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    expected_direction = torch.tensor([0.25, 0.5, 1.0])
+    expected_direction /= torch.linalg.vector_norm(expected_direction)
+    torch.testing.assert_close(result[1, :3, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
+    torch.testing.assert_close(result[1, 3:, 0, 0], expected_direction)
+
+
+def test_build_plucker_embedding_honors_shape_dtype_and_device() -> None:
+    trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
+
+    result = build_plucker_embedding(
+        trajectory,
+        height=12,
+        width=20,
+        target_height=3,
+        target_width=5,
+        device=torch.device("cpu"),
+        dtype=torch.float64,
+    )
+
+    assert result.shape == (2, 6, 3, 5)
+    assert result.dtype == torch.float64
+    assert result.device == torch.device("cpu")
+
+
+def test_build_plucker_embedding_is_deterministic() -> None:
+    trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
+    kwargs = {
+        "height": 12,
+        "width": 20,
+        "target_height": 3,
+        "target_width": 5,
+        "device": torch.device("cpu"),
+        "dtype": torch.float32,
+    }
+
+    first = build_plucker_embedding(trajectory, **kwargs)
+    second = build_plucker_embedding(trajectory, **kwargs)
+
+    torch.testing.assert_close(first, second, rtol=0.0, atol=0.0)

tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py

@@ -0,0 +1,685 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import hashlib
+import inspect
+import json
+import math
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture"
+_SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture"
+
+
+@pytest.fixture(autouse=True)
+def _inference_context():
+    with torch.inference_mode():
+        yield
+
+
+def _tiny_model(
+    module,
+    *,
+    num_layers: int = 2,
+    num_frames_per_block: int = 1,
+    sliding_window_num_frames: int = 3,
+):
+    return module.CausalLingBotWorldTransformer3DModel(
+        patch_size=(1, 2, 2),
+        num_attention_heads=2,
+        attention_head_dim=2,
+        in_channels=36,
+        out_channels=2,
+        text_dim=6,
+        freq_dim=4,
+        ffn_dim=8,
+        num_layers=num_layers,
+        cross_attn_norm=True,
+        eps=1e-6,
+        rope_max_seq_len=16,
+        sink_size=1,
+        num_frames_per_block=num_frames_per_block,
+        sliding_window_num_frames=sliding_window_num_frames,
+        local_attn_size=-1,
+    )
+
+
+def _cache(module, model, *, max_tokens: int | None = None):
+    if max_tokens is None:
+        window_frames = (
+            model.config.local_attn_size
+            if model.config.local_attn_size != -1
+            else model.config.sliding_window_num_frames
+        )
+        max_tokens = window_frames * 2 * 2
+    return module.allocate_lingbot_cache(
+        batch_size=1,
+        num_layers=len(model.blocks),
+        max_tokens=max_tokens,
+        num_local_heads=2,
+        head_dim=2,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+
+def _checkpoint_parameter_specs(model) -> dict[str, tuple[torch.Size, torch.dtype]]:
+    specs: dict[str, tuple[torch.Size, torch.dtype]] = {}
+    for name, parameter in model.named_parameters():
+        marker = ".self_attn.qkv."
+        if marker not in name:
+            specs[name] = (parameter.shape, parameter.dtype)
+            continue
+        shard_shape = torch.Size((parameter.shape[0] // 3, *parameter.shape[1:]))
+        for projection in ("q", "k", "v"):
+            specs[name.replace(marker, f".self_attn.{projection}.")] = (shard_shape, parameter.dtype)
+    return specs
+
+
+def _checkpoint_weights_for_model(model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+    checkpoint_weights: dict[str, torch.Tensor] = {}
+    expected_parameters: dict[str, torch.Tensor] = {}
+    next_value = 1
+    for name, parameter in model.named_parameters():
+        marker = ".self_attn.qkv."
+        if marker not in name:
+            value = torch.full_like(parameter, next_value)
+            checkpoint_weights[name] = value
+            expected_parameters[name] = value
+            next_value += 1
+            continue
+        shards = []
+        shard_shape = (parameter.shape[0] // 3, *parameter.shape[1:])
+        for projection in ("q", "k", "v"):
+            value = torch.full(shard_shape, next_value, dtype=parameter.dtype, device=parameter.device)
+            checkpoint_weights[name.replace(marker, f".self_attn.{projection}.")] = value
+            shards.append(value)
+            next_value += 1
+        expected_parameters[name] = torch.cat(shards)
+    return checkpoint_weights, expected_parameters
+
+
+def _self_cache_snapshot(cache) -> list[tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]]:
+    return [attention_tests._cache_snapshot(layer_cache) for layer_cache in cache.self_attention]
+
+
+def _assert_self_cache_unchanged(cache, snapshot) -> None:
+    for layer_cache, layer_snapshot in zip(cache.self_attention, snapshot, strict=True):
+        attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot)
+
+
+def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None:
+    torch.manual_seed(7)
+    module = attention_tests._load_module()
+    model = _tiny_model(module).eval()
+    cache = _cache(module, model)
+    encoder_hidden_states = torch.randn(1, 3, 6)
+    camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+    cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
+    text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
+    outputs = []
+
+    assert model.patch_embedding.in_channels == 36
+    assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2
+
+    for start_frame in range(4):
+        hidden_states = torch.randn(1, 36, 1, 4, 4)
+        timestep = torch.tensor([float(start_frame + 1)])
+        snapshot = _self_cache_snapshot(cache)
+
+        transient_output = model(
+            hidden_states,
+            timestep,
+            encoder_hidden_states,
+            camera_hidden_states,
+            cache=cache,
+            start_frame=start_frame,
+            update_cache=False,
+        )
+
+        assert transient_output.shape == (1, 2, 1, 4, 4)
+        _assert_self_cache_unchanged(cache, snapshot)
+        assert all(layer_cache is not None for layer_cache in cache.cross_attention)
+
+        committed_output = model(
+            hidden_states,
+            timestep,
+            encoder_hidden_states,
+            camera_hidden_states,
+            cache=cache,
+            start_frame=start_frame,
+            update_cache=True,
+        )
+        torch.testing.assert_close(committed_output, transient_output)
+        outputs.append(committed_output)
+
+        expected_token_end = (start_frame + 1) * 2 * 2
+        assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention)
+
+    assert len(outputs) == 4
+    assert all(len(outputs) == 1 for outputs in cross_key_outputs)
+    assert len(text_projection_outputs) == 1
+
+    alternate_cache = _cache(module, model)
+    alternate_output = model(
+        torch.zeros(1, 36, 1, 4, 4),
+        torch.tensor([1.0]),
+        encoder_hidden_states,
+        torch.zeros_like(camera_hidden_states),
+        cache=alternate_cache,
+        start_frame=0,
+        update_cache=False,
+    )
+    camera_output = model(
+        torch.zeros(1, 36, 1, 4, 4),
+        torch.tensor([1.0]),
+        encoder_hidden_states,
+        torch.ones_like(camera_hidden_states),
+        cache=_cache(module, model),
+        start_frame=0,
+        update_cache=False,
+    )
+    assert not torch.equal(camera_output, alternate_output)
+
+
+def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=2).eval()
+    cache = _cache(module, model)
+    hidden_states = torch.randn(1, 36, 1, 4, 4)
+    timestep = torch.tensor([1.0])
+    encoder_hidden_states = torch.randn(1, 3, 6)
+    camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+
+    model(
+        hidden_states,
+        timestep,
+        encoder_hidden_states,
+        camera_hidden_states,
+        cache=cache,
+        start_frame=0,
+        update_cache=False,
+    )
+    cache.cross_attention[1] = None
+    text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
+    cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
+
+    model(
+        hidden_states,
+        timestep,
+        encoder_hidden_states,
+        camera_hidden_states,
+        cache=cache,
+        start_frame=0,
+        update_cache=False,
+    )
+
+    assert len(text_projection_outputs) == 1
+    assert [len(outputs) for outputs in cross_key_outputs] == [0, 1]
+    assert all(layer_cache is not None for layer_cache in cache.cross_attention)
+
+
+@pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks"))
+def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(
+        module,
+        num_layers=1,
+        num_frames_per_block=3,
+        sliding_window_num_frames=6,
+    )
+    cache = _cache(module, model)
+    snapshot = _self_cache_snapshot(cache)
+    patch_inputs = attention_tests._record_inputs(model.patch_embedding)
+
+    with pytest.raises(ValueError, match="exactly 3 post-patch frames"):
+        model(
+            torch.randn(1, 36, frames, 4, 4),
+            torch.tensor([1.0]),
+            torch.randn(1, 3, 6),
+            torch.randn(1, 6 * 8 * 8, frames, 4, 4),
+            cache=cache,
+            start_frame=0,
+            update_cache=True,
+        )
+
+    assert patch_inputs == []
+    _assert_self_cache_unchanged(cache, snapshot)
+    assert cache.cross_attention == [None]
+
+
+def test_forward_accepts_exactly_one_configured_frame_block() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(
+        module,
+        num_layers=1,
+        num_frames_per_block=3,
+        sliding_window_num_frames=6,
+    )
+
+    output = model(
+        torch.randn(1, 36, 3, 4, 4),
+        torch.tensor([1.0]),
+        torch.randn(1, 3, 6),
+        torch.randn(1, 6 * 8 * 8, 3, 4, 4),
+        cache=_cache(module, model),
+        start_frame=0,
+        update_cache=False,
+    )
+
+    assert output.shape == (1, 2, 3, 4, 4)
+
+
+def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(
+        module,
+        num_layers=2,
+        num_frames_per_block=3,
+        sliding_window_num_frames=6,
+    )
+
+    cache = model.allocate_cache(
+        batch_size=2,
+        latent_height=8,
+        latent_width=12,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    assert len(cache.self_attention) == 2
+    assert len(cache.cross_attention) == 2
+    assert cache.cross_attention == [None, None]
+    assert all(layer.key.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
+    assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
+
+
+def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+    hidden_states = torch.zeros(1, 36, 2, 2, 4)
+    hidden_states[0, 0] = torch.tensor(
+        [
+            [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
+            [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]],
+        ]
+    )
+    with torch.no_grad():
+        model.patch_embedding.weight.zero_()
+        model.patch_embedding.bias.zero_()
+        model.patch_embedding.weight[0, 0, 0, 0, 0] = 1
+
+    tokens = model.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
+
+    torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0]))
+
+
+def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+    hidden_states = (
+        torch.stack(
+            (
+                torch.arange(0, 8),
+                torch.arange(10, 18),
+                torch.arange(20, 28),
+                torch.arange(30, 38),
+            )
+        )
+        .unsqueeze(0)
+        .float()
+    )
+    expected = torch.tensor(
+        [
+            [
+                [[0, 2, 10, 12], [4, 6, 14, 16]],
+                [[20, 22, 30, 32], [24, 26, 34, 36]],
+            ],
+            [
+                [[1, 3, 11, 13], [5, 7, 15, 17]],
+                [[21, 23, 31, 33], [25, 27, 35, 37]],
+            ],
+        ],
+        dtype=torch.float32,
+    ).unsqueeze(0)
+
+    output = model._unpatchify(
+        hidden_states,
+        batch_size=1,
+        frames=2,
+        height=1,
+        width=2,
+    )
+
+    torch.testing.assert_close(output, expected)
+
+
+def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+    model.head.norm = torch.nn.Identity()
+    with torch.no_grad():
+        model.head.modulation.zero_()
+        model.head.head.weight.zero_()
+        model.head.head.bias.zero_()
+        model.head.head.weight[:, 0] = 1
+    hidden_states = torch.ones(1, 4, model.dim)
+    timestep_embedding = torch.tensor([[[10.0, 0.0, 0.0, 0.0], [20.0, 0.0, 0.0, 0.0]]])
+    expected = torch.tensor([21.0, 21.0, 41.0, 41.0]).view(1, 4, 1).expand(1, 4, 8)
+
+    output = model.head(hidden_states, timestep_embedding)
+
+    torch.testing.assert_close(output, expected)
+
+
+def test_constructor_defaults_match_official_checkpoint_config() -> None:
+    module = attention_tests._load_module()
+    parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters
+    expected = {
+        "in_channels": 36,
+        "out_channels": 16,
+        "num_layers": 40,
+        "num_attention_heads": 40,
+        "attention_head_dim": 128,
+        "patch_size": (1, 2, 2),
+        "text_dim": 4096,
+        "ffn_dim": 13824,
+        "freq_dim": 256,
+        "rope_max_seq_len": 1024,
+        "eps": 1e-6,
+        "cross_attn_norm": True,
+        "sink_size": 9,
+        "num_frames_per_block": 3,
+        "sliding_window_num_frames": 18,
+        "local_attn_size": -1,
+    }
+
+    assert {name: parameters[name].default for name in expected} == expected
+
+
+def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(module)
+
+    assert model.dtype == next(model.parameters()).dtype
+
+
+def test_transformer_declares_regional_compile_block() -> None:
+    module = attention_tests._load_module()
+
+    assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"]
+
+
+def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
+    module = attention_tests._load_module()
+    reduced_values: list[torch.Tensor] = []
+
+    monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+
+    def all_reduce(value: torch.Tensor) -> torch.Tensor:
+        reduced_values.append(value.detach().clone())
+        return value * 2
+
+    monkeypatch.setattr(module, "tensor_model_parallel_all_reduce", all_reduce)
+    norm = module._LingBotRMSNorm(hidden_size=2, eps=0.0)
+    value = torch.tensor([[3.0, 4.0]])
+
+    output = norm(value)
+
+    torch.testing.assert_close(output, value / math.sqrt(12.5))
+    assert len(reduced_values) == 1
+    torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]]))
+
+
+def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
+    module = attention_tests._load_module()
+    monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
+
+    with pytest.raises(ValueError, match="num_heads=2.*tp_size=3"):
+        module.LingBotSelfAttention(dim=4, num_heads=2)
+
+
+def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
+    module = attention_tests._load_module()
+
+    with pytest.raises(ValueError, match="qk_norm.*rms_norm_across_heads"):
+        module.CausalLingBotWorldTransformer3DModel(qk_norm="layer_norm")
+
+
+@pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"])
+def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None:
+    module = attention_tests._load_module()
+
+    with pytest.raises(ValueError, match=field):
+        module.CausalLingBotWorldTransformer3DModel(**{field: 4})
+
+
+def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
+    module = attention_tests._load_module()
+
+    with pytest.raises(RuntimeError, match="quant_config.*not supported"):
+        module.CausalLingBotWorldTransformer3DModel(quant_config=object())
+
+
+def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None:
+    module = attention_tests._load_module()
+
+    class ConfigProbe(module.CausalLingBotWorldTransformer3DModel):
+        def __init__(self, **kwargs) -> None:
+            self.received_kwargs = kwargs
+
+    config = {
+        "_class_name": "CausalLingBotWorldTransformer3DModel",
+        "_diffusers_version": "0.35.0.dev0",
+        "patch_size": [1, 2, 2],
+        "num_attention_heads": 40,
+        "attention_head_dim": 128,
+        "in_channels": 36,
+        "out_channels": 16,
+        "text_dim": 4096,
+        "freq_dim": 256,
+        "ffn_dim": 13824,
+        "num_layers": 40,
+        "cross_attn_norm": True,
+        "eps": 1e-6,
+        "image_dim": None,
+        "added_kv_proj_dim": None,
+        "rope_max_seq_len": 1024,
+        "pos_embed_seq_len": None,
+        "qk_norm": "rms_norm_across_heads",
+        "sink_size": 9,
+        "num_frames_per_block": 3,
+        "sliding_window_num_frames": 18,
+        "local_attn_size": -1,
+    }
+
+    model = ConfigProbe.from_config(config, prefix="transformer")
+
+    assert model.received_kwargs["patch_size"] == (1, 2, 2)
+    assert model.received_kwargs["num_layers"] == 40
+    assert model.received_kwargs["prefix"] == "transformer"
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [
+        ("in_channels", 35),
+        ("num_layers", 39),
+        ("num_attention_heads", 32),
+        ("attention_head_dim", 64),
+        ("ffn_dim", 10240),
+        ("num_frames_per_block", 4),
+        ("sliding_window_num_frames", 24),
+        ("sink_size", 0),
+    ],
+)
+def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None:
+    module = attention_tests._load_module()
+    config = {
+        "_class_name": "CausalLingBotWorldTransformer3DModel",
+        "_diffusers_version": "0.35.0.dev0",
+        "patch_size": [1, 2, 2],
+        "num_attention_heads": 40,
+        "attention_head_dim": 128,
+        "in_channels": 36,
+        "out_channels": 16,
+        "text_dim": 4096,
+        "freq_dim": 256,
+        "ffn_dim": 13824,
+        "num_layers": 40,
+        "cross_attn_norm": True,
+        "eps": 1e-6,
+        "image_dim": None,
+        "added_kv_proj_dim": None,
+        "rope_max_seq_len": 1024,
+        "pos_embed_seq_len": None,
+        "qk_norm": "rms_norm_across_heads",
+        "sink_size": 9,
+        "num_frames_per_block": 3,
+        "sliding_window_num_frames": 18,
+        "local_attn_size": -1,
+    }
+    config[field] = value
+
+    with pytest.raises(ValueError, match=field):
+        module.CausalLingBotWorldTransformer3DModel.from_config(config)
+
+
+def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
+    module = attention_tests._load_module()
+    config = {
+        "_class_name": "CausalLingBotWorldTransformer3DModel",
+        "architectures": ["CausalLingBotWorldTransformer3DModel"],
+        "future_diffusers_metadata": {"version": 1},
+    }
+
+    with torch.device("meta"):
+        model = module.CausalLingBotWorldTransformer3DModel.from_config(config)
+
+    assert model.config.in_channels == 36
+    assert model.config.num_layers == 40
+
+
+def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1)
+    checkpoint_weights, expected_parameters = _checkpoint_weights_for_model(model)
+    loader_calls: list[str] = []
+    qkv_weight = model.blocks[0].self_attn.qkv.weight
+
+    def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
+        loader_calls.append(shard_id)
+        shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
+        shard_size = param.shape[0] // 3
+        param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
+
+    qkv_weight.weight_loader = record_loader
+    loaded = model.load_weights(iter(checkpoint_weights.items()))
+
+    assert set(checkpoint_weights) <= loaded
+    assert set(dict(model.named_parameters())) <= loaded
+    assert loader_calls == ["q", "k", "v"]
+    for name, param in model.named_parameters():
+        torch.testing.assert_close(param, expected_parameters[name])
+
+    first_name = next(iter(checkpoint_weights))
+    partial = _tiny_model(module, num_layers=1)
+    partial_loaded = partial.load_weights([(first_name, checkpoint_weights[first_name])])
+    expected_partial = {first_name}
+    for projection in ("q", "k", "v"):
+        marker = f".self_attn.{projection}."
+        if marker in first_name:
+            expected_partial.add(first_name.replace(marker, ".self_attn.qkv."))
+            break
+    assert partial_loaded == expected_partial
+
+    with pytest.raises(KeyError, match="unexpected_model.weight"):
+        model.load_weights([("unexpected_model.weight", torch.ones(1))])
+
+
+def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
+    module = attention_tests._load_module()
+    model = _tiny_model(module, num_layers=1)
+    first_name, first_param = next(iter(model.named_parameters()))
+    loaded_weight = torch.full_like(first_param, 17)
+
+    def weights():
+        yield first_name, loaded_weight
+        torch.testing.assert_close(first_param, loaded_weight)
+
+    assert model.load_weights(weights()) == {first_name}
+
+
+def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None:
+    utils = pytest.importorskip("vllm.model_executor.models.utils")
+    module = attention_tests._load_module()
+    transformer = _tiny_model(module, num_layers=1)
+    checkpoint_weights, _ = _checkpoint_weights_for_model(transformer)
+
+    class Pipeline(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.transformer = transformer
+
+    pipeline = Pipeline()
+    loaded = utils.AutoWeightsLoader(pipeline).load_weights(
+        (f"transformer.{name}", value) for name, value in checkpoint_weights.items()
+    )
+
+    assert set(dict(pipeline.named_parameters())) <= loaded
+
+
+def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
+    module = attention_tests._load_module()
+    fixture = json.loads(_FIXTURE_PATH.read_text())
+    model = _tiny_model(module, num_layers=1)
+    parameter_names = set(_checkpoint_parameter_specs(model))
+
+    assert {name.split(".", 1)[0] for name in parameter_names} == set(fixture["top_level_modules"])
+    assert {name for name in parameter_names if not name.startswith("blocks.")} == set(fixture["non_block_parameters"])
+    assert {name.removeprefix("blocks.0.") for name in parameter_names if name.startswith("blocks.0.")} == set(
+        fixture["block_parameter_suffixes"]
+    )
+    expanded_names = set(fixture["non_block_parameters"])
+    expanded_names.update(
+        f"blocks.{layer}.{suffix}" for layer in range(40) for suffix in fixture["block_parameter_suffixes"]
+    )
+    canonical_names = "\n".join(sorted(expanded_names))
+
+    assert len(expanded_names) == fixture["index_parameter_count"] == 1421
+    assert hashlib.sha256(canonical_names.encode()).hexdigest() == fixture["canonical_name_set_sha256"]
+    assert fixture["canonical_name_to_shard_sha256"] == (
+        "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992"
+    )
+    assert list(fixture["shards"]) == [
+        f"diffusion_pytorch_model-{index:05}-of-00008.safetensors" for index in range(1, 9)
+    ]
+    assert sum(shard["parameter_count"] for shard in fixture["shards"].values()) == 1421
+    assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values())
+
+
+def test_official_default_shapes_match_public_safetensors_header_fixture() -> None:
+    module = attention_tests._load_module()
+    fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text())
+
+    with torch.device("meta"):
+        model = module.CausalLingBotWorldTransformer3DModel()
+    parameter_specs = _checkpoint_parameter_specs(model)
+    expected_dtype = {"F32": torch.float32}[fixture["dtype"]]
+
+    assert len(fixture["shard_representatives"]) == 8
+    assert set(fixture["shard_representatives"].values()) <= set(fixture["parameter_shapes"])
+
+    for name, expected_shape in fixture["parameter_shapes"].items():
+        shape, dtype = parameter_specs[name]
+        assert tuple(shape) == tuple(expected_shape), name
+        assert dtype == expected_dtype

tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_ROOT = Path(__file__).parents[4]
+_PRODUCTION_FILES = (
+    _ROOT / "vllm_omni/diffusion/models/lingbot_world/camera.py",
+    _ROOT / "vllm_omni/diffusion/models/lingbot_world/transformer.py",
+    _ROOT / "vllm_omni/diffusion/models/lingbot_world/pipeline.py",
+)
+
+
+def test_all_lingbot_production_functions_have_complete_annotations() -> None:
+    missing: list[str] = []
+    for path in _PRODUCTION_FILES:
+        tree = ast.parse(path.read_text())
+        for node in ast.walk(tree):
+            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                continue
+            arguments = (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs)
+            for argument in arguments:
+                if argument.arg not in {"self", "cls"} and argument.annotation is None:
+                    missing.append(f"{path.name}:{node.lineno}:{node.name}:{argument.arg}")
+            if node.args.vararg is not None and node.args.vararg.annotation is None:
+                missing.append(f"{path.name}:{node.lineno}:{node.name}:*{node.args.vararg.arg}")
+            if node.args.kwarg is not None and node.args.kwarg.annotation is None:
+                missing.append(f"{path.name}:{node.lineno}:{node.name}:**{node.args.kwarg.arg}")
+            if node.returns is None:
+                missing.append(f"{path.name}:{node.lineno}:{node.name}:return")
+
+    assert missing == []