[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("")