[Diffusion] Register runtime cache for direct DLO H2D

Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup.

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/.nav.yml

@@ -116,6 +116,7 @@ nav:
         - Async Diffusion Output: design/feature/async_diffusion_output.md
         - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
         - Automatic Prefix Caching: design/feature/prefix_caching.md
+        - Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
       - Communication:
         - OmniConnector Implementations:
           - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md

docs/design/ar_diffusion_pipeline_capability.md

@@ -4,6 +4,9 @@
 > `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
 > direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
 > but it is not yet part of the stable module-design hierarchy.
+>
+> The end-to-end request, identity, failure, and memory contracts are documented
+> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
 
 `ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
 autoregressive attention KV across requests. The engine only selects
@@ -51,14 +54,13 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
 those paths requires a batch-aware state-binding contract rather than silently
 bypassing `bind_ar_diffusion_state()`.
 
-The current implementation retains one session per model instance. A request
-for a different `session_id` closes the resident session through the same LRU
-lifecycle used by explicit close and failure cleanup before allocating the new
-one. Self-attention and cross-attention memory are therefore sized for one
-resident session even if the pipeline declares a larger `session_capacity`.
-This prioritizes rollout throughput and prevents block-pool exhaustion before a
-count-based eviction can run. Multi-session residency can be added later with a
-memory-backed capacity calculation or allocation-pressure eviction.
+The implementation may retain multiple sessions per model instance. It selects
+the largest feasible capacity up to the pipeline's `session_capacity`, based on
+self-attention pages, cross-attention allocations, scratch, and declared
+model-owned state. One session is admitted whenever it fits actual free device
+memory; `gpu_memory_fraction` controls additional residency. Requests beyond
+that capacity evict the least-recently-used session through the same lifecycle
+hook used by explicit close and failure cleanup.
 
 Scratch storage is also capability-driven. `frames_per_block` reserves space for
 an uncommitted current video block, while `max_scratch_tokens_per_branch`

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -19,15 +19,17 @@ DLO weight collective.
 
 Host storage is selected separately from the transfer protocol. The loader can
 produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
-otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
-the same node can share immutable checkpoint pages when direct mmap is
-selected, while the ordinary-loader fallback still keeps a private runtime
-copy per process.
+otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B
+runtime cache can normalize those final ordinary-loader DiT tensors into an
+immutable node-local mmap entry. Consequently, replicas can share either
+proven-compatible checkpoint pages or equivalent final runtime layouts. An
+opt-in CUDA registration budget can make the complete final mapping directly
+H2D-copyable, avoiding the recurrent host packing otherwise required by mmap.
 
 The Phase A shared-mmap support boundary is TP1. TP greater than one is an
 ordinary-loader compatibility path: DLO can consume the resulting TP-local
-tensors, but those configurations do not use checkpoint mmap and must not be
-used to claim shared-mmap host-memory savings.
+tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
+can share each final TP coordinate among equivalent processes.
 
 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -81,6 +83,90 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
 transforms. Model-specific direct-layout knowledge, when required, lives in a
 checkpoint adapter beside the ordinary loader.
 
+### Post-load runtime cache
+
+The Phase B cache preserves the same ownership boundary. DLO neither derives
+cache keys nor interprets model loaders. When enabled with no-AllGather, every
+rank first completes the ordinary loader, custom callbacks, post-load
+processing, strict validation, and calibration. The loader then:
+
+1. discovers final CPU DiT parameters and persistent buffers;
+2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
+   ambiguous layouts before changing the model;
+3. hashes tensor names, metadata, and final bytes;
+4. builds or joins one content-addressed entry under a strict POSIX lock; and
+5. returns a `HostWeightPlan(backing_kind="runtime_cache")` only after file,
+   manifest, and remapped-tensor content validation succeeds.
+
+The identity excludes DP size/rank and SP rank. It includes TP world size/rank
+and conservative SP implementation/world-size guards. DP replicas and
+equivalent SP ranks can therefore share; different TP coordinates cannot.
+Final tensor content is authoritative, so an unmodeled loader/config change
+that changes bytes creates a different entry instead of risking silent
+sharing.
+
+Publication writes sharded safetensors and a JSON manifest into a same-
+filesystem temporary directory, flushes them, and atomically renames the
+directory. The manifest includes per-file SHA-256 digests and the final tensor
+content digest. Kernel `flock` ownership distinguishes a live writer from a
+stale lock file: process death releases the lock automatically. Stale temporary
+entries are removed only by the next lock owner. A bounded wait or any cache
+error retains the already-valid ordinary tensors.
+
+DLO maps this plan without moving the DiT to `meta` and without rerunning any
+post-load hook. It changes tensor storage while preserving `Parameter` objects,
+then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
+streaming hooks. The source mappings remain open through DLO disable. With the
+default zero registration budget they feed the existing bounded two-slot host
+staging path. Every streamed block is first packed from pageable mmap storage
+into one of those pinned slots and is then copied to the device. This removes
+model-sized private pinned storage, but adds a recurrent host-to-host copy and
+can amplify TP synchronization waits when several ranks compete for host
+bandwidth.
+
+This version is deliberately opt-in and steady-state-oriented. It requires
+roughly one local-disk model copy per runtime identity, has no eviction, and
+still performs ordinary loading plus full content-hash passes in every rank.
+Skip-load-on-hit is a separate follow-up after loader side effects can be
+modeled safely.
+
+### Registered runtime-cache transfer
+
+`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete
+runtime-cache registration. It also changes loader selection: no-AllGather
+uses the ordinary loader and final runtime cache even when a TP1 direct-
+checkpoint plan is available. A checkpoint binding can carry a deferred
+adapter such as grouped-QKV reordering. Registering that raw source would not
+remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache
+contains the transform-complete final tensors.
+
+After mapping and validating every final tensor, DLO groups source storages by
+backing file, page-aligns their address ranges, and coalesces overlapping or
+adjacent ranges within that mapping. The complete byte count is checked
+against the per-worker budget before any CUDA call. It then registers every
+range with the read-only CUDA host-registration flag and verifies that PyTorch
+recognizes every mapped tensor as pinned. Failure rolls back any ranges already
+registered and preserves the existing two-slot staging path. Partial direct
+transfer is deliberately unsupported. A rollback failure aborts worker
+initialization so a still-registered range is retained until process/context
+teardown rather than being unsafely unmapped.
+
+On success, the hooks copy each immutable tensor view directly into its offset
+in the existing flattened rotating device buffer on the copy stream. Parameter
+repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
+implementation currently issues more H2D copy operations than block packing,
+but avoids the much larger recurrent CPU copy. A packed cache schema should be
+considered only if launch-count measurements justify its added format and
+publication complexity.
+
+Registration is per process and CUDA context; physical file-cache pages remain
+shared and are not copied into one private host allocation per worker. The
+operator must nevertheless budget page-locked memory and satisfy OS/CUDA
+registration limits for every worker. During shutdown, DLO first synchronizes
+outstanding device work, unregisters all ranges, and then closes the safetensors
+mappings. Worker teardown invokes this cleanup before destroying the
+distributed environment and CUDA context.
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -114,11 +200,17 @@ When DP is greater than one, the engine can process one request per DP rank in
 the same denoising wave. Because AllGather is a collective, all participating
 requests must take the same execution path at every denoising step.
 
+Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
+independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
+plus shared mmap storage avoids coupling their request schedules and failure
+domains.
+
 ### Rank-local path without DLO AllGather
 
 With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
 one and streams complete blocks using H2D copies only. The host backing may be
-either a loader-approved checkpoint mapping or ordinary runtime tensors.
+a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or
+ordinary runtime tensors.
 
 For direct mmap, each process retains immutable safetensors views and uses two
 bounded pinned host staging slots. Processes on the same node that map the same
@@ -127,10 +219,14 @@ the persistent private full-model copy per pure-DP process, but each process
 still packs and transfers every complete block. Sharing is node-local; each
 node has its own page cache.
 
+For a runtime-cache mapping, a positive registration budget replaces those
+host slots with direct copies from the shared final tensor views. A zero budget
+or failed registration retains the bounded staging behavior.
+
 When direct mmap preflight fails, the regular model loader remains responsible
 for preparing each rank's weights, including TP-local tensors or HSDP-managed
-parameters. In that fallback, each pure-DP process keeps a private full runtime
-copy.
+parameters. Supported CPU layouts may then enter the opt-in runtime cache.
+Unsupported layouts keep one private runtime copy per process.
 
 This mode means:
 
@@ -141,17 +237,17 @@ This mode means:
 - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
   DLO's additional weight AllGather is disabled.
 - Pure DP deployments share one checkpoint-backed copy per node when direct
-  mmap is selected; the ordinary-loader fallback keeps one private runtime
-  copy per rank.
+  mmap is selected. The runtime cache can share equivalent ordinary-loader
+  outputs; its fallback keeps one private runtime copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
 ## Parallelism compatibility
 
 | Parallelism | DLO + AllGather | DLO without AllGather |
 |---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
-| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
-| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
+| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings removes staging-driven TP skew but does not remove ordinary TP collectives. |
 | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
 
 ### Combined dimensions
@@ -182,12 +278,10 @@ requirements.
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
 ordinary loader. Online quantization remains incompatible with DLO AllGather;
-use `--dlo-no-use-allgather` or disable online quantization.
-
-A normalized runtime mmap cache, built through the ordinary loader, is the
-proposed general mechanism for sharing transformed TP or quantized layouts.
-That cache and its publication/lifecycle protocol are intentionally outside
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+use `--dlo-no-use-allgather` or disable online quantization. Runtime-cache v1
+is no-AllGather only and rejects all quantization configurations, HSDP/DTensor,
+expert parallelism, CFG parallelism, PP, non-contiguous tensors, and
+tied/shared storage.
 
 ## Validation coverage
 
@@ -200,6 +294,13 @@ Current source-level validation includes:
 - exact loader-to-backend plan transfer and ordinary-loader fallback;
 - rank-local mmap source retention, bounded two-slot staging, and adapter
   transforms without parameter-side flags;
+- runtime-cache content identity, sharded atomic publication, concurrent
+  writer election, corruption rebuild, lock timeout, and fail-closed layout
+  rejection;
+- loader ordering through final post-load mutation and DLO remapping without
+  rerunning post-load hooks;
+- all-or-nothing page-aligned registration, budget rejection, partial-failure
+  rollback, direct-copy staging bypass, and explicit shutdown cleanup;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
@@ -259,11 +360,99 @@ from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
 because it counts a shared physical page in every process that maps it; summed
 PSS is the appropriate node-memory comparison.
 
-The highest-value missing coverage is broader end-to-end numerical and
-lifecycle comparison against ordinary layerwise offload for DP+SP,
-HSDP+SP+no-AllGather, and TP greater than one across additional models and
-target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
-change the Phase A direct-mmap TP1 support boundary.
+### L20X runtime-cache smoke matrix
+
+Phase B was also exercised on four L20X GPUs in one NVLink domain with
+MiniMax-H3 FL2VA, CUDNN attention, eager execution, 256x256 four-second
+outputs, two denoising steps, `dlo_resident_layers=0`, text-encoder TP1, and
+VAE TP1. This is a functional, memory, and profiler smoke test rather than a
+production-quality performance benchmark. Host values below are node-wide
+maxima during request waves from descendant-process `smaps_rollup`; they do
+not represent the private startup peak. HBM is the engine-reported request
+peak because sampled `nvidia-smi` data can miss short-lived allocations.
+
+One internal DP2xTP2 engine compared the three host-storage/transfer choices.
+Each measured wave contained two requests. NVLink TX is the physical per-link
+counter delta summed across the four GPUs; RX had the same value and is not
+added again here.
+
+| Mode | Two-request wave | Throughput | Host PSS | `Private_Dirty` | Engine peak HBM | NVLink TX / wave |
+|---|---:|---:|---:|---:|---:|---:|
+| AllGather, ordinary loader | 3.41 s | 0.587 req/s | 210.70 GiB | 209.72 GiB | 12,498 MB | 73.37 GiB |
+| no-AllGather, private loader weights | 6.96 s | 0.287 req/s | 312.56 GiB | 311.56 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (cold publish) | 18.94 s | 0.106 req/s | 178.71 GiB | 116.00 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (warm reuse) | 24.23 s | 0.083 req/s | 178.29 GiB | 115.58 GiB | 11,882 MB | 19.78 GiB |
+
+The no-AllGather cache preserved the no-AllGather NVLink traffic and HBM
+footprint while reducing node PSS by about 134 GiB relative to private loader
+weights. Cold and warm labels describe cache publication or reuse during
+initialization; both request phases ran after initialization. A warm hit still
+performed ordinary loading and full content hashing, so it was not a startup
+fast path. All measured no-AllGather private/cache outputs were byte-identical
+for matching seeds.
+
+The deployment shape targeted by the cache was tested separately as two
+independent DP1xTP2 engines on the same node. Both engines mapped the same two
+TP-coordinate entries, including the same file device/inode tuples; the cache
+occupied 66,281,995,903 bytes. A start barrier kept all four measured requests
+ahead of profiling.
+
+| Storage | Mean request latency, engines A / B | Combined throughput | Host PSS | `Private_Dirty` | Engine peak HBM |
+|---|---:|---:|---:|---:|---:|
+| Private loader weights | 3.50 / 3.60 s | 0.463 req/s | 366.20 GiB | 365.22 GiB | 11,882 MB |
+| Shared runtime cache | 8.96 / 12.30 s | 0.142 req/s | 232.23 GiB | 169.52 GiB | 11,882 MB |
+
+The shared cache reduced request-time PSS by 133.97 GiB (36.6%) and
+`Private_Dirty` by 195.70 GiB (53.6%), but reduced combined throughput by
+69.3% in this host-contention smoke. Across one profiled request per engine,
+aggregate GPU compute was unchanged (1.359 versus 1.357 seconds) and H2D
+payload was exactly 264.24 GiB in both modes. Aggregate CPU `aten::copy_` time
+rose from 3.75 to 32.49 seconds, while communication-kernel time rose from
+3.56 to 5.80 seconds as TP ranks waited on skewed staging. These durations are
+summed across ranks and can overlap; they identify the mechanism, not wall
+time. Matching private/cache outputs were byte-identical.
+
+The result confirms that zero-budget runtime-cache staging is a host-capacity
+option, not a latency-neutral replacement for private pinned weights. Broader
+SP, model, platform, and production-quality validation remains tracked in
+[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231). The
+historical B300 rows above predate the runtime cache and must not be treated as
+Phase B memory validation.
+
+### L20X registered-runtime-cache smoke
+
+The read-only registered path was then tested with the same MiniMax-H3
+workload. These are functional and mechanism-validation measurements, not a
+production benchmark. The TP1 pair used one request per wave. Both modes moved
+exactly 121.86 GiB H2D in the profiled request and produced byte-identical video
+and audio for both measured seeds.
+
+| TP1 no-AllGather mode | Mean wave | Throughput | Request-time host PSS | `Private_Dirty` | Engine peak HBM | CPU `aten::copy_` | GPU compute |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| Checkpoint mmap + bounded staging | 33.87 s | 0.0295 req/s | 133.65 GiB | 71.16 GiB | 13,226 MB | 21.91 s | 419.28 ms |
+| Final runtime cache + registered direct H2D | 4.88 s | 0.2055 req/s | 129.41 GiB | 66.93 GiB | 13,226 MB | 1.64 s | 419.05 ms |
+
+Registration covered 61.73 GiB in 13 page-aligned ranges. Mean wave latency
+fell by 85.6% and throughput increased 6.94x. GPU compute and H2D payload were
+unchanged, identifying recurrent host transform/packing as the removed work.
+The direct path issued 2,425 H2D events versus 1,945 for packed staging, but
+that launch increase was much smaller than the eliminated CPU cost.
+
+Two independent TP2 engines also registered the same two TP-coordinate cache
+entries across four workers. Compared with the zero-budget cache, mean wave
+latency improved from 9.21/12.55 seconds to 3.80/3.86 seconds for engines A/B;
+the private-loader reference was 3.61/3.72 seconds. Registered request-time
+host PSS was 224.26 GiB versus 232.23 GiB staged and 366.20 GiB private.
+Aggregate CPU `aten::copy_` time was 3.83 seconds registered, 32.49 seconds
+staged, and 3.75 seconds private; aggregate GPU compute and the 264.24 GiB H2D
+payload were effectively unchanged. Ordinary TP communication remained, but
+the extra wait caused by staging-driven rank skew largely disappeared.
+
+All workers explicitly unregistered before releasing their mappings, the
+system's acquired-minus-released pin counter returned to its pre-run value,
+and every GPU allocation was released. Cold runtime-cache publication remained
+expensive because it still performs ordinary loading and multiple full-content
+passes; that startup problem is outside the registered transfer scope.
 
 ## Recommendations
 
@@ -272,8 +461,12 @@ change the Phase A direct-mmap TP1 support boundary.
 - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
   not the goal.
 - Use **no-AllGather** when independent replica execution is required. TP1
-  direct-mmap deployments can share checkpoint pages per node; other layouts
-  retain the ordinary loader's private host memory behavior and are outside
-  the Phase A shared-mmap support scope.
+  direct-mmap deployments can share checkpoint pages per node. Enable the
+  runtime cache only when independent replicas have repeated final layouts and
+  benchmark the target CPU/memory topology, especially for matching
+  coordinates across multiple TP engines. On CUDA, provide an explicit
+  registration budget large enough for the complete final layout when
+  recurrent staging cost is unacceptable; leave it at zero when page-locking
+  that layout is not operationally acceptable.
 - Prefer **HSDP alone** for production HSDP deployments until the combined
   HSDP + DLO no-AllGather path has broader end-to-end coverage.

docs/design/feature/realtime_ar_diffusion.md

@@ -0,0 +1,213 @@
+# Realtime AR-Diffusion sessions
+
+> **Status:** experimental. This document describes the internal execution
+> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
+> define a public HTTP or WebSocket schema.
+
+## Scope
+
+Realtime AR-Diffusion serves world models that generate one autoregressive
+latent block per request while retaining model state across requests. The
+generic runtime owns identity, event ordering, session lifecycle, routing, and
+KV allocation. A model adapter owns model-specific controls and conditioning.
+
+LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
+same runner capability with two logical CFG branches.
+
+## Identities
+
+Four identifiers have distinct lifetimes:
+
+| Identifier | Lifetime | Owner | Contract |
+| --- | --- | --- | --- |
+| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
+| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
+| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
+| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
+
+`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
+`request_id` remains inside the immutable AR-Diffusion snapshot and the
+standard output metadata, while caller-visible `OmniRequestOutput.request_id`
+continues to use AsyncOmni's existing external-ID contract. Protocol identity
+therefore does not require a private entrypoint override.
+
+## Request lifecycle
+
+1. `ARDiffusionSession.accept_event()` validates and queues an immutable
+   `ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
+   the same event.
+2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
+   `_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
+   of the current prompt, reduced controls, ordered event IDs, and chunk index.
+3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
+   parameters, stores the typed tick in
+   `sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
+   prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
+4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
+   `OmniDiffusionRequest` through the orchestrator.
+5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
+   `_get_or_create_session()` obtains runner-owned paged KV, and
+   `bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
+   of `execute_model()`.
+6. The model pipeline validates the tick and generates exactly one latent
+   block. LingBot interprets camera controls in
+   `LingBotCameraControlReducer` and
+   `LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
+   camera or action schemas.
+7. The pipeline returns one `DiffusionOutput` with the generated payload and an
+   `ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
+   `AsyncOmni` exposes it as
+   `OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
+8. The consumer parses the standard envelope. The session commits its reducer
+   state, prompt, controls, event IDs, and next chunk index only when returned
+   metadata exactly equals the submitted tick snapshot.
+
+The metadata envelope is:
+
+```json
+{
+  "metadata": {
+    "ar_diffusion": {
+      "session_id": "world-7",
+      "request_id": "ar-world-7-3-...",
+      "chunk_index": 3,
+      "applied_event_ids": [10, 11]
+    }
+  }
+}
+```
+
+## Transaction and failure semantics
+
+Accepted events are not acknowledged as applied until the model output and
+metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
+and the session snapshot form one logical commit. A model, metadata, or reducer
+failure leaves the chunk index and queued events uncommitted and transitions the
+session to `FAILED`.
+
+After a tick failure, worker state is closed because a forward may have
+partially modified or already released KV. The same chunk is not retried in
+place. A caller must explicitly reset the session, starting again at chunk
+zero, or close it.
+
+Explicit close follows this state machine:
+
+```text
+ACTIVE or FAILED
+       |
+       v
+    CLOSING ---- worker cleanup succeeds ----> CLOSED
+       |
+       +------- worker cleanup fails --------> CLEANUP_FAILED
+                                                   |
+                                                   +-- close retry --> CLOSING
+```
+
+`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
+creation of the same `session_id`, and rejects reset. Only an idempotent close
+retry that succeeds removes the session. Transport disconnect uses the same
+path.
+
+## KV ownership and capacity
+
+`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
+declares:
+
+- TP-local self-attention geometry, AR block size, sink, and recent window;
+- logical KV branches and their worker-local slots;
+- fixed-length cross-attention caches;
+- scratch tokens required by an uncommitted forward;
+- persistent model-owned CUDA bytes per resident session; and
+- the requested resident-session capacity.
+
+The runner computes:
+
+```text
+required(capacity) =
+    scratch
+  + managed_self_attention_pages(capacity)
+  + capacity * cross_attention_bytes_per_session
+  + capacity * model_owned_state_bytes_per_session
+```
+
+`gpu_memory_fraction` is a soft expansion budget:
+
+```text
+configured_budget = available_device_bytes * gpu_memory_fraction
+effective_budget = max(configured_budget, required(1))
+```
+
+If `required(1)` exceeds actual available device memory, initialization fails.
+Otherwise the runner always admits one viable session and uses the configured
+fraction to select additional resident sessions up to the model-declared cap.
+All model-owned reservations are deducted before the paged self-attention pool
+is allocated.
+
+LingBot reports its persistent full-horizon image-condition tensor. DreamZero
+reports the measured 603 MiB per-session Wan VAE causal-convolution state.
+
+## Concurrency and routing limits
+
+The current runner executes one request at a time:
+
+- `max_num_seqs` must be one;
+- request batching and diffusion step execution are rejected;
+- one AR block is generated per engine request; and
+- an AR-Diffusion stage must have exactly one replica.
+
+The single-replica restriction is intentional. Multi-replica support requires
+session-affine routing so every tick, reset, and close reaches the worker that
+owns the session. Random or round-robin routing is not correct.
+
+Multiple resident sessions may share one runner. Capacity selection accounts
+for all persistent pools, and the runner performs LRU eviction through the same
+model lifecycle hook used by explicit close.
+
+## Model boundary
+
+The generic runtime transports controls as opaque
+`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
+and reduce those values. This keeps target/velocity, key state, pose,
+intrinsics, and other model-specific conventions out of the session protocol.
+
+For LingBot World v2:
+
+- `lingbot_world/actions.py` parses key-state and trajectory controls;
+- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
+- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
+- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
+  produces one block plus the standard metadata envelope.
+
+## Non-goals
+
+This contract does not currently provide:
+
+- a public HTTP or WebSocket request schema (the structured-interaction
+  frontend is tracked separately in
+  [#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
+- camera/action semantics shared by every world model;
+- session migration or replication across workers;
+- retry of an ambiguous partially executed chunk;
+- cross-stage KV transfer; or
+- stateful streaming VAE decode.
+
+Those features can be layered on top without adding LingBot-specific identity
+or lifecycle fields to the generic runtime.
+
+## Required regression coverage
+
+CPU contract tests cover:
+
+- separation of tick correlation IDs from internal engine routing IDs;
+- event ordering, metadata equality, and reducer fail-closed commit;
+- close/disconnect failure tombstones and cleanup retry;
+- capacity one under the default fraction for shipped LingBot geometry;
+- capacity expansion under a tuned fraction;
+- DreamZero model-owned state accounting;
+- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
+- LingBot model registration, imports, camera/action reduction, and one-block
+  output metadata.
+
+GPU validation must additionally exercise real model loading, action input,
+prompt switching, rolling-window replacement, reset, close, finite output
+latents, and exact metadata/request identity.

docs/design/index.md

@@ -22,6 +22,7 @@ implementation contract; it is not, by itself, a general support claim.
 - [Async Diffusion Output](feature/async_diffusion_output.md)
 - [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
 - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
+- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
 
 ### Communication
 

docs/models/supported_models.md

@@ -37,6 +37,7 @@ th {
 | `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -117,6 +118,10 @@ support metadata until their recipe is audited.
 
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
 
+See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
+for offline and experimental realtime usage, supported controls, validation,
+and current limitations.
+
 <sup>H3</sup> MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
 gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
 Other AMD SKUs (e.g. MI325X) are not listed until their

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -5,7 +5,11 @@ deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
 the host weights and reconstructs each layer at runtime. Without AllGather,
 each rank streams a complete block independently. Compatible TP1 deployments
 can share checkpoint-backed host pages among processes on the same node;
-otherwise DLO streams the ordinary loader's rank-local tensors.
+otherwise DLO streams the ordinary loader's rank-local tensors. The optional
+runtime-weight cache extends node-local sharing to final ordinary-loader
+layouts, including matching TP coordinates across independent engines. On
+CUDA, an explicit registration budget can make those shared mappings directly
+H2D-copyable and avoid recurrent mmap-to-pinned host packing.
 
 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -21,8 +25,10 @@ AllGather:       [AG N+1]           [AG N+2]
 Slots:           slot 0: Layer N    slot 1: Layer N+1
 ```
 
-AllGather communicates only request-independent weight shards, so data-
-parallel ranks may process different requests concurrently.
+AllGather communicates request-independent weight shards, but every member of
+the transfer group must request the same next block in the same collective
+order. A synchronized DP wave may contain different prompts only when those
+requests follow the same denoising and block-execution path.
 
 ## Usage
 
@@ -38,6 +44,15 @@ vllm serve /path/to/model --omni \
   --data-parallel-size 4 \
   --dlo-no-use-allgather
 
+# Independently scheduled replicas sharing final runtime weights on one node
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --tensor-parallel-size 2 \
+  --dlo-no-use-allgather \
+  --dlo-enable-runtime-cache \
+  --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \
+  --dlo-runtime-cache-pin-limit-gib 80
+
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
@@ -50,7 +65,10 @@ from vllm_omni import Omni
 omni = Omni(
     model="/path/to/model",
     enable_distributed_layerwise_offload=True,
-    dlo_use_allgather=True,
+    dlo_use_allgather=False,
+    dlo_enable_runtime_cache=True,
+    dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights",
+    dlo_runtime_cache_pin_limit_gib=80,
 )
 ```
 
@@ -62,6 +80,10 @@ omni = Omni(
 | `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
 | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
+| `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` |
+| `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` |
+| `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` |
+| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
 
 ## Host-weight loading
@@ -78,6 +100,50 @@ DLO may still consume those TP-local tensors, but this is a compatibility path:
 it does not share checkpoint-backed runtime weights across DP replicas and
 provides no shared-mmap host-memory guarantee.
 
+When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading
+gains a second storage opportunity. Every rank completes loader callbacks,
+post-load processing, validation, and calibration. It then hashes the final
+CPU DiT parameters and persistent buffers. One process per equivalent runtime
+layout publishes sharded safetensors plus a checksummed manifest; peers
+validate and mmap the same files. DLO replaces only tensor storage, preserves
+the final `Parameter` objects, and releases the private allocator-backed copies
+before installing its streaming hooks.
+
+A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final
+runtime layout even when a direct TP1 checkpoint plan is available. Raw
+checkpoint mappings may require deferred loader transforms such as grouped-QKV
+reordering; registering them would still leave recurrent host work. Publishing
+the fully transformed layout once makes every streamed source directly
+copyable.
+
+The runtime layout identity excludes DP rank and SP rank. It includes TP world
+size/rank and conservative SP implementation/world-size guards. Therefore two
+independent `TP=2` engines normally create two entries: both TP0 workers map
+one entry and both TP1 workers map the other. Cache sharing does not create a
+process group or an inference-time collective.
+
+With a zero registration budget, cached tensors remain pageable mmap views.
+DLO packs each streamed block into one of two bounded pinned host slots before
+the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent
+host-to-host copies; competing TP ranks or engines can turn DDR contention and
+staging skew into collective wait time.
+
+With a positive budget on CUDA, DLO instead attempts to register the complete
+page-aligned cache mapping as read-only pinned host memory. Registration is
+all-or-nothing: the total mapped range must fit the per-worker budget, and any
+CUDA registration or verification failure rolls back partial registrations and
+uses the bounded staging path. If CUDA cannot roll back a partial registration,
+worker initialization aborts instead of unmapping memory that CUDA may still
+own. On success, each immutable tensor view is copied directly into the
+existing rotating device buffers on the H2D stream, so no private host staging
+slots are allocated.
+
+CUDA registration is process-local, but it does not clone file contents. The
+same physical page-cache pages remain shared among equivalent workers; each
+worker does, however, hold its own registration reference and must satisfy the
+host's page-locking limits. Shutdown synchronizes outstanding GPU work,
+unregisters every range, and only then closes the mmap handles.
+
 The mmap plan skips only dedicated DiT weight sources. Other component sources,
 such as a text encoder loaded through the shared diffusion loader, continue to
 use their ordinary component loader. A checkpoint source that mixes DiT and
@@ -104,6 +170,34 @@ process.
 When the effective DLO group size is one, `dlo_use_allgather=True` does not
 perform a collective and uses the same rank-local transfer behavior.
 
+## Runtime-cache lifecycle
+
+Runtime-cache publication uses a per-layout POSIX file lock, a temporary
+directory on the same filesystem, file and final-tensor SHA-256 validation,
+`fsync`, and an atomic directory rename. A crashed writer releases its kernel
+lock; the next writer removes only that layout's stale temporary directories
+while holding the lock. Lock timeout, unsupported layout, insufficient disk,
+or validation failure keeps the already-valid ordinary tensors.
+
+The cache root is deliberately independent of `VLLM_CACHE_ROOT`, because
+vLLM-Omni may isolate that variable between stage replicas. All consumers that
+should share must receive the same `--dlo-runtime-cache-dir` on a local,
+disk-backed filesystem. Do not place the cache on tmpfs when the goal is lower
+host-memory PSS, and do not use a cross-node directory in this version.
+
+There is no automatic eviction. Budget roughly one final DiT copy per distinct
+TP coordinate, dtype/layout, and model version, plus temporary space while an
+entry is published. Operators may remove entries or the complete cache root
+only after all workers using its mmap files have stopped. The feature is
+opt-in for this reason.
+
+For direct H2D, set the registration limit above the complete page-aligned
+layout size reported in the worker log. Partial registration is intentionally
+unsupported because mixing direct and staged tensors would complicate buffer
+lifetime and make performance topology-dependent. Registration is CUDA-only;
+an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or
+registration error is reported and falls back to bounded staging.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -132,22 +226,37 @@ must enter each collective.
 
 ## Limitations
 
-- Direct checkpoint mmap currently requires TP1. TP greater than one is
-  outside the Phase A shared-mmap support scope and falls back before model
-  mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
-  but it provides no shared-mmap host-memory benefit or guarantee.
+- Direct checkpoint mmap currently requires TP1. TP greater than one falls
+  back to the ordinary TP-aware loader. With the optional runtime cache,
+  equivalent processes at the same TP coordinate can then share that final
+  rank-local layout in no-AllGather mode. A positive registration budget also
+  chooses the final runtime cache at TP1 so deferred checkpoint transforms are
+  not repeated during inference.
 - HSDP plus AllGather is rejected to avoid double sharding. HSDP without
   AllGather has limited end-to-end validation.
 - Online quantization uses the ordinary loader with no-AllGather. It remains
-  incompatible with DLO AllGather.
+  incompatible with DLO AllGather and runtime-cache v1. Runtime-cache v1 also
+  rejects other quantized, non-contiguous, aliased/tied, device-only, HSDP,
+  expert-parallel, CFG-parallel, and PP layouts.
 - Resident leading layers require `--dlo-no-use-allgather` and a model
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
-[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
-the direct-checkpoint path.
+The runtime cache improves steady-state host PSS, not startup peak: every rank
+loads private weights before remapping, and validation adds full content-hash
+passes. Direct H2D removes recurrent staging but does not change that startup
+contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+
+Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3
+smoke documented in the L20X runtime-cache matrix in the
+[feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix),
+two independent TP2 engines reduced request-time host PSS by 36.6% with the
+cache, but recurrent mmap-to-pinned staging reduced combined throughput by
+69.3%. The registered path recovered near-private per-engine latency while
+retaining the shared-memory benefit; a matched TP1 run improved mean wave
+latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the
+feature design for the full smoke-test caveats and metrics.
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

examples/offline_inference/aura_omni/end2end.py

@@ -24,7 +24,6 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.media.audio import load_audio
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.stage_input_processors.aura_omni import (
     DEFAULT_QWEN3_TTS_REF_TEXT,
@@ -228,7 +227,6 @@ def parse_args():
     parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
     parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
     parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/diffusion/lingbot_world_v2.py

@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Generate a LingBot-World v2 video from an image and camera trajectory.
+
+The official checkpoint is licensed separately under CC BY-NC-SA and is
+restricted to non-commercial use. This vLLM-Omni integration remains
+Apache-2.0 licensed.
+
+Example:
+    python examples/offline_inference/diffusion/lingbot_world_v2.py \
+        --prompt "The camera moves slowly forward through the scene." \
+        --image /path/to/first_frame.png \
+        --action-dir /path/to/actions/forward \
+        --output lingbot_world_v2.mp4
+
+``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
+example configures its parent as the trusted action root and sends the
+contained directory through ``extra_args["action_path"]``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
+_NUM_INFERENCE_STEPS = 4
+_MAX_SEQUENCE_LENGTH = 512
+_MAX_PIXEL_AREA = 480 * 832
+_MAX_RAW_FRAMES = 117
+_MAX_SOURCE_CAMERA_FRAMES = 4096
+_TEMPORAL_COMPRESSION = 4
+_LATENT_FRAMES_PER_BLOCK = 3
+_SPATIAL_ALIGNMENT = 16
+
+
+@dataclass(frozen=True)
+class LingBotPaths:
+    image: Path
+    action_dir: Path
+    action_root: Path
+    action_relative: Path
+    camera_frames: int
+    output: Path
+
+
+def _validate_camera_arrays(action_dir: Path) -> int:
+    """Validate the bounded camera contract before starting the engine."""
+
+    try:
+        poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
+        intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
+    except (OSError, TypeError, ValueError):
+        raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
+    if poses.ndim != 3 or poses.shape[1:] != (4, 4):
+        raise ValueError("poses.npy must have shape [frames, 4, 4].")
+    if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
+        raise ValueError("intrinsics.npy must have shape [frames, 4].")
+    if poses.shape[0] != intrinsics.shape[0]:
+        raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
+    if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
+        raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
+    if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
+        raise ValueError("Camera arrays must contain real numeric values.")
+    if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
+        raise ValueError("Camera arrays must contain only finite values.")
+    return int(poses.shape[0])
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
+    )
+    parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
+    parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
+    parser.add_argument("--image", required=True, help="Path to the first-frame image.")
+    parser.add_argument(
+        "--action-dir",
+        required=True,
+        help="Camera action directory containing poses.npy and intrinsics.npy.",
+    )
+    parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
+    parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=9,
+        help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism inside the DiT.",
+    )
+    parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
+    parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
+    parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
+    parser.add_argument(
+        "--enforce-eager",
+        action="store_true",
+        help="Disable compilation and force eager execution.",
+    )
+    return parser.parse_args(argv)
+
+
+def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
+    """Canonicalize local inputs and construct a safe trusted-root mapping."""
+
+    image = Path(args.image).expanduser().resolve()
+    if not image.is_file():
+        raise ValueError("--image must point to an existing file.")
+
+    action_dir = Path(args.action_dir).expanduser().resolve()
+    if not action_dir.is_dir():
+        raise ValueError("--action-dir must point to an existing directory.")
+    if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
+        raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
+    camera_frames = _validate_camera_arrays(action_dir)
+
+    action_root = action_dir.parent.resolve()
+    try:
+        action_relative = action_dir.relative_to(action_root)
+    except ValueError:
+        raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
+    if not action_relative.parts:
+        raise ValueError("--action-dir must name a directory contained by a trusted action root.")
+
+    output = Path(args.output).expanduser().resolve()
+    return LingBotPaths(
+        image=image,
+        action_dir=action_dir,
+        action_root=action_root,
+        action_relative=action_relative,
+        camera_frames=camera_frames,
+        output=output,
+    )
+
+
+def _positive_finite(value: float, flag: str) -> float:
+    result = float(value)
+    if not math.isfinite(result) or result <= 0:
+        raise ValueError(f"{flag} must be a positive finite number.")
+    return result
+
+
+def build_omni_kwargs(
+    args: argparse.Namespace,
+    paths: LingBotPaths,
+    *,
+    parallel_config: Any,
+) -> dict[str, Any]:
+    """Build Omni configuration without overriding checkpoint class discovery."""
+
+    if args.tensor_parallel_size <= 0:
+        raise ValueError("--tensor-parallel-size must be a positive integer.")
+    flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
+    model_path = Path(args.model).expanduser()
+    model = str(model_path.resolve()) if model_path.exists() else args.model
+    return {
+        "model": model,
+        "flow_shift": flow_shift,
+        "parallel_config": parallel_config,
+        "enforce_eager": args.enforce_eager,
+        "model_config": {"lingbot_action_root": str(paths.action_root)},
+    }
+
+
+def build_request(
+    args: argparse.Namespace,
+    paths: LingBotPaths,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Build the prompt and fixed v1 causal-DMD sampling contract."""
+
+    prompt_text = args.prompt.strip()
+    if not prompt_text:
+        raise ValueError("--prompt must contain non-whitespace text.")
+    if args.height <= 0 or args.width <= 0:
+        raise ValueError("--height and --width must be positive integers.")
+    if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
+        raise ValueError("--height and --width must each be divisible by 16.")
+    if args.height * args.width > _MAX_PIXEL_AREA:
+        raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
+    if args.num_frames <= 0:
+        raise ValueError("--num-frames must be a positive integer.")
+    if args.num_frames > _MAX_RAW_FRAMES:
+        raise ValueError("--num-frames must not exceed 117 raw frames.")
+    if args.num_frames > paths.camera_frames:
+        raise ValueError("Camera arrays must contain at least --num-frames frames.")
+    if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
+        raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
+    latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
+    if latent_frames % _LATENT_FRAMES_PER_BLOCK:
+        raise ValueError("--num-frames must map to whole three-frame latent blocks.")
+    if args.fps <= 0:
+        raise ValueError("--fps must be a positive integer.")
+    flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
+
+    prompt = {
+        "prompt": prompt_text,
+        "multi_modal_data": {"image": str(paths.image)},
+    }
+    sampling_kwargs = {
+        "height": args.height,
+        "width": args.width,
+        "num_frames": args.num_frames,
+        "num_inference_steps": _NUM_INFERENCE_STEPS,
+        "max_sequence_length": _MAX_SEQUENCE_LENGTH,
+        "seed": args.seed,
+        "fps": args.fps,
+        "extra_args": {
+            "action_path": paths.action_relative.as_posix(),
+            "flow_shift": flow_shift,
+        },
+    }
+    return prompt, sampling_kwargs
+
+
+def extract_video_array(outputs: Any) -> np.ndarray:
+    """Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
+
+    if not isinstance(outputs, list) or len(outputs) != 1:
+        raise ValueError("Expected exactly one Omni output for the single LingBot request.")
+    output = outputs[0]
+    request_output = getattr(output, "request_output", None)
+    if request_output is not None and getattr(request_output, "images", None):
+        output = request_output
+
+    images = getattr(output, "images", None)
+    if not images:
+        multimodal_output = getattr(output, "multimodal_output", None) or {}
+        images = multimodal_output.get("video")
+    if images is None:
+        raise ValueError("No video frames found in Omni output.")
+    if isinstance(images, list) and len(images) == 1:
+        images = images[0]
+
+    video: np.ndarray = np.asarray(images)
+    if video.ndim == 5:
+        if video.shape[0] != 1:
+            raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
+        video = video[0]
+    if video.ndim != 4 or video.shape[-1] not in (3, 4):
+        raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
+    return video
+
+
+def main(argv: Sequence[str] | None = None) -> Path:
+    args = parse_args(argv)
+    paths = resolve_cli_paths(args)
+    prompt, sampling_kwargs = build_request(args, paths)
+    paths.output.parent.mkdir(parents=True, exist_ok=True)
+
+    # Keep vLLM imports below pure CLI/path validation so ``--help`` and the
+    # helper tests work on machines without the native vLLM runtime.
+    from diffusers.utils import export_to_video
+
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
+    omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
+    omni = Omni(**omni_kwargs)
+    try:
+        outputs = omni.generate(
+            prompt,
+            OmniDiffusionSamplingParams(**sampling_kwargs),
+            use_tqdm=False,
+        )
+    finally:
+        omni.close()
+
+    frames = extract_video_array(outputs)
+    export_to_video(frames, str(paths.output), fps=args.fps)
+    print(f"Saved {frames.shape[0]} frames to {paths.output}")
+    return paths.output
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/diffusion/lingbot_world_v2_realtime.py

@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
+
+This example exercises the internal session API. It is not a public HTTP or
+WebSocket protocol. Each JSONL event advances the world by one three-latent
+frame AR block and writes the latent plus its identity metadata. The current
+scoped implementation supports at most ten ticks per generation epoch because
+its image condition is bounded to 117 pixel frames; reset or create a session
+to start a new world.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import math
+import time
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
+_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
+_FRAMES_PER_BLOCK = 3
+# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
+_MAX_REALTIME_TICKS = 10
+
+
+def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
+    """Build the event-side script consumed by LingBotCameraControlReducer."""
+    return {"mode": "script", "frames": frames}
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
+    parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
+    parser.add_argument("--image", required=True, help="Initial RGB image.")
+    parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
+    parser.add_argument(
+        "--events",
+        required=True,
+        help="JSONL file with one prompt/action event per AR block; at most 10 events.",
+    )
+    parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
+    parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=832)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--enforce-eager", action="store_true")
+    return parser.parse_args(argv)
+
+
+def _load_events(path: Path) -> list[dict[str, Any]]:
+    events: list[dict[str, Any]] = []
+    for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
+        line = raw_line.strip()
+        if not line:
+            continue
+        try:
+            value = json.loads(line)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
+        if not isinstance(value, dict):
+            raise ValueError(f"events line {line_number} must be a JSON object.")
+        event_id = value.get("event_id")
+        if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
+            raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
+        prompt = value.get("prompt")
+        if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
+            raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
+        frames = value.get("frames")
+        if frames is not None:
+            if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
+                raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
+            for frame in frames:
+                if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
+                    raise ValueError(f"events line {line_number} frames must contain only action strings.")
+        if prompt is None and frames is None:
+            raise ValueError(f"events line {line_number} must update prompt and/or frames.")
+        events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
+    if not events:
+        raise ValueError("events file must contain at least one event.")
+    if len(events) > _MAX_REALTIME_TICKS:
+        raise ValueError(
+            "events file must contain at most "
+            f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
+            "image-condition horizon is 117 pixel frames."
+        )
+    event_ids = [event["event_id"] for event in events]
+    if event_ids != sorted(set(event_ids)):
+        raise ValueError("event_id values must be unique and strictly increasing.")
+    return events
+
+
+def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
+    image = Path(args.image).expanduser().resolve()
+    events = Path(args.events).expanduser().resolve()
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    if not image.is_file():
+        raise ValueError("--image must point to an existing file.")
+    if not events.is_file():
+        raise ValueError("--events must point to an existing JSONL file.")
+    if not args.prompt.strip():
+        raise ValueError("--prompt must contain non-whitespace text.")
+    if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
+        raise ValueError("--height and --width must be positive multiples of 16.")
+    if args.tensor_parallel_size <= 0:
+        raise ValueError("--tensor-parallel-size must be positive.")
+    if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
+        raise ValueError("--gpu-memory-fraction must be in (0, 1].")
+    return image, events, output_dir
+
+
+async def run(argv: Sequence[str] | None = None) -> Path:
+    args = parse_args(argv)
+    image, events_path, output_dir = _validate_args(args)
+    events = _load_events(events_path)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Imports stay below pure input validation so --help and helper tests do not
+    # require a CUDA-enabled vLLM installation.
+    import torch
+
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
+    from vllm_omni.experimental.ar_diffusion.session import (
+        ARDiffusionSessionEvent,
+        ARDiffusionSessionManager,
+        ARDiffusionWorkerLifecycle,
+    )
+    from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    engine = AsyncOmni(
+        model=args.model,
+        engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
+        enforce_eager=args.enforce_eager,
+        parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
+        max_num_seqs=1,
+        model_config={
+            "ar_diffusion_height": args.height,
+            "ar_diffusion_width": args.width,
+            "ar_diffusion_kv_config": {
+                "gpu_memory_fraction": args.gpu_memory_fraction,
+                "warmup_cudagraph": True,
+            },
+        },
+    )
+    sampling = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_frames=9,
+        num_inference_steps=4,
+        max_sequence_length=512,
+        seed=args.seed,
+        output_type="latent",
+        extra_args={"flow_shift": 5.0},
+    )
+    consumer = ARDiffusionOmniTickConsumer(
+        engine,
+        prompt_provider=lambda tick: {
+            "prompt": tick.prompt,
+            "multi_modal_data": {"image": str(image)},
+        },
+        sampling_params_list=[sampling],
+        diffusion_stage_id=0,
+    )
+    manager = ARDiffusionSessionManager(
+        tick_consumer=consumer,
+        lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
+        max_pending_events=32,
+        control_reducer_factory=LingBotCameraControlReducer,
+    )
+    session = await manager.create_session(args.session_id)
+    measurements: list[dict[str, Any]] = []
+    try:
+        for chunk_index, event in enumerate(events):
+            controls = ()
+            if event["frames"] is not None:
+                controls = (
+                    ARDiffusionControlInput(
+                        track="camera",
+                        schema=_CAMERA_ACTION_SCHEMA,
+                        data=_camera_event_data(event["frames"]),
+                    ),
+                )
+            await session.accept_event(
+                ARDiffusionSessionEvent(
+                    event_id=event["event_id"],
+                    prompt=event["prompt"]
+                    if event["prompt"] is not None
+                    else (args.prompt if chunk_index == 0 else None),
+                    controls=controls,
+                )
+            )
+            started = time.perf_counter()
+            output = await session.next_chunk()
+            elapsed = time.perf_counter() - started
+            if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
+                raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
+            latent = output.images[0].detach().float().cpu()
+            metadata = consumer.chunk_metadata(output).to_dict()
+            torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
+            (output_dir / f"chunk_{chunk_index:03d}.json").write_text(
+                json.dumps(metadata, indent=2, sort_keys=True) + "\n"
+            )
+            measurements.append(
+                {
+                    "chunk_index": chunk_index,
+                    "latency_seconds": elapsed,
+                    "shape": list(latent.shape),
+                    "finite": bool(torch.isfinite(latent).all()),
+                    "metadata": metadata,
+                }
+            )
+            print(json.dumps(measurements[-1], sort_keys=True), flush=True)
+    finally:
+        try:
+            await manager.close_session(args.session_id)
+        finally:
+            engine.shutdown()
+
+    summary_path = output_dir / "summary.json"
+    summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
+    return summary_path
+
+
+def main(argv: Sequence[str] | None = None) -> Path:
+    return asyncio.run(run(argv))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/lance/end2end.py

@@ -82,9 +82,6 @@ def parse_args():
     # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
     # and ``create_default_diffusion`` materializes the stage config.
 
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(p)
     return p.parse_args()
 
 

examples/offline_inference/nemotron_voicechat/end2end.py

@@ -67,7 +67,7 @@ def main() -> None:
     parser.add_argument("--output-dir", default="results/nemotron_voicechat")
     parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
     parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
         default=None,
         help=(
             "Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
@@ -134,8 +134,8 @@ def main() -> None:
     ]
 
     omni_kwargs: dict = {}
-    if args.stage_configs_path:
-        omni_kwargs["stage_configs_path"] = args.stage_configs_path
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
     outputs = omni.generate(inputs, sampling_params_list)
 

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -38,17 +38,6 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 logger = logging.getLogger(__name__)
 
 DEFAULT_MODEL = "fishaudio/s2-pro"
-DEFAULT_DEPLOY_CONFIG = os.path.join(
-    os.path.dirname(__file__),
-    "..",
-    "..",
-    "..",
-    "..",
-    "vllm_omni",
-    "model_executor",
-    "stage_configs",
-    "fish_speech_s2_pro.yaml",
-)
 
 
 def build_prompt(
@@ -152,7 +141,8 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
-    omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     return omni_kwargs
 
 

examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py

@@ -32,11 +32,10 @@ from end2end import _estimate_prompt_len
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
 
 
-def _default_stage_config() -> str:
+def _default_deploy_config() -> str:
     repo_root = Path(__file__).resolve().parents[4]
     return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
 
@@ -109,7 +108,7 @@ def main(args: Any) -> None:
         "output_dir",
     ):
         omni_kwargs.pop(key, None)
-    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
+    omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
     omni_kwargs["log_stats"] = args.log_stats
     omni = Omni(**omni_kwargs)
 
@@ -156,7 +155,7 @@ def parse_args() -> Any:
         default=None,
         help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
     )
-    parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
+    parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
     parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
     parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
     parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
@@ -164,7 +163,6 @@ def parse_args() -> Any:
     parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
     parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
     parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
-    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

recipes/BosonAI/Higgs-Audio-V3-TTS.md

@@ -45,7 +45,7 @@ MusicGen-style delay pattern).
 **Online serving:**
 
 ```bash
-vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
+vllm serve bosonai/higgs-audio-v3-tts-4b \
     --host 0.0.0.0 --port 8095 \
     --trust-remote-code --omni
 ```

recipes/NVIDIA/Nemotron-Labs-Audex.md

@@ -51,7 +51,7 @@ NemotronH (~3B active parameters). Pick the pipeline by task:
 `model_type` auto-resolves `audex_tts.yaml`):
 
 ```bash
-vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
+vllm serve nvidia/Nemotron-Labs-Audex-2B \
     --host 0.0.0.0 --port 8097 \
     --trust-remote-code --omni
 # other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
@@ -109,7 +109,7 @@ The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
 auto-resolves to the 2B-tuned config otherwise):
 
 ```bash
-vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
+vllm serve nvidia/Nemotron-Labs-Audex-30B-A3B \
     --host 0.0.0.0 --port 8097 \
     --trust-remote-code --omni \
     --stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml

recipes/README.md

@@ -60,6 +60,7 @@ recipes/
 | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
 | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
+| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 |
 | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
 | [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
 | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |

recipes/Robbyant/LingBot-World-2.0.md

@@ -0,0 +1,122 @@
+# LingBot-World 2.0
+
+> Offline and experimental realtime interactive world generation
+
+## Summary
+
+- Vendor: Robbyant
+- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
+- Task: image-conditioned interactive world generation
+- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
+- Hardware validated: NVIDIA H200 and B200
+- Maintainer: Community
+
+The checkpoint is separately licensed under CC BY-NC-SA and restricted to
+non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
+
+## Offline generation
+
+The offline path consumes one source image and an action directory containing
+`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape
+`[frames, 4]`.
+
+```bash
+python examples/offline_inference/diffusion/lingbot_world_v2.py \
+  --prompt "The camera moves slowly forward through the scene." \
+  --image /path/to/first_frame.png \
+  --action-dir /path/to/actions/forward \
+  --num-frames 81 \
+  --output lingbot_world_v2.mp4
+```
+
+The current causal-fast checkpoint uses four DMD sampling steps. Raw frame
+counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
+longer bounded source trajectory (the official example contains 269 frames);
+the runtime consumes only the prefix needed by the request.
+
+## Realtime in-process generation
+
+The realtime example keeps AR-Diffusion KV and model-owned state across
+requests. Each JSONL line describes the prompt and/or three latent-frame camera
+actions applied at the next chunk boundary:
+
+```json
+{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]}
+{"event_id":2,"frames":[["w"],["w"],["w"]]}
+{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]}
+```
+
+Run:
+
+```bash
+python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \
+  --image /path/to/first_frame.png \
+  --events /path/to/events.jsonl \
+  --output-dir /tmp/lingbot-realtime \
+  --gpu-memory-fraction 0.6
+```
+
+The example writes one latent tensor and one metadata JSON file per chunk. It
+exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
+-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
+transport.
+
+This PR intentionally does not define a public realtime HTTP/WebSocket schema.
+An online serving client should be added together with the public transport
+API rather than exposing LingBot-specific event fields from the generic model
+runtime. The structured camera-interaction frontend and transport work is
+tracked separately in
+[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
+
+## Realtime identity and controls
+
+- `session_id` identifies the persistent world and its worker-owned state.
+- `event_id` identifies a prompt/control update and remains monotonic across reset.
+- `chunk_index` is contiguous from zero and restarts from zero after reset.
+- `request_id` correlates one chunk snapshot with its output metadata.
+- AsyncOmni uses a separate UUID-suffixed internal engine routing ID.
+
+The generic runtime transports controls as opaque snapshots. LingBot's adapter
+accepts:
+
+- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`,
+  `s`, `d`, `i`, `j`, `k`, and `l`;
+- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories.
+
+## Validation
+
+Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42.
+The exercised matrix includes:
+
+- TP=1 and TP=2 execution;
+- two interleaved resident sessions;
+- action input and prompt switching;
+- seven contiguous chunks crossing the sink plus recent rolling window;
+- direct versus paged replay;
+- CUDA Graph execution;
+- reset, close, failure cleanup, and exact metadata matching; and
+- nine-frame VAE decode.
+
+Official `generate.py` accuracy and performance numbers, together with
+generation artifacts, are recorded in the PR validation section for the exact
+tested commit.
+
+## Current limitations
+
+- Only the 14B causal-fast checkpoint is supported.
+- The realtime control plane is internal; there is no public server transport yet.
+- AR-Diffusion stages currently require one replica because session-affine
+  routing across replicas is not implemented.
+- One AR block is generated per request and `max_num_seqs` must be one.
+- Stateful streaming VAE decode is not implemented; the realtime example emits
+  latent chunks.
+- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
+  Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed.
+- No AMD GPU, Ascend NPU, or Intel GPU support is claimed.
+
+## References
+
+- Checkpoint: <https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers>
+- Official implementation: <https://github.com/robbyant/lingbot-world-v2>
+- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py)
+- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md)

recipes/StabilityAI/Stable-Audio-Open.md

@@ -78,7 +78,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \
 Start the online serving endpoint:
 
 ```bash
-vllm-omni serve /path/to/stable-audio-open-1.0 \
+vllm serve /path/to/stable-audio-open-1.0 \
   --host 0.0.0.0 \
   --port 8091 \
   --gpu-memory-utilization 0.9 \

recipes/TEMPLATE.md

@@ -39,6 +39,9 @@ platform section, document one or more tested hardware configurations.
 
 #### Command
 
+Serve commands use `vllm serve <model> --omni`. The `--omni` flag is what
+selects the Omni pipeline, so it is required and cannot be dropped.
+
 ```bash
 # Add the exact command(s) here
 ```

recipes/inclusionAI/Ming-omni-tts.md

@@ -39,7 +39,7 @@ Other hardware is welcome as community validation lands.
 Launch the two-stage talker:
 
 ```bash
-vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
+vllm serve inclusionAI/Ming-omni-tts-0.5B \
     --deploy-config vllm_omni/deploy/ming_tts.yaml \
     --omni \
     --port 8091 \

recipes/zai-org/GLM-TTS.md

@@ -44,20 +44,20 @@ conditioned on reference audio and its transcript.
 Start the server from the repository root:
 
 ```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
 ```
 
 Async chunking is enabled by default in the bundled deployment config. For
 the sync (non-streaming) path:
 
 ```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
 ```
 
 Use a custom deploy config for advanced cases:
 
 ```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
   --deploy-config /path/to/your_glm_tts_overrides.yaml
 ```
 

tests/config/test_config_factory.py

@@ -1102,6 +1102,13 @@ class TestResolveScheduler:
 
 
 class TestDeployConfigLoading:
+    def test_rejects_legacy_stage_args_schema(self, tmp_path):
+        deploy_path = tmp_path / "legacy.yaml"
+        deploy_path.write_text("stage_args:\n  - stage_id: 0\n", encoding="utf-8")
+
+        with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
+            load_deploy_config(deploy_path)
+
     def test_load_minicpmo_duplex_deploy_config(self):
         deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
 
@@ -1377,7 +1384,8 @@ stages:
         assert async_stages[1].custom_process_input_func is None
 
     def test_no_bundled_legacy_stage_config_yamls(self):
-        stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
+        repo_root = Path(__file__).resolve().parents[2]
+        stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
         assert not list(stage_config_dir.glob("*.yaml"))
 
     def test_merge_pipeline_deploy(self):

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/model_loader/test_diffusers_loader.py

@@ -12,6 +12,7 @@ import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -335,7 +336,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
         backing_kind="checkpoint_mmap",
         bindings={
             "transformer.weight": TensorBinding(
-                checkpoint_key="weight",
+                storage_key="weight",
                 file_path="unused",
             )
         },
@@ -394,6 +395,192 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     assert loader.take_host_weight_plan() is None
 
 
+def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+            ulysses_degree=1,
+            ring_degree=1,
+            allgather_degree=1,
+            ulysses_mode="strict",
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_lock_timeout=1.0,
+        model="unused",
+        revision=None,
+        model_config={},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnquantizedLinearMethod()
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    def load_weights(_model):
+        events.append("load")
+        with torch.no_grad():
+            model.transformer.weight.fill_(1)
+
+    def process_weights(_model, _device):
+        events.append("process")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    def calibrate(_model):
+        events.append("calibrate")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    original_eval = model.eval
+
+    def eval_model():
+        events.append("eval")
+        return original_eval()
+
+    def build_cache(_pipeline, **kwargs):
+        events.append("cache")
+        assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
+        assert kwargs["tensor_parallel_rank"] == 0
+        assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
+        assert kwargs["quantization_config"] is None
+        return HostWeightPlanResult(runtime_plan)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._process_weights_after_loading = process_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = calibrate  # type: ignore[method-assign]
+    model.eval = eval_model  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
+    )
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "calibrate", "eval", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_pin_limit_gib=1.0,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: events.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: events.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def unexpected_checkpoint_plan(*_args, **_kwargs):
+        pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
+
+    def build_cache(_pipeline, **_kwargs):
+        events.append("cache")
+        return HostWeightPlanResult(runtime_plan)
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object()
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
+    )
+
+    def build_cache(_pipeline, **kwargs):
+        assert kwargs["quantization_config"] == "module_quant_method"
+        return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
+
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    original_weight = model.transformer.weight.detach().clone()
+    assert loader.load_model(load_device="cpu") is model
+    assert torch.equal(model.transformer.weight, original_weight)
+    assert loader.take_host_weight_plan() is None
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     from vllm_omni.diffusion.offloader.module_collector import PipelineModules

tests/diffusion/model_loader/test_runtime_weight_cache.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import fcntl
+import multiprocessing as mp
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+    build_runtime_weight_cache_plan,
+    default_runtime_cache_root,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _LoaderMarker:
+    pass
+
+
+class _Transformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(3, 2)
+        self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
+        self.register_buffer("temporary", torch.ones(1), persistent=False)
+        with torch.no_grad():
+            self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
+            self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
+
+
+class _Pipeline(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Transformer()
+
+
+def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
+    monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
+
+    assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
+
+
+def _build(pipeline: nn.Module, cache_root: Path, **overrides):
+    kwargs = {
+        "dit_modules": (("transformer", pipeline.transformer),),
+        "loader_type": _LoaderMarker,
+        "cache_root": cache_root,
+        "lock_timeout_seconds": 2.0,
+        "max_shard_bytes": 32,
+        "model_identity": "test/model",
+        "revision": "revision",
+        "runtime_dtype": torch.float32,
+        "load_format": "default",
+        "loader_inputs": {"model_config": {"hidden_size": 3}},
+        "tensor_parallel_size": 1,
+        "tensor_parallel_rank": 0,
+        "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
+        "use_hsdp": False,
+        "enable_expert_parallel": False,
+        "quantization_config": None,
+        "cfg_parallel_size": 1,
+        "pipeline_parallel_size": 1,
+    }
+    kwargs.update(overrides)
+    return build_runtime_weight_cache_plan(pipeline, **kwargs)
+
+
+def _multiprocess_builder(cache_root: str, queue) -> None:
+    result = _build(_Pipeline(), Path(cache_root))
+    queue.put(
+        (
+            result.fallback_code,
+            result.plan.runtime_layout_key if result.plan is not None else None,
+            sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
+        )
+    )
+
+
+def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    second = _build(_Pipeline(), tmp_path)
+
+    assert first.plan is not None
+    assert second.plan is not None
+    assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
+    assert first.plan.post_load_complete
+    assert {binding.file_path for binding in first.plan.bindings.values()} == {
+        binding.file_path for binding in second.plan.bindings.values()
+    }
+    assert set(first.plan.bindings) == {
+        "transformer.persistent",
+        "transformer.proj.bias",
+        "transformer.proj.weight",
+    }
+
+
+def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
+    pipeline = _Pipeline()
+    tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
+    tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
+    sp2 = _build(
+        pipeline,
+        tmp_path,
+        sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
+    )
+
+    assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
+    assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
+
+
+def test_final_runtime_content_is_authoritative(tmp_path):
+    first_pipeline = _Pipeline()
+    second_pipeline = _Pipeline()
+    with torch.no_grad():
+        second_pipeline.transformer.proj.weight[0, 0] += 1
+
+    first = _build(first_pipeline, tmp_path)
+    second = _build(second_pipeline, tmp_path)
+
+    assert first.plan is not None and second.plan is not None
+    assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
+
+
+def test_expert_parallel_ownership_fails_closed(tmp_path):
+    result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
+
+    assert result.plan is None
+    assert result.fallback_code == "unsupported_expert_parallel"
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    assert first.plan is not None
+    shard = Path(next(iter(first.plan.bindings.values())).file_path)
+    old_inode = shard.stat().st_ino
+    with shard.open("r+b") as handle:
+        handle.seek(-1, 2)
+        byte = handle.read(1)
+        handle.seek(-1, 2)
+        handle.write(bytes([byte[0] ^ 0xFF]))
+
+    second = _build(_Pipeline(), tmp_path)
+
+    assert second.plan is not None
+    assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
+    assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
+
+
+def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
+
+    with lock_path.open("a+b") as lock_handle:
+        fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
+
+    assert result.plan is None
+    assert result.fallback_code == "lock_timeout"
+
+
+def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
+    stale_temp.mkdir(parents=True)
+    (stale_temp / "partial").write_bytes(b"partial")
+
+    result = _build(_Pipeline(), tmp_path)
+
+    assert result.plan is not None
+    assert not stale_temp.exists()
+
+
+@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
+def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
+    pipeline = _Pipeline()
+    if case == "alias":
+        pipeline.transformer.tied = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    elif case == "external_alias":
+        pipeline.encoder = nn.Module()
+        pipeline.encoder.shared = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    else:
+        pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
+        expected_code = "unsupported_tensor"
+
+    result = _build(pipeline, tmp_path)
+
+    assert result.plan is None
+    assert result.fallback_code == expected_code
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
+    context = mp.get_context("spawn")
+    queue = context.Queue()
+    processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
+    for process in processes:
+        process.start()
+    results = [queue.get(timeout=30) for _ in processes]
+    for process in processes:
+        process.join(timeout=30)
+
+    assert all(process.exitcode == 0 for process in processes)
+    assert all(code is None for code, _, _ in results)
+    assert len({key for _, key, _ in results}) == 1
+    assert len({tuple(paths) for _, _, paths in results}) == 1
+    assert len(list((tmp_path / "v1").iterdir())) == 1

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 == []

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationBudgetError,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+    _coalesce_ranges,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeStorage:
+    def __init__(self, pointer: int, size: int) -> None:
+        self._pointer = pointer
+        self._size = size
+
+    def data_ptr(self) -> int:
+        return self._pointer
+
+    def nbytes(self) -> int:
+        return self._size
+
+
+class _FakeTensor:
+    device = SimpleNamespace(type="cpu")
+
+    def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
+        self._storage = _FakeStorage(pointer, size)
+        self._pinned = pinned
+
+    def untyped_storage(self) -> _FakeStorage:
+        return self._storage
+
+    def numel(self) -> int:
+        return self._storage.nbytes()
+
+    def is_pinned(self) -> bool:
+        return self._pinned
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
+    assert _coalesce_ranges(
+        [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
+        page_size=4096,
+    ) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+        registration_module._AddressRange(0x9000, 0xA000),
+    )
+
+
+def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert runtime.registered == []
+
+
+def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7], unregister_results=[9])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+
+def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {
+            "first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)],  # type: ignore[list-item]
+            "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+        },
+        max_bytes=12288,
+    )
+
+    assert registration.total_bytes == 12288
+    assert registration.region_count == 2
+    assert registration.close() == []
+    assert runtime.unregistered == [0x9000, 0x1000]
+    assert registration.close() == []
+
+
+def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    registration = CudaHostRegistration.create(
+        {
+            "first.safetensors": [_FakeTensor(0x1000, 4096)],  # type: ignore[list-item]
+            "second.safetensors": [_FakeTensor(0x2000, 4096)],  # type: ignore[list-item]
+        },
+        max_bytes=8192,
+    )
+
+    assert registration.region_count == 2
+    assert runtime.registered == [
+        (0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+        (0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+    ]
+    assert registration.close() == []
+
+
+def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+        max_bytes=4096,
+    )
+
+    assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
+    assert registration.close() == []
+    assert runtime.unregistered == [0x1000, 0x1000]

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -473,6 +473,31 @@ class TestDistributedLayerwiseOffloadHook:
         assert staging[0][torch.float32].numel() == 12
         assert staging[1][torch.float32].numel() == 12
 
+    def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
+        current_block = nn.Linear(2, 2, bias=False)
+        next_block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        next_block.weight.data.copy_(expected)
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        hook.initialize_hook(current_block)
+        hook.registered_mmap = True
+
+        def fail_staging(_slot):
+            raise AssertionError("registered mmap must bypass host staging")
+
+        hook._stage_mmap_sources = fail_staging  # type: ignore[method-assign]
+        hook.prefetch_layer(slot=0, non_blocking=True)
+
+        assert torch.equal(next_block.weight, expected)
+
 
 class TestPinnedResidentLayerGroup:
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -563,6 +588,24 @@ class TestPinnedResidentLayerGroup:
         group.offload()
         assert all(block.weight.numel() == 0 for block in blocks)
 
+    def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
+        block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        block.weight.data.copy_(expected)
+        group = PinnedResidentLayerGroup(
+            [block],
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        group.registered_mmap = True
+        group._cpu_staging_buffers.clear()
+
+        group.load()
+
+        assert torch.equal(block.weight, expected)
+
     def test_shared_allgather_output_is_narrowed_to_current_block(
         self, dist_group, patched_offload_runtime, monkeypatch
     ):
@@ -692,6 +735,95 @@ class _MmapPostLoadPipeline(nn.Module):
 
 
 class TestMmapWeightLoading:
+    def test_runtime_cache_maps_final_weights_without_rerunning_post_load(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        class Transformer(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.proj = nn.Linear(2, 2, bias=False)
+                self.post_load_calls = 0
+
+            def post_load_weights(self):
+                self.post_load_calls += 1
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        old_parameter = pipeline.transformer.proj.weight
+        old_parameter.runtime_marker = "preserved"
+        weight_file = tmp_path / "runtime.safetensors"
+        save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file))
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={
+                "transformer.proj.weight": TensorBinding(
+                    storage_key="transformer.proj.weight",
+                    file_path=str(weight_file),
+                )
+            },
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        backend._load_weights_from_runtime_cache(pipeline, modules, plan)
+
+        assert pipeline.transformer.post_load_calls == 0
+        assert pipeline.transformer.proj.weight is old_parameter
+        assert pipeline.transformer.proj.weight.runtime_marker == "preserved"
+        assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0))
+        backend._release_mmap_handles()
+
+    def test_runtime_cache_mapping_failure_leaves_ordinary_tensors_intact(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        pipeline = nn.Module()
+        pipeline.transformer = nn.Linear(2, 2, bias=False)
+        old_parameter = pipeline.transformer.weight
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={
+                "transformer.weight": TensorBinding(
+                    storage_key="transformer.weight",
+                    file_path=str(tmp_path / "missing.safetensors"),
+                )
+            },
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        with pytest.raises(Exception, match="missing.safetensors"):
+            backend._load_weights_from_runtime_cache(pipeline, modules, plan)
+
+        assert pipeline.transformer.weight is old_parameter
+
     def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
         pipeline = _MmapPostLoadPipeline()
         weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
@@ -715,7 +847,7 @@ class TestMmapWeightLoading:
             backing_kind="checkpoint_mmap",
             bindings={
                 name: TensorBinding(
-                    checkpoint_key=name,
+                    storage_key=name,
                     file_path=str(weight_file),
                 )
                 for name in weights
@@ -786,6 +918,99 @@ class TestMmapWeightLoading:
         assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
         backend.disable()
 
+    def test_runtime_cache_enable_reclaims_private_weights_before_hook_setup(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
+        weight_file = tmp_path / "runtime.safetensors"
+        save_file(weights, str(weight_file))
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        cleanup_calls: list[None] = []
+        backend._cleanup_after_loading = lambda: cleanup_calls.append(None)  # type: ignore[method-assign]
+
+        backend.enable(pipeline)
+
+        assert backend._using_rank_local_mmap
+        # One reclamation immediately after remap, plus the existing final
+        # post-hook cleanup.
+        assert len(cleanup_calls) == 2
+        backend.disable()
+
+    def test_runtime_cache_registration_uses_budget_and_releases_before_mmap(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        events: list[str] = []
+
+        class Registration:
+            total_bytes = 4096
+            region_count = 1
+
+            @staticmethod
+            def close() -> list[str]:
+                events.append("unregister")
+                return []
+
+        create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
+
+        def create(sources, *, max_bytes):
+            create_calls.append((sources, max_bytes))
+            return Registration()
+
+        monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create))
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_runtime_cache_pin_limit_gib=1.5,
+            ),
+            torch.device("cuda"),
+        )
+        sources = {"runtime.safetensors": [torch.ones(1)]}
+        backend._runtime_cache_mapped_sources = sources
+
+        assert backend._try_register_runtime_cache_mmap()
+        assert create_calls == [(sources, int(1.5 * 1024**3))]
+
+        backend.enabled = True
+        backend._using_rank_local_mmap = True
+        backend._mmap_file_cache = {}
+        backend._release_mmap_handles = lambda: events.append("mmap")  # type: ignore[method-assign]
+        backend.disable()
+
+        assert events == ["unregister", "mmap"]
+
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
@@ -1091,7 +1316,7 @@ class TestMmapValidation:
         assert result.fallback_reason is None
         assert result.plan is not None
         assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
-            checkpoint_key="blocks.0.weight",
+            storage_key="blocks.0.weight",
             file_path=str(checkpoint_file),
         )
         assert result.plan.planned_source_prefixes == frozenset({"transformer."})
@@ -1381,7 +1606,7 @@ class TestMmapValidation:
 
         assert result.plan is not None
         assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
-            checkpoint_key="proj_in.weight",
+            storage_key="proj_in.weight",
             file_path=str(checkpoint_file),
         )
 

tests/diffusion/test_result_pump.py

@@ -418,3 +418,77 @@ class TestShutdownCleansUpFutures:
 
         assert len(executor._rpc_futures) == 0
         assert len(executor._output_futures) == 0
+
+
+class _RacyFuture(concurrent.futures.Future):
+    # done() always lies and reports False, to deterministically force the
+    # pump's check-then-act race window without needing real thread timing.
+    def done(self) -> bool:
+        return False
+
+
+def _racy_cancelled_future() -> concurrent.futures.Future:
+    fut = _RacyFuture()
+    fut.cancel()
+    assert fut.cancelled()
+    return fut
+
+
+class TestResultPumpCancelledFutureRace:
+    """Regression for #5793: a future cancelled concurrently with the pump's
+    resolve call (e.g. an asyncio.wait_for timeout, or a request abort) must
+    be dropped, not raise InvalidStateError and kill the pump thread.
+    """
+
+    def test_compute_done_racing_cancel_does_not_crash_pump(self):
+        executor = _make_executor()
+        fut = _racy_cancelled_future()
+        with executor._futures_lock:
+            executor._rpc_futures["1"] = fut
+
+        msg = AsyncDiffusionOutput(kind=AsyncOutputKind.COMPUTE_DONE, rpc_id="1", async_output_id="abc")
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.cancelled()
+
+    def test_output_ready_racing_cancel_does_not_crash_pump(self):
+        executor = _make_executor()
+        fut = _racy_cancelled_future()
+        with executor._futures_lock:
+            executor._output_futures["abc123"] = fut
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="abc123",
+            output=DiffusionOutput(output="late"),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.cancelled()
+
+    def test_batch_split_racing_cancel_does_not_crash_pump(self):
+        executor = _make_executor()
+        cancelled_fut = _racy_cancelled_future()
+        healthy_fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._output_futures["batch-1/r-aborted"] = cancelled_fut
+            executor._output_futures["batch-1/r-healthy"] = healthy_fut
+            executor._batch_split_map["batch-1"] = {
+                "batch-1/r-aborted": "r-aborted",
+                "batch-1/r-healthy": "r-healthy",
+            }
+
+        outputs = {
+            "r-aborted": DiffusionOutput(output="late"),
+            "r-healthy": DiffusionOutput(output="ok"),
+        }
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id="batch-1",
+            output=_FakeBatchOutput(outputs),
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert cancelled_fut.cancelled()
+        assert healthy_fut.done()
+        assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"]

tests/diffusion/test_worker_wrapper_base.py

@@ -12,11 +12,13 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
+import vllm_omni.diffusion.worker.diffusion_worker as worker_module
 from vllm_omni.diffusion.worker.diffusion_worker import (
     CustomPipelineWorkerExtension,
     DiffusionWorker,
@@ -204,6 +206,29 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(
+        self,
+        mocker: MockerFixture,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        worker = object.__new__(DiffusionWorker)
+        offload_backend = mocker.Mock()
+        kv_transfer_manager = mocker.Mock()
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_transfer_manager,
+        )
+        events: list[str] = []
+        offload_backend.disable.side_effect = lambda: events.append("offloader")
+        kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
+        monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
+
+        worker.shutdown()
+
+        offload_backend.disable.assert_called_once_with()
+        kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
+        assert events == ["offloader", "kv_transfer", "distributed"]
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method

tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py

@@ -3,19 +3,18 @@
 """Smoke test for HunyuanImage-3.0 Image-to-Text (I2T) pipeline."""
 
 from collections.abc import Generator
-from pathlib import Path
 
 import pytest
 import torch
 from PIL import Image
 
 from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni import Omni
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt
 
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-REPO_ROOT = Path(__file__).resolve().parents[3]
-STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_i2t.yaml"
+DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3_ar.yaml")
 
 # First 20 generated token IDs from the HF greedy reference on this input.
 # vllm-omni AR output matches this prefix bitwise; the two implementations
@@ -45,14 +44,14 @@ EXPECTED_PREFIX_TOKEN_IDS: list[int] = [
 # Decoded form, kept only for human-readable assertion messages.
 EXPECTED_PREFIX_TEXT = "The image is a solid, uniform green color with no variations, objects, or details present. It"
 
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cuda]
 
 
 @pytest.fixture(scope="module")
 def omni() -> Generator[Omni, None, None]:
     with OmniRunner(
         MODEL_NAME,
-        deploy_config=str(STAGE_CONFIG_PATH),
+        deploy_config=DEPLOY_CONFIG_PATH,
         trust_remote_code=True,
     ) as runner:
         yield runner.omni

tests/e2e/offline_inference/test_lingbot_world_v2.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end smoke test for LingBot-World v2 generation."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunnerHandler
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = os.environ.get(
+    "VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH",
+    "robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
+)
+_IMAGE_PATH = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH")
+_ACTION_DIR_VALUE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR")
+_ACTION_DIR = Path(_ACTION_DIR_VALUE).expanduser().resolve() if _ACTION_DIR_VALUE else None
+_HAS_INPUT_ASSETS = _IMAGE_PATH is not None and _ACTION_DIR is not None
+
+_OMNI_RUNNER_PARAM = (
+    MODEL,
+    None,
+    {
+        "flow_shift": 5.0,
+        "parallel_config": DiffusionParallelConfig(tensor_parallel_size=1),
+        "model_config": {
+            "lingbot_action_root": str(_ACTION_DIR.parent) if _ACTION_DIR is not None else "",
+        },
+    },
+)
+
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.diffusion,
+    pytest.mark.skipif(
+        not _HAS_INPUT_ASSETS,
+        reason=("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH and VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR are required"),
+    ),
+    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
+]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Load the real checkpoint and generate one causal block on one GPU."""
+
+    assert _IMAGE_PATH is not None
+    assert _ACTION_DIR is not None
+    response = omni_runner_handler.send_diffusion_request(
+        {
+            "model": omni_runner_handler.runner.model_name,
+            "prompt": "The camera moves slowly forward through the scene.",
+            "multi_modal_data": {"image": str(Path(_IMAGE_PATH).expanduser().resolve())},
+            "sampling_params": OmniDiffusionSamplingParams(
+                height=480,
+                width=832,
+                num_frames=9,
+                num_inference_steps=4,
+                max_sequence_length=512,
+                seed=42,
+                fps=16,
+                extra_args={
+                    "action_path": _ACTION_DIR.name,
+                    "flow_shift": 5.0,
+                },
+            ),
+        }
+    )
+
+    assert response.success
+    assert response.images is not None
+    assert len(response.images) == 1
+    video = np.asarray(response.images[0])
+    assert video.shape == (1, 9, 480, 832, 3)
+    assert np.isfinite(video).all()
+    assert float(np.ptp(video)) > 1e-6
+    assert float(np.mean(np.abs(np.diff(video, axis=1)))) > 1e-6

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -328,6 +328,13 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--omni",
             "--enable-distributed-layerwise-offload",
             "--dlo-no-use-allgather",
+            "--dlo-enable-runtime-cache",
+            "--dlo-runtime-cache-dir",
+            "/var/cache/vllm-omni-dlo",
+            "--dlo-runtime-cache-lock-timeout",
+            "30",
+            "--dlo-runtime-cache-pin-limit-gib",
+            "80",
             "--dlo-resident-layers",
             "20",
         ]
@@ -339,9 +346,17 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
+    assert args.dlo_enable_runtime_cache is True
+    assert args.dlo_runtime_cache_dir == "/var/cache/vllm-omni-dlo"
+    assert args.dlo_runtime_cache_lock_timeout == 30
+    assert args.dlo_runtime_cache_pin_limit_gib == 80
     assert args.dlo_resident_layers == 20
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
+    assert engine_args["dlo_enable_runtime_cache"] is True
+    assert engine_args["dlo_runtime_cache_dir"] == "/var/cache/vllm-omni-dlo"
+    assert engine_args["dlo_runtime_cache_lock_timeout"] == 30
+    assert engine_args["dlo_runtime_cache_pin_limit_gib"] == 80
     assert engine_args["dlo_resident_layers"] == 20
 
 
@@ -492,6 +507,24 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
     assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
 
 
+@pytest.mark.parametrize(
+    ("legacy_arg", "value"),
+    [
+        ("stage_configs_path", "legacy.yaml"),
+        ("stage_configs", [{"stage_id": 0}]),
+    ],
+)
+def test_resolve_stage_configs_rejects_legacy_config_arguments(legacy_arg, value):
+    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+
+    with pytest.raises(ValueError, match=rf"`{legacy_arg}`.*`deploy_config`"):
+        engine._resolve_stage_configs(
+            "dummy-model",
+            {legacy_arg: value},
+            trust_remote_code=False,
+        )
+
+
 def test_default_stage_config_includes_quantization_config():
     """Ensure structured quantization_config survives default diffusion-stage creation."""
     quantization_config = {

tests/entrypoints/test_utils.py

@@ -27,7 +27,7 @@ from vllm_omni.entrypoints.utils import (
     filter_dataclass_kwargs,
     filter_stages,
     load_and_resolve_stage_configs,
-    load_stage_configs_from_yaml,
+    load_stage_configs_from_model,
     resolve_model_config_path,
 )
 
@@ -417,8 +417,9 @@ class TestTryResolveOmniModelType:
 
 class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
-        """Ensure that dtype survives default stage creation."""
-        kwargs = {"dtype": torch.float32}
+        """Ensure that kwargs survive default stage creation."""
+        engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+        kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
         mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
         mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
 
@@ -431,6 +432,7 @@ class TestLoadAndResolveStageConfigs:
         assert config_path is None
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
+        assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
 
     def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
         deploy_path = tmp_path / "qwen3_multi.yaml"
@@ -518,71 +520,20 @@ class TestLoadAndResolveStageConfigs:
         assert filtered[0].final_output_type is None
 
 
-class TestLoadStageConfigsFromYaml:
-    """Regression tests for stage-config loading and merging."""
-
-    def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
-        yaml_config = create_config(
-            {
-                "async_chunk": True,
-                "stage_args": [
-                    {
-                        "stage_id": 0,
-                        "runtime": {"device": 0},
-                        "engine_args": {
-                            "parallel_config": {"tensor_parallel_size": 4},
-                        },
-                    }
-                ],
-            }
-        )
+class TestLoadStageConfigsFromModel:
+    def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog):
         mocker.patch(
-            "vllm_omni.entrypoints.utils.load_yaml_config",
-            return_value=yaml_config,
+            "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model",
+            return_value=(None, None),
         )
+        resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path")
 
-        stages = load_stage_configs_from_yaml(
-            "fake.yaml",
-            base_engine_args={
-                "parallel_config": {
-                    "tensor_parallel_size": 1,
-                    "pipeline_parallel_size": 2,
-                },
-                "model": "base-model",
-            },
-        )
+        result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False)
 
-        merged_engine_args = stages[0]["engine_args"]
-        assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
-        assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
-        assert merged_engine_args["model"] == "base-model"
-        assert merged_engine_args["async_chunk"] is True
-
-    def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
-        yaml_config = create_config(
-            {
-                "stage_args": [
-                    {
-                        "stage_id": 0,
-                        "engine_args": {
-                            "nested": {"override": 2},
-                        },
-                    }
-                ],
-            }
-        )
-        mocker.patch(
-            "vllm_omni.entrypoints.utils.load_yaml_config",
-            return_value=yaml_config,
-        )
-
-        stages = load_stage_configs_from_yaml(
-            "fake.yaml",
-            base_engine_args={"nested": {"base": 1}},
-        )
-
-        assert stages[0]["engine_args"]["nested"]["base"] == 1
-        assert stages[0]["engine_args"]["nested"]["override"] == 2
+        assert result == ([], None)
+        resolve_path.assert_not_called()
+        assert "No registered PipelineConfig resolved" in caplog.text
+        assert "deploy_config" in caplog.text
 
 
 class TestCumulativeStreamingCoercion:

tests/model_tests/diffusion/test_alignment.py

@@ -64,6 +64,9 @@ EXCLUDED_MODELS = [
     "HunyuanVideo15Pipeline",
     "HunyuanVideo15ImageToVideoPipeline",
     "LingBotVideoPipeline",
+    # Requires camera/action assets and persistent AR state; covered by its
+    # dedicated contract tests and real-checkpoint E2E smoke instead.
+    "LingBotWorldCausalDMDPipeline",
     "MagiHumanPipeline",
     "OmniVoicePipeline",
     "OmniVoice",

vllm_omni/config/omni_config.py

@@ -576,6 +576,10 @@ class _DiffusionConfigProjection:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_enable_runtime_cache: bool = False
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float = Field(default=600.0, gt=0)
+    dlo_runtime_cache_pin_limit_gib: float = Field(default=0.0, ge=0)
     dlo_resident_layers: int = Field(default=0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -638,6 +642,13 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
         from vllm_omni.quantization import build_quant_config
 
+        if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
+            raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
+        if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
+            raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
+
         if self.tf_model_config is None:
             self.tf_model_config = TransformerConfig()
         elif isinstance(self.tf_model_config, Mapping):

vllm_omni/config/stage_config.py

@@ -382,6 +382,10 @@ class StageDeployConfig:
 
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
+    dlo_enable_runtime_cache: bool | None = None
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float | None = None
+    dlo_runtime_cache_pin_limit_gib: float | None = None
     dlo_resident_layers: int | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
@@ -632,6 +636,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]:
 def load_deploy_config(path: str | Path) -> DeployConfig:
     """Load a deploy YAML (with optional base_config inheritance)."""
     raw_dict = resolve_deploy_yaml(path)
+    if "stage_args" in raw_dict:
+        raise ValueError(
+            f"Deploy config {path} uses the removed `stage_args` schema; "
+            "define topology in PipelineConfig and deployment overrides under `stages`."
+        )
 
     stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])]
 

vllm_omni/diffusion/data.py

@@ -708,6 +708,18 @@ class OmniDiffusionConfig:
     # This avoids AllGather synchronization, while host memory follows the
     # loader's existing rank-local layout instead of adding a second DP shard.
     dlo_use_allgather: bool = True
+    # Build or join a node-local mmap cache after ordinary loading. This is
+    # used when direct checkpoint mmap is unavailable, or when direct CUDA
+    # registration requires a transform-complete final layout. The first
+    # version is no-AllGather only.
+    dlo_enable_runtime_cache: bool = False
+    # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-runtime-weights.
+    dlo_runtime_cache_dir: str | None = None
+    # Maximum wait for the per-layout POSIX writer lock.
+    dlo_runtime_cache_lock_timeout: float = 600.0
+    # Maximum file-backed memory a worker may register with CUDA. Zero keeps
+    # the bounded mmap-to-pinned staging path.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
 
@@ -928,6 +940,16 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if not math.isfinite(self.dlo_runtime_cache_lock_timeout) or self.dlo_runtime_cache_lock_timeout <= 0:
+            raise ValueError("dlo_runtime_cache_lock_timeout must be positive")
+        if not math.isfinite(self.dlo_runtime_cache_pin_limit_gib) or self.dlo_runtime_cache_pin_limit_gib < 0:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib must be non-negative")
+        if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
+            raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
+        if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
+            raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, Asy
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
 from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
+from vllm_omni.diffusion.utils.future_utils import try_set_exception, try_set_result
 from vllm_omni.diffusion.worker import WorkerProc
 
 if TYPE_CHECKING:
@@ -870,9 +871,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     fut = self._rpc_futures.pop(msg.rpc_id, None) if msg.rpc_id else None
                 if fut is not None and not fut.done():
                     if msg.error:
-                        fut.set_exception(RuntimeError(msg.error))
+                        try_set_exception(fut, RuntimeError(msg.error))
                     else:
-                        fut.set_result(msg)
+                        try_set_result(fut, msg)
             elif msg.kind == AsyncOutputKind.OUTPUT_READY:
                 batch_id = msg.async_output_id
                 with self._futures_lock:
@@ -903,9 +904,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                             pending = self._output_futures.pop(batch_id, None)
                             if pending is not None and not pending.done():
                                 if exc is not None:
-                                    pending.set_exception(exc)
+                                    try_set_exception(pending, exc)
                                 else:
-                                    pending.set_result(output_result)
+                                    try_set_result(pending, output_result)
                             else:
                                 fut = concurrent.futures.Future()
                                 if exc is not None:
@@ -933,7 +934,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             with self._futures_lock:
                 pending = self._output_futures.pop(per_req_id, None)
                 if pending is not None and not pending.done():
-                    pending.set_result(per_req_result)
+                    try_set_result(pending, per_req_result)
                 else:
                     fut: concurrent.futures.Future = concurrent.futures.Future()
                     fut.set_result(per_req_result)
@@ -999,10 +1000,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             with self._futures_lock:
                 for fut in self._rpc_futures.values():
                     if not fut.done():
-                        fut.set_exception(RuntimeError("Executor shut down"))
+                        try_set_exception(fut, RuntimeError("Executor shut down"))
                 for fut in self._output_futures.values():
                     if not fut.done():
-                        fut.set_exception(RuntimeError("Executor shut down"))
+                        try_set_exception(fut, RuntimeError("Executor shut down"))
                 self._rpc_futures.clear()
                 self._output_futures.clear()
                 self._batch_split_map.clear()

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -14,6 +14,7 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
@@ -440,10 +441,16 @@ class DiffusersPipelineLoader:
                 _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1))
                 _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1))
                 _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size
+                _prefer_registered_runtime_cache = (
+                    _dist_offload
+                    and not _use_ag
+                    and bool(getattr(self.od_config, "dlo_enable_runtime_cache", False))
+                    and float(getattr(self.od_config, "dlo_runtime_cache_pin_limit_gib", 0.0)) > 0
+                )
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
-                if _dist_offload:
+                if _dist_offload and not _prefer_registered_runtime_cache:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -455,6 +462,16 @@ class DiffusersPipelineLoader:
                         online_quantization=_has_online_quant,
                     )
                     self.host_weight_plan = plan_result.plan
+                elif _prefer_registered_runtime_cache:
+                    # Direct checkpoint bindings can carry deferred transforms
+                    # (for example grouped-QKV reordering), so registering the
+                    # raw checkpoint mapping would not make every recurrent H2D
+                    # source directly copyable. Materialize once, apply every
+                    # post-load mutation, then publish/register the final layout.
+                    logger.info(
+                        "DLO runtime-cache registration requested; using the ordinary loader "
+                        "once to publish a transform-complete final mmap layout"
+                    )
 
                 _skip_load = self.host_weight_plan is not None
 
@@ -491,7 +508,8 @@ class DiffusersPipelineLoader:
                         )
                     if _dist_offload and plan_result is not None:
                         logger.info(
-                            "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
+                            "DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s",
+                            plan_result.fallback_code or "unknown",
                             plan_result.fallback_reason,
                         )
                     logger.debug("Loading weights on %s ...", load_device)
@@ -516,7 +534,107 @@ class DiffusersPipelineLoader:
                 logger.info("Quantization complete, offloaded model back to CPU")
 
         self._apply_skip_softmax_calibration(model)
-        return model.eval()
+        model.eval()
+        self._maybe_build_runtime_weight_cache(model, load_format)
+        return model
+
+    def _maybe_build_runtime_weight_cache(self, model: nn.Module, load_format: str) -> None:
+        """Publish final ordinary-loader DiT tensors for no-AllGather DLO."""
+        if self.host_weight_plan is not None:
+            return
+        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
+            return
+        if not getattr(self.od_config, "dlo_enable_runtime_cache", False):
+            return
+        if getattr(self.od_config, "dlo_use_allgather", True):
+            logger.info("DLO runtime cache is no-AllGather only; leaving AllGather storage unchanged")
+            return
+
+        parallel_config = self.parallel_config
+        tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1))
+        tp_rank = 0
+        if tp_size > 1:
+            try:
+                from vllm.distributed import get_tensor_model_parallel_rank
+
+                tp_rank = int(get_tensor_model_parallel_rank())
+            except Exception as exc:
+                logger.warning(
+                    "DLO runtime cache unavailable [parallel_rank_unavailable]: cannot resolve TP rank for TP=%d: %s",
+                    tp_size,
+                    exc,
+                )
+                return
+
+        modules = ModuleDiscovery.discover(model)
+        from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+            DEFAULT_LOCK_TIMEOUT_SECONDS,
+            DEFAULT_SHARD_SIZE_BYTES,
+            build_runtime_weight_cache_plan,
+        )
+
+        loader_inputs = {
+            "model_class_name": getattr(self.od_config, "model_class_name", None),
+            "model_config": getattr(self.od_config, "model_config", None),
+            "diffusers_load_kwargs": getattr(self.od_config, "diffusers_load_kwargs", None),
+            "custom_pipeline_args": getattr(self.od_config, "custom_pipeline_args", None),
+            "lora_path": getattr(self.od_config, "lora_path", None),
+            "lora_scale": getattr(self.od_config, "lora_scale", None),
+            "lora_backend": getattr(self.od_config, "lora_backend", None),
+        }
+        sp_guard = {
+            "sequence_parallel_size": int(getattr(parallel_config, "sequence_parallel_size", 1)),
+            "ulysses_degree": int(getattr(parallel_config, "ulysses_degree", 1)),
+            "ring_degree": int(getattr(parallel_config, "ring_degree", 1)),
+            "allgather_degree": int(getattr(parallel_config, "allgather_degree", 1)),
+            "ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"),
+        }
+        model_identity = getattr(self.od_config, "model", None)
+        runtime_quantization = self.quant_config
+        if runtime_quantization is None and any(
+            (quant_method := getattr(module, "quant_method", None)) is not None
+            and not isinstance(quant_method, UnquantizedLinearMethod)
+            for module in model.modules()
+        ):
+            runtime_quantization = "module_quant_method"
+        result = build_runtime_weight_cache_plan(
+            model,
+            dit_modules=tuple(zip(modules.dit_names, modules.dits)),
+            loader_type=type(self),
+            cache_root=getattr(self.od_config, "dlo_runtime_cache_dir", None),
+            lock_timeout_seconds=float(
+                getattr(self.od_config, "dlo_runtime_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS)
+            ),
+            max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES,
+            model_identity=str(model_identity) if model_identity is not None else None,
+            revision=getattr(self.od_config, "revision", None),
+            runtime_dtype=getattr(self.od_config, "dtype", None),
+            load_format=load_format,
+            loader_inputs=loader_inputs,
+            tensor_parallel_size=tp_size,
+            tensor_parallel_rank=tp_rank,
+            sequence_parallel_guard=sp_guard,
+            use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)),
+            enable_expert_parallel=bool(getattr(parallel_config, "enable_expert_parallel", False)),
+            quantization_config=runtime_quantization,
+            cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)),
+            pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)),
+        )
+        if result.plan is not None:
+            self.host_weight_plan = result.plan
+            logger.info(
+                "DLO runtime cache selected (layout=%s, TP=%d/%d, SP=%d)",
+                result.plan.runtime_layout_key[:12] if result.plan.runtime_layout_key else "unknown",
+                tp_rank,
+                tp_size,
+                sp_guard["sequence_parallel_size"],
+            )
+        else:
+            logger.warning(
+                "DLO runtime cache unavailable [%s]; retaining ordinary host tensors: %s",
+                result.fallback_code or "unknown",
+                result.fallback_reason or "unspecified reason",
+            )
 
     @staticmethod
     def _request_offload_after_quant(model: nn.Module) -> int:

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -28,7 +28,7 @@ TensorTransform = Callable[[torch.Tensor], torch.Tensor]
 class TensorBinding:
     """One runtime tensor backed by one safetensors entry."""
 
-    checkpoint_key: str
+    storage_key: str
     file_path: str
     transform: TensorTransform | None = None
 
@@ -40,6 +40,8 @@ class HostWeightPlan:
     backing_kind: str
     bindings: dict[str, TensorBinding]
     planned_source_prefixes: frozenset[str] = frozenset()
+    runtime_layout_key: str | None = None
+    post_load_complete: bool = False
 
 
 @dataclass(frozen=True)
@@ -48,6 +50,7 @@ class HostWeightPlanResult:
 
     plan: HostWeightPlan | None
     fallback_reason: str | None = None
+    fallback_code: str | None = None
 
 
 class _PlanIncompatibleError(RuntimeError):
@@ -258,11 +261,11 @@ def _validate_source_metadata(
         with safe_open(file_path, framework="pt", device="cpu") as handle:
             available = set(handle.keys())
             for runtime_name, target, binding in entries:
-                if binding.checkpoint_key not in available:
+                if binding.storage_key not in available:
                     raise _PlanIncompatibleError(
-                        f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
+                        f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}"
                     )
-                source = handle.get_slice(binding.checkpoint_key)
+                source = handle.get_slice(binding.storage_key)
                 source_shape = tuple(source.get_shape())
                 source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
                 if source_shape != tuple(target.shape):
@@ -287,11 +290,19 @@ def build_checkpoint_mmap_plan(
 ) -> HostWeightPlanResult:
     """Build a complete direct-checkpoint plan or return a fallback reason."""
     if tensor_parallel_size != 1:
-        return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
+        return HostWeightPlanResult(
+            None,
+            f"TP={tensor_parallel_size} requires the ordinary loader",
+            "unsupported_tp",
+        )
     if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp")
     if online_quantization:
-        return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
+        return HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+            "unsupported_online_quantization",
+        )
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
     if not callable(remap_fn):
@@ -328,14 +339,14 @@ def build_checkpoint_mmap_plan(
                 )
             checkpoint_key, file_path = model_to_ckpt[runtime_name]
             bindings[runtime_name] = TensorBinding(
-                checkpoint_key=checkpoint_key,
+                storage_key=checkpoint_key,
                 file_path=file_path,
                 transform=policy.transform if policy is not None else None,
             )
 
         _validate_source_metadata(required, bindings)
     except (OSError, ValueError, _PlanIncompatibleError) as exc:
-        return HostWeightPlanResult(None, str(exc))
+        return HostWeightPlanResult(None, str(exc), "checkpoint_plan_incompatible")
 
     return HostWeightPlanResult(
         HostWeightPlan(

vllm_omni/diffusion/model_loader/runtime_weight_cache.py

@@ -0,0 +1,700 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Node-local, content-addressed runtime-weight cache for diffusion DLO."""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import errno
+import fcntl
+import hashlib
+import inspect
+import json
+import math
+import os
+import shutil
+import time
+import uuid
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    HostWeightPlan,
+    HostWeightPlanResult,
+    TensorBinding,
+)
+
+logger = init_logger(__name__)
+
+RUNTIME_CACHE_SCHEMA_VERSION = 1
+DEFAULT_RUNTIME_CACHE_ROOT = "~/.cache/vllm-omni/dlo-runtime-weights"
+DEFAULT_LOCK_TIMEOUT_SECONDS = 600.0
+DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3
+_HASH_CHUNK_BYTES = 64 * 1024**2
+_FILE_HASH_CHUNK_BYTES = 8 * 1024**2
+
+_SUPPORTED_DTYPES = {
+    torch.bool,
+    torch.uint8,
+    torch.int8,
+    torch.int16,
+    torch.int32,
+    torch.int64,
+    torch.float16,
+    torch.bfloat16,
+    torch.float32,
+    torch.float64,
+}
+for _dtype_name in ("float8_e4m3fn", "float8_e5m2"):
+    if (_dtype := getattr(torch, _dtype_name, None)) is not None:
+        _SUPPORTED_DTYPES.add(_dtype)
+
+
+@dataclass(frozen=True)
+class _RuntimeTensor:
+    name: str
+    tensor: torch.Tensor
+    kind: str
+
+
+class _RuntimeCacheError(RuntimeError):
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+
+
+def default_runtime_cache_root() -> str:
+    """Return the default root, intentionally independent of VLLM_CACHE_ROOT."""
+    return os.path.expanduser(DEFAULT_RUNTIME_CACHE_ROOT)
+
+
+def _persistent_buffer(module: nn.Module, local_name: str) -> bool:
+    parent_path, _, leaf_name = local_name.rpartition(".")
+    owner = module.get_submodule(parent_path)
+    return leaf_name not in owner._non_persistent_buffers_set
+
+
+def _named_parameters_with_duplicates(module: nn.Module):
+    try:
+        return module.named_parameters(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_parameters()
+
+
+def _named_buffers_with_duplicates(module: nn.Module):
+    try:
+        return module.named_buffers(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_buffers()
+
+
+def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None:
+    parent_path, _, leaf_name = runtime_name.rpartition(".")
+    parent = pipeline.get_submodule(parent_path)
+    tensor = parent._parameters.get(leaf_name)
+    if tensor is None:
+        tensor = parent._buffers.get(leaf_name)
+    return tensor
+
+
+def _collect_runtime_tensors(
+    pipeline: nn.Module,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+) -> list[_RuntimeTensor]:
+    records: dict[str, _RuntimeTensor] = {}
+    for dit_name, dit_module in dit_modules:
+        for local_name, tensor in _named_parameters_with_duplicates(dit_module):
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = _RuntimeTensor(runtime_name, tensor, "parameter")
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise _RuntimeCacheError(
+                    "ambiguous_ownership",
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+        for local_name, tensor in _named_buffers_with_duplicates(dit_module):
+            if not _persistent_buffer(dit_module, local_name):
+                continue
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = _RuntimeTensor(runtime_name, tensor, "buffer")
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise _RuntimeCacheError(
+                    "ambiguous_ownership",
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+
+    if not records:
+        raise _RuntimeCacheError(
+            "ambiguous_ownership",
+            "no DiT parameters or persistent buffers were discovered",
+        )
+
+    storage_owners: dict[tuple[int, int], str] = {}
+    for record in records.values():
+        tensor = record.tensor
+        try:
+            pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name)
+        except AttributeError as exc:
+            raise _RuntimeCacheError(
+                "ambiguous_ownership",
+                f"DiT tensor {record.name!r} is not owned by the pipeline",
+            ) from exc
+        if pipeline_tensor is not tensor:
+            raise _RuntimeCacheError(
+                "ambiguous_ownership",
+                f"DiT tensor {record.name!r} does not resolve to the discovered object",
+            )
+        if tensor.device.type != "cpu" or tensor.is_meta:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}",
+            )
+        if tensor.layout != torch.strided:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} uses unsupported layout {tensor.layout}",
+            )
+        if not tensor.is_contiguous():
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is non-contiguous with stride {tensor.stride()}",
+            )
+        if tensor.dtype not in _SUPPORTED_DTYPES:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} uses unsupported dtype {tensor.dtype}",
+            )
+        if hasattr(tensor, "to_local"):
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is a distributed tensor",
+            )
+
+        tensor_nbytes = tensor.numel() * tensor.element_size()
+        storage = tensor.untyped_storage()
+        if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is a view into a larger storage",
+            )
+        if tensor_nbytes:
+            storage_id = (storage.data_ptr(), storage.nbytes())
+            owner = storage_owners.setdefault(storage_id, record.name)
+            if owner != record.name:
+                raise _RuntimeCacheError(
+                    "unsupported_alias",
+                    f"{record.name!r} shares storage with {owner!r}",
+                )
+
+    # The ownership boundary is the complete pipeline, not only the discovered
+    # DiT traversal. Reject a DiT storage that is also registered by an encoder,
+    # VAE, resident component, or non-persistent buffer.
+    for pipeline_name, tensor in chain(
+        _named_parameters_with_duplicates(pipeline),
+        _named_buffers_with_duplicates(pipeline),
+    ):
+        if tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0:
+            continue
+        storage = tensor.untyped_storage()
+        owner = storage_owners.get((storage.data_ptr(), storage.nbytes()))
+        if owner is not None and owner != pipeline_name:
+            raise _RuntimeCacheError(
+                "unsupported_alias",
+                f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}",
+            )
+
+    return [records[name] for name in sorted(records)]
+
+
+def _json_normalize(value: Any) -> Any:
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, torch.dtype):
+        return str(value)
+    if isinstance(value, os.PathLike):
+        return os.fspath(value)
+    if dataclasses.is_dataclass(value):
+        return _json_normalize(dataclasses.asdict(value))
+    model_dump = getattr(value, "model_dump", None)
+    if callable(model_dump):
+        try:
+            dumped = model_dump(mode="json")
+        except (TypeError, ValueError):
+            dumped = model_dump()
+        return _json_normalize(dumped)
+    if isinstance(value, Mapping):
+        return {str(key): _json_normalize(item) for key, item in sorted(value.items(), key=lambda item: str(item[0]))}
+    if isinstance(value, (set, frozenset)):
+        return sorted((_json_normalize(item) for item in value), key=lambda item: _canonical_json(item))
+    if isinstance(value, (list, tuple)):
+        return [_json_normalize(item) for item in value]
+    if inspect.isclass(value):
+        return f"{value.__module__}.{value.__qualname__}"
+    return {
+        "type": f"{type(value).__module__}.{type(value).__qualname__}",
+        "value": str(value),
+    }
+
+
+def _canonical_json(value: Any) -> bytes:
+    return json.dumps(
+        _json_normalize(value),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=True,
+    ).encode()
+
+
+def _implementation_fingerprint(
+    loader_type: type, pipeline: nn.Module, dit_modules: Sequence[tuple[str, nn.Module]]
+) -> str:
+    objects: list[Any] = [loader_type, type(pipeline)]
+    load_weights = getattr(type(pipeline), "load_weights", None)
+    if load_weights is not None:
+        objects.append(load_weights)
+    for _, dit_module in dit_modules:
+        objects.append(type(dit_module))
+        post_load = getattr(type(dit_module), "post_load_weights", None)
+        if post_load is not None:
+            objects.append(post_load)
+
+    digest = hashlib.sha256()
+    identities: set[str] = set()
+    for obj in objects:
+        identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}"
+        if identity in identities:
+            continue
+        identities.add(identity)
+        digest.update(identity.encode())
+        try:
+            digest.update(inspect.getsource(obj).encode())
+        except (OSError, TypeError):
+            pass
+    return digest.hexdigest()
+
+
+def _tensor_metadata(record: _RuntimeTensor) -> dict[str, Any]:
+    tensor = record.tensor
+    return {
+        "kind": record.kind,
+        "shape": list(tensor.shape),
+        "dtype": str(tensor.dtype),
+        "nbytes": tensor.numel() * tensor.element_size(),
+        "layout": "contiguous",
+    }
+
+
+def _update_hash_with_tensor(digest: Any, tensor: torch.Tensor) -> None:
+    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
+    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
+        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
+        digest.update(memoryview(chunk.numpy()))
+
+
+def _runtime_content_digest(records: Sequence[_RuntimeTensor]) -> str:
+    digest = hashlib.sha256()
+    for record in records:
+        digest.update(_canonical_json({"name": record.name, **_tensor_metadata(record)}))
+        _update_hash_with_tensor(digest, record.tensor)
+    return digest.hexdigest()
+
+
+def _file_digest(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(_FILE_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+@contextlib.contextmanager
+def _exclusive_lock(path: Path, timeout_seconds: float):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    deadline = time.monotonic() + timeout_seconds
+    with path.open("a+b") as handle:
+        while True:
+            try:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+                break
+            except OSError as exc:
+                if exc.errno not in (errno.EACCES, errno.EAGAIN):
+                    raise _RuntimeCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise _RuntimeCacheError(
+                        "lock_timeout",
+                        f"timed out after {timeout_seconds:g}s waiting for runtime-cache writer {path.name}",
+                    ) from exc
+                time.sleep(min(0.2, remaining))
+        try:
+            yield
+        finally:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+
+
+def _fsync_file(path: Path) -> None:
+    with path.open("rb") as handle:
+        os.fsync(handle.fileno())
+
+
+def _fsync_directory(path: Path) -> None:
+    fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def _split_shards(records: Sequence[_RuntimeTensor], max_shard_bytes: int) -> list[list[_RuntimeTensor]]:
+    shards: list[list[_RuntimeTensor]] = []
+    current: list[_RuntimeTensor] = []
+    current_bytes = 0
+    for record in records:
+        nbytes = record.tensor.numel() * record.tensor.element_size()
+        if current and current_bytes + nbytes > max_shard_bytes:
+            shards.append(current)
+            current = []
+            current_bytes = 0
+        current.append(record)
+        current_bytes += nbytes
+    if current:
+        shards.append(current)
+    return shards
+
+
+def _publish_entry(
+    *,
+    cache_root: Path,
+    entry_dir: Path,
+    cache_key: str,
+    layout_identity: dict[str, Any],
+    content_digest: str,
+    records: Sequence[_RuntimeTensor],
+    max_shard_bytes: int,
+) -> None:
+    total_nbytes = sum(record.tensor.numel() * record.tensor.element_size() for record in records)
+    free_bytes = shutil.disk_usage(cache_root).free
+    required_bytes = total_nbytes + max(64 * 1024**2, total_nbytes // 100)
+    if free_bytes < required_bytes:
+        raise _RuntimeCacheError(
+            "insufficient_disk",
+            f"runtime cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}",
+        )
+
+    tmp_root = cache_root / ".tmp"
+    tmp_root.mkdir(parents=True, exist_ok=True)
+    temp_dir = tmp_root / f"{cache_key}.{os.getpid()}.{uuid.uuid4().hex}"
+    temp_dir.mkdir()
+    try:
+        shards = _split_shards(records, max_shard_bytes)
+        tensor_manifest: dict[str, dict[str, Any]] = {}
+        file_manifest: dict[str, dict[str, Any]] = {}
+        shard_count = len(shards)
+        for index, shard in enumerate(shards, start=1):
+            filename = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
+            path = temp_dir / filename
+            save_file(
+                {record.name: record.tensor.detach() for record in shard},
+                str(path),
+                metadata={"format": "pt", "dlo_runtime_cache_schema": str(RUNTIME_CACHE_SCHEMA_VERSION)},
+            )
+            _fsync_file(path)
+            file_manifest[filename] = {
+                "size": path.stat().st_size,
+                "sha256": _file_digest(path),
+            }
+            for record in shard:
+                tensor_manifest[record.name] = {
+                    **_tensor_metadata(record),
+                    "file": filename,
+                    "storage_key": record.name,
+                }
+
+        manifest = {
+            "schema_version": RUNTIME_CACHE_SCHEMA_VERSION,
+            "cache_key": cache_key,
+            "layout_identity": layout_identity,
+            "runtime_content_sha256": content_digest,
+            "total_tensor_bytes": total_nbytes,
+            "created_at_unix_seconds": time.time(),
+            "writer_pid": os.getpid(),
+            "files": file_manifest,
+            "tensors": tensor_manifest,
+        }
+        manifest_path = temp_dir / "manifest.json"
+        with manifest_path.open("w", encoding="utf-8") as handle:
+            json.dump(manifest, handle, sort_keys=True, separators=(",", ":"))
+            handle.flush()
+            os.fsync(handle.fileno())
+        _fsync_directory(temp_dir)
+
+        invalid_dir: Path | None = None
+        if entry_dir.exists():
+            invalid_dir = tmp_root / f"{cache_key}.invalid.{os.getpid()}.{uuid.uuid4().hex}"
+            os.replace(entry_dir, invalid_dir)
+        os.replace(temp_dir, entry_dir)
+        _fsync_directory(entry_dir.parent)
+        if invalid_dir is not None:
+            shutil.rmtree(invalid_dir, ignore_errors=True)
+    except Exception:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        raise
+
+
+def _validate_entry(
+    *,
+    entry_dir: Path,
+    cache_key: str,
+    layout_identity: dict[str, Any],
+    content_digest: str,
+    expected_records: Sequence[_RuntimeTensor],
+) -> HostWeightPlan:
+    manifest_path = entry_dir / "manifest.json"
+    try:
+        with manifest_path.open(encoding="utf-8") as handle:
+            manifest = json.load(handle)
+    except (OSError, ValueError) as exc:
+        raise _RuntimeCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc
+
+    if manifest.get("schema_version") != RUNTIME_CACHE_SCHEMA_VERSION:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache schema mismatch")
+    if manifest.get("cache_key") != cache_key:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache key mismatch")
+    if manifest.get("layout_identity") != layout_identity:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache layout identity mismatch")
+    if manifest.get("runtime_content_sha256") != content_digest:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache content identity mismatch")
+
+    expected_metadata = {record.name: _tensor_metadata(record) for record in expected_records}
+    tensor_manifest = manifest.get("tensors")
+    file_manifest = manifest.get("files")
+    if not isinstance(tensor_manifest, dict) or not isinstance(file_manifest, dict):
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache manifest has invalid sections")
+    if set(tensor_manifest) != set(expected_metadata):
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache tensor names differ from the loaded DiT")
+
+    for filename, metadata in file_manifest.items():
+        if Path(filename).name != filename or not isinstance(metadata, dict):
+            raise _RuntimeCacheError("entry_rejected", f"invalid cache filename {filename!r}")
+        path = entry_dir / filename
+        try:
+            stat = path.stat()
+        except OSError as exc:
+            raise _RuntimeCacheError("entry_rejected", f"missing cache shard {path}: {exc}") from exc
+        if stat.st_size != metadata.get("size") or _file_digest(path) != metadata.get("sha256"):
+            raise _RuntimeCacheError("entry_rejected", f"content digest mismatch for cache shard {path}")
+
+    bindings: dict[str, TensorBinding] = {}
+    mapped_records: list[_RuntimeTensor] = []
+    with contextlib.ExitStack() as stack:
+        handles = {
+            filename: stack.enter_context(safe_open(entry_dir / filename, framework="pt", device="cpu"))
+            for filename in file_manifest
+        }
+        expected_keys_by_file: dict[str, set[str]] = {filename: set() for filename in file_manifest}
+        for record in expected_records:
+            stored = tensor_manifest.get(record.name)
+            if not isinstance(stored, dict):
+                raise _RuntimeCacheError("entry_rejected", f"invalid tensor metadata for {record.name!r}")
+            filename = stored.get("file")
+            storage_key = stored.get("storage_key")
+            if filename not in handles or not isinstance(storage_key, str):
+                raise _RuntimeCacheError("entry_rejected", f"invalid storage binding for {record.name!r}")
+            actual_metadata = {key: stored.get(key) for key in expected_metadata[record.name]}
+            if actual_metadata != expected_metadata[record.name]:
+                raise _RuntimeCacheError("entry_rejected", f"metadata mismatch for {record.name!r}")
+            expected_keys_by_file[filename].add(storage_key)
+            try:
+                tensor = handles[filename].get_tensor(storage_key)
+            except Exception as exc:
+                raise _RuntimeCacheError(
+                    "entry_rejected",
+                    f"cannot map {record.name!r} from {filename}: {exc}",
+                ) from exc
+            if tuple(tensor.shape) != tuple(record.tensor.shape) or tensor.dtype != record.tensor.dtype:
+                raise _RuntimeCacheError("entry_rejected", f"mapped metadata mismatch for {record.name!r}")
+            mapped_records.append(_RuntimeTensor(record.name, tensor, record.kind))
+            bindings[record.name] = TensorBinding(storage_key=storage_key, file_path=str(entry_dir / filename))
+
+        for filename, expected_keys in expected_keys_by_file.items():
+            if set(handles[filename].keys()) != expected_keys:
+                raise _RuntimeCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}")
+
+        if _runtime_content_digest(mapped_records) != content_digest:
+            raise _RuntimeCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher")
+
+    return HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings=bindings,
+        runtime_layout_key=cache_key,
+        post_load_complete=True,
+    )
+
+
+def _remove_stale_temps(cache_root: Path, cache_key: str) -> None:
+    temp_root = cache_root / ".tmp"
+    if not temp_root.is_dir():
+        return
+    for path in temp_root.glob(f"{cache_key}.*"):
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)
+
+
+def build_runtime_weight_cache_plan(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    loader_type: type,
+    cache_root: str | os.PathLike[str] | None,
+    lock_timeout_seconds: float,
+    max_shard_bytes: int,
+    model_identity: str | None,
+    revision: str | None,
+    runtime_dtype: Any,
+    load_format: str,
+    loader_inputs: Mapping[str, Any],
+    tensor_parallel_size: int,
+    tensor_parallel_rank: int,
+    sequence_parallel_guard: Mapping[str, Any],
+    use_hsdp: bool,
+    enable_expert_parallel: bool,
+    quantization_config: Any,
+    cfg_parallel_size: int,
+    pipeline_parallel_size: int,
+) -> HostWeightPlanResult:
+    """Publish or join one immutable cache entry for final CPU DiT tensors."""
+    if load_format == "dummy":
+        return HostWeightPlanResult(
+            None,
+            "dummy/random weights do not have a reusable runtime identity",
+            "unsupported_load_format",
+        )
+    if use_hsdp:
+        return HostWeightPlanResult(None, "HSDP/DTensor layouts are not runtime-cache compatible", "unsupported_hsdp")
+    if enable_expert_parallel:
+        return HostWeightPlanResult(
+            None,
+            "expert-parallel weight ownership is outside runtime-cache v1",
+            "unsupported_expert_parallel",
+        )
+    if quantization_config is not None:
+        return HostWeightPlanResult(
+            None,
+            "quantized runtime layouts are not supported by the first runtime-cache version",
+            "unsupported_quantization",
+        )
+    if cfg_parallel_size != 1:
+        return HostWeightPlanResult(None, "CFG-parallel component ownership is not proven", "unsupported_cfg_parallel")
+    if pipeline_parallel_size != 1:
+        return HostWeightPlanResult(None, "pipeline-parallel ownership is outside runtime-cache v1", "unsupported_pp")
+    if not math.isfinite(lock_timeout_seconds) or lock_timeout_seconds <= 0:
+        return HostWeightPlanResult(None, "runtime-cache lock timeout must be positive", "invalid_configuration")
+    if max_shard_bytes <= 0:
+        return HostWeightPlanResult(None, "runtime-cache shard size must be positive", "invalid_configuration")
+
+    try:
+        records = _collect_runtime_tensors(pipeline, dit_modules)
+        content_digest = _runtime_content_digest(records)
+        layout_identity = {
+            "schema_version": RUNTIME_CACHE_SCHEMA_VERSION,
+            "model": model_identity,
+            "revision": revision,
+            "runtime_dtype": str(runtime_dtype),
+            "load_format": load_format,
+            "loader_inputs": _json_normalize(loader_inputs),
+            "loader_implementation_sha256": _implementation_fingerprint(loader_type, pipeline, dit_modules),
+            "components": [name for name, _ in dit_modules],
+            "tensor_parallel_size": tensor_parallel_size,
+            "tensor_parallel_rank": tensor_parallel_rank,
+            "sequence_parallel_guard": _json_normalize(sequence_parallel_guard),
+            "runtime_content_sha256": content_digest,
+        }
+        cache_key = hashlib.sha256(_canonical_json(layout_identity)).hexdigest()
+        root = Path(cache_root or default_runtime_cache_root()).expanduser().resolve()
+        entries_root = root / f"v{RUNTIME_CACHE_SCHEMA_VERSION}"
+        entry_dir = entries_root / cache_key
+        entries_root.mkdir(parents=True, exist_ok=True)
+
+        if entry_dir.is_dir():
+            try:
+                plan = _validate_entry(
+                    entry_dir=entry_dir,
+                    cache_key=cache_key,
+                    layout_identity=layout_identity,
+                    content_digest=content_digest,
+                    expected_records=records,
+                )
+                logger.info("Reusing DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
+                return HostWeightPlanResult(plan)
+            except _RuntimeCacheError as exc:
+                logger.warning("DLO runtime cache entry %s will be rebuilt: %s", cache_key[:12], exc)
+
+        lock_path = root / ".locks" / f"{cache_key}.lock"
+        with _exclusive_lock(lock_path, lock_timeout_seconds):
+            _remove_stale_temps(root, cache_key)
+            if entry_dir.is_dir():
+                try:
+                    plan = _validate_entry(
+                        entry_dir=entry_dir,
+                        cache_key=cache_key,
+                        layout_identity=layout_identity,
+                        content_digest=content_digest,
+                        expected_records=records,
+                    )
+                    logger.info("Joined DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
+                    return HostWeightPlanResult(plan)
+                except _RuntimeCacheError:
+                    pass
+
+            _publish_entry(
+                cache_root=root,
+                entry_dir=entry_dir,
+                cache_key=cache_key,
+                layout_identity=layout_identity,
+                content_digest=content_digest,
+                records=records,
+                max_shard_bytes=max_shard_bytes,
+            )
+            plan = _validate_entry(
+                entry_dir=entry_dir,
+                cache_key=cache_key,
+                layout_identity=layout_identity,
+                content_digest=content_digest,
+                expected_records=records,
+            )
+            logger.info(
+                "Published DLO runtime cache entry %s with %d tensors at %s",
+                cache_key[:12],
+                len(records),
+                entry_dir,
+            )
+            return HostWeightPlanResult(plan)
+    except _RuntimeCacheError as exc:
+        return HostWeightPlanResult(None, str(exc), exc.code)
+    except Exception as exc:
+        return HostWeightPlanResult(None, f"runtime-cache operation failed: {exc}", "cache_operation_failed")
+
+
+__all__ = [
+    "DEFAULT_LOCK_TIMEOUT_SECONDS",
+    "DEFAULT_RUNTIME_CACHE_ROOT",
+    "DEFAULT_SHARD_SIZE_BYTES",
+    "RUNTIME_CACHE_SCHEMA_VERSION",
+    "build_runtime_weight_cache_plan",
+    "default_runtime_cache_root",
+]

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -70,6 +70,9 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 logger = logging.getLogger(__name__)
 MAX_DREAMZERO_SESSIONS = 64
+# Shipped DreamZero geometry retains 24 Wan VAE causal-convolution cache
+# entries. This is the measured persistent CUDA upper bound per live session.
+DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION = 603 * 1024 * 1024
 
 # The pipeline's per-session state is a bespoke ``DreamZeroState`` by default, or
 # a ``DreamZeroStateAdapter`` view when the opt-in session manager is enabled.
@@ -149,6 +152,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             session_capacity=MAX_DREAMZERO_SESSIONS,
             cross_attention=tuple(cross_attention),
             max_scratch_tokens_per_branch=int(transformer.num_action_per_block + transformer.num_state_per_block),
+            model_owned_state_bytes_per_session=DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
         )
 
     @contextmanager

vllm_omni/diffusion/models/lingbot_world/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .actions import LingBotCameraControlReducer
+from .pipeline import (
+    LingBotWorldCausalDMDPipeline,
+    get_lingbot_world_post_process_func,
+    get_lingbot_world_pre_process_func,
+)
+from .transformer import CausalLingBotWorldTransformer3DModel
+
+__all__ = [
+    "CausalLingBotWorldTransformer3DModel",
+    "LingBotCameraControlReducer",
+    "LingBotWorldCausalDMDPipeline",
+    "get_lingbot_world_post_process_func",
+    "get_lingbot_world_pre_process_func",
+]

vllm_omni/diffusion/models/lingbot_world/actions.py

@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Realtime keyboard controls for LingBot World 2.0."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.models.lingbot_world.camera import (
+    CameraTrajectory,
+)
+from vllm_omni.experimental.ar_diffusion.session import (
+    ARDiffusionPreparedControls,
+    ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionControlInput,
+)
+
+LINGBOT_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
+LINGBOT_CAMERA_TRAJECTORY_SCHEMA = "lingbot.camera_trajectory.v1"
+_ACTION_ORDER = ("w", "a", "s", "d", "i", "j", "k", "l")
+_VALID_ACTIONS = frozenset(_ACTION_ORDER)
+_REFERENCE_HEIGHT = 480
+_REFERENCE_WIDTH = 832
+
+
+def _normalize_actions(value: object, *, field: str) -> tuple[str, ...]:
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        raise ValueError(f"{field} must be a sequence of LingBot action keys.")
+    actions: set[str] = set()
+    for item in value:
+        if not isinstance(item, str) or item.lower() not in _VALID_ACTIONS:
+            raise ValueError(f"{field} supports only W/A/S/D/I/J/K/L action keys.")
+        actions.add(item.lower())
+    return tuple(action for action in _ACTION_ORDER if action in actions)
+
+
+def _normalize_frames(value: object, *, field: str) -> tuple[tuple[str, ...], ...]:
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        raise ValueError(f"{field} must be a sequence of per-frame action lists.")
+    return tuple(_normalize_actions(actions, field=f"{field}[{index}]") for index, actions in enumerate(value))
+
+
+def parse_lingbot_camera_action_frames(
+    data: Mapping[str, Any],
+    *,
+    expected_frames: int,
+) -> tuple[tuple[str, ...], ...]:
+    """Validate the chunk-sized payload emitted by the session reducer."""
+
+    if data.get("mode") != "frames":
+        raise ValueError("lingbot.camera_actions.v1 model ticks require mode='frames'.")
+    frames = _normalize_frames(data.get("frames"), field="camera action frames")
+    if len(frames) != expected_frames:
+        raise ValueError(
+            "lingbot.camera_actions.v1 must contain exactly one action list per "
+            f"latent frame; expected {expected_frames}, got {len(frames)}."
+        )
+    return frames
+
+
+@dataclass(frozen=True)
+class _LingBotCameraReducerState:
+    mode: str | None = None
+    held_actions: tuple[str, ...] = ()
+    script_frames: tuple[tuple[str, ...], ...] = ()
+
+
+class LingBotCameraControlReducer:
+    """Sample live state/script controls into one LingBot AR block.
+
+    State mode preserves held keys across chunks and guarantees a one-frame
+    pulse when a press and release both arrive before the next chunk. Script
+    mode is a finite FIFO padded with neutral frames after exhaustion.
+    """
+
+    def __init__(self, *, frames_per_block: int = 3) -> None:
+        if isinstance(frames_per_block, bool) or not isinstance(frames_per_block, int) or frames_per_block <= 0:
+            raise ValueError("frames_per_block must be a positive integer.")
+        self._frames_per_block = frames_per_block
+        self._state = _LingBotCameraReducerState()
+
+    @staticmethod
+    def _state_transitions(data: Mapping[str, Any]) -> tuple[tuple[str, ...], ...]:
+        transitions = data.get("transitions")
+        if not isinstance(transitions, Sequence) or isinstance(transitions, (str, bytes)) or not transitions:
+            raise ValueError("LingBot state-mode camera actions require non-empty transitions.")
+        normalized: list[tuple[str, ...]] = []
+        previous_timestamp: int | None = None
+        for index, transition in enumerate(transitions):
+            if not isinstance(transition, Mapping):
+                raise ValueError("Each LingBot camera state transition must be a mapping.")
+            timestamp = transition.get("client_ts_ms")
+            if timestamp is not None:
+                if isinstance(timestamp, bool) or not isinstance(timestamp, int) or timestamp < 0:
+                    raise ValueError("client_ts_ms must be a non-negative integer when provided.")
+                if previous_timestamp is not None and timestamp < previous_timestamp:
+                    raise ValueError("client_ts_ms values must be monotonic within one event.")
+                previous_timestamp = timestamp
+            normalized.append(
+                _normalize_actions(
+                    transition.get("actions"),
+                    field=f"camera state transitions[{index}].actions",
+                )
+            )
+        return tuple(normalized)
+
+    @staticmethod
+    def _script(data: Mapping[str, Any]) -> tuple[tuple[str, ...], ...]:
+        return _normalize_frames(data.get("frames"), field="camera action script frames")
+
+    def prepare(
+        self,
+        *,
+        current_controls: Mapping[str, ARDiffusionControlInput],
+        events: Sequence[ARDiffusionSessionEvent],
+        chunk_index: int,
+    ) -> ARDiffusionPreparedControls:
+        del chunk_index
+        controls = {
+            track: control
+            for track, control in current_controls.items()
+            if not (track == "camera" and control.schema == LINGBOT_CAMERA_ACTION_SCHEMA)
+        }
+        state = self._state
+        pending_transitions: tuple[tuple[str, ...], ...] = ()
+
+        for event in events:
+            for control in event.controls:
+                if control.track != "camera" or control.schema != LINGBOT_CAMERA_ACTION_SCHEMA:
+                    controls[control.track] = control
+                    if control.track == "camera":
+                        state = _LingBotCameraReducerState()
+                        pending_transitions = ()
+                    continue
+
+                mode = control.data.get("mode")
+                controls.pop("camera", None)
+                if mode == "script":
+                    state = _LingBotCameraReducerState(
+                        mode="script",
+                        script_frames=self._script(control.data),
+                    )
+                    pending_transitions = ()
+                elif mode == "state":
+                    transitions = self._state_transitions(control.data)
+                    if state.mode != "state":
+                        state = _LingBotCameraReducerState(mode="state")
+                    else:
+                        state = _LingBotCameraReducerState(
+                            mode="state",
+                            held_actions=state.held_actions,
+                        )
+                    pending_transitions += transitions
+                else:
+                    raise ValueError("lingbot.camera_actions.v1 events require mode='state' or mode='script'.")
+
+        if state.mode == "script":
+            frames = state.script_frames[: self._frames_per_block]
+            frames += ((),) * (self._frames_per_block - len(frames))
+            state = _LingBotCameraReducerState(
+                mode="script",
+                script_frames=state.script_frames[self._frames_per_block :],
+            )
+        elif state.mode == "state":
+            final_actions = pending_transitions[-1] if pending_transitions else state.held_actions
+            pulse = next(
+                (actions for actions in reversed(pending_transitions) if actions),
+                None,
+            )
+            if pulse is not None and pulse != final_actions:
+                frames = (pulse,) + (final_actions,) * (self._frames_per_block - 1)
+            else:
+                frames = (final_actions,) * self._frames_per_block
+            state = _LingBotCameraReducerState(
+                mode="state",
+                held_actions=final_actions,
+            )
+        else:
+            frames = ()
+
+        if frames:
+            controls["camera"] = ARDiffusionControlInput(
+                track="camera",
+                schema=LINGBOT_CAMERA_ACTION_SCHEMA,
+                data={
+                    "mode": "frames",
+                    "frames": [list(actions) for actions in frames],
+                },
+            )
+        return ARDiffusionPreparedControls(
+            controls=tuple(controls[track] for track in sorted(controls)),
+            private_state=state,
+        )
+
+    def commit(self, prepared: ARDiffusionPreparedControls) -> None:
+        if not isinstance(prepared.private_state, _LingBotCameraReducerState):
+            raise TypeError("LingBot reducer prepared state has an unexpected type.")
+        self._state = prepared.private_state
+
+    def reset(self) -> None:
+        self._state = _LingBotCameraReducerState()
+
+
+def _rotation_matrix(axis: str, angle: float) -> np.ndarray:
+    cosine = math.cos(angle)
+    sine = math.sin(angle)
+    if axis == "x":
+        return np.array(
+            ((1.0, 0.0, 0.0), (0.0, cosine, -sine), (0.0, sine, cosine)),
+            dtype=np.float64,
+        )
+    return np.array(
+        ((cosine, 0.0, sine), (0.0, 1.0, 0.0), (-sine, 0.0, cosine)),
+        dtype=np.float64,
+    )
+
+
+def integrate_lingbot_camera_actions(
+    frames: Sequence[Sequence[str]],
+    *,
+    width: int,
+    height: int,
+    initial_pose: torch.Tensor | None = None,
+    initial_pitch: float = 0.0,
+) -> tuple[CameraTrajectory, float]:
+    """Convert latent-frame WASD/IJKL controls to cumulative C2W poses.
+
+    Calibration and motion constants intentionally match SGLang's LingBot
+    adapter: movement 0.05, pitch 4 degrees, yaw 6 degrees, pitch clamp 85
+    degrees, and target-resolution focal lengths of 500 pixels.
+    """
+
+    normalized_frames = tuple(
+        _normalize_actions(actions, field=f"camera action frames[{index}]") for index, actions in enumerate(frames)
+    )
+    if not normalized_frames:
+        raise ValueError("camera action frames must not be empty.")
+    if width <= 0 or height <= 0:
+        raise ValueError("camera action resolution must be positive.")
+
+    if initial_pose is None:
+        current_pose = np.eye(4, dtype=np.float64)
+    else:
+        if initial_pose.shape != (4, 4) or not torch.isfinite(initial_pose).all():
+            raise ValueError("initial_pose must be one finite 4x4 camera-to-world matrix.")
+        current_pose = initial_pose.detach().cpu().double().numpy().copy()
+    current_pitch = float(initial_pitch)
+    pitch_limit = math.radians(85.0)
+    poses: list[np.ndarray] = []
+
+    for actions in normalized_frames:
+        rotation = current_pose[:3, :3]
+        translation = current_pose[:3, 3]
+        pitch_delta = math.radians(4.0) * (("i" in actions) - ("k" in actions))
+        if not -pitch_limit <= current_pitch + pitch_delta <= pitch_limit:
+            pitch_delta = 0.0
+        else:
+            current_pitch += pitch_delta
+        yaw_delta = math.radians(6.0) * (("l" in actions) - ("j" in actions))
+        new_rotation = _rotation_matrix("y", yaw_delta) @ rotation @ _rotation_matrix("x", pitch_delta)
+
+        forward = np.array((new_rotation[0, 2], 0.0, new_rotation[2, 2]))
+        right = np.array((new_rotation[0, 0], 0.0, new_rotation[2, 0]))
+        forward_norm = np.linalg.norm(forward)
+        right_norm = np.linalg.norm(right)
+        if forward_norm > 0:
+            forward /= forward_norm + 1e-6
+        if right_norm > 0:
+            right /= right_norm + 1e-6
+
+        movement = np.zeros(3, dtype=np.float64)
+        movement += forward * 0.05 * (("w" in actions) - ("s" in actions))
+        movement += right * 0.05 * (("d" in actions) - ("a" in actions))
+        current_pose = np.eye(4, dtype=np.float64)
+        current_pose[:3, :3] = new_rotation
+        current_pose[:3, 3] = translation + movement
+        poses.append(current_pose)
+
+    # build_plucker_embedding expects intrinsics in the 832x480 reference
+    # coordinate system. These values become SGLang's [500, 500, W/2, H/2]
+    # after that function scales them to the requested resolution.
+    reference_intrinsics = (
+        500.0 * _REFERENCE_WIDTH / width,
+        500.0 * _REFERENCE_HEIGHT / height,
+        _REFERENCE_WIDTH / 2,
+        _REFERENCE_HEIGHT / 2,
+    )
+    trajectory = CameraTrajectory(
+        # Keep the cumulative controller pose in FP64 between chunks, matching
+        # SGLang's NumPy integration and avoiding long-session drift.
+        poses=torch.from_numpy(np.stack(poses)),
+        intrinsics=torch.tensor(reference_intrinsics, dtype=torch.float32).repeat(len(poses), 1),
+    )
+    return trajectory, current_pitch

vllm_omni/diffusion/models/lingbot_world/camera.py

@@ -0,0 +1,481 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Camera trajectory loading and geometry for LingBot World v2 conditioning."""
+
+from __future__ import annotations
+
+import io
+import math
+import os
+import stat
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import BinaryIO
+
+import numpy as np
+import torch
+
+_REFERENCE_HEIGHT = 480
+_REFERENCE_WIDTH = 832
+_MAX_ACTION_FRAMES = 117
+_MAX_SOURCE_ACTION_FRAMES = 4096
+_MAX_NPY_HEADER_BYTES = 10_000
+_DIRECTORY_OPEN_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
+_FILE_OPEN_FLAGS = os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK
+
+
+@dataclass(frozen=True)
+class CameraTrajectory:
+    """Camera-to-world poses and ``(fx, fy, cx, cy)`` intrinsics by frame."""
+
+    poses: torch.Tensor
+    intrinsics: torch.Tensor
+
+
+@dataclass(frozen=True)
+class TrustedActionDirectory:
+    """Canonical action location tied to the trusted root's filesystem identity."""
+
+    root: Path
+    relative: Path
+    root_device: int
+    root_inode: int
+
+
+def resolve_trusted_action_directory(
+    action_path: str | os.PathLike[str],
+    trusted_root: str | os.PathLike[str],
+) -> TrustedActionDirectory:
+    """Resolve an action directory beneath a trusted root without exposing paths."""
+
+    try:
+        root = Path(trusted_root).expanduser().resolve(strict=True)
+    except (OSError, RuntimeError):
+        raise ValueError("The configured LingBot trusted action root is unavailable.") from None
+    if not root.is_dir():
+        raise ValueError("The configured LingBot trusted action root must be a directory.")
+
+    try:
+        candidate = Path(action_path).expanduser()
+        if not candidate.is_absolute():
+            candidate = root / candidate
+        candidate = candidate.resolve(strict=True)
+        relative = candidate.relative_to(root)
+    except (OSError, RuntimeError, ValueError):
+        raise ValueError(
+            "sampling_params.extra_args.action_path must be contained by the trusted action root."
+        ) from None
+    if not candidate.is_dir():
+        raise ValueError("sampling_params.extra_args.action_path must identify a directory in the trusted action root.")
+
+    try:
+        root_stat = root.stat()
+    except OSError:
+        raise ValueError("The configured LingBot trusted action root is unavailable.") from None
+    return TrustedActionDirectory(
+        root=root,
+        relative=relative,
+        root_device=root_stat.st_dev,
+        root_inode=root_stat.st_ino,
+    )
+
+
+def _validate_trajectory(
+    trajectory: CameraTrajectory,
+    *,
+    file_suffix: str = "",
+) -> None:
+    poses = trajectory.poses
+    intrinsics = trajectory.intrinsics
+    if poses.ndim != 3 or poses.shape[1:] != (4, 4):
+        raise ValueError(f"poses{file_suffix} must have shape [frames, 4, 4], got {tuple(poses.shape)}")
+    if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
+        raise ValueError(
+            f"intrinsics{file_suffix} must have shape [frames, 4] ordered as fx, fy, cx, cy, "
+            f"got {tuple(intrinsics.shape)}"
+        )
+    if poses.shape[0] == 0:
+        raise ValueError("camera trajectory must contain at least one frame")
+    if poses.shape[0] != intrinsics.shape[0]:
+        raise ValueError("poses and intrinsics must contain the same number of frames")
+    if not torch.isfinite(poses).all() or not torch.isfinite(intrinsics).all():
+        raise ValueError("camera trajectory values must all be finite")
+
+
+@contextmanager
+def _open_action_directory(action_directory: TrustedActionDirectory) -> Iterator[int]:
+    root_fd: int | None = None
+    action_fd: int | None = None
+    try:
+        root_fd = os.open(action_directory.root, _DIRECTORY_OPEN_FLAGS)
+        root_stat = os.fstat(root_fd)
+        if (root_stat.st_dev, root_stat.st_ino) != (
+            action_directory.root_device,
+            action_directory.root_inode,
+        ):
+            raise ValueError("The configured LingBot trusted action root changed before camera files were opened.")
+
+        action_fd = os.dup(root_fd)
+        for component in action_directory.relative.parts:
+            next_fd = os.open(component, _DIRECTORY_OPEN_FLAGS, dir_fd=action_fd)
+            os.close(action_fd)
+            action_fd = next_fd
+    except OSError:
+        if action_fd is not None:
+            os.close(action_fd)
+        if root_fd is not None:
+            os.close(root_fd)
+        raise ValueError("Unable to securely open the LingBot action directory within the trusted root.") from None
+    except ValueError:
+        if action_fd is not None:
+            os.close(action_fd)
+        if root_fd is not None:
+            os.close(root_fd)
+        raise
+    assert action_fd is not None
+    assert root_fd is not None
+    try:
+        yield action_fd
+    finally:
+        os.close(action_fd)
+        os.close(root_fd)
+
+
+@contextmanager
+def _open_camera_file(action_fd: int, filename: str) -> Iterator[tuple[BinaryIO, os.stat_result]]:
+    file_fd = -1
+    try:
+        file_fd = os.open(filename, _FILE_OPEN_FLAGS, dir_fd=action_fd)
+        file_stat = os.fstat(file_fd)
+        if not stat.S_ISREG(file_stat.st_mode):
+            raise ValueError(f"{filename} must be a regular NPY file.")
+        with os.fdopen(file_fd, "rb", closefd=True) as file_handle:
+            file_fd = -1
+            yield file_handle, file_stat
+    except FileNotFoundError:
+        raise FileNotFoundError(f"camera trajectory file not found: {filename}") from None
+    except OSError:
+        raise ValueError(f"{filename} could not be opened securely from the trusted action directory.") from None
+    finally:
+        if file_fd >= 0:
+            os.close(file_fd)
+
+
+def _load_bounded_npy(
+    action_fd: int,
+    filename: str,
+    *,
+    trailing_shape: tuple[int, ...],
+) -> np.ndarray:
+    with _open_camera_file(action_fd, filename) as (file_handle, file_stat):
+        max_payload_bytes = _MAX_SOURCE_ACTION_FRAMES * math.prod(trailing_shape) * 8
+        max_file_bytes = _MAX_NPY_HEADER_BYTES + max_payload_bytes + 32
+        if file_stat.st_size > max_file_bytes:
+            raise ValueError(f"{filename} byte size exceeds the bounded camera-array limit.")
+        snapshot = file_handle.read(max_file_bytes + 1)
+        if len(snapshot) != file_stat.st_size:
+            raise ValueError(f"{filename} changed while its bounded contents were read.")
+        snapshot_file = io.BytesIO(snapshot)
+        try:
+            version = np.lib.format.read_magic(snapshot_file)
+            if version == (1, 0):
+                shape, _fortran_order, dtype = np.lib.format.read_array_header_1_0(snapshot_file)
+            elif version == (2, 0):
+                shape, _fortran_order, dtype = np.lib.format.read_array_header_2_0(snapshot_file)
+            else:
+                raise ValueError(f"{filename} must use NPY format version 1.0 or 2.0.")
+        except (EOFError, TypeError, ValueError):
+            raise ValueError(f"{filename} has an invalid or oversized NPY header.") from None
+
+        if len(shape) != len(trailing_shape) + 1 or tuple(shape[1:]) != trailing_shape:
+            expected = ", ".join(str(value) for value in trailing_shape)
+            raise ValueError(f"{filename} must have shape [frames, {expected}], got {shape}.")
+        frames = shape[0]
+        if isinstance(frames, bool) or not isinstance(frames, int) or not 1 <= frames <= _MAX_SOURCE_ACTION_FRAMES:
+            raise ValueError(f"{filename} must contain between 1 and {_MAX_SOURCE_ACTION_FRAMES} source frames.")
+        if dtype.kind not in "fiu" or dtype.hasobject or dtype.itemsize > 8:
+            raise ValueError(f"{filename} must contain real numeric values with at most 8 bytes per element.")
+
+        data_offset = snapshot_file.tell()
+        expected_bytes = math.prod(shape) * dtype.itemsize
+        if len(snapshot) != data_offset + expected_bytes:
+            raise ValueError(f"{filename} byte size does not match its NPY header.")
+        snapshot_file.seek(0)
+        try:
+            array = np.load(snapshot_file, allow_pickle=False)
+        except (EOFError, TypeError, ValueError):
+            raise ValueError(f"{filename} does not contain a valid numeric NPY array.") from None
+    # Official LingBot trajectories may be longer than one request. Only
+    # materialize the bounded prefix that any supported request can consume.
+    result: np.ndarray = np.asarray(array[:_MAX_ACTION_FRAMES], dtype=np.float32)
+    return result
+
+
+def load_camera_trajectory(action_directory: TrustedActionDirectory) -> CameraTrajectory:
+    """Load bounded numeric camera arrays through a trusted directory handle."""
+
+    if not isinstance(action_directory, TrustedActionDirectory):
+        raise TypeError("action_directory must be resolved against a trusted action root before loading.")
+    with _open_action_directory(action_directory) as action_fd:
+        poses_array = _load_bounded_npy(action_fd, "poses.npy", trailing_shape=(4, 4))
+        intrinsics_array = _load_bounded_npy(action_fd, "intrinsics.npy", trailing_shape=(4,))
+
+    trajectory = CameraTrajectory(
+        poses=torch.from_numpy(poses_array),
+        intrinsics=torch.from_numpy(intrinsics_array),
+    )
+    _validate_trajectory(trajectory, file_suffix=".npy")
+    return trajectory
+
+
+def _linear_interpolate(values: torch.Tensor, target_times: np.ndarray) -> torch.Tensor:
+    source_times = np.linspace(0.0, 1.0, values.shape[0])
+    values_array = values.detach().cpu().double().numpy()
+    columns = [
+        np.interp(target_times, source_times, values_array[:, column]) for column in range(values_array.shape[1])
+    ]
+    interpolated = torch.from_numpy(np.stack(columns, axis=1))
+    return interpolated.to(device=values.device, dtype=values.dtype)
+
+
+def _rotation_matrices_to_quaternions(matrices: np.ndarray) -> np.ndarray:
+    """Convert rotation matrices to normalized ``(w, x, y, z)`` quaternions."""
+
+    quaternions: np.ndarray = np.empty((matrices.shape[0], 4), dtype=np.float64)
+    for index, matrix in enumerate(matrices):
+        trace = float(np.trace(matrix))
+        if trace > 0.0:
+            scale = math.sqrt(trace + 1.0) * 2.0
+            quaternion = np.array(
+                [
+                    0.25 * scale,
+                    (matrix[2, 1] - matrix[1, 2]) / scale,
+                    (matrix[0, 2] - matrix[2, 0]) / scale,
+                    (matrix[1, 0] - matrix[0, 1]) / scale,
+                ]
+            )
+        else:
+            axis = int(np.argmax(np.diag(matrix)))
+            if axis == 0:
+                scale = math.sqrt(max(0.0, 1.0 + matrix[0, 0] - matrix[1, 1] - matrix[2, 2])) * 2.0
+                quaternion = np.array(
+                    [
+                        (matrix[2, 1] - matrix[1, 2]) / scale,
+                        0.25 * scale,
+                        (matrix[0, 1] + matrix[1, 0]) / scale,
+                        (matrix[0, 2] + matrix[2, 0]) / scale,
+                    ]
+                )
+            elif axis == 1:
+                scale = math.sqrt(max(0.0, 1.0 + matrix[1, 1] - matrix[0, 0] - matrix[2, 2])) * 2.0
+                quaternion = np.array(
+                    [
+                        (matrix[0, 2] - matrix[2, 0]) / scale,
+                        (matrix[0, 1] + matrix[1, 0]) / scale,
+                        0.25 * scale,
+                        (matrix[1, 2] + matrix[2, 1]) / scale,
+                    ]
+                )
+            else:
+                scale = math.sqrt(max(0.0, 1.0 + matrix[2, 2] - matrix[0, 0] - matrix[1, 1])) * 2.0
+                quaternion = np.array(
+                    [
+                        (matrix[1, 0] - matrix[0, 1]) / scale,
+                        (matrix[0, 2] + matrix[2, 0]) / scale,
+                        (matrix[1, 2] + matrix[2, 1]) / scale,
+                        0.25 * scale,
+                    ]
+                )
+        norm = np.linalg.norm(quaternion)
+        if not np.isfinite(norm) or norm <= np.finfo(np.float64).eps:
+            raise ValueError("camera rotation matrices must describe valid rotations")
+        quaternions[index] = quaternion / norm
+    return quaternions
+
+
+def _quaternions_to_rotation_matrices(quaternions: np.ndarray) -> np.ndarray:
+    quaternions = quaternions / np.linalg.norm(quaternions, axis=1, keepdims=True)
+    w, x, y, z = quaternions.T
+    matrices: np.ndarray = np.stack(
+        (
+            1 - 2 * (y * y + z * z),
+            2 * (x * y - z * w),
+            2 * (x * z + y * w),
+            2 * (x * y + z * w),
+            1 - 2 * (x * x + z * z),
+            2 * (y * z - x * w),
+            2 * (x * z - y * w),
+            2 * (y * z + x * w),
+            1 - 2 * (x * x + y * y),
+        ),
+        axis=1,
+    ).reshape(-1, 3, 3)
+    return matrices
+
+
+def _slerp_rotations(matrices: np.ndarray, target_times: np.ndarray) -> np.ndarray:
+    quaternions = _rotation_matrices_to_quaternions(matrices)
+    interpolated = np.empty((target_times.shape[0], 4), dtype=np.float64)
+    for target_index, target_time in enumerate(target_times):
+        scaled_time = float(target_time) * (matrices.shape[0] - 1)
+        left = min(max(math.floor(scaled_time), 0), matrices.shape[0] - 2)
+        fraction = scaled_time - left
+        first = quaternions[left]
+        second = quaternions[left + 1]
+        cosine = float(np.dot(first, second))
+        if cosine < 0.0:
+            second = -second
+            cosine = -cosine
+        cosine = min(cosine, 1.0)
+        if cosine > 0.9995:
+            quaternion = first + fraction * (second - first)
+            interpolated[target_index] = quaternion / np.linalg.norm(quaternion)
+            continue
+        angle = math.acos(cosine)
+        denominator = math.sin(angle)
+        interpolated[target_index] = (
+            math.sin((1.0 - fraction) * angle) / denominator * first + math.sin(fraction * angle) / denominator * second
+        )
+    return _quaternions_to_rotation_matrices(interpolated)
+
+
+def interpolate_camera_trajectory(
+    trajectory: CameraTrajectory,
+    num_frames: int,
+) -> CameraTrajectory:
+    """Resample translations/intrinsics linearly and rotations with quaternion SLERP."""
+
+    _validate_trajectory(trajectory)
+    if num_frames <= 0:
+        raise ValueError(f"num_frames must be positive, got {num_frames}")
+
+    source_frames = trajectory.poses.shape[0]
+    if num_frames == source_frames:
+        return trajectory
+    if source_frames == 1:
+        return CameraTrajectory(
+            poses=trajectory.poses.repeat(num_frames, 1, 1),
+            intrinsics=trajectory.intrinsics.repeat(num_frames, 1),
+        )
+
+    # Rotations use quaternion SLERP; translations and intrinsics are linear.
+    target_times = np.linspace(0.0, 1.0, num_frames)
+    rotation_matrices = trajectory.poses[:, :3, :3].detach().cpu().double().numpy()
+    interpolated_rotations = _slerp_rotations(rotation_matrices, target_times)
+    translations = _linear_interpolate(trajectory.poses[:, :3, 3], target_times)
+    intrinsics = _linear_interpolate(trajectory.intrinsics, target_times)
+
+    poses = torch.eye(
+        4,
+        device=trajectory.poses.device,
+        dtype=trajectory.poses.dtype,
+    ).repeat(num_frames, 1, 1)
+    poses[:, :3, :3] = torch.from_numpy(interpolated_rotations).to(
+        device=trajectory.poses.device,
+        dtype=trajectory.poses.dtype,
+    )
+    poses[:, :3, 3] = translations
+    return CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+
+def _invert_camera_poses(poses: torch.Tensor) -> torch.Tensor:
+    """Invert rigid transforms using R^-1 = R^T and t^-1 = -R^T t."""
+
+    rotations = poses[:, :3, :3]
+    translations = poses[:, :3, 3:]
+    inverse_rotations = rotations.transpose(-1, -2)
+    inverse = torch.eye(4, device=poses.device, dtype=poses.dtype).repeat(poses.shape[0], 1, 1)
+    inverse[:, :3, :3] = inverse_rotations
+    inverse[:, :3, 3:] = -torch.bmm(inverse_rotations, translations)
+    return inverse
+
+
+def _prepare_framewise_poses(raw_poses: torch.Tensor) -> torch.Tensor:
+    """Convert raw C2W poses to normalized first-relative framewise deltas."""
+
+    # Remove global placement, then convert the sequence to framewise deltas.
+    first_world_to_camera = _invert_camera_poses(raw_poses[:1])
+    relative_poses = torch.matmul(first_world_to_camera, raw_poses).clone()
+    relative_poses[0] = torch.eye(4, device=raw_poses.device, dtype=raw_poses.dtype)
+    if relative_poses.shape[0] > 1:
+        relative_poses[1:] = torch.bmm(
+            _invert_camera_poses(relative_poses[:-1]),
+            relative_poses[1:],
+        )
+
+    # Only relative direction/magnitude matters to the checkpoint; normalizing
+    # by the largest step removes the trajectory's arbitrary translation scale.
+    translations = relative_poses[:, :3, 3]
+    max_translation_norm = torch.linalg.vector_norm(translations, dim=-1).max()
+    if max_translation_norm > 0:
+        relative_poses[:, :3, 3] = translations / max_translation_norm
+    return relative_poses
+
+
+def build_plucker_embedding(
+    trajectory: CameraTrajectory,
+    *,
+    height: int,
+    width: int,
+    target_height: int,
+    target_width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Build checkpoint-ordered ``(ray origin, ray direction)`` channels."""
+
+    _validate_trajectory(trajectory)
+    dimensions = {
+        "height": height,
+        "width": width,
+        "target_height": target_height,
+        "target_width": target_width,
+    }
+    for name, value in dimensions.items():
+        if value <= 0:
+            raise ValueError(f"{name} must be positive, got {value}")
+    if dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
+        raise ValueError(f"dtype must be floating point, got {dtype}")
+
+    # Compute geometry in at least FP32, then cast at the model boundary.
+    compute_dtype = torch.float32 if dtype in (torch.float16, torch.bfloat16) else dtype
+    poses = _prepare_framewise_poses(trajectory.poses.to(device=device, dtype=compute_dtype))
+    intrinsics = trajectory.intrinsics.to(device=device, dtype=compute_dtype).clone()
+    intrinsics[:, (0, 2)] *= width / _REFERENCE_WIDTH
+    intrinsics[:, (1, 3)] *= height / _REFERENCE_HEIGHT
+    if torch.any(intrinsics[:, :2] == 0):
+        raise ValueError("camera focal lengths fx and fy must be non-zero")
+
+    # Sample requested-pixel coordinates directly at the conditioning resolution.
+    x_coordinates = (torch.arange(target_width, device=device, dtype=compute_dtype) + 0.5) * (width / target_width)
+    y_coordinates = (torch.arange(target_height, device=device, dtype=compute_dtype) + 0.5) * (height / target_height)
+    grid_y, grid_x = torch.meshgrid(y_coordinates, x_coordinates, indexing="ij")
+    grid_x = grid_x.unsqueeze(0)
+    grid_y = grid_y.unsqueeze(0)
+
+    # Pinhole back-projection: pixel (x, y) becomes the camera-space direction
+    # ((x-cx)/fx, (y-cy)/fy, 1), which is normalized before rotation.
+    fx, fy, cx, cy = intrinsics.unbind(dim=1)
+    camera_x = (grid_x - cx[:, None, None]) / fx[:, None, None]
+    camera_y = (grid_y - cy[:, None, None]) / fy[:, None, None]
+    camera_directions = torch.stack(
+        (camera_x, camera_y, torch.ones_like(camera_x)),
+        dim=-1,
+    )
+    camera_directions = camera_directions / torch.linalg.vector_norm(camera_directions, dim=-1, keepdim=True)
+
+    # A ray is represented by its camera center (origin) and unit direction in
+    # the prepared relative coordinate frame. Channel order is checkpoint
+    # data, not an implementation preference: [origin_xyz, direction_xyz].
+    world_directions = torch.einsum(
+        "fij,fhwj->fhwi",
+        poses[:, :3, :3],
+        camera_directions,
+    )
+    world_directions = world_directions / torch.linalg.vector_norm(world_directions, dim=-1, keepdim=True)
+    ray_origins = poses[:, None, None, :3, 3].expand_as(world_directions)
+
+    embedding = torch.cat((ray_origins, world_directions), dim=-1)
+    return embedding.permute(0, 3, 1, 2).contiguous().to(dtype=dtype)

vllm_omni/diffusion/offloader/__init__.py

@@ -71,8 +71,9 @@ def get_offload_backend(
     Args:
         od_config: OmniDiffusionConfig with offload settings
         device: Target device (auto-detected if None)
-        host_weight_plan: Exact loader-produced backing plan, if ordinary
-            weight materialization was skipped.
+        host_weight_plan: Exact loader-produced backing plan. Checkpoint plans
+            may have skipped ordinary materialization; runtime-cache plans
+            describe already-valid ordinary tensors that can be remapped.
 
     Returns:
         OffloadBackend instance or None if offloading disabled
@@ -85,10 +86,14 @@ def get_offload_backend(
     # Extract and validate configuration
     config = OffloadConfig.from_od_config(od_config)
 
+    plan_requires_backend = host_weight_plan is not None and not host_weight_plan.post_load_complete
     if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE:
-        raise RuntimeError(
-            "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
-        )
+        if plan_requires_backend:
+            raise RuntimeError(
+                "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
+            )
+        logger.warning("Ignoring optional runtime-cache plan because distributed layerwise offload is not selected")
+        host_weight_plan = None
 
     # Return None if no offloading requested
     if config.strategy == OffloadStrategy.NONE:
@@ -96,7 +101,7 @@ def get_offload_backend(
 
     # Validate platform (CUDA required for now)
     if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1:
-        if host_weight_plan is not None:
+        if plan_requires_backend:
             raise RuntimeError(
                 "The loader skipped ordinary weight materialization for DLO, "
                 "but this platform cannot create the required offload backend"
@@ -112,7 +117,7 @@ def get_offload_backend(
         try:
             device = current_omni_platform.get_torch_device()
         except (NotImplementedError, AttributeError) as exc:
-            if host_weight_plan is not None:
+            if plan_requires_backend:
                 raise RuntimeError(
                     "The loader skipped ordinary weight materialization for DLO, "
                     "but the target device could not be resolved"

vllm_omni/diffusion/offloader/base.py

@@ -31,6 +31,9 @@ class OffloadConfig:
     # blocks from the loader-selected host backing with H2D only.
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0  # leading DiT layers kept on device
+    # Per-worker budget for registering shared runtime-cache mappings with CUDA.
+    # Zero retains bounded host staging.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -101,6 +104,7 @@ class OffloadConfig:
         # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
         dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        dlo_runtime_cache_pin_limit_gib = float(getattr(od_config, "dlo_runtime_cache_pin_limit_gib", 0.0))
         if dlo_resident_layers < 0:
             raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
         if dlo_resident_layers and dlo_use_allgather:
@@ -136,6 +140,7 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
+            dlo_runtime_cache_pin_limit_gib=dlo_runtime_cache_pin_limit_gib,
         )
 
 

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bounded CUDA registration for existing read-only host mappings."""
+
+from __future__ import annotations
+
+import mmap
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+_CUDA_HOST_REGISTER_READ_ONLY = 0x08
+
+
+class CudaHostRegistrationError(RuntimeError):
+    """Raised when an existing host range cannot be registered safely."""
+
+
+class CudaHostRegistrationBudgetError(CudaHostRegistrationError):
+    """Raised before mutation when the complete mapping exceeds its budget."""
+
+
+class CudaHostRegistrationCleanupError(CudaHostRegistrationError):
+    """Raised when partial registration could not be rolled back safely."""
+
+
+@dataclass(frozen=True)
+class _AddressRange:
+    start: int
+    end: int
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start
+
+
+def _coalesce_ranges(
+    ranges: Sequence[tuple[int, int]],
+    page_size: int = mmap.PAGESIZE,
+) -> tuple[_AddressRange, ...]:
+    """Page-align and merge overlapping ranges from one backing mapping."""
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    aligned: list[_AddressRange] = []
+    for start, size in ranges:
+        if start <= 0 or size < 0:
+            raise ValueError(f"invalid host range start={start}, size={size}")
+        if size == 0:
+            continue
+        aligned_start = start - start % page_size
+        end = start + size
+        aligned_end = ((end + page_size - 1) // page_size) * page_size
+        aligned.append(_AddressRange(aligned_start, aligned_end))
+
+    merged: list[_AddressRange] = []
+    for region in sorted(aligned, key=lambda item: (item.start, item.end)):
+        if merged and region.start <= merged[-1].end:
+            previous = merged[-1]
+            merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
+        else:
+            merged.append(region)
+    return tuple(merged)
+
+
+def _tensor_regions(
+    sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+) -> tuple[_AddressRange, ...]:
+    """Resolve storage spans without merging unrelated file mappings."""
+    regions: list[_AddressRange] = []
+    for mapping_name, tensors in sources_by_mapping.items():
+        ranges: list[tuple[int, int]] = []
+        for tensor in tensors:
+            if tensor.device.type != "cpu":
+                raise CudaHostRegistrationError(
+                    f"CUDA host registration requires CPU storage, but {mapping_name!r} contains {tensor.device}"
+                )
+            storage = tensor.untyped_storage()
+            ranges.append((storage.data_ptr(), storage.nbytes()))
+        regions.extend(_coalesce_ranges(ranges))
+    return tuple(regions)
+
+
+def _error_message(runtime: Any, error: Any) -> str:
+    try:
+        message = runtime.cudaGetErrorString(error)
+    except Exception:
+        return str(error)
+    if isinstance(message, bytes):
+        return message.decode(errors="replace")
+    return str(message)
+
+
+class CudaHostRegistration:
+    """Own CUDA registrations for already-existing file-backed CPU tensors."""
+
+    def __init__(
+        self,
+        runtime: Any,
+        regions: tuple[_AddressRange, ...],
+    ) -> None:
+        self._runtime = runtime
+        self._regions = regions
+        self._closed = False
+
+    @classmethod
+    def create(
+        cls,
+        sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+        *,
+        max_bytes: int,
+    ) -> CudaHostRegistration:
+        if max_bytes <= 0:
+            raise CudaHostRegistrationBudgetError("CUDA host-registration budget is disabled")
+        regions = _tensor_regions(sources_by_mapping)
+        total_bytes = sum(region.size for region in regions)
+        if total_bytes > max_bytes:
+            raise CudaHostRegistrationBudgetError(
+                f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
+            )
+        if not regions:
+            raise CudaHostRegistrationError("no non-empty host ranges were available for registration")
+        if not torch.cuda.is_available():
+            raise CudaHostRegistrationError("CUDA is not available")
+
+        try:
+            runtime = torch.cuda.cudart()
+        except Exception as exc:
+            raise CudaHostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
+
+        registered: list[_AddressRange] = []
+        try:
+            for region in regions:
+                error = runtime.cudaHostRegister(
+                    region.start,
+                    region.size,
+                    _CUDA_HOST_REGISTER_READ_ONLY,
+                )
+                if int(error) != 0:
+                    raise CudaHostRegistrationError(
+                        "cudaHostRegister(read-only) failed for "
+                        f"[{region.start:#x}, {region.end:#x}): {_error_message(runtime, error)}"
+                    )
+                registered.append(region)
+
+            unpinned = [
+                mapping_name
+                for mapping_name, tensors in sources_by_mapping.items()
+                if any(tensor.numel() and not tensor.is_pinned() for tensor in tensors)
+            ]
+            if unpinned:
+                raise CudaHostRegistrationError(
+                    f"CUDA registration succeeded but PyTorch did not recognize pinned storage for {unpinned[:3]}"
+                )
+        except Exception as exc:
+            rollback_errors: list[str] = []
+            for region in reversed(registered):
+                try:
+                    error = runtime.cudaHostUnregister(region.start)
+                    if int(error) != 0:
+                        rollback_errors.append(
+                            f"cudaHostUnregister({region.start:#x}) failed: {_error_message(runtime, error)}"
+                        )
+                except Exception as rollback_exc:
+                    rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
+            if rollback_errors:
+                raise CudaHostRegistrationCleanupError(
+                    f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}"
+                ) from exc
+            if isinstance(exc, CudaHostRegistrationError):
+                raise
+            raise CudaHostRegistrationError(f"CUDA host registration raised: {exc}") from exc
+
+        return cls(runtime, regions)
+
+    @property
+    def total_bytes(self) -> int:
+        return sum(region.size for region in self._regions)
+
+    @property
+    def region_count(self) -> int:
+        return len(self._regions)
+
+    def close(self) -> list[str]:
+        """Unregister every range, returning errors after best-effort cleanup."""
+        if self._closed:
+            return []
+        errors: list[str] = []
+        failed: list[_AddressRange] = []
+        for region in reversed(self._regions):
+            try:
+                error = self._runtime.cudaHostUnregister(region.start)
+                if int(error) != 0:
+                    errors.append(
+                        f"cudaHostUnregister({region.start:#x}) failed: {_error_message(self._runtime, error)}"
+                    )
+                    failed.append(region)
+            except Exception as exc:
+                errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
+                failed.append(region)
+        self._regions = tuple(reversed(failed))
+        self._closed = not self._regions
+        return errors

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -18,6 +18,7 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
 from __future__ import annotations
 
+import time
 from itertools import chain
 from typing import Any
 
@@ -34,6 +35,11 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
+from .cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+)
 from .module_collector import ModuleDiscovery
 from .offload_plan import (
     OffloadPlan,
@@ -93,6 +99,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.rank = rank
         self.pin_memory = pin_memory
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.tensor_transforms = tensor_transforms or {}
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -444,6 +451,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _resolve_mmap_source(
+        source_info: dict[str, Any],
+        meta: dict[str, Any],
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        source = source_info["tensor"]
+        transform = source_info["transform"]
+        if callable(transform):
+            source = transform(source)
+        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
+            raise ValueError(
+                "mmap weight transform changed tensor layout for "
+                f"{source_info['name']!r}: expected dtype={dtype}, "
+                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
+                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
+                f"stride={source.stride()}"
+            )
+        return source
+
     @staticmethod
     def _pack_mmap_sources(
         cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -456,18 +483,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = source_info["tensor"]
-                transform = source_info["transform"]
-                if callable(transform):
-                    source = transform(source)
-                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
-                    raise ValueError(
-                        "mmap weight transform changed tensor layout for "
-                        f"{source_info['name']!r}: expected dtype={dtype}, "
-                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
-                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
-                        f"stride={source.stride()}"
-                    )
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -481,6 +497,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             staged[dtype] = destination
         return staged
 
+    @staticmethod
+    def _copy_mmap_sources_to_device(
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
+        metadata: dict[torch.dtype, list[dict[str, Any]]],
+        gpu_buffers: dict[torch.dtype, torch.Tensor],
+        *,
+        non_blocking: bool,
+    ) -> None:
+        """Copy registered source views directly into flattened device buffers."""
+        for dtype, metas in metadata.items():
+            destination = gpu_buffers[dtype]
+            sources = cpu_sources[dtype]
+            for source_info, meta in zip(sources, metas, strict=True):
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                start = meta["offset"]
+                physical_storage = destination[start : start + meta["numel"]]
+                async_copy = non_blocking and source.is_pinned()
+                if source.is_contiguous():
+                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
+                else:
+                    torch.as_strided(
+                        physical_storage,
+                        size=source.shape,
+                        stride=source.stride(),
+                    ).copy_(source, non_blocking=async_copy)
+
     def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
         """Pack this block's mmap views into one bounded host staging slot."""
         previous_copy = self.cpu_staging_events[slot]
@@ -514,17 +556,27 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
-            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
-            with current_omni_platform.stream(self.copy_stream):
-                for dtype, cpu_shard in cpu_weights.items():
-                    gw = gpu_weights[dtype]
-                    async_copy = non_blocking and cpu_shard.is_pinned()
-                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
-                evt.record(self.copy_stream)
-            if self.rank_local_mmap:
-                # The CPU slot may be overwritten only after this H2D copy has
-                # finished.  The shared event protects reuse by another hook.
-                self.cpu_staging_events[slot] = evt
+            if self.rank_local_mmap and self.registered_mmap:
+                with current_omni_platform.stream(self.copy_stream):
+                    self._copy_mmap_sources_to_device(
+                        self.cpu_sources,
+                        self.metadata,
+                        gpu_weights,
+                        non_blocking=non_blocking,
+                    )
+                    evt.record(self.copy_stream)
+            else:
+                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+                with current_omni_platform.stream(self.copy_stream):
+                    for dtype, cpu_shard in cpu_weights.items():
+                        gw = gpu_weights[dtype]
+                        async_copy = non_blocking and cpu_shard.is_pinned()
+                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                    evt.record(self.copy_stream)
+                if self.rank_local_mmap:
+                    # The CPU slot may be overwritten only after this H2D copy has
+                    # finished.  The shared event protects reuse by another hook.
+                    self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -734,6 +786,7 @@ class PinnedResidentLayerGroup:
         self.copy_stream = copy_stream
         self.loaded = False
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -804,6 +857,14 @@ class PinnedResidentLayerGroup:
         ready = current_omni_platform.Event()
         with current_omni_platform.stream(self.copy_stream):
             for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
+                if self.rank_local_mmap and self.registered_mmap:
+                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
+                        state["cpu_sources"],
+                        state["metadata"],
+                        block_buffers,
+                        non_blocking=True,
+                    )
+                    continue
                 if self.rank_local_mmap:
                     slot = index % 2
                     previous_copy = self._cpu_staging_events[slot]
@@ -901,8 +962,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.host_weight_plan = host_weight_plan
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
+        self._runtime_cache_mapped_sources: dict[str, list[torch.Tensor]] = {}
+        self._cuda_host_registration: CudaHostRegistration | None = None
 
     def load_resident_layers(self) -> None:
         """Load the model-declared leading blocks for the denoise stage."""
@@ -1023,7 +1087,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     framework="pt",
                     device="cpu",
                 )
-            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
+            tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
             if is_parameter:
                 replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                 parent._parameters[leaf_name] = replacement
@@ -1103,6 +1167,154 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 f"(first 5: {remaining_meta[:5]})."
             )
 
+    def _load_weights_from_runtime_cache(
+        self,
+        pipeline: nn.Module,
+        modules,
+        plan: HostWeightPlan,
+    ) -> None:
+        """Replace final ordinary-loader tensors with validated cache views."""
+        from safetensors import safe_open
+
+        if not plan.post_load_complete:
+            raise RuntimeError("A runtime-cache plan must describe fully post-processed weights")
+        if any(binding.transform is not None for binding in plan.bindings.values()):
+            raise RuntimeError("Runtime-cache bindings cannot carry deferred checkpoint transforms")
+
+        required_names: set[str] = set()
+        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
+            required_names.update(f"{dit_name}.{name}" for name, _ in dit_module.named_parameters())
+            for name, _ in dit_module.named_buffers():
+                parent_path, _, leaf_name = name.rpartition(".")
+                owner = dit_module.get_submodule(parent_path)
+                if leaf_name not in owner._non_persistent_buffers_set:
+                    required_names.add(f"{dit_name}.{name}")
+        if required_names != set(plan.bindings):
+            missing = sorted(required_names - set(plan.bindings))
+            unexpected = sorted(set(plan.bindings) - required_names)
+            raise RuntimeError(
+                "Runtime-cache plan coverage changed after loader validation: "
+                f"missing={missing[:5]}, unexpected={unexpected[:5]}"
+            )
+
+        file_cache: dict[str, Any] = {}
+        replacements: list[tuple[torch.Tensor, torch.Tensor]] = []
+        mapped_sources: dict[str, list[torch.Tensor]] = {}
+        try:
+            # Validate and map every tensor before mutating the pipeline. A
+            # missing/corrupt shard therefore leaves the ordinary tensors intact.
+            for runtime_name, binding in plan.bindings.items():
+                parent_path, _, leaf_name = runtime_name.rpartition(".")
+                try:
+                    parent = pipeline.get_submodule(parent_path)
+                except AttributeError as exc:
+                    raise RuntimeError(f"Runtime-cache target module {parent_path!r} no longer exists") from exc
+
+                target = parent._parameters.get(leaf_name)
+                if target is None:
+                    target = parent._buffers.get(leaf_name)
+                if target is None:
+                    raise RuntimeError(f"Runtime-cache target tensor {runtime_name!r} no longer exists")
+                if target.device.type != "cpu" or target.is_meta:
+                    raise RuntimeError(
+                        f"Runtime-cache target {runtime_name!r} is not an ordinary CPU tensor: {target.device}"
+                    )
+
+                if binding.file_path not in file_cache:
+                    file_cache[binding.file_path] = safe_open(
+                        binding.file_path,
+                        framework="pt",
+                        device="cpu",
+                    )
+                tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
+                if tensor.shape != target.shape or tensor.dtype != target.dtype:
+                    raise RuntimeError(
+                        f"Runtime-cache metadata changed for {runtime_name!r}: "
+                        f"cache={tuple(tensor.shape)}/{tensor.dtype}, "
+                        f"runtime={tuple(target.shape)}/{target.dtype}"
+                    )
+                replacements.append((target, tensor))
+                mapped_sources.setdefault(binding.file_path, []).append(tensor)
+        except Exception:
+            file_cache.clear()
+            raise
+
+        # Preserve Parameter objects and their loader/runtime attributes; only
+        # their storage changes. This also keeps any references held by model
+        # code valid while dropping the private allocator-backed storage.
+        for target, tensor in replacements:
+            set_tensor_storage(target, tensor)
+
+        self._mmap_transforms_by_tensor_id.clear()
+        self._mmap_file_cache = file_cache
+        self._runtime_cache_mapped_sources = mapped_sources
+        logger.info(
+            "Remapped %d final DiT tensors from runtime cache %s",
+            len(replacements),
+            plan.runtime_layout_key[:12] if plan.runtime_layout_key else "unknown",
+        )
+
+    def _try_register_runtime_cache_mmap(self) -> bool:
+        """Register complete cache mappings when the operator budget permits."""
+        pin_limit_gib = self.config.dlo_runtime_cache_pin_limit_gib
+        if pin_limit_gib <= 0:
+            return False
+        if self.device.type != "cuda":
+            logger.warning(
+                "DLO runtime-cache registration is CUDA-only; using bounded host staging on %s",
+                self.device,
+            )
+            return False
+        if not self.config.pin_cpu_memory:
+            logger.warning("DLO runtime-cache registration requires pin_cpu_memory=True; using bounded host staging")
+            return False
+        if not self._runtime_cache_mapped_sources:
+            logger.warning("DLO runtime-cache registration found no mapped sources; using bounded host staging")
+            return False
+
+        max_bytes = int(pin_limit_gib * 1024**3)
+        started = time.perf_counter()
+        try:
+            registration = CudaHostRegistration.create(
+                self._runtime_cache_mapped_sources,
+                max_bytes=max_bytes,
+            )
+        except CudaHostRegistrationCleanupError:
+            # Falling back could unmap a range that CUDA still owns. Abort the
+            # worker so process/context teardown provides the final cleanup.
+            logger.exception("DLO runtime-cache registration rollback failed")
+            raise
+        except CudaHostRegistrationError as exc:
+            logger.warning(
+                "DLO runtime-cache direct H2D unavailable (%s); using bounded host staging",
+                exc,
+            )
+            return False
+
+        self._cuda_host_registration = registration
+        logger.info(
+            "Registered %.2f GiB of shared runtime-cache mappings in %d range(s) for direct H2D in %.3f s",
+            registration.total_bytes / 1024**3,
+            registration.region_count,
+            time.perf_counter() - started,
+        )
+        return True
+
+    def _release_registered_mmap(self) -> bool:
+        registration = self._cuda_host_registration
+        if registration is None:
+            return True
+        errors = registration.close()
+        if errors:
+            logger.error(
+                "DLO runtime-cache host unregistration failed; retaining mmap handles for retry/process exit: %s",
+                errors[:3],
+            )
+            return False
+        self._cuda_host_registration = None
+        logger.info("Unregistered DLO runtime-cache host mappings")
+        return True
+
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.
 
@@ -1443,17 +1655,31 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._using_mmap = self.host_weight_plan is not None
         self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
         if self._using_mmap:
-            if self.host_weight_plan.backing_kind != "checkpoint_mmap":
+            if self.host_weight_plan.backing_kind == "checkpoint_mmap":
+                self._load_weights_via_mmap(
+                    pipeline,
+                    modules,
+                    self.host_weight_plan,
+                )
+            elif self.host_weight_plan.backing_kind == "runtime_cache":
+                if self.config.dlo_use_allgather:
+                    raise ValueError("Runtime-cache backing is no-AllGather only in the first implementation")
+                self._using_rank_local_mmap = True
+                self._load_weights_from_runtime_cache(
+                    pipeline,
+                    modules,
+                    self.host_weight_plan,
+                )
+                # Assignments above dropped the model's private tensor
+                # references. Reclaim those allocations before hook setup so
+                # steady-state PSS reflects the shared mmap backing.
+                self._cleanup_after_loading()
+            else:
                 raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
-            self._load_weights_via_mmap(
-                pipeline,
-                modules,
-                self.host_weight_plan,
-            )
             if self._using_rank_local_mmap:
                 logger.info(
-                    "DLO rank-local mmap storage enabled: checkpoint pages are "
-                    "node-shared; each worker owns only two bounded host staging slots"
+                    "DLO rank-local mmap storage enabled: file-backed pages are node-shared; "
+                    "transfer setup will select registered direct H2D or bounded host staging"
                 )
         else:
             remaining_meta = [
@@ -1625,6 +1851,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
+            if self.enabled and self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+                self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+                self._runtime_cache_mapped_sources.clear()
+                assert self._resident_layer_group is not None
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+                if self._using_registered_mmap:
+                    self._resident_layer_group._cpu_staging_buffers.clear()
             if self._using_mmap and not self.enabled:
                 self._release_mmap_handles()
             return
@@ -1640,9 +1873,18 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
+
+        if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+            self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+            self._runtime_cache_mapped_sources.clear()
+            for hook in all_hooks:
+                hook.registered_mmap = self._using_registered_mmap
+            if self._resident_layer_group is not None:
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+
         unified_cpu_staging = None
         cpu_staging_events = None
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap and not self._using_registered_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1724,6 +1966,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     def _release_mmap_handles(self) -> None:
         """Release safetensors mmap file handles."""
         self._mmap_transforms_by_tensor_id.clear()
+        self._runtime_cache_mapped_sources.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
             del self._mmap_file_cache
@@ -1748,6 +1991,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if self._using_rank_local_mmap:
             current_omni_platform.synchronize()
+        registration_released = self._release_registered_mmap()
 
         for blocks in self._blocks:
             for block in blocks:
@@ -1763,9 +2007,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
-        self._release_mmap_handles()
+        if registration_released:
+            self._release_mmap_handles()
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/registry.py

@@ -136,6 +136,11 @@ _DIFFUSION_MODELS = {
         "pipeline_wan2_2_i2v",
         "WanI2VDMD2Pipeline",
     ),
+    "LingBotWorldCausalDMDPipeline": (
+        "lingbot_world",
+        "pipeline",
+        "LingBotWorldCausalDMDPipeline",
+    ),
     "LongCatImagePipeline": (
         "longcat_image",
         "pipeline_longcat_image",
@@ -353,6 +358,7 @@ _NO_CACHE_ACCELERATION = {
     # Pipelines that do not support cache acceleration (cache_dit / tea_cache).
     "NextStep11Pipeline",
     "AudioXPipeline",
+    "LingBotWorldCausalDMDPipeline",
 }
 
 
@@ -547,6 +553,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanS2VPipeline": "get_wan22_s2v_post_process_func",
     "WanT2VDMD2Pipeline": "get_wan22_post_process_func",
     "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
+    "LingBotWorldCausalDMDPipeline": "get_lingbot_world_post_process_func",
     "LongCatImagePipeline": "get_longcat_image_post_process_func",
     "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_post_process_func",
     "BagelPipeline": "get_bagel_post_process_func",
@@ -607,6 +614,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "WanS2VPipeline": "get_wan22_s2v_pre_process_func",
     "WanT2VDMD2Pipeline": "get_wan22_pre_process_func",
     "WanI2VDMD2Pipeline": "get_wan22_i2v_pre_process_func",
+    "LingBotWorldCausalDMDPipeline": "get_lingbot_world_pre_process_func",
     "OmniGen2Pipeline": "get_omnigen2_pre_process_func",
     "HeliosPipeline": "get_helios_pre_process_func",
     "HeliosPyramidPipeline": "get_helios_pre_process_func",

vllm_omni/diffusion/utils/future_utils.py

@@ -0,0 +1,27 @@
+"""Helpers for resolving ``concurrent.futures.Future`` objects shared across threads."""
+
+import concurrent.futures
+from typing import Any, TypeVar
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_ResultT = TypeVar("_ResultT")
+
+
+def try_set_result(fut: concurrent.futures.Future[_ResultT], result: _ResultT) -> None:
+    # fut may be cancelled concurrently (e.g. asyncio.wait_for timeout) between
+    # the caller's fut.done() check and this call; drop the late result instead
+    # of crashing the pump thread.
+    try:
+        fut.set_result(result)
+    except concurrent.futures.InvalidStateError:
+        logger.debug("Dropping late result for already-resolved/cancelled future")
+
+
+def try_set_exception(fut: concurrent.futures.Future[Any], exc: BaseException) -> None:
+    try:
+        fut.set_exception(exc)
+    except concurrent.futures.InvalidStateError:
+        logger.debug("Dropping late exception for already-resolved/cancelled future")

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -444,6 +444,25 @@ class DiffusionWorker:
         else:
             profiler.stop()
 
+    def _run_ar_diffusion_session_lifecycle(self, method: str, session_id: str) -> bool:
+        """Delegate an optional AR session lifecycle call to the model runner."""
+        if not isinstance(session_id, str) or not session_id.strip():
+            raise ValueError("session_id must be a non-empty string.")
+        assert self.model_runner is not None, "Model runner not initialized"
+        lifecycle_method = getattr(self.model_runner, method, None)
+        if not callable(lifecycle_method):
+            return False
+        lifecycle_method(session_id)
+        return True
+
+    def reset_ar_diffusion_session(self, session_id: str) -> bool:
+        """Reset runner-owned AR state through the collective RPC boundary."""
+        return self._run_ar_diffusion_session_lifecycle("reset_session", session_id)
+
+    def close_ar_diffusion_session(self, session_id: str) -> bool:
+        """Close runner-owned AR state through the collective RPC boundary."""
+        return self._run_ar_diffusion_session_lifecycle("close_session", session_id)
+
     def execute_model(
         self,
         req: OmniDiffusionRequest | list[OmniDiffusionRequest],
@@ -786,11 +805,18 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        if self.model_runner is not None:
-            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
-            if mgr is not None:
-                mgr.shutdown_prefetch()
-        destroy_distributed_env()
+        try:
+            if self.model_runner is not None:
+                mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+                try:
+                    offload_backend = getattr(self.model_runner, "offload_backend", None)
+                    if offload_backend is not None:
+                        offload_backend.disable()
+                finally:
+                    if mgr is not None:
+                        mgr.shutdown_prefetch()
+        finally:
+            destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -543,6 +543,10 @@ class OrchestratorArgs:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_enable_runtime_cache: bool = False
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float = 600.0
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     dlo_resident_layers: int = 0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -982,6 +982,9 @@ class AsyncOmniEngine:
         stage_engine_args = {
             "max_num_seqs": kwargs.get("max_num_seqs") or 1,
             "parallel_config": parallel_config,
+            # Default-stage construction bypasses the structured projection.
+            # Runner selection remains owned by the selected engine/platform.
+            "engine_backend": kwargs.get("engine_backend", "default"),
             "model_class_name": kwargs.get("model_class_name", None),
             "task_type": kwargs.get("task_type", None),
             "model_config": kwargs.get("model_config", None),
@@ -997,6 +1000,10 @@ class AsyncOmniEngine:
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
+            "dlo_enable_runtime_cache": kwargs.get("dlo_enable_runtime_cache", False),
+            "dlo_runtime_cache_dir": kwargs.get("dlo_runtime_cache_dir", None),
+            "dlo_runtime_cache_lock_timeout": kwargs.get("dlo_runtime_cache_lock_timeout", 600.0),
+            "dlo_runtime_cache_pin_limit_gib": kwargs.get("dlo_runtime_cache_pin_limit_gib", 0.0),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (
@@ -1115,15 +1122,13 @@ class AsyncOmniEngine:
     ) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
 
+        for legacy_arg in ("stage_configs_path", "stage_configs"):
+            if legacy_arg in kwargs:
+                raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.")
+
         deploy_config_path = kwargs.pop("deploy_config", None)
         strategy_config_path = kwargs.pop("strategy_config", None)
         stage_overrides_json = kwargs.pop("stage_overrides", None)
-        explicit_stage_configs = kwargs.pop("stage_configs", None)
-        if explicit_stage_configs is not None:
-            logger.warning(
-                "`stage_configs` is not part of the public API. "
-                "Ignoring it and resolving stages from deploy_config/model factory."
-            )
 
         # Parse --stage-overrides JSON string if provided
         stage_overrides = parse_stage_overrides(stage_overrides_json)

vllm_omni/entrypoints/cli/serve.py

@@ -722,6 +722,39 @@ class OmniServeCommand(CLISubcommand):
                 "independently (including existing TP shards)."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-enable-runtime-cache",
+            action="store_true",
+            help=(
+                "After ordinary loading, publish or join a node-local mmap cache "
+                "for final DiT weights. This first version requires "
+                "--dlo-no-use-allgather."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-dir",
+            type=str,
+            default=None,
+            help=(
+                "Shared local-disk directory for DLO runtime weights. Defaults "
+                "to ~/.cache/vllm-omni/dlo-runtime-weights."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-lock-timeout",
+            type=float,
+            default=600.0,
+            help="Seconds to wait for another process publishing the same runtime layout.",
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-pin-limit-gib",
+            type=float,
+            default=0.0,
+            help=(
+                "Maximum GiB of read-only runtime-cache mappings each CUDA worker "
+                "may register for direct H2D. Zero keeps bounded pinned staging."
+            ),
+        )
         omni_config_group.add_argument(
             "--dlo-resident-layers",
             type=int,

vllm_omni/entrypoints/utils.py

@@ -19,7 +19,7 @@ from vllm_omni.config.config_factory import (
 )
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import _DEPLOY_DIR
-from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
+from vllm_omni.config.yaml_util import create_config, load_yaml_config
 from vllm_omni.diffusion.utils.hf_utils import (
     _looks_like_dreamzero,
     get_diffusion_model_index,
@@ -369,11 +369,12 @@ def load_stage_configs_from_model(
 ) -> tuple[list, str | None]:
     """Load stage configurations from model's default config file.
 
-    For models registered in the pipeline registry (new path), uses
+    For models registered in the pipeline registry, uses
     ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges
     PipelineConfig + DeployConfig + CLI overrides.
 
-    For other models (legacy path), loads stage configs from YAML.
+    Models that cannot be resolved through the registry return no stages so the
+    caller can apply its default stage configuration, when available.
 
     Args:
         model: Model name or path (used to determine model_type)
@@ -383,7 +384,7 @@ def load_stage_configs_from_model(
         stage_overrides: Per-stage overrides from --stage-overrides.
         strategy_config_path: Optional path to a composable-parallel
             ``strategy.yaml`` whose derived sizing is overlaid onto the
-            registry-merged stages (opt-in; ignored on the legacy YAML path).
+            registry-merged stages.
 
     Returns:
         ``(stage_configs, omni_lb_policy)``: the list of stage configuration
@@ -427,69 +428,17 @@ def load_stage_configs_from_model(
         # Convert StageConfig objects to OmegaConf for backward compat
         return [stage.to_omegaconf() for stage in stages], omni_lb_policy
 
-    # Legacy fallback: load from YAML. A composable-parallel strategy cannot be
-    # applied here (it overlays onto registry-merged stages), so warn rather than
-    # silently dropping the operator's --strategy-config.
+    strategy_note = ""
     if strategy_config_path is not None:
-        logger.warning(
-            "--strategy-config (%s) was provided but model %r resolves via the "
-            "legacy stage_configs YAML path, which does not support "
-            "composable-parallel strategies; the strategy is ignored. Use a "
-            "registry-based model to apply it.",
-            strategy_config_path,
-            model,
-        )
-    stage_config_path = resolve_model_config_path(model)
-    if stage_config_path is None:
-        return [], None
-    stage_configs = load_stage_configs_from_yaml(
-        config_path=stage_config_path,
-        base_engine_args=base_engine_args,
-        prefer_stage_engine_args=True,
+        strategy_note = f" Strategy config {strategy_config_path!r} was not applied."
+    logger.warning(
+        "No registered PipelineConfig resolved for model %r. Legacy `stage_args` "
+        "YAMLs are no longer supported; register the pipeline and provide deployment "
+        "overrides through `deploy_config`.%s",
+        model,
+        strategy_note,
     )
-    return stage_configs, None
-
-
-def load_stage_configs_from_yaml(
-    config_path: str,
-    base_engine_args: dict | None = None,
-    prefer_stage_engine_args: bool = True,
-) -> list:
-    """Load stage configurations from a YAML file (legacy OmegaConf path).
-
-    TODO(@lishunyang12): remove once all models use PipelineConfig + DeployConfig.
-
-    Args:
-        config_path: Path to the YAML configuration file
-        base_engine_args: Engine args supplied by the caller.
-        prefer_stage_engine_args: When True, YAML stage args override caller
-            engine args. When False, caller engine args override YAML defaults.
-
-    Returns:
-        List of stage configuration dictionaries from the file's stage_args
-    """
-    if base_engine_args is None:
-        base_engine_args = {}
-    config_data = load_yaml_config(config_path)
-    stage_args = config_data.stage_args
-    global_async_chunk = config_data.get("async_chunk", False)
-    # Convert any nested dataclass objects to dicts before creating DictConfig
-    base_engine_args = _convert_dataclasses_to_dict(base_engine_args)
-    base_engine_args = create_config(base_engine_args)
-    for stage_arg in stage_args:
-        base_engine_args_tmp = base_engine_args.copy()
-        # Update base_engine_args with stage-specific engine_args if they exist
-        if hasattr(stage_arg, "engine_args") and stage_arg.engine_args is not None:
-            if prefer_stage_engine_args:
-                merged_engine_args = merge_configs(base_engine_args_tmp, stage_arg.engine_args)
-            else:
-                merged_engine_args = merge_configs(stage_arg.engine_args, base_engine_args_tmp)
-            base_engine_args_tmp = create_config(merged_engine_args)
-        stage_type = getattr(stage_arg, "stage_type", "llm")
-        if hasattr(stage_arg, "runtime") and stage_arg.runtime is not None and stage_type != "diffusion":
-            base_engine_args_tmp.async_chunk = global_async_chunk
-        stage_arg.engine_args = base_engine_args_tmp
-    return stage_args
+    return [], None
 
 
 def filter_stages(

vllm_omni/experimental/ar_diffusion/__init__.py

@@ -13,12 +13,32 @@ from vllm_omni.experimental.ar_diffusion.capability import (
     SupportsARDiffusionPipeline,
     SupportsARDiffusionWarmup,
 )
+from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
 from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+from vllm_omni.experimental.ar_diffusion.session import (
+    ARDiffusionSession,
+    ARDiffusionSessionManager,
+    ARDiffusionWorkerLifecycle,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    AR_DIFFUSION_TICK_KEY,
+    ARDiffusionChunkMetadata,
+    ARDiffusionControlInput,
+    ARDiffusionTickRequest,
+)
 
 __all__ = [
     "ARDiffusionKVBranchSpec",
     "ARDiffusionCrossAttentionKVSpec",
     "ARDiffusionEngine",
+    "ARDiffusionOmniTickConsumer",
+    "ARDiffusionSession",
+    "ARDiffusionSessionManager",
+    "ARDiffusionWorkerLifecycle",
+    "AR_DIFFUSION_TICK_KEY",
+    "ARDiffusionChunkMetadata",
+    "ARDiffusionControlInput",
+    "ARDiffusionTickRequest",
     "ARDiffusionKVCacheSpec",
     "SupportsARDiffusionPipeline",
     "SupportsARDiffusionWarmup",

vllm_omni/experimental/ar_diffusion/capability.py

@@ -58,6 +58,9 @@ class ARDiffusionKVCacheSpec:
 
     ``max_scratch_tokens_per_branch`` is the maximum non-video KV (for example,
     action/state registers) that must coexist with an uncommitted video block.
+    ``model_owned_state_bytes_per_session`` is a conservative upper bound for
+    persistent per-session CUDA state outside runner-owned self/cross-attention
+    KV and scratch pools.
     ``session_capacity`` is the pipeline's upper bound; a runner may retain
     fewer sessions when its current memory and execution policy is stricter.
 
@@ -79,6 +82,7 @@ class ARDiffusionKVCacheSpec:
     cross_attention: tuple[ARDiffusionCrossAttentionKVSpec, ...] = ()
     max_model_len: int = 1 << 20
     max_scratch_tokens_per_branch: int = 0
+    model_owned_state_bytes_per_session: int = 0
 
     def __post_init__(self) -> None:
         positive_fields = {
@@ -101,6 +105,11 @@ class ARDiffusionKVCacheSpec:
                 "AR-Diffusion max_scratch_tokens_per_branch must be non-negative, "
                 f"got {self.max_scratch_tokens_per_branch}"
             )
+        if self.model_owned_state_bytes_per_session < 0:
+            raise ValueError(
+                "AR-Diffusion model_owned_state_bytes_per_session must be non-negative, "
+                f"got {self.model_owned_state_bytes_per_session}"
+            )
         if not self.kv_branches:
             raise ValueError("AR-Diffusion requires at least one KV branch")
         kv_branch_names = [kv_branch.name for kv_branch in self.kv_branches]

vllm_omni/experimental/ar_diffusion/consumer.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AsyncOmni consumer for typed realtime AR-Diffusion ticks."""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
+from typing import Any, Protocol
+
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionChunkMetadata,
+    ARDiffusionTickRequest,
+)
+from vllm_omni.inputs.data import (
+    OmniDiffusionSamplingParams,
+    OmniPromptType,
+    OmniSamplingParams,
+)
+from vllm_omni.outputs import OmniRequestOutput
+
+AR_DIFFUSION_OUTPUT_METADATA_KEY = "ar_diffusion"
+
+
+class ARDiffusionGenerateClient(Protocol):
+    """Internal subset of AsyncOmni used to execute one tick."""
+
+    def generate(
+        self,
+        prompt: OmniPromptType,
+        *,
+        request_id: str,
+        sampling_params_list: Sequence[OmniSamplingParams],
+        output_modalities: list[str] | None = None,
+    ) -> AsyncIterator[OmniRequestOutput]: ...
+
+
+def _metadata_mapping(output: OmniRequestOutput) -> Mapping[str, Any]:
+    multimodal_output = output.multimodal_output
+    if not isinstance(multimodal_output, Mapping):
+        raise ValueError("AR-Diffusion output must contain a multimodal output mapping.")
+    metadata = multimodal_output.get("metadata")
+    if not isinstance(metadata, Mapping):
+        raise ValueError("AR-Diffusion output must contain multimodal_output.metadata.")
+    value = metadata.get(AR_DIFFUSION_OUTPUT_METADATA_KEY)
+    if not isinstance(value, Mapping):
+        raise ValueError(f"AR-Diffusion output metadata must contain {AR_DIFFUSION_OUTPUT_METADATA_KEY!r}.")
+    return value
+
+
+def _parse_chunk_metadata(output: OmniRequestOutput) -> ARDiffusionChunkMetadata:
+    value = _metadata_mapping(output)
+    session_id = value.get("session_id")
+    request_id = value.get("request_id")
+    chunk_index = value.get("chunk_index")
+    applied_event_ids = value.get("applied_event_ids", ())
+    if not isinstance(session_id, str) or not session_id.strip():
+        raise ValueError("AR-Diffusion output metadata session_id must be a non-empty string.")
+    if not isinstance(request_id, str) or not request_id.strip():
+        raise ValueError("AR-Diffusion output metadata request_id must be a non-empty string.")
+    if isinstance(chunk_index, bool) or not isinstance(chunk_index, int) or chunk_index < 0:
+        raise ValueError("AR-Diffusion output metadata chunk_index must be a non-negative integer.")
+    if not isinstance(applied_event_ids, Sequence) or isinstance(applied_event_ids, (str, bytes, bytearray)):
+        raise ValueError("AR-Diffusion output metadata applied_event_ids must be a sequence.")
+    event_ids = tuple(applied_event_ids)
+    if any(isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0 for event_id in event_ids):
+        raise ValueError("AR-Diffusion output metadata applied_event_ids must contain non-negative integers.")
+    if event_ids != tuple(sorted(set(event_ids))):
+        raise ValueError("AR-Diffusion output metadata applied_event_ids must be unique and strictly increasing.")
+    return ARDiffusionChunkMetadata(
+        session_id=session_id,
+        request_id=request_id,
+        chunk_index=chunk_index,
+        applied_event_ids=event_ids,
+    )
+
+
+class ARDiffusionOmniTickConsumer:
+    """Converts a typed tick into one AsyncOmni diffusion request.
+
+    ``prompt_provider`` packages session initialization data such as an initial
+    image into the normal Omni prompt. Control ``track/schema/data`` remains in
+    the typed tick and is not interpreted by this class.
+    """
+
+    def __init__(
+        self,
+        engine_client: ARDiffusionGenerateClient,
+        *,
+        prompt_provider: Callable[[ARDiffusionTickRequest], OmniPromptType],
+        sampling_params_list: Sequence[OmniSamplingParams],
+        diffusion_stage_id: int,
+        output_modalities: Sequence[str] | None = None,
+    ) -> None:
+        templates = tuple(sampling_params_list)
+        if not templates:
+            raise ValueError("sampling_params_list must not be empty.")
+        if (
+            isinstance(diffusion_stage_id, bool)
+            or not isinstance(diffusion_stage_id, int)
+            or not 0 <= diffusion_stage_id < len(templates)
+        ):
+            raise ValueError("diffusion_stage_id must index sampling_params_list.")
+        if not isinstance(templates[diffusion_stage_id], OmniDiffusionSamplingParams):
+            raise TypeError("The selected diffusion stage must use OmniDiffusionSamplingParams.")
+        stage_pools = getattr(getattr(engine_client, "engine", None), "stage_pools", None)
+        if stage_pools is None or not 0 <= diffusion_stage_id < len(stage_pools):
+            raise ValueError(
+                "AR-Diffusion realtime execution requires an initialized AsyncOmni "
+                "engine with a matching diffusion stage pool."
+            )
+        replica_count = getattr(stage_pools[diffusion_stage_id], "num_replicas", None)
+        if isinstance(replica_count, bool) or not isinstance(replica_count, int) or replica_count < 1:
+            raise ValueError("The selected AR-Diffusion stage must expose a positive num_replicas.")
+        if replica_count != 1:
+            raise ValueError(
+                "AR-Diffusion realtime sessions currently require exactly one "
+                "replica for the selected diffusion stage; session-affine "
+                f"routing is not implemented (got num_replicas={replica_count})."
+            )
+        self._engine_client = engine_client
+        self._prompt_provider = prompt_provider
+        self._sampling_params_templates = templates
+        self._diffusion_stage_id = diffusion_stage_id
+        self._output_modalities = list(output_modalities) if output_modalities is not None else None
+
+    def _sampling_params_for_tick(
+        self,
+        tick: ARDiffusionTickRequest,
+    ) -> list[OmniSamplingParams]:
+        params = [copy.copy(template) for template in self._sampling_params_templates]
+        diffusion_params = params[self._diffusion_stage_id]
+        assert isinstance(diffusion_params, OmniDiffusionSamplingParams)
+        diffusion_params.extra_args = {
+            **(diffusion_params.extra_args or {}),
+            **tick.to_extra_args(),
+        }
+        return params
+
+    async def execute_tick(self, tick: ARDiffusionTickRequest) -> OmniRequestOutput:
+        prompt = self._prompt_provider(tick)
+        final_output: OmniRequestOutput | None = None
+        async for output in self._engine_client.generate(
+            prompt,
+            request_id=tick.request_id,
+            sampling_params_list=self._sampling_params_for_tick(tick),
+            output_modalities=self._output_modalities,
+        ):
+            if output.stage_id != self._diffusion_stage_id:
+                continue
+            if output.finished:
+                final_output = output
+
+        if final_output is None:
+            raise RuntimeError(
+                f"AR-Diffusion tick {tick.request_id!r} completed without a final "
+                f"output from stage {self._diffusion_stage_id}."
+            )
+        if final_output.error:
+            raise RuntimeError(f"AR-Diffusion tick {tick.request_id!r} failed: {final_output.error}")
+        if final_output.request_id != tick.request_id:
+            raise ValueError("AR-Diffusion output request_id does not match the submitted tick request_id.")
+        # Parse here so malformed/missing standard metadata fails at the engine
+        # boundary, before the session attempts to commit accepted events.
+        self.chunk_metadata(final_output)
+        return final_output
+
+    def chunk_metadata(self, output: OmniRequestOutput) -> ARDiffusionChunkMetadata:
+        return _parse_chunk_metadata(output)

vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py

@@ -32,7 +32,6 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
     ARDiffusionPagedLayerInputs,
     ar_diffusion_paged_attention,
     paged_write_attn,
-    set_current_paged_kv_cache,
 )
 
 __all__ = [
@@ -53,5 +52,4 @@ __all__ = [
     "paged_write_attn",
     "pool_write_chunk",
     "resident_block_ids",
-    "set_current_paged_kv_cache",
 ]

vllm_omni/experimental/ar_diffusion/kv_cache/config.py

@@ -23,7 +23,8 @@ class ARDiffusionKVConfig:
     sink_chunks: int = 0
     # Boundary reset vs. sliding replacement.
     reset_at_boundary: bool = False
-    # Fraction of free device memory budgeted for the AR-Diffusion KV pool.
+    # Fraction of free device memory used to admit additional resident
+    # sessions. One session is admitted whenever it fits actual free memory.
     gpu_memory_fraction: float = 0.1
     # When CUDA graph / torch.compile is on (not enforce_eager), pre-capture the
     # DiT graphs for every window-fill shape at load time via a synthetic rollout,

vllm_omni/experimental/ar_diffusion/kv_cache/manager.py

@@ -174,6 +174,7 @@ class ARDiffusionKVCache:
         device: torch.device | None = None,
         frames_per_block: int = 1,
         max_scratch_tokens_per_branch: int = 0,
+        model_owned_state_bytes_per_session: int = 0,
     ) -> None:
         if not config.enable:
             raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig")
@@ -197,13 +198,17 @@ class ARDiffusionKVCache:
 
         self.config = config
         self.kv_branches = kv_branches
-        self.session_capacity = session_capacity
+        self.requested_session_capacity = session_capacity
         self._kv_branch_local_indices = {kv_branch.name: kv_branch.local_index for kv_branch in kv_branches}
         self.num_local_kv_branches = max(local_indices) + 1
         if frames_per_block <= 0:
             raise ValueError(f"frames_per_block must be positive, got {frames_per_block}")
         if max_scratch_tokens_per_branch < 0:
             raise ValueError(f"max_scratch_tokens_per_branch must be non-negative, got {max_scratch_tokens_per_branch}")
+        if model_owned_state_bytes_per_session < 0:
+            raise ValueError(
+                f"model_owned_state_bytes_per_session must be non-negative, got {model_owned_state_bytes_per_session}"
+            )
         self.frames_per_block = int(frames_per_block)
         self.block_size = block_size
         self.num_layers = num_layers
@@ -217,28 +222,11 @@ class ARDiffusionKVCache:
         self.device = device or torch.device("cpu")
         self._allocate_tensors = device is not None
         self._adapters: dict[str, ARDiffusionRequestAdapter] = {}
-
-        # Named cross-attention K/V is allocated lazily per session and released
-        # with that session. Reserve its worst-case capacity when sizing the
-        # paged self-attention pool so the two stores share one memory budget.
         self._cross_sessions: dict[
             str,
             dict[str, dict[str, tuple[list[torch.Tensor], list[torch.Tensor]]]],
         ] = {}
 
-        def _cross_pool_bytes(length: int) -> int:
-            return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers)
-
-        cross_bytes_per_session = sum(_cross_pool_bytes(length) for length in self.cross_attention_lengths.values())
-        cross_total_bytes = cross_bytes_per_session * session_capacity
-        if cross_total_bytes:
-            _log.info(
-                "AR-Diffusion cross-attn reservation: %.1f MiB/session × %d sessions = %.1f MiB",
-                cross_bytes_per_session / (1024 * 1024),
-                session_capacity,
-                cross_total_bytes / (1024 * 1024),
-            )
-
         self.spec = ChunkWindowSpec(
             block_size=block_size,
             num_kv_heads=num_kv_heads,
@@ -250,41 +238,7 @@ class ARDiffusionKVCache:
             sink_chunks=config.sink_chunks,
             reset_at_boundary=config.reset_at_boundary,
         )
-        # Each pool block spans all layers' K/V, so size against the per-layer
-        # page size times the layer count.
-        # Size the self-attn pool against memory reserved for the maximum number
-        # of lazily allocated cross-attention sessions.
-        num_blocks = compute_num_blocks(
-            max(0, available_bytes - cross_total_bytes),
-            config.gpu_memory_fraction,
-            self.spec.page_size_bytes * num_layers,
-        )
-        # Floor: one forward needs the resident window plus the in-flight chunk
-        # (frames_per_block frame-blocks) for every KV branch THIS rank runs,
-        # with a little eviction-transient headroom. The memory-fraction
-        # heuristic can under-size this once block_size grows — e.g. frame-granular
-        # paging at the true frame_seqlen makes each block larger and the pool
-        # fewer-blocks — so guarantee the minimum the rollout cannot run without,
-        # otherwise allocate_chunk hits an exhausted pool mid-forward.
-        resident_blocks = config.sink_chunks + config.window_chunks
-        min_blocks = self.num_local_kv_branches * (resident_blocks + self.frames_per_block) + 2
-        if num_blocks < min_blocks:
-            _log.warning(
-                "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
-                "floor (%d local KV branch(es) x (sink_chunks=%d + window_chunks=%d "
-                "+ frames_per_block=%d) + 2 headroom)",
-                num_blocks,
-                min_blocks,
-                self.num_local_kv_branches,
-                config.sink_chunks,
-                config.window_chunks,
-                self.frames_per_block,
-            )
-            num_blocks = min_blocks
-        layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
-        self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
-        self.managed_num_blocks = num_blocks
-        self.num_blocks = num_blocks
+
         # Scratch blocks are outside KVCacheManager ownership. A non-committing
         # forward needs one block per current frame plus space for any
         # model-declared action/state tokens that coexist with video KV.
@@ -297,6 +251,101 @@ class ARDiffusionKVCache:
         scratch_per_kv_branch = max(minimum_scratch_blocks, override_blocks)
         self.scratch_blocks_per_kv_branch = scratch_per_kv_branch
         self.scratch_num_blocks = self.num_local_kv_branches * scratch_per_kv_branch
+
+        # The self-attention pool, scratch pool, and lazily materialized
+        # cross-attention caches share one hard memory budget. Select the largest
+        # feasible resident-session count rather than raising a block-count floor
+        # past that budget. All resident sessions need a complete sink + window;
+        # only one request can be in flight, so frames_per_block is counted once.
+        page_size_bytes = self.spec.page_size_bytes * num_layers
+        self.available_memory_bytes = available_bytes
+        self.configured_memory_budget_bytes = int(available_bytes * config.gpu_memory_fraction)
+        # Reuse the public helper's validation for the memory fraction/page size.
+        compute_num_blocks(available_bytes, config.gpu_memory_fraction, page_size_bytes)
+        self.scratch_reserved_bytes = self.scratch_num_blocks * page_size_bytes
+        self.model_owned_state_bytes_per_session = model_owned_state_bytes_per_session
+
+        def _cross_pool_bytes(length: int) -> int:
+            return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers)
+
+        self.cross_attention_bytes_per_session = sum(
+            _cross_pool_bytes(length) for length in self.cross_attention_lengths.values()
+        )
+
+        def _required_managed_blocks(capacity: int) -> int:
+            resident_per_session = config.sink_chunks + config.window_chunks
+            return self.num_local_kv_branches * (capacity * resident_per_session + self.frames_per_block) + 2
+
+        def _required_bytes(capacity: int) -> int:
+            return (
+                self.scratch_reserved_bytes
+                + capacity * self.cross_attention_bytes_per_session
+                + capacity * self.model_owned_state_bytes_per_session
+                + _required_managed_blocks(capacity) * page_size_bytes
+            )
+
+        one_session_bytes = _required_bytes(1)
+        if one_session_bytes > available_bytes:
+            raise ValueError(
+                "AR-Diffusion available device memory cannot fit one session: "
+                f"available={available_bytes} bytes, required={one_session_bytes} bytes "
+                "(managed self-attention + cross-attention + scratch + model-owned state)."
+            )
+        self.memory_budget_bytes = max(self.configured_memory_budget_bytes, one_session_bytes)
+        if self.memory_budget_bytes > self.configured_memory_budget_bytes:
+            _log.warning(
+                "AR-Diffusion raised the configured memory budget from %d to %d bytes "
+                "to admit one session that fits actual free device memory",
+                self.configured_memory_budget_bytes,
+                self.memory_budget_bytes,
+            )
+
+        effective_capacity = 0
+        required_managed_blocks = 0
+        for candidate in range(session_capacity, 0, -1):
+            candidate_managed_blocks = _required_managed_blocks(candidate)
+            if _required_bytes(candidate) <= self.memory_budget_bytes:
+                effective_capacity = candidate
+                required_managed_blocks = candidate_managed_blocks
+                break
+        assert effective_capacity > 0
+
+        self.session_capacity = effective_capacity
+        self.cross_attention_reserved_bytes = self.cross_attention_bytes_per_session * effective_capacity
+        self.model_owned_state_reserved_bytes = self.model_owned_state_bytes_per_session * effective_capacity
+        self_attn_budget_bytes = (
+            self.memory_budget_bytes
+            - self.scratch_reserved_bytes
+            - self.cross_attention_reserved_bytes
+            - self.model_owned_state_reserved_bytes
+        )
+        num_blocks = self_attn_budget_bytes // page_size_bytes
+        assert num_blocks >= required_managed_blocks
+        if effective_capacity < session_capacity:
+            _log.warning(
+                "AR-Diffusion resident session capacity reduced from %d to %d by the KV memory budget",
+                session_capacity,
+                effective_capacity,
+            )
+        if self.cross_attention_reserved_bytes:
+            _log.info(
+                "AR-Diffusion cross-attn reservation: %.1f MiB/session × %d sessions = %.1f MiB",
+                self.cross_attention_bytes_per_session / (1024 * 1024),
+                effective_capacity,
+                self.cross_attention_reserved_bytes / (1024 * 1024),
+            )
+        if self.model_owned_state_reserved_bytes:
+            _log.info(
+                "AR-Diffusion model-owned state reservation: %.1f MiB/session × %d sessions = %.1f MiB",
+                self.model_owned_state_bytes_per_session / (1024 * 1024),
+                effective_capacity,
+                self.model_owned_state_reserved_bytes / (1024 * 1024),
+            )
+
+        layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
+        self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
+        self.managed_num_blocks = num_blocks
+        self.num_blocks = num_blocks
         self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks
         self.null_block_id = self.manager.block_pool.null_block.block_id
 
@@ -561,6 +610,11 @@ class ARDiffusionKVCache:
         """
         _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
         adapter.on_chunk_committed()
+        self.manager.remove_skipped_blocks(
+            adapter.request_id,
+            adapter.num_computed_tokens,
+            num_prompt_tokens=adapter.num_prompt_tokens,
+        )
         _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
 
     # -- pool-backed K/V access --------------------------------------------

vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py

@@ -10,19 +10,6 @@ import torch
 
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
 
-# Set by ARDiffusionPagedForwardContext.prepare() before each KV branch forward and
-# read by the fused write+attend custom op below. The pools are process-lifetime
-# allocations owned by one ARDiffusionKVCache per worker process (CFG-parallel /
-# TP ranks are separate processes), mirroring vLLM's forward-context pattern for
-# keeping multi-GiB cache tensors out of the compiled graph's inputs.
-_CURRENT_PAGED_KV_CACHE: Any = None
-
-
-def set_current_paged_kv_cache(kv_cache: Any) -> None:
-    global _CURRENT_PAGED_KV_CACHE
-    _CURRENT_PAGED_KV_CACHE = kv_cache
-
-
 _LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
 
 
@@ -50,6 +37,9 @@ class ARDiffusionPagedLayerInputs(NamedTuple):
     """
 
     layer_idx: torch.Tensor
+    key_pool: torch.Tensor
+    value_pool: torch.Tensor
+    block_size: int
     seq_len: int
     video_slots: torch.Tensor
     action_slots: torch.Tensor
@@ -228,7 +218,6 @@ class ARDiffusionPagedForwardContext:
         ) = self.build_block_table(action_len=action_len, query_len=query_len, device=device)
         if self.action_slot_mapping is None:
             self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
-        set_current_paged_kv_cache(self.kv_cache)
         self._prepared = True
 
     def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
@@ -236,6 +225,9 @@ class ARDiffusionPagedForwardContext:
             raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()")
         return ARDiffusionPagedLayerInputs(
             layer_idx=_layer_idx_tensor(layer_idx),
+            key_pool=self.kv_cache._k_pools[layer_idx],
+            value_pool=self.kv_cache._v_pools[layer_idx],
+            block_size=int(self.kv_cache.block_size),
             seq_len=int(self.seq_len),
             video_slots=self.current_video_slot_mapping,
             action_slots=self.action_slot_mapping,
@@ -476,17 +468,18 @@ def ar_diffusion_paged_attention(
 # One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats
 # it as a single graph node (no eager island, no graph breaks), and the K/V slot
 # writes happen inside the op so write→read ordering with the block-table kernel
-# is internal — no hidden-mutation ordering hazards against a separate reader op
-# and no multi-GiB pool tensors as graph inputs (the flat and paged pool views
-# alias the same storage, which AOTAutograd handles poorly as inputs). Pattern
-# follows sage_attn3.py in-repo and vLLM's own unified attention ops.
+# is internal. The flat pools are explicit mutable inputs: Inductor/CUDA Graph
+# must track their storage lifetime instead of observing an undeclared mutation
+# through a process-global registry.
 def _paged_write_attn_impl(
     query: torch.Tensor,
     k_curr: torch.Tensor,
     v_curr: torch.Tensor,
     k_act: torch.Tensor | None,
     v_act: torch.Tensor | None,
-    layer_idx: torch.Tensor,
+    key_pool: torch.Tensor,
+    value_pool: torch.Tensor,
+    block_size: int,
     video_slots: torch.Tensor,
     action_slots: torch.Tensor,
     block_table: torch.Tensor,
@@ -496,21 +489,17 @@ def _paged_write_attn_impl(
     max_seq_len: int,
     softmax_scale: float,
 ) -> torch.Tensor:
-    kv = _CURRENT_PAGED_KV_CACHE
-    if kv is None:
-        raise RuntimeError("ar_diffusion_paged_write_attn called before prepare() set the KV pool registry")
-    layer_idx = int(layer_idx)
-    k_pool = kv._k_pools[layer_idx]
-    v_pool = kv._v_pools[layer_idx]
-    k_pool[video_slots] = k_curr.to(k_pool.dtype)
-    v_pool[video_slots] = v_curr.to(v_pool.dtype)
+    key_pool[video_slots] = k_curr.to(key_pool.dtype)
+    value_pool[video_slots] = v_curr.to(value_pool.dtype)
     if k_act is not None and v_act is not None and k_act.shape[0] > 0:
-        k_pool[action_slots] = k_act.to(k_pool.dtype)
-        v_pool[action_slots] = v_act.to(v_pool.dtype)
+        key_pool[action_slots] = k_act.to(key_pool.dtype)
+        value_pool[action_slots] = v_act.to(value_pool.dtype)
+    key_cache = key_pool.unflatten(0, (-1, block_size))
+    value_cache = value_pool.unflatten(0, (-1, block_size))
     return ar_diffusion_paged_attention(
         query,
-        kv.key_cache(layer_idx),
-        kv.value_cache(layer_idx),
+        key_cache,
+        value_cache,
         block_table=block_table,
         query_start_loc=query_start_loc,
         seq_lens=seq_lens,
@@ -525,14 +514,19 @@ def _paged_write_attn_impl(
 # the module from sys.modules (same as sage_attn3.py).
 if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
 
-    @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=())
+    @torch.library.custom_op(
+        "vllm_omni::ar_diffusion_paged_write_attn",
+        mutates_args=("key_pool", "value_pool"),
+    )
     def _paged_write_attn_op(
         query: torch.Tensor,
         k_curr: torch.Tensor,
         v_curr: torch.Tensor,
         k_act: torch.Tensor | None,
         v_act: torch.Tensor | None,
-        layer_idx: torch.Tensor,
+        key_pool: torch.Tensor,
+        value_pool: torch.Tensor,
+        block_size: int,
         video_slots: torch.Tensor,
         action_slots: torch.Tensor,
         block_table: torch.Tensor,
@@ -548,7 +542,9 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
             v_curr,
             k_act,
             v_act,
-            layer_idx,
+            key_pool,
+            value_pool,
+            block_size,
             video_slots,
             action_slots,
             block_table,
@@ -566,7 +562,9 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
         v_curr,
         k_act,
         v_act,
-        layer_idx,
+        key_pool,
+        value_pool,
+        block_size,
         video_slots,
         action_slots,
         block_table,
@@ -589,7 +587,9 @@ def paged_write_attn(
         v_curr,
         k_act,
         v_act,
-        inputs.layer_idx,
+        inputs.key_pool,
+        inputs.value_pool,
+        inputs.block_size,
         inputs.video_slots,
         inputs.action_slots,
         inputs.block_table,

vllm_omni/experimental/ar_diffusion/kv_cache/state.py

@@ -130,7 +130,7 @@ class ARDiffusionKVState:
         if ctx.commit_current and ctx._allocated_video:
             n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size
             for _ in range(n_chunks):
-                ctx.adapter.on_chunk_committed()
+                self.kv_cache.commit_chunk(ctx.adapter)
             self._committed[kv_branch] += ctx.seq_len
             _log.debug(
                 "AR-Diffusion COMMIT [%s] new_tokens=%d chunks=%d resident=%d/%d",
@@ -173,6 +173,17 @@ class ARDiffusionKVState:
             for i in range(self.num_layers)
         ]
 
+    def clear_cross_attention(self) -> None:
+        """Invalidate all named cross-attention caches without dropping self-KV.
+
+        Prompt changes may preserve the autoregressive world history while
+        requiring fresh text projections. This operation deliberately leaves
+        branch adapters and their resident self-attention blocks untouched.
+        """
+        if self._closed:
+            raise RuntimeError(f"AR-Diffusion session {self.session_id!r} is closed")
+        self.kv_cache.release_cross_attention(self.session_id)
+
     def close(self) -> None:
         """Release all self- and cross-attention storage owned by this session."""
         if self._closed:

vllm_omni/experimental/ar_diffusion/runner.py

@@ -21,6 +21,7 @@ from vllm_omni.experimental.ar_diffusion.capability import (
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
 from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
 
 logger = init_logger(__name__)
 
@@ -51,7 +52,6 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
     """
 
     _WARMUP_SID = "__ardiffusion_warmup__"
-    _MAX_RESIDENT_SESSIONS = 1
 
     def __init__(self, vllm_config: object, od_config: OmniDiffusionConfig, device: torch.device) -> None:
         super().__init__(vllm_config, od_config, device)
@@ -123,7 +123,6 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
         self.ar_diffusion_kv_config = config
         self._ar_diffusion_capability = capability
         self._ar_diffusion_kv_cache_spec = spec
-        self._session_capacity = min(spec.session_capacity, self._MAX_RESIDENT_SESSIONS)
         self.kv_cache = ARDiffusionKVCache(
             config,
             num_layers=spec.num_layers,
@@ -134,12 +133,14 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
             max_model_len=spec.max_model_len,
             available_bytes=self._available_memory_bytes() if available_bytes is None else available_bytes,
             kv_branches=spec.kv_branches,
-            session_capacity=self._session_capacity,
+            session_capacity=spec.session_capacity,
             cross_attention_lengths=spec.cross_attention_lengths,
             frames_per_block=spec.frames_per_block,
             max_scratch_tokens_per_branch=spec.max_scratch_tokens_per_branch,
+            model_owned_state_bytes_per_session=spec.model_owned_state_bytes_per_session,
             device=self.device,
         )
+        self._session_capacity = self.kv_cache.session_capacity
         logger.info(
             "AR-Diffusion KV cache: blocks=%d layers=%d local_kv_heads=%d head_size=%d "
             "tokens/frame=%d frames/block=%d window=%d sink=%d kv_branches=%s cross=%s "
@@ -230,9 +231,14 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
         return state
 
     @staticmethod
-    def _request_session(req: OmniDiffusionRequest) -> tuple[str, dict]:
+    def _request_session(
+        req: OmniDiffusionRequest,
+    ) -> tuple[str, dict, ARDiffusionTickRequest | None]:
         extra_args = req.sampling_params.extra_args or {}
-        return str(extra_args.get("session_id") or "default"), extra_args
+        tick = ARDiffusionTickRequest.from_extra_args(extra_args)
+        if tick is not None:
+            return tick.session_id, extra_args, tick
+        return str(extra_args.get("session_id") or "default"), extra_args, None
 
     def execute_model(
         self,
@@ -245,8 +251,10 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
         if capability is None:
             raise RuntimeError("AR-Diffusion capability missing after KV cache initialization")
 
-        session_id, extra_args = self._request_session(req)
-        if extra_args.get("reset", False):
+        session_id, extra_args, tick = self._request_session(req)
+        reset = tick.reset if tick is not None else bool(extra_args.get("reset", False))
+        close_session = tick.close_session if tick is not None else bool(extra_args.get("close_session", False))
+        if reset:
             self.reset_session(session_id)
         state = self._get_or_create_session(session_id)
         started = time.perf_counter()
@@ -268,7 +276,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
             )
             raise
         self._perf_e2e_times.append(time.perf_counter() - started)
-        if extra_args.get("close_session", False):
+        if close_session:
             self.close_session(session_id)
         return output
 

vllm_omni/experimental/ar_diffusion/session.py

@@ -0,0 +1,578 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Internal realtime session control plane for AR-Diffusion.
+
+This module owns transport-level identity, event ordering, backpressure, and
+chunk-boundary snapshots. Model-defined control payloads remain opaque here.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from collections import OrderedDict
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Generic, Protocol, TypeVar
+
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+    ARDiffusionChunkMetadata,
+    ARDiffusionControlInput,
+    ARDiffusionTickRequest,
+)
+
+
+class ARDiffusionSessionError(RuntimeError):
+    """Base error for realtime AR-Diffusion session operations."""
+
+
+class ARDiffusionSessionStateError(ARDiffusionSessionError):
+    """Raised when an operation is invalid for the current session state."""
+
+
+class ARDiffusionEventQueueFullError(ARDiffusionSessionError):
+    """Raised when accepting an event would exceed the bounded queue."""
+
+
+class ARDiffusionStaleEventError(ARDiffusionSessionError):
+    """Raised when an event is older than the current monotonic boundary."""
+
+
+class ARDiffusionChunkMetadataError(ARDiffusionSessionError):
+    """Raised when a tick consumer returns metadata for a different chunk."""
+
+
+class ARDiffusionEventAcceptanceStatus(str, Enum):
+    ACCEPTED = "accepted"
+    DUPLICATE = "duplicate"
+
+
+class ARDiffusionSessionStatus(str, Enum):
+    ACTIVE = "active"
+    RESETTING = "resetting"
+    FAILED = "failed"
+    CLOSING = "closing"
+    CLEANUP_FAILED = "cleanup_failed"
+    CLOSED = "closed"
+
+
+def _copy_control(control: ARDiffusionControlInput) -> ARDiffusionControlInput:
+    return ARDiffusionControlInput(
+        track=control.track,
+        schema=control.schema,
+        data=control.data,
+    )
+
+
+@dataclass(frozen=True)
+class ARDiffusionSessionEvent:
+    """One atomically accepted prompt/control update.
+
+    ``prompt=None`` means that this event does not update the prompt. Controls
+    are keyed by ``track`` when the chunk snapshot is assembled; their schema
+    and data are otherwise opaque to this layer.
+    """
+
+    event_id: int
+    prompt: str | None = None
+    controls: tuple[ARDiffusionControlInput, ...] = ()
+
+    def __post_init__(self) -> None:
+        if isinstance(self.event_id, bool) or not isinstance(self.event_id, int) or self.event_id < 0:
+            raise ValueError("event_id must be a non-negative integer.")
+        if self.prompt is not None and not isinstance(self.prompt, str):
+            raise ValueError("prompt must be a string or None.")
+        controls = tuple(self.controls)
+        if any(not isinstance(control, ARDiffusionControlInput) for control in controls):
+            raise ValueError("controls must contain ARDiffusionControlInput values.")
+        tracks = [control.track for control in controls]
+        if len(tracks) != len(set(tracks)):
+            raise ValueError("A composite event may update each control track at most once.")
+        if self.prompt is None and not controls:
+            raise ValueError("An event must update the prompt, one or more controls, or both.")
+        object.__setattr__(self, "controls", controls)
+
+
+@dataclass(frozen=True)
+class ARDiffusionEventAcceptance:
+    """Immediate queue acceptance result; this does not mean applied."""
+
+    session_id: str
+    event_id: int
+    status: ARDiffusionEventAcceptanceStatus
+    pending_event_count: int
+
+
+TickOutputT = TypeVar("TickOutputT")
+
+
+class ARDiffusionTickConsumer(Protocol[TickOutputT]):
+    """Consumes one immutable chunk-boundary tick."""
+
+    async def execute_tick(self, tick: ARDiffusionTickRequest) -> TickOutputT: ...
+
+    def chunk_metadata(self, output: TickOutputT) -> ARDiffusionChunkMetadata: ...
+
+
+@dataclass(frozen=True)
+class ARDiffusionPreparedControls:
+    """One reducer-owned, transactional control snapshot.
+
+    ``private_state`` is opaque to the session layer and becomes durable only
+    after the corresponding model tick succeeds.
+    """
+
+    controls: tuple[ARDiffusionControlInput, ...]
+    private_state: Any = None
+
+
+class ARDiffusionControlReducer(Protocol):
+    """Optional model adapter for stateful controls sampled per AR block."""
+
+    def prepare(
+        self,
+        *,
+        current_controls: Mapping[str, ARDiffusionControlInput],
+        events: Sequence[ARDiffusionSessionEvent],
+        chunk_index: int,
+    ) -> ARDiffusionPreparedControls: ...
+
+    def commit(self, prepared: ARDiffusionPreparedControls) -> None: ...
+
+    def reset(self) -> None: ...
+
+
+class ARDiffusionSessionLifecycle(Protocol):
+    """Worker-facing session lifecycle operations."""
+
+    async def reset_session(self, session_id: str) -> None: ...
+
+    async def close_session(self, session_id: str) -> None: ...
+
+
+class ARDiffusionCollectiveRPCClient(Protocol):
+    """Subset of the engine control RPC surface used by session lifecycle."""
+
+    async def collective_rpc(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        stage_ids: list[int] | None = None,
+    ) -> list[Any]: ...
+
+
+def _collect_rpc_support_results(value: Any, *, errors: list[str], supported: list[bool]) -> None:
+    if isinstance(value, bool):
+        supported.append(value)
+        return
+    if isinstance(value, (list, tuple)):
+        for item in value:
+            _collect_rpc_support_results(item, errors=errors, supported=supported)
+        return
+    if isinstance(value, Mapping):
+        if value.get("supported") is False:
+            errors.append(str(value.get("error") or "target stage does not support this lifecycle RPC"))
+            return
+        if "result" in value:
+            _collect_rpc_support_results(value["result"], errors=errors, supported=supported)
+
+
+class ARDiffusionWorkerLifecycle:
+    """Routes reset/close through the existing orchestrator collective RPC."""
+
+    def __init__(
+        self,
+        rpc_client: ARDiffusionCollectiveRPCClient,
+        *,
+        stage_ids: Sequence[int],
+        timeout: float | None = None,
+    ) -> None:
+        normalized_stage_ids = tuple(stage_ids)
+        if not normalized_stage_ids:
+            raise ValueError("stage_ids must select at least one AR-Diffusion stage.")
+        if any(isinstance(stage_id, bool) or not isinstance(stage_id, int) or stage_id < 0 for stage_id in stage_ids):
+            raise ValueError("stage_ids must contain non-negative integers.")
+        self._rpc_client = rpc_client
+        self._stage_ids = normalized_stage_ids
+        self._timeout = timeout
+
+    async def _invoke(self, method: str, session_id: str) -> None:
+        results = await self._rpc_client.collective_rpc(
+            method=method,
+            timeout=self._timeout,
+            args=(session_id,),
+            stage_ids=list(self._stage_ids),
+        )
+        errors: list[str] = []
+        supported: list[bool] = []
+        _collect_rpc_support_results(results, errors=errors, supported=supported)
+        if errors:
+            raise ARDiffusionSessionError(f"{method} failed: {'; '.join(errors)}")
+        if not supported or not all(supported):
+            raise ARDiffusionSessionError(f"{method} is not supported by every selected worker.")
+
+    async def reset_session(self, session_id: str) -> None:
+        await self._invoke("reset_ar_diffusion_session", session_id)
+
+    async def close_session(self, session_id: str) -> None:
+        await self._invoke("close_ar_diffusion_session", session_id)
+
+
+def _default_request_id(session_id: str, chunk_index: int) -> str:
+    return f"ar-{session_id}-{chunk_index}-{uuid.uuid4().hex}"
+
+
+class ARDiffusionSession(Generic[TickOutputT]):
+    """Transactional event queue and chunk snapshot for one realtime world."""
+
+    def __init__(
+        self,
+        session_id: str,
+        *,
+        tick_consumer: ARDiffusionTickConsumer[TickOutputT],
+        lifecycle: ARDiffusionSessionLifecycle,
+        max_pending_events: int,
+        request_id_factory: Callable[[str, int], str] | None = None,
+        control_reducer: ARDiffusionControlReducer | None = None,
+    ) -> None:
+        if not isinstance(session_id, str) or not session_id.strip():
+            raise ValueError("session_id must be a non-empty string.")
+        if isinstance(max_pending_events, bool) or not isinstance(max_pending_events, int) or max_pending_events <= 0:
+            raise ValueError("max_pending_events must be a positive integer.")
+        self.session_id = session_id
+        self._tick_consumer = tick_consumer
+        self._lifecycle = lifecycle
+        self._max_pending_events = max_pending_events
+        self._request_id_factory = request_id_factory or _default_request_id
+        self._control_reducer = control_reducer
+
+        self._status = ARDiffusionSessionStatus.ACTIVE
+        self._pending_events: dict[int, ARDiffusionSessionEvent] = {}
+        self._recent_event_ids: OrderedDict[int, None] = OrderedDict()
+        self._event_floor = -1
+        self._inflight_event_floor: int | None = None
+        self._chunk_index = 0
+        self._prompt: str | None = None
+        self._controls: dict[str, ARDiffusionControlInput] = {}
+
+        self._state_lock = asyncio.Lock()
+        self._tick_lock = asyncio.Lock()
+
+    @property
+    def status(self) -> ARDiffusionSessionStatus:
+        return self._status
+
+    @property
+    def chunk_index(self) -> int:
+        return self._chunk_index
+
+    @property
+    def pending_event_count(self) -> int:
+        return len(self._pending_events)
+
+    def _require_active(self) -> None:
+        if self._status is not ARDiffusionSessionStatus.ACTIVE:
+            raise ARDiffusionSessionStateError(f"AR-Diffusion session {self.session_id!r} is {self._status.value}.")
+
+    def _require_resettable(self) -> None:
+        if self._status not in {
+            ARDiffusionSessionStatus.ACTIVE,
+            ARDiffusionSessionStatus.FAILED,
+        }:
+            raise ARDiffusionSessionStateError(f"AR-Diffusion session {self.session_id!r} is {self._status.value}.")
+
+    async def accept_event(self, event: ARDiffusionSessionEvent) -> ARDiffusionEventAcceptance:
+        """Atomically enqueue an event or return an idempotent duplicate result."""
+        if not isinstance(event, ARDiffusionSessionEvent):
+            raise TypeError("event must be an ARDiffusionSessionEvent.")
+        async with self._state_lock:
+            self._require_active()
+            if event.event_id in self._pending_events:
+                return ARDiffusionEventAcceptance(
+                    session_id=self.session_id,
+                    event_id=event.event_id,
+                    status=ARDiffusionEventAcceptanceStatus.DUPLICATE,
+                    pending_event_count=len(self._pending_events),
+                )
+            if event.event_id in self._recent_event_ids:
+                return ARDiffusionEventAcceptance(
+                    session_id=self.session_id,
+                    event_id=event.event_id,
+                    status=ARDiffusionEventAcceptanceStatus.DUPLICATE,
+                    pending_event_count=len(self._pending_events),
+                )
+
+            event_floor = self._event_floor
+            if self._inflight_event_floor is not None:
+                event_floor = max(event_floor, self._inflight_event_floor)
+            if event.event_id <= event_floor:
+                raise ARDiffusionStaleEventError(
+                    f"event_id {event.event_id} is stale; the session event floor is {event_floor}."
+                )
+            if len(self._pending_events) >= self._max_pending_events:
+                raise ARDiffusionEventQueueFullError(
+                    f"AR-Diffusion session {self.session_id!r} event queue is full "
+                    f"({self._max_pending_events} pending)."
+                )
+
+            stored_event = ARDiffusionSessionEvent(
+                event_id=event.event_id,
+                prompt=event.prompt,
+                controls=tuple(_copy_control(control) for control in event.controls),
+            )
+            self._pending_events[event.event_id] = stored_event
+            self._recent_event_ids[event.event_id] = None
+            while len(self._recent_event_ids) > self._max_pending_events:
+                self._recent_event_ids.popitem(last=False)
+            return ARDiffusionEventAcceptance(
+                session_id=self.session_id,
+                event_id=event.event_id,
+                status=ARDiffusionEventAcceptanceStatus.ACCEPTED,
+                pending_event_count=len(self._pending_events),
+            )
+
+    def _build_tick_locked(
+        self,
+    ) -> tuple[
+        ARDiffusionTickRequest,
+        str | None,
+        dict[str, ARDiffusionControlInput],
+        ARDiffusionPreparedControls | None,
+    ]:
+        event_ids = tuple(sorted(self._pending_events))
+        events = tuple(self._pending_events[event_id] for event_id in event_ids)
+        prompt = self._prompt
+        for event in events:
+            if event.prompt is not None:
+                prompt = event.prompt
+
+        prepared_controls: ARDiffusionPreparedControls | None = None
+        if self._control_reducer is None:
+            controls = {track: _copy_control(control) for track, control in self._controls.items()}
+            for event in events:
+                for control in event.controls:
+                    controls[control.track] = _copy_control(control)
+        else:
+            prepared_controls = self._control_reducer.prepare(
+                current_controls={track: _copy_control(control) for track, control in self._controls.items()},
+                events=events,
+                chunk_index=self._chunk_index,
+            )
+            if not isinstance(prepared_controls, ARDiffusionPreparedControls):
+                raise TypeError("control reducer prepare() must return ARDiffusionPreparedControls.")
+            controls = {}
+            for control in prepared_controls.controls:
+                if not isinstance(control, ARDiffusionControlInput):
+                    raise TypeError("prepared controls must contain ARDiffusionControlInput values.")
+                if control.track in controls:
+                    raise ValueError("prepared controls must contain at most one value per track.")
+                controls[control.track] = _copy_control(control)
+
+        chunk_index = self._chunk_index
+        request_id = self._request_id_factory(self.session_id, chunk_index)
+        tick = ARDiffusionTickRequest(
+            session_id=self.session_id,
+            request_id=request_id,
+            chunk_index=chunk_index,
+            applied_event_ids=event_ids,
+            prompt=prompt,
+            controls=tuple(_copy_control(controls[track]) for track in sorted(controls)),
+        )
+        self._inflight_event_floor = event_ids[-1] if event_ids else None
+        return tick, prompt, controls, prepared_controls
+
+    async def _commit_tick(
+        self,
+        tick: ARDiffusionTickRequest,
+        prompt: str | None,
+        controls: dict[str, ARDiffusionControlInput],
+        prepared_controls: ARDiffusionPreparedControls | None,
+    ) -> None:
+        async with self._state_lock:
+            # Reducer state and the session snapshot form one logical commit.
+            # Run the only fallible operation before mutating session fields so
+            # a reducer failure cannot acknowledge events or advance the chunk.
+            if prepared_controls is not None:
+                assert self._control_reducer is not None
+                self._control_reducer.commit(prepared_controls)
+            for event_id in tick.applied_event_ids:
+                self._pending_events.pop(event_id, None)
+            if tick.applied_event_ids:
+                self._event_floor = tick.applied_event_ids[-1]
+            self._prompt = prompt
+            self._controls = controls
+            self._chunk_index += 1
+            self._inflight_event_floor = None
+
+    async def _fail_tick(self, tick_error: BaseException) -> None:
+        async with self._state_lock:
+            self._inflight_event_floor = None
+            self._status = ARDiffusionSessionStatus.FAILED
+        try:
+            # A model forward may already have released its state. This
+            # idempotent close also covers successful forwards whose metadata
+            # or local reducer commit could not be accepted.
+            await self._lifecycle.close_session(self.session_id)
+        except Exception as cleanup_error:
+            raise ARDiffusionSessionError(
+                "AR-Diffusion tick failed "
+                f"({type(tick_error).__name__}: {tick_error}) and worker "
+                f"session cleanup also failed: {cleanup_error}"
+            ) from tick_error
+
+    async def next_chunk(self) -> TickOutputT:
+        """Execute and atomically commit one chunk-boundary snapshot."""
+        async with self._tick_lock:
+            async with self._state_lock:
+                self._require_active()
+                tick, prompt, controls, prepared_controls = self._build_tick_locked()
+
+            expected_metadata = ARDiffusionChunkMetadata.from_tick(tick)
+            try:
+                output = await self._tick_consumer.execute_tick(tick)
+                metadata = self._tick_consumer.chunk_metadata(output)
+                if not isinstance(metadata, ARDiffusionChunkMetadata) or metadata != expected_metadata:
+                    raise ARDiffusionChunkMetadataError(
+                        "Tick consumer metadata does not match the submitted AR-Diffusion tick."
+                    )
+            except BaseException as tick_error:
+                await self._fail_tick(tick_error)
+                raise
+
+            commit_task = asyncio.create_task(
+                self._commit_tick(
+                    tick,
+                    prompt,
+                    controls,
+                    prepared_controls,
+                )
+            )
+            try:
+                await asyncio.shield(commit_task)
+            except asyncio.CancelledError:
+                # Once a consumer reports success, the local chunk/event state
+                # must commit even if the transport task is being cancelled.
+                try:
+                    await commit_task
+                except BaseException as commit_error:
+                    await self._fail_tick(commit_error)
+                    raise
+                raise
+            except BaseException as commit_error:
+                await self._fail_tick(commit_error)
+                raise
+            return output
+
+    async def reset(self) -> None:
+        """Reset worker state and discard queued/current transport state."""
+        async with self._tick_lock:
+            async with self._state_lock:
+                self._require_resettable()
+                self._status = ARDiffusionSessionStatus.RESETTING
+            try:
+                await self._lifecycle.reset_session(self.session_id)
+            except BaseException:
+                async with self._state_lock:
+                    self._status = ARDiffusionSessionStatus.FAILED
+                raise
+            async with self._state_lock:
+                if self._pending_events:
+                    self._event_floor = max(self._event_floor, max(self._pending_events))
+                self._pending_events.clear()
+                self._prompt = None
+                self._controls.clear()
+                if self._control_reducer is not None:
+                    self._control_reducer.reset()
+                self._inflight_event_floor = None
+                self._chunk_index = 0
+                self._status = ARDiffusionSessionStatus.ACTIVE
+
+    async def close(self) -> None:
+        """Close worker state, retaining a tombstone until cleanup succeeds."""
+        async with self._tick_lock:
+            async with self._state_lock:
+                if self._status is ARDiffusionSessionStatus.CLOSED:
+                    return
+                self._status = ARDiffusionSessionStatus.CLOSING
+            try:
+                await self._lifecycle.close_session(self.session_id)
+            except BaseException:
+                async with self._state_lock:
+                    self._status = ARDiffusionSessionStatus.CLEANUP_FAILED
+                raise
+            else:
+                async with self._state_lock:
+                    self._pending_events.clear()
+                    self._prompt = None
+                    self._controls.clear()
+                    self._inflight_event_floor = None
+                    self._status = ARDiffusionSessionStatus.CLOSED
+
+    async def disconnect(self) -> None:
+        """Disconnect uses the same owning-worker cleanup as explicit close."""
+        await self.close()
+
+
+class ARDiffusionSessionManager(Generic[TickOutputT]):
+    """Creates and owns transport sessions sharing one tick/lifecycle backend."""
+
+    def __init__(
+        self,
+        *,
+        tick_consumer: ARDiffusionTickConsumer[TickOutputT],
+        lifecycle: ARDiffusionSessionLifecycle,
+        max_pending_events: int,
+        request_id_factory: Callable[[str, int], str] | None = None,
+        session_id_factory: Callable[[], str] | None = None,
+        control_reducer_factory: Callable[[], ARDiffusionControlReducer] | None = None,
+    ) -> None:
+        self._tick_consumer = tick_consumer
+        self._lifecycle = lifecycle
+        self._max_pending_events = max_pending_events
+        self._request_id_factory = request_id_factory
+        self._session_id_factory = session_id_factory or (lambda: uuid.uuid4().hex)
+        self._control_reducer_factory = control_reducer_factory
+        self._sessions: dict[str, ARDiffusionSession[TickOutputT]] = {}
+        self._lock = asyncio.Lock()
+
+    async def create_session(self, session_id: str | None = None) -> ARDiffusionSession[TickOutputT]:
+        resolved_session_id = session_id or self._session_id_factory()
+        async with self._lock:
+            if resolved_session_id in self._sessions:
+                raise ARDiffusionSessionStateError(f"AR-Diffusion session {resolved_session_id!r} already exists.")
+            session = ARDiffusionSession(
+                resolved_session_id,
+                tick_consumer=self._tick_consumer,
+                lifecycle=self._lifecycle,
+                max_pending_events=self._max_pending_events,
+                request_id_factory=self._request_id_factory,
+                control_reducer=(
+                    self._control_reducer_factory() if self._control_reducer_factory is not None else None
+                ),
+            )
+            self._sessions[resolved_session_id] = session
+            return session
+
+    async def get_session(self, session_id: str) -> ARDiffusionSession[TickOutputT]:
+        async with self._lock:
+            try:
+                return self._sessions[session_id]
+            except KeyError as exc:
+                raise ARDiffusionSessionStateError(f"AR-Diffusion session {session_id!r} does not exist.") from exc
+
+    async def reset_session(self, session_id: str) -> None:
+        session = await self.get_session(session_id)
+        await session.reset()
+
+    async def close_session(self, session_id: str) -> None:
+        session = await self.get_session(session_id)
+        await session.close()
+        async with self._lock:
+            if self._sessions.get(session_id) is session:
+                self._sessions.pop(session_id)
+
+    async def disconnect(self, session_id: str) -> None:
+        """Release the owning worker route when its transport disconnects."""
+        await self.close_session(session_id)

vllm_omni/experimental/ar_diffusion/tick_protocol.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Typed request boundary shared by AR-Diffusion models and serving APIs.
+
+The transport layer owns event ordering and snapshots. Model adapters consume
+the resulting immutable tick without learning about WebSocket messages.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Any
+
+AR_DIFFUSION_TICK_KEY = "ar_diffusion_tick"
+
+
+def _non_empty_string(value: object, *, field: str) -> str:
+    if not isinstance(value, str) or not value.strip():
+        raise ValueError(f"{field} must be a non-empty string.")
+    return value
+
+
+def _non_negative_int(value: object, *, field: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+        raise ValueError(f"{field} must be a non-negative integer.")
+    return value
+
+
+def _freeze_transport_value(value: Any, *, path: str) -> Any:
+    """Copy JSON-like transport data into recursively immutable containers."""
+    if value is None or isinstance(value, (str, bool, int, float)):
+        return value
+    if isinstance(value, Mapping):
+        frozen: dict[str, Any] = {}
+        for key, item in value.items():
+            if not isinstance(key, str):
+                raise ValueError(f"{path} mapping keys must be strings.")
+            frozen[key] = _freeze_transport_value(item, path=f"{path}.{key}")
+        return MappingProxyType(frozen)
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
+        return tuple(_freeze_transport_value(item, path=f"{path}[]") for item in value)
+    raise ValueError(
+        f"{path} must contain only transport-safe mappings, sequences, and scalar values; got {type(value).__name__}."
+    )
+
+
+def _thaw_transport_value(value: Any) -> Any:
+    """Return a detached JSON-serializable copy of an immutable snapshot."""
+    if isinstance(value, Mapping):
+        return {key: _thaw_transport_value(item) for key, item in value.items()}
+    if isinstance(value, tuple):
+        return [_thaw_transport_value(item) for item in value]
+    return value
+
+
+@dataclass(frozen=True)
+class ARDiffusionControlInput:
+    """One model-defined control track in a chunk-boundary snapshot."""
+
+    track: str
+    schema: str
+    data: Mapping[str, Any]
+
+    def __post_init__(self) -> None:
+        _non_empty_string(self.track, field="control.track")
+        _non_empty_string(self.schema, field="control.schema")
+        if not isinstance(self.data, Mapping):
+            raise ValueError("control.data must be a mapping.")
+        object.__setattr__(
+            self,
+            "data",
+            _freeze_transport_value(self.data, path=f"control[{self.track!r}].data"),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "track": self.track,
+            "schema": self.schema,
+            "data": _thaw_transport_value(self.data),
+        }
+
+    @classmethod
+    def from_dict(cls, value: object) -> ARDiffusionControlInput:
+        if not isinstance(value, Mapping):
+            raise ValueError("Each control must be a mapping.")
+        return cls(
+            track=value.get("track"),  # type: ignore[arg-type]
+            schema=value.get("schema"),  # type: ignore[arg-type]
+            data=value.get("data"),  # type: ignore[arg-type]
+        )
+
+
+@dataclass(frozen=True)
+class ARDiffusionTickRequest:
+    """A single atomic AR-block computation.
+
+    ``request_id`` identifies this computation, ``session_id`` identifies the
+    persistent rollout, and ``applied_event_ids`` records which accepted
+    interaction events are represented by this immutable snapshot.
+    """
+
+    session_id: str
+    request_id: str
+    chunk_index: int
+    applied_event_ids: tuple[int, ...] = ()
+    prompt: str | None = None
+    controls: tuple[ARDiffusionControlInput, ...] = ()
+    reset: bool = False
+    close_session: bool = False
+
+    def __post_init__(self) -> None:
+        _non_empty_string(self.session_id, field="session_id")
+        _non_empty_string(self.request_id, field="request_id")
+        _non_negative_int(self.chunk_index, field="chunk_index")
+        if self.prompt is not None and not isinstance(self.prompt, str):
+            raise ValueError("prompt must be a string or None.")
+        if not isinstance(self.reset, bool):
+            raise ValueError("reset must be a boolean.")
+        if not isinstance(self.close_session, bool):
+            raise ValueError("close_session must be a boolean.")
+
+        event_ids = tuple(
+            _non_negative_int(event_id, field="applied_event_ids item") for event_id in self.applied_event_ids
+        )
+        if event_ids != tuple(sorted(set(event_ids))):
+            raise ValueError("applied_event_ids must be unique and strictly increasing.")
+        controls = tuple(self.controls)
+        if any(not isinstance(control, ARDiffusionControlInput) for control in controls):
+            raise ValueError("controls must contain ARDiffusionControlInput values.")
+        object.__setattr__(self, "applied_event_ids", event_ids)
+        object.__setattr__(self, "controls", controls)
+
+    def to_extra_args(self) -> dict[str, Any]:
+        """Serialize the tick under one namespaced ``extra_args`` key."""
+        return {
+            AR_DIFFUSION_TICK_KEY: {
+                "session_id": self.session_id,
+                "request_id": self.request_id,
+                "chunk_index": self.chunk_index,
+                "applied_event_ids": list(self.applied_event_ids),
+                "prompt": self.prompt,
+                "controls": [control.to_dict() for control in self.controls],
+                "reset": self.reset,
+                "close_session": self.close_session,
+            }
+        }
+
+    @classmethod
+    def from_extra_args(
+        cls,
+        extra_args: Mapping[str, Any] | None,
+    ) -> ARDiffusionTickRequest | None:
+        """Parse a typed tick, returning ``None`` for legacy AR requests."""
+        if extra_args is None or AR_DIFFUSION_TICK_KEY not in extra_args:
+            return None
+        value = extra_args[AR_DIFFUSION_TICK_KEY]
+        if not isinstance(value, Mapping):
+            raise ValueError(f"{AR_DIFFUSION_TICK_KEY} must be a mapping.")
+        raw_event_ids = value.get("applied_event_ids", ())
+        raw_controls = value.get("controls", ())
+        if not isinstance(raw_event_ids, Sequence) or isinstance(raw_event_ids, (str, bytes)):
+            raise ValueError("applied_event_ids must be a sequence.")
+        if not isinstance(raw_controls, Sequence) or isinstance(raw_controls, (str, bytes)):
+            raise ValueError("controls must be a sequence.")
+        return cls(
+            session_id=value.get("session_id"),  # type: ignore[arg-type]
+            request_id=value.get("request_id"),  # type: ignore[arg-type]
+            chunk_index=value.get("chunk_index"),  # type: ignore[arg-type]
+            applied_event_ids=tuple(raw_event_ids),
+            prompt=value.get("prompt"),  # type: ignore[arg-type]
+            controls=tuple(ARDiffusionControlInput.from_dict(control) for control in raw_controls),
+            reset=value.get("reset", False),  # type: ignore[arg-type]
+            close_session=value.get("close_session", False),  # type: ignore[arg-type]
+        )
+
+
+@dataclass(frozen=True)
+class ARDiffusionChunkMetadata:
+    """Identity metadata returned with one generated AR block."""
+
+    session_id: str
+    request_id: str
+    chunk_index: int
+    applied_event_ids: tuple[int, ...] = ()
+
+    @classmethod
+    def from_tick(cls, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+        return cls(
+            session_id=tick.session_id,
+            request_id=tick.request_id,
+            chunk_index=tick.chunk_index,
+            applied_event_ids=tick.applied_event_ids,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "session_id": self.session_id,
+            "request_id": self.request_id,
+            "chunk_index": self.chunk_index,
+            "applied_event_ids": list(self.applied_event_ids),
+        }