[Docs] Clarify TP1 DLO mmap support scope

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

docs/.nav.yml

@@ -116,7 +116,6 @@ 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,9 +4,6 @@
 > `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
@@ -54,13 +51,14 @@ 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 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.
+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.
 
 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,17 +19,15 @@ 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. 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.
+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.
 
 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. Direct checkpoint mmap still does not cover TP, but the Phase B cache
-can share each final TP coordinate among equivalent processes.
+tensors, but those configurations do not use checkpoint mmap and must not be
+used to claim shared-mmap host-memory savings.
 
 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -83,90 +81,6 @@ 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
@@ -200,17 +114,11 @@ 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
-a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or
-ordinary runtime tensors.
+either a loader-approved checkpoint 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
@@ -219,14 +127,10 @@ 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. Supported CPU layouts may then enter the opt-in runtime cache.
-Unsupported layouts keep one private runtime copy per process.
+parameters. In that fallback, each pure-DP process keeps a private full runtime
+copy.
 
 This mode means:
 
@@ -237,17 +141,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 runtime cache can share equivalent ordinary-loader
-  outputs; its fallback keeps one private runtime copy per rank.
+  mmap is selected; the ordinary-loader 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; 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. |
+| **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. |
 | **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
@@ -278,10 +182,12 @@ 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. 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.
+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).
 
 ## Validation coverage
 
@@ -294,13 +200,6 @@ 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
@@ -360,99 +259,11 @@ 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.
 
-### 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.
+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.
 
 ## Recommendations
 
@@ -461,12 +272,8 @@ passes; that startup problem is outside the registered transfer scope.
 - 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. 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.
+  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.
 - 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

@@ -1,213 +0,0 @@
-# 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,7 +22,6 @@ 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,7 +37,6 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -118,10 +117,6 @@ 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,11 +5,7 @@ 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. 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.
+otherwise DLO streams the ordinary loader's rank-local tensors.
 
 See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
 for the implementation contract and compatibility matrix.
@@ -25,10 +21,8 @@ AllGather:       [AG N+1]           [AG N+2]
 Slots:           slot 0: Layer N    slot 1: Layer N+1
 ```
 
-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.
+AllGather communicates only request-independent weight shards, so data-
+parallel ranks may process different requests concurrently.
 
 ## Usage
 
@@ -44,15 +38,6 @@ 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 \
@@ -65,10 +50,7 @@ from vllm_omni import Omni
 omni = Omni(
     model="/path/to/model",
     enable_distributed_layerwise_offload=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,
+    dlo_use_allgather=True,
 )
 ```
 
@@ -80,10 +62,6 @@ 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
@@ -100,50 +78,6 @@ 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
@@ -170,34 +104,6 @@ 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:
@@ -226,37 +132,22 @@ must enter each collective.
 
 ## Limitations
 
-- 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.
+- 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.
 - 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 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.
+  incompatible with DLO AllGather.
 - 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.
 
-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.
+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.
 
 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,6 +24,7 @@ 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,
@@ -227,6 +228,7 @@ 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

@@ -1,288 +0,0 @@
-# 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

@@ -1,239 +0,0 @@
-# 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,6 +82,9 @@ 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(
-        "--deploy-config",
+        "--stage-configs-path",
         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.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
     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,6 +38,17 @@ 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(
@@ -141,8 +152,7 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
+    omni_kwargs["deploy_config"] = args.deploy_config or DEFAULT_DEPLOY_CONFIG
     return omni_kwargs
 
 

examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py

@@ -32,10 +32,11 @@ 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_deploy_config() -> str:
+def _default_stage_config() -> str:
     repo_root = Path(__file__).resolve().parents[4]
     return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
 
@@ -108,7 +109,7 @@ def main(args: Any) -> None:
         "output_dir",
     ):
         omni_kwargs.pop(key, None)
-    omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
+    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
     omni_kwargs["log_stats"] = args.log_stats
     omni = Omni(**omni_kwargs)
 
@@ -155,7 +156,7 @@ def parse_args() -> Any:
         default=None,
         help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
     )
-    parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
+    parser.add_argument("--stage-configs-path", 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")
@@ -163,6 +164,7 @@ 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 serve bosonai/higgs-audio-v3-tts-4b \
+vllm-omni 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 serve nvidia/Nemotron-Labs-Audex-2B \
+vllm-omni 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 serve nvidia/Nemotron-Labs-Audex-30B-A3B \
+vllm-omni 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,7 +60,6 @@ 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

@@ -1,122 +0,0 @@
-# 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 serve /path/to/stable-audio-open-1.0 \
+vllm-omni serve /path/to/stable-audio-open-1.0 \
   --host 0.0.0.0 \
   --port 8091 \
   --gpu-memory-utilization 0.9 \

recipes/TEMPLATE.md

@@ -39,9 +39,6 @@ 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 serve inclusionAI/Ming-omni-tts-0.5B \
+vllm-omni 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 serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+vllm-omni 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 serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
+vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
 ```
 
 Use a custom deploy config for advanced cases:
 
 ```bash
-vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
+vllm-omni 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,13 +1102,6 @@ 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"))
 
@@ -1384,8 +1377,7 @@ stages:
         assert async_stages[1].custom_process_input_func is None
 
     def test_no_bundled_legacy_stage_config_yamls(self):
-        repo_root = Path(__file__).resolve().parents[2]
-        stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
+        stage_config_dir = Path(__file__).parent.parent / "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,7 +18,6 @@ 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"
@@ -57,21 +56,6 @@ 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
@@ -119,7 +103,6 @@ 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(
@@ -130,10 +113,7 @@ def make_runner(
     )
     runner.device = torch.device("cpu")
     runner.pipeline = pipeline
-    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
-        enable=True,
-        gpu_memory_fraction=gpu_memory_fraction,
-    )
+    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
     runner.kv_cache = None
     runner._ar_diffusion_capability = None
     runner._ar_diffusion_kv_cache_spec = None
@@ -152,16 +132,6 @@ 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)
@@ -170,43 +140,6 @@ 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)
@@ -237,31 +170,6 @@ 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
@@ -290,67 +198,24 @@ def test_dreamzero_like_two_branches_are_independent():
 
 
 def test_lru_eviction_releases_blocks_and_notifies_pipeline():
-    pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
+    pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
     runner = make_runner(pipeline)
     kv = runner.kv_cache
     assert kv is not None
-    assert runner._session_capacity == 2
-    assert kv.session_capacity == 2
+    assert runner._session_capacity == 1
+    assert kv.session_capacity == 1
     free_total = kv.manager.block_pool.get_num_free_blocks()
     old = commit_one_frame(runner, "old", "main")
     k = torch.randn(1, 8, 4, 64)
     old.populate_cross_attention("main", "text", [(k, k)] * old.num_layers)
     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", "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
+    assert "old" not in kv._cross_sessions
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
 
 
 def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
@@ -367,10 +232,7 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
         raise RuntimeError("layer exploded")
 
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
-    request = SimpleNamespace(
-        request_id="broken-request",
-        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
-    )
+    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
 
     with pytest.raises(RuntimeError, match="layer exploded"):
         runner.execute_model(request)
@@ -401,10 +263,7 @@ 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(
-        request_id="broken-request",
-        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
-    )
+    request = SimpleNamespace(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

@@ -1,231 +0,0 @@
-# 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,8 +30,6 @@ 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(
@@ -97,6 +95,7 @@ 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
 
 
@@ -202,149 +201,12 @@ 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},
-    )
-    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
-    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 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
+    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)
+    assert kv.num_blocks == expected
+    assert expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested
 
 
 def _make_kv(
@@ -359,24 +221,15 @@ 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,
-            gpu_memory_fraction=1.0,
-        ),
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
         num_layers=1,
         num_kv_heads=4,
         head_size=64,
         dtype=torch.float32,
         block_size=BLOCK,
         max_model_len=4096,
-        available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
+        available_bytes=1 << 16,  # tiny -> the floor binds
         device=torch.device("cpu"),
         kv_branches=kv_branches,
         session_capacity=1,

tests/diffusion/ar_diffusion/test_kv_orchestration.py

@@ -54,15 +54,13 @@ 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)
-        # Allocation may transiently add one in-flight chunk.
+        # Resident window stays bounded by window + the in-flight chunk.
         assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
-        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())
+        kv.commit_chunk(adapter)
 
     # Pool memory plateaus once the window is full (eviction recycles blocks).
-    assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
 
     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")
-    kv, st = make_state(num_layers=2)
+    _, 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,10 +349,6 @@ 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,19 +174,6 @@ 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

@@ -1,505 +0,0 @@
-# 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

@@ -1,108 +0,0 @@
-# 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

@@ -1,53 +0,0 @@
-# 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,7 +12,6 @@ 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
@@ -336,7 +335,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
         backing_kind="checkpoint_mmap",
         bindings={
             "transformer.weight": TensorBinding(
-                storage_key="weight",
+                checkpoint_key="weight",
                 file_path="unused",
             )
         },
@@ -395,192 +394,6 @@ 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

@@ -1,222 +0,0 @@
-# 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,10 +6,7 @@ from types import SimpleNamespace
 
 import pytest
 
-from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
-    DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
-    DreamZeroPipeline,
-)
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -35,37 +32,6 @@ 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

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

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

@@ -1,13 +0,0 @@
-{
-  "_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

@@ -1,42 +0,0 @@
-{
-  "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

@@ -1,30 +0,0 @@
-{
-  "_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

@@ -1,112 +0,0 @@
-{
-  "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

@@ -1,294 +0,0 @@
-# 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

@@ -1,543 +0,0 @@
-# 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

@@ -1,525 +0,0 @@
-# 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

@@ -1,685 +0,0 @@
-# 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

@@ -1,39 +0,0 @@
-# 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

@@ -1,205 +0,0 @@
-# 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,31 +473,6 @@ 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):
@@ -588,24 +563,6 @@ 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
     ):
@@ -735,95 +692,6 @@ 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()}
@@ -847,7 +715,7 @@ class TestMmapWeightLoading:
             backing_kind="checkpoint_mmap",
             bindings={
                 name: TensorBinding(
-                    storage_key=name,
+                    checkpoint_key=name,
                     file_path=str(weight_file),
                 )
                 for name in weights
@@ -918,99 +786,6 @@ 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):
@@ -1316,7 +1091,7 @@ class TestMmapValidation:
         assert result.fallback_reason is None
         assert result.plan is not None
         assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
-            storage_key="blocks.0.weight",
+            checkpoint_key="blocks.0.weight",
             file_path=str(checkpoint_file),
         )
         assert result.plan.planned_source_prefixes == frozenset({"transformer."})
@@ -1606,7 +1381,7 @@ class TestMmapValidation:
 
         assert result.plan is not None
         assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
-            storage_key="proj_in.weight",
+            checkpoint_key="proj_in.weight",
             file_path=str(checkpoint_file),
         )
 

tests/diffusion/test_result_pump.py

@@ -418,77 +418,3 @@ 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,13 +12,11 @@ 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,
@@ -206,29 +204,6 @@ 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,18 +3,19 @@
 """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"
-DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3_ar.yaml")
+REPO_ROOT = Path(__file__).resolve().parents[3]
+STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_i2t.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
@@ -44,14 +45,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, pytest.mark.cuda]
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
 
 @pytest.fixture(scope="module")
 def omni() -> Generator[Omni, None, None]:
     with OmniRunner(
         MODEL_NAME,
-        deploy_config=DEPLOY_CONFIG_PATH,
+        deploy_config=str(STAGE_CONFIG_PATH),
         trust_remote_code=True,
     ) as runner:
         yield runner.omni

tests/e2e/offline_inference/test_lingbot_world_v2.py

@@ -1,84 +0,0 @@
-# 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,13 +328,6 @@ 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",
         ]
@@ -346,17 +339,9 @@ 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
 
 
@@ -507,24 +492,6 @@ 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_model,
+    load_stage_configs_from_yaml,
     resolve_model_config_path,
 )
 
@@ -417,9 +417,8 @@ class TestTryResolveOmniModelType:
 
 class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
-        """Ensure that kwargs survive default stage creation."""
-        engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
-        kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
+        """Ensure that dtype survives default stage creation."""
+        kwargs = {"dtype": torch.float32}
         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))
 
@@ -432,7 +431,6 @@ 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"
@@ -520,20 +518,71 @@ class TestLoadAndResolveStageConfigs:
         assert filtered[0].final_output_type is None
 
 
-class TestLoadStageConfigsFromModel:
-    def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog):
+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},
+                        },
+                    }
+                ],
+            }
+        )
         mocker.patch(
-            "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model",
-            return_value=(None, None),
+            "vllm_omni.entrypoints.utils.load_yaml_config",
+            return_value=yaml_config,
         )
-        resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path")
 
-        result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False)
+        stages = load_stage_configs_from_yaml(
+            "fake.yaml",
+            base_engine_args={
+                "parallel_config": {
+                    "tensor_parallel_size": 1,
+                    "pipeline_parallel_size": 2,
+                },
+                "model": "base-model",
+            },
+        )
 
-        assert result == ([], None)
-        resolve_path.assert_not_called()
-        assert "No registered PipelineConfig resolved" in caplog.text
-        assert "deploy_config" in caplog.text
+        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
 
 
 class TestCumulativeStreamingCoercion:

tests/model_tests/diffusion/test_alignment.py

@@ -64,9 +64,6 @@ 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,10 +576,6 @@ 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"
@@ -642,13 +638,6 @@ 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,10 +382,6 @@ 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
@@ -636,11 +632,6 @@ 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,18 +708,6 @@ 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
 
@@ -940,16 +928,6 @@ 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,7 +24,6 @@ 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:
@@ -871,9 +870,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:
-                        try_set_exception(fut, RuntimeError(msg.error))
+                        fut.set_exception(RuntimeError(msg.error))
                     else:
-                        try_set_result(fut, msg)
+                        fut.set_result(msg)
             elif msg.kind == AsyncOutputKind.OUTPUT_READY:
                 batch_id = msg.async_output_id
                 with self._futures_lock:
@@ -904,9 +903,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:
-                                    try_set_exception(pending, exc)
+                                    pending.set_exception(exc)
                                 else:
-                                    try_set_result(pending, output_result)
+                                    pending.set_result(output_result)
                             else:
                                 fut = concurrent.futures.Future()
                                 if exc is not None:
@@ -934,7 +933,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():
-                    try_set_result(pending, per_req_result)
+                    pending.set_result(per_req_result)
                 else:
                     fut: concurrent.futures.Future = concurrent.futures.Future()
                     fut.set_result(per_req_result)
@@ -1000,10 +999,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             with self._futures_lock:
                 for fut in self._rpc_futures.values():
                     if not fut.done():
-                        try_set_exception(fut, RuntimeError("Executor shut down"))
+                        fut.set_exception(RuntimeError("Executor shut down"))
                 for fut in self._output_futures.values():
                     if not fut.done():
-                        try_set_exception(fut, RuntimeError("Executor shut down"))
+                        fut.set_exception(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,7 +14,6 @@ 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,
@@ -441,16 +440,10 @@ 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 and not _prefer_registered_runtime_cache:
+                if _dist_offload:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -462,16 +455,6 @@ 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
 
@@ -508,8 +491,7 @@ class DiffusersPipelineLoader:
                         )
                     if _dist_offload and plan_result is not None:
                         logger.info(
-                            "DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s",
-                            plan_result.fallback_code or "unknown",
+                            "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
                             plan_result.fallback_reason,
                         )
                     logger.debug("Loading weights on %s ...", load_device)
@@ -534,107 +516,7 @@ class DiffusersPipelineLoader:
                 logger.info("Quantization complete, offloaded model back to CPU")
 
         self._apply_skip_softmax_calibration(model)
-        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",
-            )
+        return model.eval()
 
     @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."""
 
-    storage_key: str
+    checkpoint_key: str
     file_path: str
     transform: TensorTransform | None = None
 
@@ -40,8 +40,6 @@ 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)
@@ -50,7 +48,6 @@ class HostWeightPlanResult:
 
     plan: HostWeightPlan | None
     fallback_reason: str | None = None
-    fallback_code: str | None = None
 
 
 class _PlanIncompatibleError(RuntimeError):
@@ -261,11 +258,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.storage_key not in available:
+                if binding.checkpoint_key not in available:
                     raise _PlanIncompatibleError(
-                        f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}"
+                        f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
                     )
-                source = handle.get_slice(binding.storage_key)
+                source = handle.get_slice(binding.checkpoint_key)
                 source_shape = tuple(source.get_shape())
                 source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
                 if source_shape != tuple(target.shape):
@@ -290,19 +287,11 @@ 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",
-            "unsupported_tp",
-        )
+        return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
     if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp")
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
     if online_quantization:
-        return HostWeightPlanResult(
-            None,
-            "online quantization requires the ordinary loader",
-            "unsupported_online_quantization",
-        )
+        return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
     if not callable(remap_fn):
@@ -339,14 +328,14 @@ def build_checkpoint_mmap_plan(
                 )
             checkpoint_key, file_path = model_to_ckpt[runtime_name]
             bindings[runtime_name] = TensorBinding(
-                storage_key=checkpoint_key,
+                checkpoint_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), "checkpoint_plan_incompatible")
+        return HostWeightPlanResult(None, str(exc))
 
     return HostWeightPlanResult(
         HostWeightPlan(

vllm_omni/diffusion/model_loader/runtime_weight_cache.py

@@ -1,700 +0,0 @@
-# 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,9 +70,6 @@ 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.
@@ -152,7 +149,6 @@ 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

@@ -1,18 +0,0 @@
-# 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

@@ -1,302 +0,0 @@
-# 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

@@ -1,481 +0,0 @@
-# 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,9 +71,8 @@ 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. Checkpoint plans
-            may have skipped ordinary materialization; runtime-cache plans
-            describe already-valid ordinary tensors that can be remapped.
+        host_weight_plan: Exact loader-produced backing plan, if ordinary
+            weight materialization was skipped.
 
     Returns:
         OffloadBackend instance or None if offloading disabled
@@ -86,14 +85,10 @@ 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:
-        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
+        raise RuntimeError(
+            "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
+        )
 
     # Return None if no offloading requested
     if config.strategy == OffloadStrategy.NONE:
@@ -101,7 +96,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 plan_requires_backend:
+        if host_weight_plan is not None:
             raise RuntimeError(
                 "The loader skipped ordinary weight materialization for DLO, "
                 "but this platform cannot create the required offload backend"
@@ -117,7 +112,7 @@ def get_offload_backend(
         try:
             device = current_omni_platform.get_torch_device()
         except (NotImplementedError, AttributeError) as exc:
-            if plan_requires_backend:
+            if host_weight_plan is not None:
                 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,9 +31,6 @@ 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":
@@ -104,7 +101,6 @@ 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:
@@ -140,7 +136,6 @@ 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

@@ -1,205 +0,0 @@
-# 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,7 +18,6 @@ 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
 
@@ -35,11 +34,6 @@ 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,
@@ -99,7 +93,6 @@ 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()
@@ -451,26 +444,6 @@ 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]]],
@@ -483,7 +456,18 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                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()}"
+                    )
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -497,32 +481,6 @@ 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]
@@ -556,27 +514,17 @@ 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:
-            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
+            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]
@@ -786,7 +734,6 @@ 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]] = []
@@ -857,14 +804,6 @@ 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]
@@ -962,11 +901,8 @@ 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."""
@@ -1087,7 +1023,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     framework="pt",
                     device="cpu",
                 )
-            tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
+            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
             if is_parameter:
                 replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                 parent._parameters[leaf_name] = replacement
@@ -1167,154 +1103,6 @@ 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.
 
@@ -1655,31 +1443,17 @@ 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":
-                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:
+            if self.host_weight_plan.backing_kind != "checkpoint_mmap":
                 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: file-backed pages are node-shared; "
-                    "transfer setup will select registered direct H2D or bounded host staging"
+                    "DLO rank-local mmap storage enabled: checkpoint pages are "
+                    "node-shared; each worker owns only two bounded host staging slots"
                 )
         else:
             remaining_meta = [
@@ -1851,13 +1625,6 @@ 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
@@ -1873,18 +1640,9 @@ 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 and not self._using_registered_mmap:
+        if self._using_rank_local_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1966,7 +1724,6 @@ 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
@@ -1991,7 +1748,6 @@ 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:
@@ -2007,11 +1763,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
-        if registration_released:
-            self._release_mmap_handles()
+        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,11 +136,6 @@ _DIFFUSION_MODELS = {
         "pipeline_wan2_2_i2v",
         "WanI2VDMD2Pipeline",
     ),
-    "LingBotWorldCausalDMDPipeline": (
-        "lingbot_world",
-        "pipeline",
-        "LingBotWorldCausalDMDPipeline",
-    ),
     "LongCatImagePipeline": (
         "longcat_image",
         "pipeline_longcat_image",
@@ -358,7 +353,6 @@ _NO_CACHE_ACCELERATION = {
     # Pipelines that do not support cache acceleration (cache_dit / tea_cache).
     "NextStep11Pipeline",
     "AudioXPipeline",
-    "LingBotWorldCausalDMDPipeline",
 }
 
 
@@ -553,7 +547,6 @@ _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",
@@ -614,7 +607,6 @@ _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

@@ -1,27 +0,0 @@
-"""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,25 +444,6 @@ 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],
@@ -805,18 +786,11 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        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()
+        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()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -543,10 +543,6 @@ 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,9 +982,6 @@ 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),
@@ -1000,10 +997,6 @@ 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": (
@@ -1122,13 +1115,15 @@ 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,39 +722,6 @@ 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
+from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs
 from vllm_omni.diffusion.utils.hf_utils import (
     _looks_like_dreamzero,
     get_diffusion_model_index,
@@ -369,12 +369,11 @@ 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, uses
+    For models registered in the pipeline registry (new path), uses
     ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges
     PipelineConfig + DeployConfig + CLI overrides.
 
-    Models that cannot be resolved through the registry return no stages so the
-    caller can apply its default stage configuration, when available.
+    For other models (legacy path), loads stage configs from YAML.
 
     Args:
         model: Model name or path (used to determine model_type)
@@ -384,7 +383,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.
+            registry-merged stages (opt-in; ignored on the legacy YAML path).
 
     Returns:
         ``(stage_configs, omni_lb_policy)``: the list of stage configuration
@@ -428,17 +427,69 @@ 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
 
-    strategy_note = ""
+    # 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.
     if strategy_config_path is not None:
-        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,
+        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,
     )
-    return [], None
+    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
 
 
 def filter_stages(

vllm_omni/experimental/ar_diffusion/__init__.py

@@ -13,32 +13,12 @@ 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,9 +58,6 @@ 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.
 
@@ -82,7 +79,6 @@ 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 = {
@@ -105,11 +101,6 @@ 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

@@ -1,168 +0,0 @@
-# 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,6 +32,7 @@ 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__ = [
@@ -52,4 +53,5 @@ __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,8 +23,7 @@ class ARDiffusionKVConfig:
     sink_chunks: int = 0
     # Boundary reset vs. sliding replacement.
     reset_at_boundary: bool = False
-    # Fraction of free device memory used to admit additional resident
-    # sessions. One session is admitted whenever it fits actual free memory.
+    # Fraction of free device memory budgeted for the AR-Diffusion KV pool.
     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,7 +174,6 @@ 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")
@@ -198,17 +197,13 @@ class ARDiffusionKVCache:
 
         self.config = config
         self.kv_branches = kv_branches
-        self.requested_session_capacity = session_capacity
+        self.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
@@ -222,11 +217,28 @@ 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,
@@ -238,7 +250,41 @@ 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.
@@ -251,101 +297,6 @@ 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
 
@@ -610,11 +561,6 @@ 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,6 +10,19 @@ 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] = {}
 
 
@@ -37,9 +50,6 @@ 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
@@ -218,6 +228,7 @@ 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:
@@ -225,9 +236,6 @@ 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,
@@ -468,18 +476,17 @@ 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. 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.
+# 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.
 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,
-    key_pool: torch.Tensor,
-    value_pool: torch.Tensor,
-    block_size: int,
+    layer_idx: torch.Tensor,
     video_slots: torch.Tensor,
     action_slots: torch.Tensor,
     block_table: torch.Tensor,
@@ -489,17 +496,21 @@ def _paged_write_attn_impl(
     max_seq_len: int,
     softmax_scale: float,
 ) -> torch.Tensor:
-    key_pool[video_slots] = k_curr.to(key_pool.dtype)
-    value_pool[video_slots] = v_curr.to(value_pool.dtype)
+    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)
     if k_act is not None and v_act is not None and k_act.shape[0] > 0:
-        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))
+        k_pool[action_slots] = k_act.to(k_pool.dtype)
+        v_pool[action_slots] = v_act.to(v_pool.dtype)
     return ar_diffusion_paged_attention(
         query,
-        key_cache,
-        value_cache,
+        kv.key_cache(layer_idx),
+        kv.value_cache(layer_idx),
         block_table=block_table,
         query_start_loc=query_start_loc,
         seq_lens=seq_lens,
@@ -514,19 +525,14 @@ 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=("key_pool", "value_pool"),
-    )
+    @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=())
     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,
-        key_pool: torch.Tensor,
-        value_pool: torch.Tensor,
-        block_size: int,
+        layer_idx: torch.Tensor,
         video_slots: torch.Tensor,
         action_slots: torch.Tensor,
         block_table: torch.Tensor,
@@ -542,9 +548,7 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
             v_curr,
             k_act,
             v_act,
-            key_pool,
-            value_pool,
-            block_size,
+            layer_idx,
             video_slots,
             action_slots,
             block_table,
@@ -562,9 +566,7 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
         v_curr,
         k_act,
         v_act,
-        key_pool,
-        value_pool,
-        block_size,
+        layer_idx,
         video_slots,
         action_slots,
         block_table,
@@ -587,9 +589,7 @@ def paged_write_attn(
         v_curr,
         k_act,
         v_act,
-        inputs.key_pool,
-        inputs.value_pool,
-        inputs.block_size,
+        inputs.layer_idx,
         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):
-                self.kv_cache.commit_chunk(ctx.adapter)
+                ctx.adapter.on_chunk_committed()
             self._committed[kv_branch] += ctx.seq_len
             _log.debug(
                 "AR-Diffusion COMMIT [%s] new_tokens=%d chunks=%d resident=%d/%d",
@@ -173,17 +173,6 @@ 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,7 +21,6 @@ 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__)
 
@@ -52,6 +51,7 @@ 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,6 +123,7 @@ 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,
@@ -133,14 +134,12 @@ 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=spec.session_capacity,
+            session_capacity=self._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 "
@@ -231,14 +230,9 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
         return state
 
     @staticmethod
-    def _request_session(
-        req: OmniDiffusionRequest,
-    ) -> tuple[str, dict, ARDiffusionTickRequest | None]:
+    def _request_session(req: OmniDiffusionRequest) -> tuple[str, dict]:
         extra_args = req.sampling_params.extra_args or {}
-        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
+        return str(extra_args.get("session_id") or "default"), extra_args
 
     def execute_model(
         self,
@@ -251,10 +245,8 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
         if capability is None:
             raise RuntimeError("AR-Diffusion capability missing after KV cache initialization")
 
-        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:
+        session_id, extra_args = self._request_session(req)
+        if extra_args.get("reset", False):
             self.reset_session(session_id)
         state = self._get_or_create_session(session_id)
         started = time.perf_counter()
@@ -276,7 +268,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
             )
             raise
         self._perf_e2e_times.append(time.perf_counter() - started)
-        if close_session:
+        if extra_args.get("close_session", False):
             self.close_session(session_id)
         return output
 

vllm_omni/experimental/ar_diffusion/session.py

@@ -1,578 +0,0 @@
-# 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

@@ -1,202 +0,0 @@
-# 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),
-        }