diff --git a/docs/.nav.yml b/docs/.nav.yml
index ffaa46368..d367ded05 100644
--- a/docs/.nav.yml
+++ b/docs/.nav.yml
@@ -116,6 +116,7 @@ nav:
- Async Diffusion Output: design/feature/async_diffusion_output.md
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
- Automatic Prefix Caching: design/feature/prefix_caching.md
+ - Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
- Communication:
- OmniConnector Implementations:
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
diff --git a/docs/design/ar_diffusion_pipeline_capability.md b/docs/design/ar_diffusion_pipeline_capability.md
index c9e05c8e8..3c1b4b9cd 100644
--- a/docs/design/ar_diffusion_pipeline_capability.md
+++ b/docs/design/ar_diffusion_pipeline_capability.md
@@ -4,6 +4,9 @@
> `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
> direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
> but it is not yet part of the stable module-design hierarchy.
+>
+> The end-to-end request, identity, failure, and memory contracts are documented
+> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
`ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
autoregressive attention KV across requests. The engine only selects
@@ -51,14 +54,13 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
those paths requires a batch-aware state-binding contract rather than silently
bypassing `bind_ar_diffusion_state()`.
-The current implementation retains one session per model instance. A request
-for a different `session_id` closes the resident session through the same LRU
-lifecycle used by explicit close and failure cleanup before allocating the new
-one. Self-attention and cross-attention memory are therefore sized for one
-resident session even if the pipeline declares a larger `session_capacity`.
-This prioritizes rollout throughput and prevents block-pool exhaustion before a
-count-based eviction can run. Multi-session residency can be added later with a
-memory-backed capacity calculation or allocation-pressure eviction.
+The implementation may retain multiple sessions per model instance. It selects
+the largest feasible capacity up to the pipeline's `session_capacity`, based on
+self-attention pages, cross-attention allocations, scratch, and declared
+model-owned state. One session is admitted whenever it fits actual free device
+memory; `gpu_memory_fraction` controls additional residency. Requests beyond
+that capacity evict the least-recently-used session through the same lifecycle
+hook used by explicit close and failure cleanup.
Scratch storage is also capability-driven. `frames_per_block` reserves space for
an uncommitted current video block, while `max_scratch_tokens_per_branch`
diff --git a/docs/design/feature/realtime_ar_diffusion.md b/docs/design/feature/realtime_ar_diffusion.md
new file mode 100644
index 000000000..5963befe7
--- /dev/null
+++ b/docs/design/feature/realtime_ar_diffusion.md
@@ -0,0 +1,213 @@
+# Realtime AR-Diffusion sessions
+
+> **Status:** experimental. This document describes the internal execution
+> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
+> define a public HTTP or WebSocket schema.
+
+## Scope
+
+Realtime AR-Diffusion serves world models that generate one autoregressive
+latent block per request while retaining model state across requests. The
+generic runtime owns identity, event ordering, session lifecycle, routing, and
+KV allocation. A model adapter owns model-specific controls and conditioning.
+
+LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
+same runner capability with two logical CFG branches.
+
+## Identities
+
+Four identifiers have distinct lifetimes:
+
+| Identifier | Lifetime | Owner | Contract |
+| --- | --- | --- | --- |
+| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
+| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
+| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
+| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
+
+`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
+`request_id` remains inside the immutable AR-Diffusion snapshot and the
+standard output metadata, while caller-visible `OmniRequestOutput.request_id`
+continues to use AsyncOmni's existing external-ID contract. Protocol identity
+therefore does not require a private entrypoint override.
+
+## Request lifecycle
+
+1. `ARDiffusionSession.accept_event()` validates and queues an immutable
+ `ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
+ the same event.
+2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
+ `_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
+ of the current prompt, reduced controls, ordered event IDs, and chunk index.
+3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
+ parameters, stores the typed tick in
+ `sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
+ prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
+4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
+ `OmniDiffusionRequest` through the orchestrator.
+5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
+ `_get_or_create_session()` obtains runner-owned paged KV, and
+ `bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
+ of `execute_model()`.
+6. The model pipeline validates the tick and generates exactly one latent
+ block. LingBot interprets camera controls in
+ `LingBotCameraControlReducer` and
+ `LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
+ camera or action schemas.
+7. The pipeline returns one `DiffusionOutput` with the generated payload and an
+ `ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
+ `AsyncOmni` exposes it as
+ `OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
+8. The consumer parses the standard envelope. The session commits its reducer
+ state, prompt, controls, event IDs, and next chunk index only when returned
+ metadata exactly equals the submitted tick snapshot.
+
+The metadata envelope is:
+
+```json
+{
+ "metadata": {
+ "ar_diffusion": {
+ "session_id": "world-7",
+ "request_id": "ar-world-7-3-...",
+ "chunk_index": 3,
+ "applied_event_ids": [10, 11]
+ }
+ }
+}
+```
+
+## Transaction and failure semantics
+
+Accepted events are not acknowledged as applied until the model output and
+metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
+and the session snapshot form one logical commit. A model, metadata, or reducer
+failure leaves the chunk index and queued events uncommitted and transitions the
+session to `FAILED`.
+
+After a tick failure, worker state is closed because a forward may have
+partially modified or already released KV. The same chunk is not retried in
+place. A caller must explicitly reset the session, starting again at chunk
+zero, or close it.
+
+Explicit close follows this state machine:
+
+```text
+ACTIVE or FAILED
+ |
+ v
+ CLOSING ---- worker cleanup succeeds ----> CLOSED
+ |
+ +------- worker cleanup fails --------> CLEANUP_FAILED
+ |
+ +-- close retry --> CLOSING
+```
+
+`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
+creation of the same `session_id`, and rejects reset. Only an idempotent close
+retry that succeeds removes the session. Transport disconnect uses the same
+path.
+
+## KV ownership and capacity
+
+`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
+declares:
+
+- TP-local self-attention geometry, AR block size, sink, and recent window;
+- logical KV branches and their worker-local slots;
+- fixed-length cross-attention caches;
+- scratch tokens required by an uncommitted forward;
+- persistent model-owned CUDA bytes per resident session; and
+- the requested resident-session capacity.
+
+The runner computes:
+
+```text
+required(capacity) =
+ scratch
+ + managed_self_attention_pages(capacity)
+ + capacity * cross_attention_bytes_per_session
+ + capacity * model_owned_state_bytes_per_session
+```
+
+`gpu_memory_fraction` is a soft expansion budget:
+
+```text
+configured_budget = available_device_bytes * gpu_memory_fraction
+effective_budget = max(configured_budget, required(1))
+```
+
+If `required(1)` exceeds actual available device memory, initialization fails.
+Otherwise the runner always admits one viable session and uses the configured
+fraction to select additional resident sessions up to the model-declared cap.
+All model-owned reservations are deducted before the paged self-attention pool
+is allocated.
+
+LingBot reports its persistent full-horizon image-condition tensor. DreamZero
+reports the measured 603 MiB per-session Wan VAE causal-convolution state.
+
+## Concurrency and routing limits
+
+The current runner executes one request at a time:
+
+- `max_num_seqs` must be one;
+- request batching and diffusion step execution are rejected;
+- one AR block is generated per engine request; and
+- an AR-Diffusion stage must have exactly one replica.
+
+The single-replica restriction is intentional. Multi-replica support requires
+session-affine routing so every tick, reset, and close reaches the worker that
+owns the session. Random or round-robin routing is not correct.
+
+Multiple resident sessions may share one runner. Capacity selection accounts
+for all persistent pools, and the runner performs LRU eviction through the same
+model lifecycle hook used by explicit close.
+
+## Model boundary
+
+The generic runtime transports controls as opaque
+`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
+and reduce those values. This keeps target/velocity, key state, pose,
+intrinsics, and other model-specific conventions out of the session protocol.
+
+For LingBot World v2:
+
+- `lingbot_world/actions.py` parses key-state and trajectory controls;
+- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
+- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
+- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
+ produces one block plus the standard metadata envelope.
+
+## Non-goals
+
+This contract does not currently provide:
+
+- a public HTTP or WebSocket request schema (the structured-interaction
+ frontend is tracked separately in
+ [#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
+- camera/action semantics shared by every world model;
+- session migration or replication across workers;
+- retry of an ambiguous partially executed chunk;
+- cross-stage KV transfer; or
+- stateful streaming VAE decode.
+
+Those features can be layered on top without adding LingBot-specific identity
+or lifecycle fields to the generic runtime.
+
+## Required regression coverage
+
+CPU contract tests cover:
+
+- separation of tick correlation IDs from internal engine routing IDs;
+- event ordering, metadata equality, and reducer fail-closed commit;
+- close/disconnect failure tombstones and cleanup retry;
+- capacity one under the default fraction for shipped LingBot geometry;
+- capacity expansion under a tuned fraction;
+- DreamZero model-owned state accounting;
+- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
+- LingBot model registration, imports, camera/action reduction, and one-block
+ output metadata.
+
+GPU validation must additionally exercise real model loading, action input,
+prompt switching, rolling-window replacement, reset, close, finite output
+latents, and exact metadata/request identity.
diff --git a/docs/design/index.md b/docs/design/index.md
index 0db152335..acac90458 100644
--- a/docs/design/index.md
+++ b/docs/design/index.md
@@ -22,6 +22,7 @@ implementation contract; it is not, by itself, a general support claim.
- [Async Diffusion Output](feature/async_diffusion_output.md)
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
+- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
### Communication
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 5097c83fa..62645aae6 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -37,6 +37,7 @@ th {
| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -117,6 +118,10 @@ support metadata until their recipe is audited.
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
+See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
+for offline and experimental realtime usage, supported controls, validation,
+and current limitations.
+
H3 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
diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py
new file mode 100644
index 000000000..145f5f641
--- /dev/null
+++ b/examples/offline_inference/diffusion/lingbot_world_v2.py
@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Generate a LingBot-World v2 video from an image and camera trajectory.
+
+The official checkpoint is licensed separately under CC BY-NC-SA and is
+restricted to non-commercial use. This vLLM-Omni integration remains
+Apache-2.0 licensed.
+
+Example:
+ python examples/offline_inference/diffusion/lingbot_world_v2.py \
+ --prompt "The camera moves slowly forward through the scene." \
+ --image /path/to/first_frame.png \
+ --action-dir /path/to/actions/forward \
+ --output lingbot_world_v2.mp4
+
+``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
+example configures its parent as the trusted action root and sends the
+contained directory through ``extra_args["action_path"]``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
+_NUM_INFERENCE_STEPS = 4
+_MAX_SEQUENCE_LENGTH = 512
+_MAX_PIXEL_AREA = 480 * 832
+_MAX_RAW_FRAMES = 117
+_MAX_SOURCE_CAMERA_FRAMES = 4096
+_TEMPORAL_COMPRESSION = 4
+_LATENT_FRAMES_PER_BLOCK = 3
+_SPATIAL_ALIGNMENT = 16
+
+
+@dataclass(frozen=True)
+class LingBotPaths:
+ image: Path
+ action_dir: Path
+ action_root: Path
+ action_relative: Path
+ camera_frames: int
+ output: Path
+
+
+def _validate_camera_arrays(action_dir: Path) -> int:
+ """Validate the bounded camera contract before starting the engine."""
+
+ try:
+ poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
+ intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
+ except (OSError, TypeError, ValueError):
+ raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
+ if poses.ndim != 3 or poses.shape[1:] != (4, 4):
+ raise ValueError("poses.npy must have shape [frames, 4, 4].")
+ if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
+ raise ValueError("intrinsics.npy must have shape [frames, 4].")
+ if poses.shape[0] != intrinsics.shape[0]:
+ raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
+ if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
+ raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
+ if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
+ raise ValueError("Camera arrays must contain real numeric values.")
+ if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
+ raise ValueError("Camera arrays must contain only finite values.")
+ return int(poses.shape[0])
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
+ )
+ parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
+ parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
+ parser.add_argument("--image", required=True, help="Path to the first-frame image.")
+ parser.add_argument(
+ "--action-dir",
+ required=True,
+ help="Camera action directory containing poses.npy and intrinsics.npy.",
+ )
+ parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
+ parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
+ parser.add_argument(
+ "--num-frames",
+ type=int,
+ default=9,
+ help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
+ )
+ parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+ parser.add_argument(
+ "--tensor-parallel-size",
+ type=int,
+ default=1,
+ help="Number of GPUs used for tensor parallelism inside the DiT.",
+ )
+ parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
+ parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
+ parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
+ parser.add_argument(
+ "--enforce-eager",
+ action="store_true",
+ help="Disable compilation and force eager execution.",
+ )
+ return parser.parse_args(argv)
+
+
+def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
+ """Canonicalize local inputs and construct a safe trusted-root mapping."""
+
+ image = Path(args.image).expanduser().resolve()
+ if not image.is_file():
+ raise ValueError("--image must point to an existing file.")
+
+ action_dir = Path(args.action_dir).expanduser().resolve()
+ if not action_dir.is_dir():
+ raise ValueError("--action-dir must point to an existing directory.")
+ if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
+ raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
+ camera_frames = _validate_camera_arrays(action_dir)
+
+ action_root = action_dir.parent.resolve()
+ try:
+ action_relative = action_dir.relative_to(action_root)
+ except ValueError:
+ raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
+ if not action_relative.parts:
+ raise ValueError("--action-dir must name a directory contained by a trusted action root.")
+
+ output = Path(args.output).expanduser().resolve()
+ return LingBotPaths(
+ image=image,
+ action_dir=action_dir,
+ action_root=action_root,
+ action_relative=action_relative,
+ camera_frames=camera_frames,
+ output=output,
+ )
+
+
+def _positive_finite(value: float, flag: str) -> float:
+ result = float(value)
+ if not math.isfinite(result) or result <= 0:
+ raise ValueError(f"{flag} must be a positive finite number.")
+ return result
+
+
+def build_omni_kwargs(
+ args: argparse.Namespace,
+ paths: LingBotPaths,
+ *,
+ parallel_config: Any,
+) -> dict[str, Any]:
+ """Build Omni configuration without overriding checkpoint class discovery."""
+
+ if args.tensor_parallel_size <= 0:
+ raise ValueError("--tensor-parallel-size must be a positive integer.")
+ flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
+ model_path = Path(args.model).expanduser()
+ model = str(model_path.resolve()) if model_path.exists() else args.model
+ return {
+ "model": model,
+ "flow_shift": flow_shift,
+ "parallel_config": parallel_config,
+ "enforce_eager": args.enforce_eager,
+ "model_config": {"lingbot_action_root": str(paths.action_root)},
+ }
+
+
+def build_request(
+ args: argparse.Namespace,
+ paths: LingBotPaths,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+ """Build the prompt and fixed v1 causal-DMD sampling contract."""
+
+ prompt_text = args.prompt.strip()
+ if not prompt_text:
+ raise ValueError("--prompt must contain non-whitespace text.")
+ if args.height <= 0 or args.width <= 0:
+ raise ValueError("--height and --width must be positive integers.")
+ if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
+ raise ValueError("--height and --width must each be divisible by 16.")
+ if args.height * args.width > _MAX_PIXEL_AREA:
+ raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
+ if args.num_frames <= 0:
+ raise ValueError("--num-frames must be a positive integer.")
+ if args.num_frames > _MAX_RAW_FRAMES:
+ raise ValueError("--num-frames must not exceed 117 raw frames.")
+ if args.num_frames > paths.camera_frames:
+ raise ValueError("Camera arrays must contain at least --num-frames frames.")
+ if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
+ raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
+ latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
+ if latent_frames % _LATENT_FRAMES_PER_BLOCK:
+ raise ValueError("--num-frames must map to whole three-frame latent blocks.")
+ if args.fps <= 0:
+ raise ValueError("--fps must be a positive integer.")
+ flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
+
+ prompt = {
+ "prompt": prompt_text,
+ "multi_modal_data": {"image": str(paths.image)},
+ }
+ sampling_kwargs = {
+ "height": args.height,
+ "width": args.width,
+ "num_frames": args.num_frames,
+ "num_inference_steps": _NUM_INFERENCE_STEPS,
+ "max_sequence_length": _MAX_SEQUENCE_LENGTH,
+ "seed": args.seed,
+ "fps": args.fps,
+ "extra_args": {
+ "action_path": paths.action_relative.as_posix(),
+ "flow_shift": flow_shift,
+ },
+ }
+ return prompt, sampling_kwargs
+
+
+def extract_video_array(outputs: Any) -> np.ndarray:
+ """Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
+
+ if not isinstance(outputs, list) or len(outputs) != 1:
+ raise ValueError("Expected exactly one Omni output for the single LingBot request.")
+ output = outputs[0]
+ request_output = getattr(output, "request_output", None)
+ if request_output is not None and getattr(request_output, "images", None):
+ output = request_output
+
+ images = getattr(output, "images", None)
+ if not images:
+ multimodal_output = getattr(output, "multimodal_output", None) or {}
+ images = multimodal_output.get("video")
+ if images is None:
+ raise ValueError("No video frames found in Omni output.")
+ if isinstance(images, list) and len(images) == 1:
+ images = images[0]
+
+ video: np.ndarray = np.asarray(images)
+ if video.ndim == 5:
+ if video.shape[0] != 1:
+ raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
+ video = video[0]
+ if video.ndim != 4 or video.shape[-1] not in (3, 4):
+ raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
+ return video
+
+
+def main(argv: Sequence[str] | None = None) -> Path:
+ args = parse_args(argv)
+ paths = resolve_cli_paths(args)
+ prompt, sampling_kwargs = build_request(args, paths)
+ paths.output.parent.mkdir(parents=True, exist_ok=True)
+
+ # Keep vLLM imports below pure CLI/path validation so ``--help`` and the
+ # helper tests work on machines without the native vLLM runtime.
+ from diffusers.utils import export_to_video
+
+ from vllm_omni.diffusion.data import DiffusionParallelConfig
+ from vllm_omni.entrypoints.omni import Omni
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ parallel_config = DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size)
+ omni_kwargs = build_omni_kwargs(args, paths, parallel_config=parallel_config)
+ omni = Omni(**omni_kwargs)
+ try:
+ outputs = omni.generate(
+ prompt,
+ OmniDiffusionSamplingParams(**sampling_kwargs),
+ use_tqdm=False,
+ )
+ finally:
+ omni.close()
+
+ frames = extract_video_array(outputs)
+ export_to_video(frames, str(paths.output), fps=args.fps)
+ print(f"Saved {frames.shape[0]} frames to {paths.output}")
+ return paths.output
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py
new file mode 100644
index 000000000..3d6d7a7e1
--- /dev/null
+++ b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py
@@ -0,0 +1,239 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
+
+This example exercises the internal session API. It is not a public HTTP or
+WebSocket protocol. Each JSONL event advances the world by one three-latent
+frame AR block and writes the latent plus its identity metadata. The current
+scoped implementation supports at most ten ticks per generation epoch because
+its image condition is bounded to 117 pixel frames; reset or create a session
+to start a new world.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import math
+import time
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
+_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
+_FRAMES_PER_BLOCK = 3
+# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
+_MAX_REALTIME_TICKS = 10
+
+
+def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
+ """Build the event-side script consumed by LingBotCameraControlReducer."""
+ return {"mode": "script", "frames": frames}
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
+ parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
+ parser.add_argument("--image", required=True, help="Initial RGB image.")
+ parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
+ parser.add_argument(
+ "--events",
+ required=True,
+ help="JSONL file with one prompt/action event per AR block; at most 10 events.",
+ )
+ parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
+ parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
+ parser.add_argument("--height", type=int, default=480)
+ parser.add_argument("--width", type=int, default=832)
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
+ parser.add_argument("--tensor-parallel-size", type=int, default=1)
+ parser.add_argument("--enforce-eager", action="store_true")
+ return parser.parse_args(argv)
+
+
+def _load_events(path: Path) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+ for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
+ line = raw_line.strip()
+ if not line:
+ continue
+ try:
+ value = json.loads(line)
+ except json.JSONDecodeError as exc:
+ raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
+ if not isinstance(value, dict):
+ raise ValueError(f"events line {line_number} must be a JSON object.")
+ event_id = value.get("event_id")
+ if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
+ raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
+ prompt = value.get("prompt")
+ if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
+ raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
+ frames = value.get("frames")
+ if frames is not None:
+ if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
+ raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
+ for frame in frames:
+ if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
+ raise ValueError(f"events line {line_number} frames must contain only action strings.")
+ if prompt is None and frames is None:
+ raise ValueError(f"events line {line_number} must update prompt and/or frames.")
+ events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
+ if not events:
+ raise ValueError("events file must contain at least one event.")
+ if len(events) > _MAX_REALTIME_TICKS:
+ raise ValueError(
+ "events file must contain at most "
+ f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
+ "image-condition horizon is 117 pixel frames."
+ )
+ event_ids = [event["event_id"] for event in events]
+ if event_ids != sorted(set(event_ids)):
+ raise ValueError("event_id values must be unique and strictly increasing.")
+ return events
+
+
+def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
+ image = Path(args.image).expanduser().resolve()
+ events = Path(args.events).expanduser().resolve()
+ output_dir = Path(args.output_dir).expanduser().resolve()
+ if not image.is_file():
+ raise ValueError("--image must point to an existing file.")
+ if not events.is_file():
+ raise ValueError("--events must point to an existing JSONL file.")
+ if not args.prompt.strip():
+ raise ValueError("--prompt must contain non-whitespace text.")
+ if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
+ raise ValueError("--height and --width must be positive multiples of 16.")
+ if args.tensor_parallel_size <= 0:
+ raise ValueError("--tensor-parallel-size must be positive.")
+ if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
+ raise ValueError("--gpu-memory-fraction must be in (0, 1].")
+ return image, events, output_dir
+
+
+async def run(argv: Sequence[str] | None = None) -> Path:
+ args = parse_args(argv)
+ image, events_path, output_dir = _validate_args(args)
+ events = _load_events(events_path)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Imports stay below pure input validation so --help and helper tests do not
+ # require a CUDA-enabled vLLM installation.
+ import torch
+
+ from vllm_omni.diffusion.data import DiffusionParallelConfig
+ from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
+ from vllm_omni.entrypoints.async_omni import AsyncOmni
+ from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
+ from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionSessionEvent,
+ ARDiffusionSessionManager,
+ ARDiffusionWorkerLifecycle,
+ )
+ from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ engine = AsyncOmni(
+ model=args.model,
+ engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
+ enforce_eager=args.enforce_eager,
+ parallel_config=DiffusionParallelConfig(tensor_parallel_size=args.tensor_parallel_size),
+ max_num_seqs=1,
+ model_config={
+ "ar_diffusion_height": args.height,
+ "ar_diffusion_width": args.width,
+ "ar_diffusion_kv_config": {
+ "gpu_memory_fraction": args.gpu_memory_fraction,
+ "warmup_cudagraph": True,
+ },
+ },
+ )
+ sampling = OmniDiffusionSamplingParams(
+ height=args.height,
+ width=args.width,
+ num_frames=9,
+ num_inference_steps=4,
+ max_sequence_length=512,
+ seed=args.seed,
+ output_type="latent",
+ extra_args={"flow_shift": 5.0},
+ )
+ consumer = ARDiffusionOmniTickConsumer(
+ engine,
+ prompt_provider=lambda tick: {
+ "prompt": tick.prompt,
+ "multi_modal_data": {"image": str(image)},
+ },
+ sampling_params_list=[sampling],
+ diffusion_stage_id=0,
+ )
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
+ max_pending_events=32,
+ control_reducer_factory=LingBotCameraControlReducer,
+ )
+ session = await manager.create_session(args.session_id)
+ measurements: list[dict[str, Any]] = []
+ try:
+ for chunk_index, event in enumerate(events):
+ controls = ()
+ if event["frames"] is not None:
+ controls = (
+ ARDiffusionControlInput(
+ track="camera",
+ schema=_CAMERA_ACTION_SCHEMA,
+ data=_camera_event_data(event["frames"]),
+ ),
+ )
+ await session.accept_event(
+ ARDiffusionSessionEvent(
+ event_id=event["event_id"],
+ prompt=event["prompt"]
+ if event["prompt"] is not None
+ else (args.prompt if chunk_index == 0 else None),
+ controls=controls,
+ )
+ )
+ started = time.perf_counter()
+ output = await session.next_chunk()
+ elapsed = time.perf_counter() - started
+ if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
+ raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
+ latent = output.images[0].detach().float().cpu()
+ metadata = consumer.chunk_metadata(output).to_dict()
+ torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
+ (output_dir / f"chunk_{chunk_index:03d}.json").write_text(
+ json.dumps(metadata, indent=2, sort_keys=True) + "\n"
+ )
+ measurements.append(
+ {
+ "chunk_index": chunk_index,
+ "latency_seconds": elapsed,
+ "shape": list(latent.shape),
+ "finite": bool(torch.isfinite(latent).all()),
+ "metadata": metadata,
+ }
+ )
+ print(json.dumps(measurements[-1], sort_keys=True), flush=True)
+ finally:
+ try:
+ await manager.close_session(args.session_id)
+ finally:
+ engine.shutdown()
+
+ summary_path = output_dir / "summary.json"
+ summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
+ return summary_path
+
+
+def main(argv: Sequence[str] | None = None) -> Path:
+ return asyncio.run(run(argv))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/recipes/README.md b/recipes/README.md
index 6a3b23770..5fee4fdae 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -60,6 +60,7 @@ recipes/
| [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
+| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 |
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
| [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md
new file mode 100644
index 000000000..851aeb930
--- /dev/null
+++ b/recipes/Robbyant/LingBot-World-2.0.md
@@ -0,0 +1,122 @@
+# LingBot-World 2.0
+
+> Offline and experimental realtime interactive world generation
+
+## Summary
+
+- Vendor: Robbyant
+- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
+- Task: image-conditioned interactive world generation
+- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
+- Hardware validated: NVIDIA H200 and B200
+- Maintainer: Community
+
+The checkpoint is separately licensed under CC BY-NC-SA and restricted to
+non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
+
+## Offline generation
+
+The offline path consumes one source image and an action directory containing
+`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape
+`[frames, 4]`.
+
+```bash
+python examples/offline_inference/diffusion/lingbot_world_v2.py \
+ --prompt "The camera moves slowly forward through the scene." \
+ --image /path/to/first_frame.png \
+ --action-dir /path/to/actions/forward \
+ --num-frames 81 \
+ --output lingbot_world_v2.mp4
+```
+
+The current causal-fast checkpoint uses four DMD sampling steps. Raw frame
+counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
+longer bounded source trajectory (the official example contains 269 frames);
+the runtime consumes only the prefix needed by the request.
+
+## Realtime in-process generation
+
+The realtime example keeps AR-Diffusion KV and model-owned state across
+requests. Each JSONL line describes the prompt and/or three latent-frame camera
+actions applied at the next chunk boundary:
+
+```json
+{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]}
+{"event_id":2,"frames":[["w"],["w"],["w"]]}
+{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]}
+```
+
+Run:
+
+```bash
+python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \
+ --image /path/to/first_frame.png \
+ --events /path/to/events.jsonl \
+ --output-dir /tmp/lingbot-realtime \
+ --gpu-memory-fraction 0.6
+```
+
+The example writes one latent tensor and one metadata JSON file per chunk. It
+exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
+-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
+transport.
+
+This PR intentionally does not define a public realtime HTTP/WebSocket schema.
+An online serving client should be added together with the public transport
+API rather than exposing LingBot-specific event fields from the generic model
+runtime. The structured camera-interaction frontend and transport work is
+tracked separately in
+[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
+
+## Realtime identity and controls
+
+- `session_id` identifies the persistent world and its worker-owned state.
+- `event_id` identifies a prompt/control update and remains monotonic across reset.
+- `chunk_index` is contiguous from zero and restarts from zero after reset.
+- `request_id` correlates one chunk snapshot with its output metadata.
+- AsyncOmni uses a separate UUID-suffixed internal engine routing ID.
+
+The generic runtime transports controls as opaque snapshots. LingBot's adapter
+accepts:
+
+- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`,
+ `s`, `d`, `i`, `j`, `k`, and `l`;
+- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories.
+
+## Validation
+
+Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42.
+The exercised matrix includes:
+
+- TP=1 and TP=2 execution;
+- two interleaved resident sessions;
+- action input and prompt switching;
+- seven contiguous chunks crossing the sink plus recent rolling window;
+- direct versus paged replay;
+- CUDA Graph execution;
+- reset, close, failure cleanup, and exact metadata matching; and
+- nine-frame VAE decode.
+
+Official `generate.py` accuracy and performance numbers, together with
+generation artifacts, are recorded in the PR validation section for the exact
+tested commit.
+
+## Current limitations
+
+- Only the 14B causal-fast checkpoint is supported.
+- The realtime control plane is internal; there is no public server transport yet.
+- AR-Diffusion stages currently require one replica because session-affine
+ routing across replicas is not implemented.
+- One AR block is generated per request and `max_num_seqs` must be one.
+- Stateful streaming VAE decode is not implemented; the realtime example emits
+ latent chunks.
+- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
+ Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed.
+- No AMD GPU, Ascend NPU, or Intel GPU support is claimed.
+
+## References
+
+- Checkpoint:
+- Official implementation:
+- 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)
diff --git a/tests/diffusion/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py
index 7b94fe80f..10451c134 100644
--- a/tests/diffusion/ar_diffusion/test_capability_runner.py
+++ b/tests/diffusion/ar_diffusion/test_capability_runner.py
@@ -18,6 +18,7 @@ from vllm_omni.experimental.ar_diffusion.capability import (
)
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
BLOCK = 16
POS = "positive"
@@ -56,6 +57,21 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
)
+def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
+ return ARDiffusionKVCacheSpec(
+ num_layers=1,
+ num_kv_heads=1,
+ head_size=1,
+ tokens_per_frame=1,
+ frames_per_block=3,
+ window_frames=3,
+ sink_frames=3,
+ kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+ session_capacity=capacity,
+ cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
+ )
+
+
class CapablePipeline:
def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
self.spec = spec
@@ -103,6 +119,7 @@ def make_runner(
*,
available_bytes: int = 1 << 28,
step_execution: bool = False,
+ gpu_memory_fraction: float = 0.1,
) -> ARDiffusionModelRunner:
runner = object.__new__(ARDiffusionModelRunner)
runner.od_config = SimpleNamespace(
@@ -113,7 +130,10 @@ def make_runner(
)
runner.device = torch.device("cpu")
runner.pipeline = pipeline
- runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
+ runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
+ enable=True,
+ gpu_memory_fraction=gpu_memory_fraction,
+ )
runner.kv_cache = None
runner._ar_diffusion_capability = None
runner._ar_diffusion_kv_cache_spec = None
@@ -132,6 +152,16 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
return state
+def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
+ state = runner._get_or_create_session(session_id)
+ assert runner.kv_cache is not None
+ seq_len = BLOCK * runner.kv_cache.frames_per_block
+ ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
+ ctx.ensure_video_slots(torch.device("cpu"))
+ state.commit_paged_context(kv_branch)
+ return state
+
+
def test_ar_runner_rejects_pipeline_without_capability():
runner = object.__new__(ARDiffusionModelRunner)
runner.od_config = SimpleNamespace(max_num_seqs=1)
@@ -140,6 +170,43 @@ def test_ar_runner_rejects_pipeline_without_capability():
runner._preallocate_kv_cache(available_bytes=1 << 20)
+def test_runner_uses_typed_tick_as_authoritative_session_contract():
+ tick = ARDiffusionTickRequest(
+ session_id="world-7",
+ request_id="request-3",
+ chunk_index=3,
+ reset=True,
+ )
+ req = SimpleNamespace(
+ request_id="request-3",
+ sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
+ )
+
+ session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
+
+ assert session_id == "world-7"
+ assert extra_args == tick.to_extra_args()
+ assert parsed == tick
+
+
+def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
+ tick = ARDiffusionTickRequest(
+ session_id="world-7",
+ request_id="client-request-3",
+ chunk_index=3,
+ )
+ req = SimpleNamespace(
+ request_id="engine-request-uuid",
+ sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
+ )
+
+ session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
+
+ assert session_id == "world-7"
+ assert req.request_id == "engine-request-uuid"
+ assert parsed.request_id == "client-request-3"
+
+
def test_lingbot_like_single_branch_session_reuse_reset_and_close():
pipeline = CapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline)
@@ -170,6 +237,31 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
assert pipeline.closes == ["s1"]
+def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
+ runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
+ kv = runner.kv_cache
+ assert kv is not None
+
+ session_a = commit_one_block(runner, "world-a", "main")
+ a0_blocks = kv.window_block_ids(session_a.adapter("main"))
+ session_b = commit_one_block(runner, "world-b", "main")
+ b0_blocks = kv.window_block_ids(session_b.adapter("main"))
+ session_a_again = commit_one_block(runner, "world-a", "main")
+ a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
+
+ assert session_a_again is session_a
+ assert session_a.adapter("main").request_id == "ar::world-a::main"
+ assert session_b.adapter("main").request_id == "ar::world-b::main"
+ assert session_a.adapter("main").completed_chunks == 6
+ assert session_b.adapter("main").completed_chunks == 3
+ assert len(a0_blocks) == 3
+ assert len(b0_blocks) == 3
+ assert len(a1_blocks) == 6
+ assert set(a0_blocks) <= set(a1_blocks)
+ assert set(a1_blocks).isdisjoint(b0_blocks)
+ assert tuple(runner._sessions) == ("world-b", "world-a")
+
+
def test_lingbot_like_sink_survives_sliding_window_eviction():
runner = make_runner(CapablePipeline(lingbot_like_spec()))
kv = runner.kv_cache
@@ -198,12 +290,12 @@ def test_dreamzero_like_two_branches_are_independent():
def test_lru_eviction_releases_blocks_and_notifies_pipeline():
- pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
+ pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
runner = make_runner(pipeline)
kv = runner.kv_cache
assert kv is not None
- assert runner._session_capacity == 1
- assert kv.session_capacity == 1
+ assert runner._session_capacity == 2
+ assert kv.session_capacity == 2
free_total = kv.manager.block_pool.get_num_free_blocks()
old = commit_one_frame(runner, "old", "main")
k = torch.randn(1, 8, 4, 64)
@@ -211,13 +303,56 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
assert kv.manager.block_pool.get_num_free_blocks() < free_total
runner._get_or_create_session("new")
+ assert tuple(runner._sessions) == ("old", "new")
+ runner._get_or_create_session("newest")
- assert tuple(runner._sessions) == ("new",)
+ assert tuple(runner._sessions) == ("new", "newest")
assert pipeline.closes == ["old"]
assert "old" not in kv._cross_sessions
assert kv.manager.block_pool.get_num_free_blocks() == free_total
+def test_budget_reduced_capacity_drives_runner_lru():
+ pipeline = CapablePipeline(tiny_spec(capacity=2))
+ # One tiny LingBot-like session requires 128 bytes; two require 192.
+ runner = make_runner(
+ pipeline,
+ available_bytes=128,
+ gpu_memory_fraction=1.0,
+ )
+ kv = runner.kv_cache
+ assert kv is not None
+ assert kv.requested_session_capacity == 2
+ assert kv.session_capacity == 1
+ assert runner._session_capacity == 1
+
+ runner._get_or_create_session("old")
+ runner._get_or_create_session("new")
+
+ assert tuple(runner._sessions) == ("new",)
+ assert pipeline.closes == ["old"]
+
+
+def test_dreamzero_like_requested_capacity_is_capped_by_budget():
+ spec = dreamzero_like_spec(capacity=64)
+ # Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
+ # managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
+ # cross-attention=1 page/session, for 44 pages total.
+ page_bytes = 65_536
+ pipeline = CapablePipeline(spec)
+ runner = make_runner(
+ pipeline,
+ available_bytes=44 * page_bytes,
+ gpu_memory_fraction=1.0,
+ )
+ kv = runner.kv_cache
+ assert kv is not None
+ assert kv.requested_session_capacity == 64
+ assert kv.session_capacity == 2
+ assert runner._session_capacity == 2
+ assert kv.cross_attention_reserved_bytes == 2 * page_bytes
+
+
def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
pipeline = CapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline)
@@ -232,7 +367,10 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
raise RuntimeError("layer exploded")
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
- request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+ request = SimpleNamespace(
+ request_id="broken-request",
+ sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+ )
with pytest.raises(RuntimeError, match="layer exploded"):
runner.execute_model(request)
@@ -263,7 +401,10 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
runner.device = torch.device("cuda")
- request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+ request = SimpleNamespace(
+ request_id="broken-request",
+ sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+ )
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
runner.execute_model(request)
diff --git a/tests/diffusion/ar_diffusion/test_consumer.py b/tests/diffusion/ar_diffusion/test_consumer.py
new file mode 100644
index 000000000..d1197113c
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_consumer.py
@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionSession,
+ ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+from vllm_omni.inputs.data import (
+ OmniDiffusionSamplingParams,
+ OmniPromptType,
+ OmniSamplingParams,
+)
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeStagePool:
+ def __init__(self, num_replicas: int = 1) -> None:
+ self.num_replicas = num_replicas
+
+
+class FakeEngine:
+ def __init__(self, num_replicas: int = 1) -> None:
+ self.stage_pools = [FakeStagePool(num_replicas)]
+
+
+class FakeGenerateClient:
+ def __init__(
+ self,
+ *,
+ include_metadata: bool = True,
+ applied_event_ids: list[int] | None = None,
+ output_error: str | None = None,
+ num_replicas: int = 1,
+ ) -> None:
+ self.engine = FakeEngine(num_replicas)
+ self.include_metadata = include_metadata
+ self.applied_event_ids = applied_event_ids
+ self.output_error = output_error
+ self.calls: list[dict[str, Any]] = []
+
+ async def generate(
+ self,
+ prompt: OmniPromptType,
+ *,
+ request_id: str,
+ sampling_params_list: Sequence[OmniSamplingParams],
+ output_modalities: list[str] | None = None,
+ ) -> AsyncIterator[OmniRequestOutput]:
+ self.calls.append(
+ {
+ "prompt": prompt,
+ "request_id": request_id,
+ "sampling_params_list": sampling_params_list,
+ "output_modalities": output_modalities,
+ }
+ )
+ params = sampling_params_list[0]
+ assert isinstance(params, OmniDiffusionSamplingParams)
+ tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
+ assert tick is not None
+ chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
+ if self.applied_event_ids is not None:
+ chunk_metadata["applied_event_ids"] = self.applied_event_ids
+ metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
+ yield OmniRequestOutput(
+ request_id=request_id,
+ stage_id=0,
+ final_output_type="video",
+ images=["chunk"], # type: ignore[list-item]
+ _multimodal_output=metadata,
+ finished=True,
+ error=self.output_error,
+ )
+
+
+class FakeLifecycle:
+ async def reset_session(self, session_id: str) -> None:
+ pass
+
+ async def close_session(self, session_id: str) -> None:
+ pass
+
+
+def make_tick() -> ARDiffusionTickRequest:
+ return ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="world-1-chunk-3",
+ chunk_index=3,
+ applied_event_ids=(7, 8),
+ prompt="turn left",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"samples": [{"x": 1.0}]},
+ ),
+ ),
+ )
+
+
+def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
+ client = FakeGenerateClient(num_replicas=2)
+
+ with pytest.raises(
+ ValueError,
+ match=r"require exactly one replica.*num_replicas=2",
+ ):
+ ARDiffusionOmniTickConsumer(
+ client,
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ assert client.calls == []
+
+
+@pytest.mark.asyncio
+async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
+ client = FakeGenerateClient()
+ template = OmniDiffusionSamplingParams(
+ output_type="latent",
+ extra_args={"deployment_option": "kept"},
+ )
+ consumer = ARDiffusionOmniTickConsumer(
+ client,
+ prompt_provider=lambda tick: {
+ "prompt": tick.prompt,
+ "multi_modal_data": {"image": "initial-image"},
+ },
+ sampling_params_list=[template],
+ diffusion_stage_id=0,
+ output_modalities=["video"],
+ )
+ tick = make_tick()
+
+ output = await consumer.execute_tick(tick)
+
+ assert output.images == ["chunk"]
+ assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
+ assert len(client.calls) == 1
+ call = client.calls[0]
+ assert call["request_id"] == tick.request_id
+ assert call["output_modalities"] == ["video"]
+ assert call["prompt"] == {
+ "prompt": "turn left",
+ "multi_modal_data": {"image": "initial-image"},
+ }
+ submitted_params = call["sampling_params_list"][0]
+ assert submitted_params.extra_args["deployment_option"] == "kept"
+ assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
+ assert "ar_diffusion_tick" not in template.extra_args
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_missing_chunk_metadata() -> None:
+ consumer = ARDiffusionOmniTickConsumer(
+ FakeGenerateClient(include_metadata=False),
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ with pytest.raises(ValueError, match="multimodal_output.metadata"):
+ await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
+ consumer = ARDiffusionOmniTickConsumer(
+ FakeGenerateClient(applied_event_ids=[8, 7]),
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ with pytest.raises(ValueError, match="unique and strictly increasing"):
+ await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
+ consumer = ARDiffusionOmniTickConsumer(
+ FakeGenerateClient(include_metadata=False, output_error="model failed"),
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ with pytest.raises(RuntimeError, match="model failed"):
+ await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
+ client = FakeGenerateClient()
+ consumer = ARDiffusionOmniTickConsumer(
+ client,
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+ session = ARDiffusionSession(
+ "world-1",
+ tick_consumer=consumer,
+ lifecycle=FakeLifecycle(),
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
+ )
+ await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
+
+ output = await session.next_chunk()
+
+ assert output.request_id == "world-1-chunk-0"
+ assert session.chunk_index == 1
+ assert session.pending_event_count == 0
+ assert client.calls[0]["request_id"] == "world-1-chunk-0"
diff --git a/tests/diffusion/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py
index d09c1c1b6..086666063 100644
--- a/tests/diffusion/ar_diffusion/test_kv_cache.py
+++ b/tests/diffusion/ar_diffusion/test_kv_cache.py
@@ -30,6 +30,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
BLOCK = 16
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
return ChunkWindowSpec(
@@ -95,7 +97,6 @@ def test_sliding_replace_snaps_to_chunk_boundary():
def test_sink_chunks_protected():
# sink = 1 chunk (16 tokens) is never skipped.
skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
- # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
assert skip == 32
@@ -201,12 +202,149 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
session_capacity=1,
cross_attention_lengths={"text": L},
- device=torch.device("cpu"),
)
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
- expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+ page_bytes = kv.spec.page_size_bytes * 2
+ scratch_bytes = kv.scratch_num_blocks * page_bytes
+ expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
assert kv.num_blocks == expected
- assert expected > 1 * (2 + 1) + 2 # above the local-slot floor, so the cross deduction is what's tested
+ assert kv.cross_attention_reserved_bytes == cross_bytes
+ assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
+ assert expected > 1 * (2 + 1) + 2 # above the local-slot minimum, so the cross deduction is what's tested
+
+
+def _make_tiny_capacity_kv(
+ *,
+ requested_capacity: int,
+ available_bytes: int,
+ gpu_memory_fraction: float = 1.0,
+ model_owned_state_bytes_per_session: int = 0,
+) -> ARDiffusionKVCache:
+ return ARDiffusionKVCache(
+ ARDiffusionKVConfig(
+ enable=True,
+ chunk_size=1,
+ window_chunks=3,
+ sink_chunks=3,
+ gpu_memory_fraction=gpu_memory_fraction,
+ ),
+ num_layers=1,
+ num_kv_heads=1,
+ head_size=1,
+ dtype=torch.float32,
+ block_size=1,
+ max_model_len=128,
+ available_bytes=available_bytes,
+ kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+ session_capacity=requested_capacity,
+ cross_attention_lengths={"text": 2},
+ frames_per_block=3,
+ model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
+ device=torch.device("cpu"),
+ )
+
+
+def test_capacity_two_retains_both_windows_and_allocates_next_block():
+ # LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
+ # 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
+ # reservations: 192 bytes exactly.
+ kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
+ assert kv.requested_session_capacity == 2
+ assert kv.session_capacity == 2
+ assert kv.managed_num_blocks == 17
+
+ adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
+ for adapter in adapters:
+ for _ in range(6): # three sink + three window blocks
+ kv.allocate_token_slots(adapter, 1)
+ kv.commit_chunk(adapter)
+ assert len(kv.window_block_ids(adapter)) == 6
+
+ # Only one request is in flight, so its three-frame block is counted once.
+ kv.allocate_token_slots(adapters[0], 3)
+
+
+def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
+ kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
+
+ assert kv.requested_session_capacity == 64
+ assert kv.session_capacity == 2
+ assert kv.cross_attention_bytes_per_session == 16
+ assert kv.cross_attention_reserved_bytes == 32
+ assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
+
+
+def test_capacity_rejects_budget_that_cannot_fit_one_session():
+ # One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
+ # 14 * 8 + 16 = 128 bytes.
+ with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
+ _make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
+
+
+def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
+ kv = _make_tiny_capacity_kv(
+ requested_capacity=2,
+ available_bytes=1280,
+ gpu_memory_fraction=0.05,
+ )
+
+ assert kv.configured_memory_budget_bytes == 64
+ assert kv.memory_budget_bytes == 128
+ assert kv.session_capacity == 1
+
+
+def test_model_owned_state_reduces_effective_session_capacity():
+ kv = _make_tiny_capacity_kv(
+ requested_capacity=2,
+ available_bytes=200,
+ model_owned_state_bytes_per_session=16,
+ )
+
+ assert kv.session_capacity == 1
+ assert kv.model_owned_state_reserved_bytes == 16
+ assert (
+ kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
+ <= kv.memory_budget_bytes
+ )
+
+
+def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
+ return ARDiffusionKVCache(
+ ARDiffusionKVConfig(
+ enable=True,
+ chunk_size=1560,
+ window_chunks=9,
+ sink_chunks=9,
+ gpu_memory_fraction=gpu_memory_fraction,
+ ),
+ num_layers=40,
+ num_kv_heads=40,
+ head_size=128,
+ dtype=torch.bfloat16,
+ block_size=1560,
+ max_model_len=1 << 20,
+ available_bytes=100 * (1 << 30),
+ kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+ session_capacity=2,
+ cross_attention_lengths={"text": 512},
+ frames_per_block=3,
+ model_owned_state_bytes_per_session=7_488_000,
+ )
+
+
+def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
+ kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
+
+ assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
+ assert kv.memory_budget_bytes == 33_653_670_400
+ assert kv.session_capacity == 1
+
+
+def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
+ kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
+
+ assert kv.memory_budget_bytes == 60 * (1 << 30)
+ assert kv.session_capacity == 2
def _make_kv(
@@ -221,15 +359,24 @@ def _make_kv(
if local_branches == 1
else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
)
+ page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
+ declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
+ scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
+ managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
return ARDiffusionKVCache(
- ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+ ARDiffusionKVConfig(
+ enable=True,
+ chunk_size=BLOCK,
+ window_chunks=window_chunks,
+ gpu_memory_fraction=1.0,
+ ),
num_layers=1,
num_kv_heads=4,
head_size=64,
dtype=torch.float32,
block_size=BLOCK,
max_model_len=4096,
- available_bytes=1 << 16, # tiny -> the floor binds
+ available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
device=torch.device("cpu"),
kv_branches=kv_branches,
session_capacity=1,
diff --git a/tests/diffusion/ar_diffusion/test_kv_orchestration.py b/tests/diffusion/ar_diffusion/test_kv_orchestration.py
index 32603bf88..4482fa6fb 100644
--- a/tests/diffusion/ar_diffusion/test_kv_orchestration.py
+++ b/tests/diffusion/ar_diffusion/test_kv_orchestration.py
@@ -54,13 +54,15 @@ def test_full_request_lifecycle_and_eviction():
used = {int(s) // kv.block_size for s in slots}
assert kv.null_block_id not in used
assert used <= set(block_table)
- # Resident window stays bounded by window + the in-flight chunk.
+ # Allocation may transiently add one in-flight chunk.
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
- free_after.append(kv.manager.block_pool.get_num_free_blocks())
kv.commit_chunk(adapter)
+ # A successful commit immediately prunes physical ownership.
+ assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
+ free_after.append(kv.manager.block_pool.get_num_free_blocks())
# Pool memory plateaus once the window is full (eviction recycles blocks).
- assert free_after[-1] == free_after[kv.spec.window_chunks]
+ assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
kv.end_request(adapter)
assert kv.manager.block_pool.get_num_free_blocks() == free_total
diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py
index 27135874f..603a2a175 100644
--- a/tests/diffusion/ar_diffusion/test_paged_attention.py
+++ b/tests/diffusion/ar_diffusion/test_paged_attention.py
@@ -332,7 +332,7 @@ def test_block_table_padded_to_fixed_width():
def test_prepare_is_idempotent_and_layers_share_metadata():
device = torch.device("cpu")
- _, st = make_state(num_layers=2)
+ kv, st = make_state(num_layers=2)
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
fctx = contexts[0].forward_ctx
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
@@ -349,6 +349,10 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
assert i0.block_table is i1.block_table
assert i0.seq_lens is i1.seq_lens
assert i0.video_slots is i1.video_slots
+ assert i0.key_pool is kv._k_pools[0]
+ assert i0.value_pool is kv._v_pools[0]
+ assert i1.key_pool is kv._k_pools[1]
+ assert i1.value_pool is kv._v_pools[1]
def test_layer_inputs_before_prepare_raises():
diff --git a/tests/diffusion/ar_diffusion/test_pipeline_bridge.py b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
index ec60290b9..ed5233800 100644
--- a/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
@@ -174,6 +174,19 @@ def test_window_reset_can_keep_named_cross_attention_cache():
assert not st.is_cross_attention_populated(POS, "text")
+def test_clear_cross_attention_preserves_self_attention_history():
+ kv, st = make_state(num_layers=1, cross_attn_length=8)
+ _prepare_and_commit(st, POS, 2)
+ k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
+ st.populate_cross_attention(POS, "text", [(k, k)])
+ resident_before = kv.window_block_ids(st.adapter(POS))
+
+ st.clear_cross_attention()
+
+ assert kv.window_block_ids(st.adapter(POS)) == resident_before
+ assert not st.is_cross_attention_populated(POS, "text")
+
+
def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
_, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
assert st.kv_cache.num_local_kv_branches == 1
diff --git a/tests/diffusion/ar_diffusion/test_session.py b/tests/diffusion/ar_diffusion/test_session.py
new file mode 100644
index 000000000..a31be388b
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_session.py
@@ -0,0 +1,505 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionChunkMetadataError,
+ ARDiffusionEventAcceptanceStatus,
+ ARDiffusionEventQueueFullError,
+ ARDiffusionPreparedControls,
+ ARDiffusionSession,
+ ARDiffusionSessionError,
+ ARDiffusionSessionEvent,
+ ARDiffusionSessionManager,
+ ARDiffusionSessionStateError,
+ ARDiffusionSessionStatus,
+ ARDiffusionStaleEventError,
+ ARDiffusionWorkerLifecycle,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeTickConsumer:
+ def __init__(self) -> None:
+ self.ticks: list[ARDiffusionTickRequest] = []
+ self.failures_remaining = 0
+ self.started: asyncio.Event | None = None
+ self.release: asyncio.Event | None = None
+ self.metadata_override: ARDiffusionChunkMetadata | None = None
+
+ async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+ self.ticks.append(tick)
+ if self.started is not None:
+ self.started.set()
+ if self.release is not None:
+ await self.release.wait()
+ if self.failures_remaining:
+ self.failures_remaining -= 1
+ raise RuntimeError("fake tick failed")
+ return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
+
+ def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
+ return output
+
+
+class FakeLifecycle:
+ def __init__(self) -> None:
+ self.resets: list[str] = []
+ self.closes: list[str] = []
+ self.close_failures_remaining = 0
+
+ async def reset_session(self, session_id: str) -> None:
+ self.resets.append(session_id)
+
+ async def close_session(self, session_id: str) -> None:
+ self.closes.append(session_id)
+ if self.close_failures_remaining:
+ self.close_failures_remaining -= 1
+ raise RuntimeError("fake close failed")
+
+
+class FailingCommitReducer:
+ def __init__(self) -> None:
+ self.commit_calls = 0
+ self.reset_calls = 0
+
+ def prepare(
+ self,
+ *,
+ current_controls: dict[str, ARDiffusionControlInput],
+ events: tuple[ARDiffusionSessionEvent, ...],
+ chunk_index: int,
+ ) -> ARDiffusionPreparedControls:
+ del current_controls, events, chunk_index
+ return ARDiffusionPreparedControls(controls=())
+
+ def commit(self, prepared: ARDiffusionPreparedControls) -> None:
+ del prepared
+ self.commit_calls += 1
+ raise RuntimeError("fake reducer commit failed")
+
+ def reset(self) -> None:
+ self.reset_calls += 1
+
+
+def make_session(
+ *,
+ max_pending_events: int = 8,
+ consumer: FakeTickConsumer | None = None,
+ lifecycle: FakeLifecycle | None = None,
+) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
+ resolved_consumer = consumer or FakeTickConsumer()
+ resolved_lifecycle = lifecycle or FakeLifecycle()
+ session = ARDiffusionSession(
+ "world-1",
+ tick_consumer=resolved_consumer,
+ lifecycle=resolved_lifecycle,
+ max_pending_events=max_pending_events,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+ )
+ return session, resolved_consumer, resolved_lifecycle
+
+
+@pytest.mark.asyncio
+async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
+ session, consumer, _ = make_session()
+ event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
+ event_one = ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"samples": [{"x": 1.0}]},
+ ),
+ ),
+ )
+
+ accepted_two = await session.accept_event(event_two)
+ accepted_one = await session.accept_event(event_one)
+ duplicate = await session.accept_event(event_two)
+
+ assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+ assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+ assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+ assert consumer.ticks == []
+
+ metadata = await session.next_chunk()
+
+ assert metadata.applied_event_ids == (1, 2)
+ assert consumer.ticks[0] == ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="world-1-request-0",
+ chunk_index=0,
+ applied_event_ids=(1, 2),
+ prompt="turn",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"samples": [{"x": 1.0}]},
+ ),
+ ),
+ )
+ assert session.pending_event_count == 0
+
+ applied_duplicate = await session.accept_event(event_two)
+ assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+
+ next_metadata = await session.next_chunk()
+ assert next_metadata.chunk_index == 1
+ assert next_metadata.applied_event_ids == ()
+ assert consumer.ticks[1].prompt == "turn"
+ assert consumer.ticks[1].controls == consumer.ticks[0].controls
+
+
+@pytest.mark.asyncio
+async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
+ session, consumer, _ = make_session(max_pending_events=1)
+ payload = {"samples": [{"x": 1.0}]}
+ composite = ARDiffusionSessionEvent(
+ event_id=5,
+ prompt="new prompt",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data=payload,
+ ),
+ ARDiffusionControlInput(
+ track="action",
+ schema="example.action.v1",
+ data={"keys": ["left"]},
+ ),
+ ),
+ )
+ await session.accept_event(composite)
+ payload["samples"][0]["x"] = 99.0
+
+ duplicate = await session.accept_event(composite)
+ assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+ with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
+ await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
+
+ await session.next_chunk()
+ assert consumer.ticks[0].applied_event_ids == (5,)
+ assert consumer.ticks[0].prompt == "new prompt"
+ assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
+ assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
+ assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
+ assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
+
+
+def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
+ with pytest.raises(ValueError, match="at most once"):
+ ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ ARDiffusionControlInput(track="camera", schema="one", data={}),
+ ARDiffusionControlInput(track="camera", schema="two", data={}),
+ ),
+ )
+
+ with pytest.raises(ValueError, match="transport-safe"):
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"object": object()},
+ )
+
+
+@pytest.mark.asyncio
+async def test_failed_tick_closes_worker_and_requires_reset() -> None:
+ consumer = FakeTickConsumer()
+ consumer.failures_remaining = 1
+ lifecycle = FakeLifecycle()
+ session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
+ await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
+
+ with pytest.raises(RuntimeError, match="fake tick failed"):
+ await session.next_chunk()
+
+ assert session.status is ARDiffusionSessionStatus.FAILED
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 1
+ assert lifecycle.closes == ["world-1"]
+ with pytest.raises(ARDiffusionSessionStateError, match="failed"):
+ await session.next_chunk()
+
+ await session.reset()
+ assert session.status is ARDiffusionSessionStatus.ACTIVE
+ assert session.pending_event_count == 0
+ assert lifecycle.resets == ["world-1"]
+ await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
+ metadata = await session.next_chunk()
+ assert metadata.chunk_index == 0
+ assert metadata.applied_event_ids == (4,)
+ assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
+
+
+@pytest.mark.asyncio
+async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
+ consumer = FakeTickConsumer()
+ consumer.started = asyncio.Event()
+ consumer.release = asyncio.Event()
+ session, _, _ = make_session(consumer=consumer)
+ await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
+
+ first_chunk = asyncio.create_task(session.next_chunk())
+ await consumer.started.wait()
+ accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
+ assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+ assert consumer.ticks[0].prompt == "first"
+ assert consumer.ticks[0].applied_event_ids == (10,)
+
+ with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
+ await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
+
+ consumer.release.set()
+ first_metadata = await first_chunk
+ second_metadata = await session.next_chunk()
+ assert first_metadata.applied_event_ids == (10,)
+ assert second_metadata.applied_event_ids == (11,)
+ assert consumer.ticks[1].prompt == "second"
+
+
+@pytest.mark.asyncio
+async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
+ consumer = FakeTickConsumer()
+ consumer.metadata_override = ARDiffusionChunkMetadata(
+ session_id="another-world",
+ request_id="wrong",
+ chunk_index=0,
+ )
+ lifecycle = FakeLifecycle()
+ session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
+ await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
+
+ with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
+ await session.next_chunk()
+
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 1
+ assert session.status is ARDiffusionSessionStatus.FAILED
+ assert lifecycle.closes == ["world-1"]
+
+
+@pytest.mark.asyncio
+async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
+ consumer = FakeTickConsumer()
+ lifecycle = FakeLifecycle()
+ reducer = FailingCommitReducer()
+ session = ARDiffusionSession(
+ "world-1",
+ tick_consumer=consumer,
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+ control_reducer=reducer,
+ )
+ await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
+
+ with pytest.raises(RuntimeError, match="fake reducer commit failed"):
+ await session.next_chunk()
+
+ assert reducer.commit_calls == 1
+ assert session.status is ARDiffusionSessionStatus.FAILED
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 1
+ assert lifecycle.closes == ["world-1"]
+ with pytest.raises(ARDiffusionSessionStateError, match="failed"):
+ await session.next_chunk()
+
+ await session.reset()
+ assert reducer.reset_calls == 1
+ assert session.status is ARDiffusionSessionStatus.ACTIVE
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 0
+
+
+@pytest.mark.asyncio
+async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
+ consumer = FakeTickConsumer()
+ lifecycle = FakeLifecycle()
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
+ )
+ session = await manager.create_session("world-reset")
+ await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
+ await session.next_chunk()
+ await manager.reset_session("world-reset")
+
+ assert lifecycle.resets == ["world-reset"]
+ assert session.status is ARDiffusionSessionStatus.ACTIVE
+ assert session.chunk_index == 0
+ duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
+ assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+
+ reset_tick = await session.next_chunk()
+ assert reset_tick.chunk_index == 0
+ assert consumer.ticks[-1].prompt is None
+ assert consumer.ticks[-1].controls == ()
+
+ await manager.close_session("world-reset")
+ assert lifecycle.closes == ["world-reset"]
+ assert session.status is ARDiffusionSessionStatus.CLOSED
+ with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
+ await manager.get_session("world-reset")
+
+ disconnected = await manager.create_session("world-disconnect")
+ await manager.disconnect("world-disconnect")
+ assert disconnected.status is ARDiffusionSessionStatus.CLOSED
+ assert lifecycle.closes[-1] == "world-disconnect"
+
+
+@pytest.mark.asyncio
+async def test_manager_keeps_interleaved_sessions_independent() -> None:
+ consumer = FakeTickConsumer()
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=FakeLifecycle(),
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+ )
+ session_a = await manager.create_session("world-a")
+ session_b = await manager.create_session("world-b")
+
+ await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
+ await session_a.next_chunk()
+ await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
+ await session_b.next_chunk()
+ await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
+ await session_a.next_chunk()
+
+ assert [
+ (
+ tick.session_id,
+ tick.request_id,
+ tick.chunk_index,
+ tick.applied_event_ids,
+ )
+ for tick in consumer.ticks
+ ] == [
+ ("world-a", "world-a-request-0", 0, (10,)),
+ ("world-b", "world-b-request-0", 0, (1,)),
+ ("world-a", "world-a-request-1", 1, (11,)),
+ ]
+ assert session_a.chunk_index == 2
+ assert session_b.chunk_index == 1
+ with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
+ await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
+ assert (
+ await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
+ ).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+
+
+@pytest.mark.asyncio
+async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
+ consumer = FakeTickConsumer()
+ lifecycle = FakeLifecycle()
+ lifecycle.close_failures_remaining = 1
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ )
+ session = await manager.create_session("world-close-retry")
+
+ with pytest.raises(RuntimeError, match="fake close failed"):
+ await manager.close_session("world-close-retry")
+
+ assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
+ assert await manager.get_session("world-close-retry") is session
+ with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
+ await manager.create_session("world-close-retry")
+ with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
+ await manager.reset_session("world-close-retry")
+
+ await manager.close_session("world-close-retry")
+
+ assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
+ assert session.status is ARDiffusionSessionStatus.CLOSED
+ with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
+ await manager.get_session("world-close-retry")
+ recreated = await manager.create_session("world-close-retry")
+ assert recreated is not session
+
+
+@pytest.mark.asyncio
+async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
+ lifecycle = FakeLifecycle()
+ lifecycle.close_failures_remaining = 1
+ manager = ARDiffusionSessionManager(
+ tick_consumer=FakeTickConsumer(),
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ )
+ session = await manager.create_session("world-disconnect-retry")
+
+ with pytest.raises(RuntimeError, match="fake close failed"):
+ await manager.disconnect("world-disconnect-retry")
+
+ assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
+ assert await manager.get_session("world-disconnect-retry") is session
+
+ await manager.close_session("world-disconnect-retry")
+ assert session.status is ARDiffusionSessionStatus.CLOSED
+
+
+class FakeRPCClient:
+ def __init__(self, result: list[Any]) -> None:
+ self.result = result
+ self.calls: list[dict[str, Any]] = []
+
+ async def collective_rpc(self, **kwargs: Any) -> list[Any]:
+ self.calls.append(kwargs)
+ return self.result
+
+
+@pytest.mark.asyncio
+async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
+ client = FakeRPCClient([[True]])
+ lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
+
+ await lifecycle.reset_session("world-1")
+ await lifecycle.close_session("world-1")
+
+ assert client.calls == [
+ {
+ "method": "reset_ar_diffusion_session",
+ "timeout": 3.0,
+ "args": ("world-1",),
+ "stage_ids": [2],
+ },
+ {
+ "method": "close_ar_diffusion_session",
+ "timeout": 3.0,
+ "args": ("world-1",),
+ "stage_ids": [2],
+ },
+ ]
+
+
+@pytest.mark.asyncio
+async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
+ lifecycle = ARDiffusionWorkerLifecycle(
+ FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
+ stage_ids=[1],
+ )
+
+ with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
+ await lifecycle.close_session("world-1")
diff --git a/tests/diffusion/ar_diffusion/test_tick_protocol.py b/tests/diffusion/ar_diffusion/test_tick_protocol.py
new file mode 100644
index 000000000..70e71ba80
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_tick_protocol.py
@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
+ tick = ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="chunk-request-7",
+ chunk_index=7,
+ applied_event_ids=(10, 11),
+ prompt="turn left",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_trajectory.v1",
+ data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
+ ),
+ ),
+ )
+
+ restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
+
+ assert restored == tick
+ assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
+ "session_id": "world-1",
+ "request_id": "chunk-request-7",
+ "chunk_index": 7,
+ "applied_event_ids": [10, 11],
+ }
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "message"),
+ [
+ ("session_id", "", "session_id"),
+ ("request_id", "", "request_id"),
+ ("chunk_index", -1, "chunk_index"),
+ ("applied_event_ids", (2, 1), "strictly increasing"),
+ ("applied_event_ids", (1, 1), "strictly increasing"),
+ ],
+)
+def test_tick_rejects_invalid_identity(
+ field: str,
+ value: object,
+ message: str,
+) -> None:
+ values = {
+ "session_id": "world-1",
+ "request_id": "request-1",
+ "chunk_index": 0,
+ "applied_event_ids": (),
+ }
+ values[field] = value
+
+ with pytest.raises(ValueError, match=message):
+ ARDiffusionTickRequest(**values)
+
+
+def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
+ payload = {"samples": [{"keys": ["w"]}]}
+ event_ids = [1, 2]
+ controls = [
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data=payload,
+ )
+ ]
+ tick = ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="request-a",
+ chunk_index=0,
+ applied_event_ids=event_ids, # type: ignore[arg-type]
+ controls=controls, # type: ignore[arg-type]
+ )
+ serialized = tick.to_extra_args()
+
+ payload["samples"][0]["keys"].append("j")
+ event_ids.append(3)
+ controls.clear()
+ serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
+
+ assert tick.applied_event_ids == (1, 2)
+ assert isinstance(tick.controls, tuple)
+ assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
+ {
+ "track": "camera",
+ "schema": "example.camera.v1",
+ "data": {"samples": [{"keys": ["w"]}]},
+ }
+ ]
+ with pytest.raises(TypeError):
+ tick.controls[0].data["samples"] = () # type: ignore[index]
+ with pytest.raises(AttributeError):
+ tick.controls[0].data["samples"][0]["keys"].append("j")
+
+
+def test_legacy_request_has_no_typed_tick() -> None:
+ assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None
diff --git a/tests/diffusion/ar_diffusion/test_worker_lifecycle.py b/tests/diffusion/ar_diffusion/test_worker_lifecycle.py
new file mode 100644
index 000000000..42204049c
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_worker_lifecycle.py
@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeARDiffusionRunner:
+ def __init__(self) -> None:
+ self.resets: list[str] = []
+ self.closes: list[str] = []
+
+ def reset_session(self, session_id: str) -> None:
+ self.resets.append(session_id)
+
+ def close_session(self, session_id: str) -> None:
+ self.closes.append(session_id)
+
+
+class FakeRegularRunner:
+ pass
+
+
+def make_worker(runner: object) -> DiffusionWorker:
+ worker = object.__new__(DiffusionWorker)
+ worker.model_runner = runner
+ return worker
+
+
+def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
+ runner = FakeARDiffusionRunner()
+ worker = make_worker(runner)
+
+ assert worker.reset_ar_diffusion_session("world-1") is True
+ assert worker.close_ar_diffusion_session("world-1") is True
+ assert runner.resets == ["world-1"]
+ assert runner.closes == ["world-1"]
+
+
+def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
+ worker = make_worker(FakeRegularRunner())
+
+ assert worker.reset_ar_diffusion_session("world-1") is False
+ assert worker.close_ar_diffusion_session("world-1") is False
+
+
+def test_worker_rejects_empty_session_id() -> None:
+ worker = make_worker(FakeARDiffusionRunner())
+
+ with pytest.raises(ValueError, match="session_id"):
+ worker.close_ar_diffusion_session("")
diff --git a/tests/diffusion/models/dreamzero/test_pipeline_state.py b/tests/diffusion/models/dreamzero/test_pipeline_state.py
index 366876ef7..ff349d6af 100644
--- a/tests/diffusion/models/dreamzero/test_pipeline_state.py
+++ b/tests/diffusion/models/dreamzero/test_pipeline_state.py
@@ -6,7 +6,10 @@ from types import SimpleNamespace
import pytest
-from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
+ DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
+ DreamZeroPipeline,
+)
from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,6 +35,37 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None:
assert session_b.call_count == 3
+def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None:
+ pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+ pipeline.transformer = SimpleNamespace(
+ frame_seqlen=16,
+ blocks=[
+ SimpleNamespace(
+ self_attn=SimpleNamespace(
+ max_attention_size=96,
+ tp_num_heads=4,
+ )
+ )
+ ],
+ text_len=8,
+ model_type="t2v",
+ num_layers=2,
+ num_heads=4,
+ dim=256,
+ num_frame_per_block=4,
+ num_action_per_block=2,
+ num_state_per_block=1,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size",
+ lambda: 1,
+ )
+
+ spec = pipeline.ar_diffusion_kv_cache_spec()
+
+ assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION
+
+
def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
pipeline = _empty_pipeline()
diff --git a/tests/diffusion/models/lingbot_world/__init__.py b/tests/diffusion/models/lingbot_world/__init__.py
new file mode 100644
index 000000000..988131360
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture
new file mode 100644
index 000000000..015a5aa91
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture
@@ -0,0 +1,13 @@
+{
+ "_class_name": "LingBotWorldCausalDMDPipeline",
+ "_diffusers_version": "0.35.0.dev0",
+ "boundary_ratio": 0.9,
+ "image_encoder": [null, null],
+ "image_processor": [null, null],
+ "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+ "text_encoder": ["transformers", "UMT5EncoderModel"],
+ "tokenizer": ["transformers", "T5TokenizerFast"],
+ "transformer": ["diffusers", "CausalLingBotWorldTransformer3DModel"],
+ "transformer_2": [null, null],
+ "vae": ["diffusers", "AutoencoderKLWan"]
+}
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture
new file mode 100644
index 000000000..d07408a04
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture
@@ -0,0 +1,42 @@
+{
+ "sources": {
+ "config": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/config.json",
+ "shard_1_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00001-of-00008.safetensors",
+ "shard_8_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00008-of-00008.safetensors"
+ },
+ "dtype": "F32",
+ "shard_representatives": {
+ "diffusion_pytorch_model-00001-of-00008.safetensors": "blocks.0.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00002-of-00008.safetensors": "blocks.4.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00003-of-00008.safetensors": "blocks.10.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00004-of-00008.safetensors": "blocks.14.cam_scale_layer.bias",
+ "diffusion_pytorch_model-00005-of-00008.safetensors": "blocks.20.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00006-of-00008.safetensors": "blocks.25.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00007-of-00008.safetensors": "blocks.30.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00008-of-00008.safetensors": "blocks.35.cam_injector_layer1.bias"
+ },
+ "parameter_shapes": {
+ "blocks.0.cam_injector_layer1.bias": [5120],
+ "blocks.4.cam_injector_layer1.bias": [5120],
+ "blocks.10.cam_injector_layer1.bias": [5120],
+ "blocks.14.cam_scale_layer.bias": [5120],
+ "blocks.20.cam_injector_layer1.bias": [5120],
+ "blocks.25.cam_injector_layer1.bias": [5120],
+ "blocks.30.cam_injector_layer1.bias": [5120],
+ "blocks.35.cam_injector_layer1.bias": [5120],
+ "patch_embedding.weight": [5120, 36, 1, 2, 2],
+ "patch_embedding_wancamctrl.weight": [5120, 1536],
+ "c2ws_hidden_states_layer1.weight": [5120, 5120],
+ "c2ws_hidden_states_layer2.weight": [5120, 5120],
+ "blocks.0.modulation": [1, 6, 5120],
+ "blocks.0.self_attn.q.weight": [5120, 5120],
+ "blocks.0.ffn.0.weight": [13824, 5120],
+ "blocks.0.ffn.2.weight": [5120, 13824],
+ "blocks.0.cam_injector_layer1.weight": [5120, 5120],
+ "blocks.0.cam_injector_layer2.weight": [5120, 5120],
+ "blocks.0.cam_scale_layer.weight": [5120, 5120],
+ "blocks.0.cam_shift_layer.weight": [5120, 5120],
+ "head.modulation": [1, 2, 5120],
+ "head.head.weight": [64, 5120]
+ }
+}
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture
new file mode 100644
index 000000000..950d26fae
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture
@@ -0,0 +1,30 @@
+{
+ "_class_name": "UniPCMultistepScheduler",
+ "_diffusers_version": "0.35.0.dev0",
+ "beta_end": 0.02,
+ "beta_schedule": "linear",
+ "beta_start": 0.0001,
+ "disable_corrector": [],
+ "dynamic_thresholding_ratio": 0.995,
+ "final_sigmas_type": "zero",
+ "flow_shift": 5.0,
+ "lower_order_final": true,
+ "num_train_timesteps": 1000,
+ "predict_x0": true,
+ "prediction_type": "flow_prediction",
+ "rescale_betas_zero_snr": false,
+ "sample_max_value": 1.0,
+ "solver_order": 2,
+ "solver_p": null,
+ "solver_type": "bh2",
+ "steps_offset": 0,
+ "thresholding": false,
+ "time_shift_type": "exponential",
+ "timestep_spacing": "linspace",
+ "trained_betas": null,
+ "use_beta_sigmas": false,
+ "use_dynamic_shifting": false,
+ "use_exponential_sigmas": false,
+ "use_flow_sigmas": true,
+ "use_karras_sigmas": false
+}
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture
new file mode 100644
index 000000000..2a2c93efd
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture
@@ -0,0 +1,112 @@
+{
+ "source": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model.safetensors.index.json",
+ "index_parameter_count": 1421,
+ "canonical_name_set_sha256": "5d31b7f64fd5e5e271b789b2dbe50fb7bbf146bd8b750646bdc9abb4e5b8139d",
+ "canonical_name_to_shard_sha256": "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992",
+ "canonical_serialization": "sorted lines formatted as parameter_name=shard_filename with LF separators and no trailing LF",
+ "shards": {
+ "diffusion_pytorch_model-00001-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "81a3d4778de604e75609aae7dac3d8dcbc081811e7748267fe953bbeb375e79d"
+ },
+ "diffusion_pytorch_model-00002-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "4e36e2596a6468f29179fa5795396f2e9a7f7d8debc912b34a6f6c2c7c7c610d"
+ },
+ "diffusion_pytorch_model-00003-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "6f74a7b2cf7444c2daa2b3fc32bd3787589b35543c230c5c13b89f0f0f2a0998"
+ },
+ "diffusion_pytorch_model-00004-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "ec5c967ed5271ba7834337c4ee93a73f8ed0eead4492d4920624251c65c058ac"
+ },
+ "diffusion_pytorch_model-00005-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "ad1fdab5830dd9ce43ea007f112b4be223327b7c8197c766d9cf9219e74d619b"
+ },
+ "diffusion_pytorch_model-00006-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "fd58f0d86be42d8afcf55d3e07d8a22f7f6ef5064da931a600c890b71b6dfbae"
+ },
+ "diffusion_pytorch_model-00007-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "305a1714b5f4fb54d5877c593b451273ab1490d90fd0896a299d96e140295047"
+ },
+ "diffusion_pytorch_model-00008-of-00008.safetensors": {
+ "parameter_count": 161,
+ "parameter_names_sha256": "4c18d94d96168682eb212f93e3cdf7313928776d8d9e6e4de0255fe899decb25"
+ }
+ },
+ "top_level_modules": [
+ "blocks",
+ "c2ws_hidden_states_layer1",
+ "c2ws_hidden_states_layer2",
+ "head",
+ "patch_embedding",
+ "patch_embedding_wancamctrl",
+ "text_embedding",
+ "time_embedding",
+ "time_projection"
+ ],
+ "non_block_parameters": [
+ "c2ws_hidden_states_layer1.bias",
+ "c2ws_hidden_states_layer1.weight",
+ "c2ws_hidden_states_layer2.bias",
+ "c2ws_hidden_states_layer2.weight",
+ "head.head.bias",
+ "head.head.weight",
+ "head.modulation",
+ "patch_embedding.bias",
+ "patch_embedding.weight",
+ "patch_embedding_wancamctrl.bias",
+ "patch_embedding_wancamctrl.weight",
+ "text_embedding.0.bias",
+ "text_embedding.0.weight",
+ "text_embedding.2.bias",
+ "text_embedding.2.weight",
+ "time_embedding.0.bias",
+ "time_embedding.0.weight",
+ "time_embedding.2.bias",
+ "time_embedding.2.weight",
+ "time_projection.1.bias",
+ "time_projection.1.weight"
+ ],
+ "block_parameter_suffixes": [
+ "cam_injector_layer1.bias",
+ "cam_injector_layer1.weight",
+ "cam_injector_layer2.bias",
+ "cam_injector_layer2.weight",
+ "cam_scale_layer.bias",
+ "cam_scale_layer.weight",
+ "cam_shift_layer.bias",
+ "cam_shift_layer.weight",
+ "cross_attn.k.bias",
+ "cross_attn.k.weight",
+ "cross_attn.norm_k.weight",
+ "cross_attn.norm_q.weight",
+ "cross_attn.o.bias",
+ "cross_attn.o.weight",
+ "cross_attn.q.bias",
+ "cross_attn.q.weight",
+ "cross_attn.v.bias",
+ "cross_attn.v.weight",
+ "ffn.0.bias",
+ "ffn.0.weight",
+ "ffn.2.bias",
+ "ffn.2.weight",
+ "modulation",
+ "norm3.bias",
+ "norm3.weight",
+ "self_attn.k.bias",
+ "self_attn.k.weight",
+ "self_attn.norm_k.weight",
+ "self_attn.norm_q.weight",
+ "self_attn.o.bias",
+ "self_attn.o.weight",
+ "self_attn.q.bias",
+ "self_attn.q.weight",
+ "self_attn.v.bias",
+ "self_attn.v.weight"
+ ]
+}
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py
new file mode 100644
index 000000000..df012777f
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py
@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.lingbot_world.actions import (
+ LINGBOT_CAMERA_ACTION_SCHEMA,
+ LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
+ LingBotCameraControlReducer,
+ integrate_lingbot_camera_actions,
+ parse_lingbot_camera_action_frames,
+)
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionSession,
+ ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _action_control(mode: str, **data) -> ARDiffusionControlInput:
+ return ARDiffusionControlInput(
+ track="camera",
+ schema=LINGBOT_CAMERA_ACTION_SCHEMA,
+ data={"mode": mode, **data},
+ )
+
+
+def _prepared_frames(prepared) -> tuple[tuple[str, ...], ...]:
+ assert len(prepared.controls) == 1
+ return prepared.controls[0].data["frames"]
+
+
+def test_state_controls_hold_release_and_preserve_short_pulse() -> None:
+ reducer = LingBotCameraControlReducer()
+ pressed = reducer.prepare(
+ current_controls={},
+ events=(
+ ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ _action_control(
+ "state",
+ transitions=[{"actions": ["W"], "client_ts_ms": 10}],
+ ),
+ ),
+ ),
+ ),
+ chunk_index=0,
+ )
+ assert _prepared_frames(pressed) == (("w",), ("w",), ("w",))
+ assert pressed.controls[0].to_dict()["data"]["frames"] == [["w"], ["w"], ["w"]]
+ reducer.commit(pressed)
+
+ held = reducer.prepare(
+ current_controls={"camera": pressed.controls[0]},
+ events=(),
+ chunk_index=1,
+ )
+ assert _prepared_frames(held) == (("w",), ("w",), ("w",))
+ reducer.commit(held)
+
+ pulse = reducer.prepare(
+ current_controls={"camera": held.controls[0]},
+ events=(
+ ARDiffusionSessionEvent(
+ event_id=2,
+ controls=(
+ _action_control(
+ "state",
+ transitions=[
+ {"actions": ["j"], "client_ts_ms": 20},
+ {"actions": [], "client_ts_ms": 21},
+ ],
+ ),
+ ),
+ ),
+ ),
+ chunk_index=2,
+ )
+ assert _prepared_frames(pulse) == (("j",), (), ())
+ reducer.commit(pulse)
+
+ released = reducer.prepare(
+ current_controls={"camera": pulse.controls[0]},
+ events=(),
+ chunk_index=3,
+ )
+ assert _prepared_frames(released) == ((), (), ())
+
+
+def test_script_is_transactional_consumed_once_and_neutral_padded() -> None:
+ reducer = LingBotCameraControlReducer()
+ event = ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ _action_control(
+ "script",
+ frames=[["w"], ["d"], ["l"], ["i"]],
+ ),
+ ),
+ )
+ first_attempt = reducer.prepare(
+ current_controls={},
+ events=(event,),
+ chunk_index=0,
+ )
+ assert _prepared_frames(first_attempt) == (("w",), ("d",), ("l",))
+
+ # An uncommitted failed attempt must not consume the reducer's script.
+ second_attempt = reducer.prepare(
+ current_controls={},
+ events=(event,),
+ chunk_index=0,
+ )
+ assert _prepared_frames(second_attempt) == (("w",), ("d",), ("l",))
+ reducer.commit(second_attempt)
+
+ tail = reducer.prepare(
+ current_controls={"camera": second_attempt.controls[0]},
+ events=(),
+ chunk_index=1,
+ )
+ assert _prepared_frames(tail) == (("i",), (), ())
+ reducer.commit(tail)
+ exhausted = reducer.prepare(
+ current_controls={"camera": tail.controls[0]},
+ events=(),
+ chunk_index=2,
+ )
+ assert _prepared_frames(exhausted) == ((), (), ())
+
+
+def test_trajectory_control_remains_an_expert_override() -> None:
+ reducer = LingBotCameraControlReducer()
+ trajectory = ARDiffusionControlInput(
+ track="camera",
+ schema=LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
+ data={"poses": [[[1.0]]], "intrinsics": [[1.0]]},
+ )
+ prepared = reducer.prepare(
+ current_controls={},
+ events=(
+ ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(trajectory,),
+ ),
+ ),
+ chunk_index=0,
+ )
+ assert prepared.controls == (trajectory,)
+
+
+def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
+ trajectory, pitch = integrate_lingbot_camera_actions(
+ [["w"], ["d"], ["l"]],
+ width=832,
+ height=480,
+ )
+
+ assert pitch == 0.0
+ torch.testing.assert_close(
+ trajectory.poses[0, :3, 3],
+ torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64),
+ atol=1e-6,
+ rtol=0,
+ )
+ torch.testing.assert_close(
+ trajectory.poses[1, :3, 3],
+ torch.tensor([0.05, 0.0, 0.05], dtype=torch.float64),
+ atol=1e-6,
+ rtol=0,
+ )
+ expected_yaw = torch.tensor(
+ [
+ [0.9945219, 0.0, 0.10452846],
+ [0.0, 1.0, 0.0],
+ [-0.10452846, 0.0, 0.9945219],
+ ],
+ dtype=torch.float64,
+ )
+ torch.testing.assert_close(
+ trajectory.poses[2, :3, :3],
+ expected_yaw,
+ atol=1e-6,
+ rtol=0,
+ )
+ torch.testing.assert_close(
+ trajectory.intrinsics[0],
+ torch.tensor([500.0, 500.0, 416.0, 240.0]),
+ )
+
+ resized, _ = integrate_lingbot_camera_actions(
+ [[]],
+ width=416,
+ height=240,
+ )
+ scaled = resized.intrinsics[0] * torch.tensor([416 / 832, 240 / 480, 416 / 832, 240 / 480])
+ torch.testing.assert_close(
+ scaled,
+ torch.tensor([500.0, 500.0, 208.0, 120.0]),
+ )
+
+
+@pytest.mark.parametrize(
+ ("data", "message"),
+ [
+ ({"mode": "frames", "frames": [["x"], [], []]}, "W/A/S/D/I/J/K/L"),
+ ({"mode": "frames", "frames": [["w"]]}, "exactly one action list"),
+ ({"mode": "state", "frames": [[], [], []]}, "mode='frames'"),
+ ],
+)
+def test_model_tick_action_payload_validation(data, message: str) -> None:
+ with pytest.raises(ValueError, match=message):
+ parse_lingbot_camera_action_frames(data, expected_frames=3)
+
+
+class _TickConsumer:
+ def __init__(self) -> None:
+ self.ticks: list[ARDiffusionTickRequest] = []
+
+ async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+ self.ticks.append(tick)
+ return ARDiffusionChunkMetadata.from_tick(tick)
+
+ def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
+ return output
+
+
+class _Lifecycle:
+ async def reset_session(self, session_id: str) -> None:
+ del session_id
+
+ async def close_session(self, session_id: str) -> None:
+ del session_id
+
+
+def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
+ async def exercise() -> None:
+ consumer = _TickConsumer()
+ reducer = LingBotCameraControlReducer()
+ session = ARDiffusionSession(
+ "world-actions",
+ tick_consumer=consumer,
+ lifecycle=_Lifecycle(),
+ max_pending_events=8,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
+ control_reducer=reducer,
+ )
+ await session.accept_event(
+ ARDiffusionSessionEvent(
+ event_id=1,
+ prompt="turn toward the window",
+ controls=(
+ _action_control(
+ "state",
+ transitions=[{"actions": ["d"]}],
+ ),
+ ),
+ )
+ )
+ await session.next_chunk()
+
+ tick = consumer.ticks[0]
+ assert tick.prompt == "turn toward the window"
+ assert tick.applied_event_ids == (1,)
+ assert tick.controls[0].data["frames"] == (("d",), ("d",), ("d",))
+
+ await session.reset()
+ await session.accept_event(
+ ARDiffusionSessionEvent(
+ event_id=2,
+ controls=(
+ _action_control(
+ "state",
+ transitions=[{"actions": []}],
+ ),
+ ),
+ )
+ )
+ await session.next_chunk()
+ assert consumer.ticks[1].chunk_index == 0
+ assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
+
+ asyncio.run(exercise())
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py
new file mode 100644
index 000000000..f5567a706
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py
@@ -0,0 +1,543 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import sys
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/transformer.py"
+_MODULE_NAME = "_lingbot_world_attention_under_test"
+_MISSING_MODULE = object()
+_STUBBED_MODULE_NAMES = (
+ "vllm",
+ "vllm.distributed",
+ "vllm.model_executor",
+ "vllm.model_executor.layers",
+ "vllm.model_executor.layers.conv",
+ "vllm.model_executor.layers.linear",
+ "vllm.model_executor.model_loader",
+ "vllm.model_executor.model_loader.weight_utils",
+ "vllm.model_executor.utils",
+ "vllm_omni",
+ "vllm_omni.diffusion",
+ "vllm_omni.diffusion.attention",
+ "vllm_omni.diffusion.attention.layer",
+ "vllm_omni.diffusion.layers",
+ "vllm_omni.diffusion.layers.norm",
+ "vllm_omni.diffusion.layers.rope",
+ "vllm_omni.experimental",
+ "vllm_omni.experimental.ar_diffusion",
+ "vllm_omni.experimental.ar_diffusion.kv_cache",
+ "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention",
+)
+
+
+@pytest.fixture(autouse=True)
+def _inference_context():
+ with torch.inference_mode():
+ yield
+
+
+def _install_vllm_stubs() -> None:
+ for name in _STUBBED_MODULE_NAMES:
+ sys.modules[name] = ModuleType(name)
+
+ distributed = sys.modules["vllm.distributed"]
+ distributed.get_tensor_model_parallel_rank = lambda: 0
+ distributed.get_tensor_model_parallel_world_size = lambda: 1
+ distributed.tensor_model_parallel_all_reduce = lambda value: value
+
+ def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
+ for name, value in attrs.items():
+ setattr(weight, name, value)
+
+ sys.modules["vllm.model_executor.utils"].set_weight_attrs = set_weight_attrs
+ sys.modules["vllm.model_executor.model_loader.weight_utils"].default_weight_loader = lambda param, loaded_weight: (
+ param.data.copy_(loaded_weight)
+ )
+
+ class _Linear(nn.Module):
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ bias: bool = True,
+ return_bias: bool = False,
+ **kwargs,
+ ) -> None:
+ super().__init__()
+ del kwargs
+ self.return_bias = return_bias
+ self.weight = nn.Parameter(torch.empty(output_size, input_size))
+ self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+ if self.bias is not None:
+ nn.init.zeros_(self.bias)
+ self.calls = 0
+
+ def forward(self, value: torch.Tensor):
+ self.calls += 1
+ output = F.linear(value, self.weight, self.bias)
+ return (output, self.bias) if self.return_bias else output
+
+ linear = sys.modules["vllm.model_executor.layers.linear"]
+ linear.ColumnParallelLinear = _Linear
+ linear.RowParallelLinear = _Linear
+
+ class _QKVLinear(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ head_size: int,
+ total_num_heads: int,
+ total_num_kv_heads: int | None = None,
+ *,
+ bias: bool = True,
+ **kwargs,
+ ) -> None:
+ super().__init__()
+ del kwargs
+ self.num_heads = total_num_heads
+ self.num_kv_heads = total_num_kv_heads or total_num_heads
+ output_size = (self.num_heads + 2 * self.num_kv_heads) * head_size
+ self.weight = nn.Parameter(torch.empty(output_size, hidden_size))
+ self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+ if self.bias is not None:
+ nn.init.zeros_(self.bias)
+ self.calls = 0
+
+ def load_shard(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
+ shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
+ shard_size = param.shape[0] // 3
+ param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
+
+ self.weight.weight_loader = load_shard
+ if self.bias is not None:
+ self.bias.weight_loader = load_shard
+
+ def forward(self, value: torch.Tensor):
+ self.calls += 1
+ return F.linear(value, self.weight, self.bias), None
+
+ linear.QKVParallelLinear = _QKVLinear
+
+ class _Conv3dLayer(nn.Conv3d):
+ def __init__(self, in_channels: int, out_channels: int, kernel_size, stride, **kwargs) -> None:
+ super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs)
+
+ sys.modules["vllm.model_executor.layers.conv"].Conv3dLayer = _Conv3dLayer
+
+ class _LayerNorm(nn.LayerNorm):
+ def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True) -> None:
+ super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
+
+ sys.modules["vllm_omni.diffusion.layers.norm"].LayerNorm = _LayerNorm
+
+ class _Attention(nn.Module):
+ def __init__(self, *args, softmax_scale: float, **kwargs) -> None:
+ super().__init__()
+ del args
+ self.softmax_scale = softmax_scale
+ self.causal = kwargs["causal"]
+ self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata=None,
+ ) -> torch.Tensor:
+ del attn_metadata
+ self.calls.append((query.detach().clone(), key.detach().clone(), value.detach().clone()))
+ scores = torch.einsum("bqhd,bkhd->bhqk", query, key) * self.softmax_scale
+ weights = scores.softmax(dim=-1)
+ return torch.einsum("bhqk,bkhd->bqhd", weights, value)
+
+ sys.modules["vllm_omni.diffusion.attention.layer"].Attention = _Attention
+
+ class _RotaryEmbeddingWan(nn.Module):
+ def __init__(self, *args, **kwargs) -> None:
+ super().__init__()
+ del args, kwargs
+ self.calls: list[tuple[torch.Tensor, torch.Tensor]] = []
+
+ def forward(self, value: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+ self.calls.append((cos.detach().clone(), sin.detach().clone()))
+ return value
+
+ sys.modules["vllm_omni.diffusion.layers.rope"].RotaryEmbeddingWan = _RotaryEmbeddingWan
+
+ paged_attention = sys.modules["vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"]
+
+ class _PagedLayerContext:
+ pass
+
+ class _PagedLayerInputs:
+ pass
+
+ paged_attention.ARDiffusionPagedLayerContext = _PagedLayerContext
+ paged_attention.ARDiffusionPagedLayerInputs = _PagedLayerInputs
+ paged_attention.ar_diffusion_paged_attention = lambda *args, **kwargs: (_ for _ in ()).throw(
+ AssertionError("CUDA paged attention was not expected in this CPU test")
+ )
+ paged_attention.paged_write_attn = lambda *args, **kwargs: (_ for _ in ()).throw(
+ AssertionError("paged path was not expected in this direct-cache test")
+ )
+
+
+def _load_module():
+ assert _MODULE_PATH.exists(), "LingBot attention module has not been implemented"
+ module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
+ previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
+ try:
+ _install_vllm_stubs()
+ spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[spec.name] = module
+ spec.loader.exec_module(module)
+ return module
+ finally:
+ for name, original in previous.items():
+ if original is _MISSING_MODULE:
+ sys.modules.pop(name, None)
+ else:
+ sys.modules[name] = original
+
+
+def test_load_module_restores_process_module_state() -> None:
+ module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
+ previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
+
+ _load_module()
+
+ for name, original in previous.items():
+ assert sys.modules.get(name, _MISSING_MODULE) is original
+
+
+def test_allocate_lingbot_cache_creates_request_local_layer_storage() -> None:
+ module = _load_module()
+
+ cache = module.allocate_lingbot_cache(
+ batch_size=2,
+ num_layers=3,
+ max_tokens=7,
+ num_local_heads=2,
+ head_dim=4,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ assert len(cache.self_attention) == 3
+ assert len(cache.cross_attention) == 3
+ assert cache.cross_attention == [None, None, None]
+ assert cache.self_attention[0].key.shape == (2, 7, 2, 4)
+ assert cache.self_attention[0].value.shape == (2, 7, 2, 4)
+ assert cache.self_attention[0].key.dtype == torch.float32
+ assert cache.self_attention[0].end == 0
+ assert cache.self_attention[0].key.data_ptr() != cache.self_attention[1].key.data_ptr()
+
+
+def _allocate_single_layer(module, *, max_tokens: int):
+ return module.allocate_lingbot_cache(
+ batch_size=1,
+ num_layers=1,
+ max_tokens=max_tokens,
+ num_local_heads=1,
+ head_dim=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+
+def _set_identity_attention(attention: nn.Module) -> None:
+ qkv = getattr(attention, "qkv", None)
+ if qkv is not None:
+ projections = (qkv, attention.o)
+ else:
+ projections = (attention.q, attention.k, attention.v, attention.o)
+ for projection in projections:
+ with torch.no_grad():
+ identity = torch.eye(2)
+ projection.weight.copy_(torch.cat((identity, identity, identity)) if projection is qkv else identity)
+ if projection.bias is not None:
+ projection.bias.zero_()
+ attention.norm_q = nn.Identity()
+ attention.norm_k = nn.Identity()
+
+
+def _tokens(*values: float) -> torch.Tensor:
+ return torch.tensor([[[value, 0.0] for value in values]])
+
+
+def _record_inputs(module: nn.Module) -> list[tuple]:
+ records: list[tuple] = []
+ module.register_forward_pre_hook(lambda _module, args: records.append(args))
+ return records
+
+
+def _record_outputs(module: nn.Module) -> list[torch.Tensor]:
+ records: list[torch.Tensor] = []
+
+ def record(_module, _args, output) -> None:
+ records.append(output[0] if isinstance(output, tuple) else output)
+
+ module.register_forward_hook(record)
+ return records
+
+
+def _cache_snapshot(cache) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]:
+ return (
+ cache.key.detach().view(torch.uint8).clone(),
+ cache.value.detach().view(torch.uint8).clone(),
+ (cache.end, cache.absolute_end, cache.last_start, cache.sink_end),
+ )
+
+
+def _assert_cache_unchanged(cache, snapshot) -> None:
+ key, value, metadata = snapshot
+ assert torch.equal(cache.key.detach().view(torch.uint8), key)
+ assert torch.equal(cache.value.detach().view(torch.uint8), value)
+ assert (cache.end, cache.absolute_end, cache.last_start, cache.sink_end) == metadata
+
+
+def test_self_attention_repeated_offset_overwrites_then_later_offset_appends() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ cache = _allocate_single_layer(module, max_tokens=6).self_attention[0]
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+ attention(_tokens(10, 20), cache=cache, current_start=0, sink_tokens=0)
+
+ assert cache.end == 2
+ assert cache.absolute_end == 2
+ torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0]))
+
+ attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
+
+ assert cache.end == 4
+ assert cache.absolute_end == 4
+ torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0, 3.0, 4.0]))
+
+
+def test_self_attention_uses_one_fused_qkv_projection() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+
+ assert attention.qkv.calls == 1
+
+
+def test_self_attention_routes_paged_inputs_through_engine_owned_cache() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ paged_inputs = module.ARDiffusionPagedLayerInputs()
+ calls = []
+
+ def paged_write(inputs, query, key, value, k_act, v_act, scale):
+ calls.append((inputs, query.clone(), key.clone(), value.clone(), k_act, v_act, scale))
+ return query
+
+ module.paged_write_attn = paged_write
+ output = attention(
+ _tokens(1, 2),
+ cache=paged_inputs,
+ current_start=0,
+ sink_tokens=0,
+ update_cache=False,
+ )
+
+ assert output.shape == (1, 2, 2)
+ assert len(calls) == 1
+ assert calls[0][0] is paged_inputs
+ assert calls[0][4:6] == (None, None)
+ assert attention.attn.calls == []
+ assert not any(hasattr(attention, name) for name in ("q", "k", "v"))
+
+
+def test_self_attention_is_chunk_causal_without_masking_inside_current_chunk() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ cache = _allocate_single_layer(module, max_tokens=8).self_attention[0]
+ attention_inputs = _record_inputs(attention.attn)
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+ first_keys = attention_inputs[-1][1].detach().clone()
+ attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
+ second_keys = attention_inputs[-1][1].detach().clone()
+
+ assert attention.attn.causal is False
+ torch.testing.assert_close(first_keys[0, :, 0, 0], torch.tensor([1.0, 2.0]))
+ torch.testing.assert_close(second_keys[0, :, 0, 0], torch.tensor([1.0, 2.0, 3.0, 4.0]))
+
+
+def test_self_attention_retains_sink_and_latest_local_history_after_eviction() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ cache = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+ attention_inputs = _record_inputs(attention.attn)
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=1)
+ attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=1)
+ attention(_tokens(5, 6), cache=cache, current_start=4, sink_tokens=1)
+
+ visible_keys = attention_inputs[-1][1].detach().clone()
+ assert cache.end == 4
+ assert cache.absolute_end == 6
+ torch.testing.assert_close(visible_keys[0, :, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
+ torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
+
+
+def test_cross_attention_projects_encoder_kv_once_per_request() -> None:
+ module = _load_module()
+ attention = module.LingBotCrossAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ key_outputs = _record_outputs(attention.k)
+ value_outputs = _record_outputs(attention.v)
+ attention_inputs = _record_inputs(attention.attn)
+
+ output, cache = attention(_tokens(1), _tokens(2, 3), cache=None)
+ second_output, reused_cache = attention(_tokens(4), _tokens(20, 30), cache=cache)
+
+ assert output.shape == second_output.shape == (1, 1, 2)
+ assert reused_cache is cache
+ assert len(key_outputs) == 1
+ assert len(value_outputs) == 1
+ assert cache.end == 2
+ torch.testing.assert_close(cache.key[0, :, 0, 0], torch.tensor([2.0, 3.0]))
+ torch.testing.assert_close(attention_inputs[-1][1], cache.key)
+
+
+def test_self_attention_cache_is_isolated_between_requests() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ first = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+ second = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+
+ attention(_tokens(1, 2), cache=first, current_start=0, sink_tokens=0)
+
+ assert first.end == 2
+ assert second.end == 0
+ assert torch.count_nonzero(second.key) == 0
+ attention(_tokens(8), cache=second, current_start=0, sink_tokens=0)
+ torch.testing.assert_close(first.key[0, : first.end, 0, 0], torch.tensor([1.0, 2.0]))
+ torch.testing.assert_close(second.key[0, : second.end, 0, 0], torch.tensor([8.0]))
+
+
+def test_cross_attention_keeps_checkpoint_qkvo_and_norm_parameter_names() -> None:
+ module = _load_module()
+ attention = module.LingBotCrossAttention(dim=4, num_heads=2)
+
+ names = set(attention.state_dict())
+
+ assert {
+ "q.weight",
+ "k.weight",
+ "v.weight",
+ "o.weight",
+ "norm_q.weight",
+ "norm_k.weight",
+ } <= names
+ assert callable(attention.norm_q.weight.weight_loader)
+
+
+def test_tp_world_size_one_attention_output_shapes() -> None:
+ module = _load_module()
+ self_attention = module.LingBotSelfAttention(dim=4, num_heads=2)
+ cross_attention = module.LingBotCrossAttention(dim=4, num_heads=2)
+ cache = module.allocate_lingbot_cache(
+ batch_size=2,
+ num_layers=1,
+ max_tokens=5,
+ num_local_heads=2,
+ head_dim=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+ hidden_states = torch.randn(2, 3, 4)
+ encoder_hidden_states = torch.randn(2, 5, 4)
+
+ self_output = self_attention(hidden_states, cache=cache.self_attention[0], current_start=0, sink_tokens=0)
+ cross_output, cache.cross_attention[0] = cross_attention(
+ hidden_states,
+ encoder_hidden_states,
+ cache=cache.cross_attention[0],
+ )
+
+ assert self_output.shape == hidden_states.shape
+ assert cross_output.shape == hidden_states.shape
+ assert cache.cross_attention[0] is not None
+
+
+def test_self_attention_applies_rotary_embedding_to_current_query_and_key() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
+ cos = torch.ones(2, 1)
+ sin = torch.zeros(2, 1)
+ rotary_inputs = _record_inputs(attention.rotary_embedding)
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, rotary_emb=(cos, sin), sink_tokens=0)
+
+ assert len(rotary_inputs) == 2
+ torch.testing.assert_close(rotary_inputs[0][1], cos)
+ torch.testing.assert_close(rotary_inputs[1][2], sin)
+
+
+def test_attention_block_owns_checkpoint_parent_namespaces() -> None:
+ module = _load_module()
+ block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+
+ names = set(block.state_dict())
+ expected = {
+ *(f"self_attn.{child}.weight" for child in ("qkv", "o", "norm_q", "norm_k")),
+ *(f"cross_attn.{child}.weight" for child in ("q", "k", "v", "o", "norm_q", "norm_k")),
+ }
+
+ assert expected <= names
+
+
+def test_cross_attention_caches_are_isolated_between_requests() -> None:
+ module = _load_module()
+ attention = module.LingBotCrossAttention(dim=2, num_heads=1)
+
+ _, first = attention(_tokens(1), _tokens(2, 3), cache=None)
+ _, second = attention(_tokens(1), _tokens(8, 9), cache=None)
+
+ assert first.key.untyped_storage().data_ptr() != second.key.untyped_storage().data_ptr()
+ assert first.value.untyped_storage().data_ptr() != second.value.untyped_storage().data_ptr()
+ second_key = second.key.clone()
+ first.key.zero_()
+ torch.testing.assert_close(second.key, second_key)
+
+
+def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyPatch) -> None:
+ module = _load_module()
+ monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+ monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
+ norm = module._LingBotRMSNorm(2, eps=1e-6)
+
+ norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
+
+ torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py
new file mode 100644
index 000000000..8b30237c4
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py
@@ -0,0 +1,525 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/camera.py"
+_SPEC = importlib.util.spec_from_file_location("_lingbot_world_camera_under_test", _MODULE_PATH)
+assert _SPEC is not None and _SPEC.loader is not None
+_CAMERA_MODULE = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = _CAMERA_MODULE
+_SPEC.loader.exec_module(_CAMERA_MODULE)
+
+CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
+build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding
+interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory
+load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory
+resolve_trusted_action_directory = _CAMERA_MODULE.resolve_trusted_action_directory
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _identity_poses(num_frames: int) -> np.ndarray:
+ return np.repeat(np.eye(4, dtype=np.float64)[None], num_frames, axis=0)
+
+
+def _write_trajectory(
+ directory: Path,
+ poses: np.ndarray,
+ intrinsics: np.ndarray,
+) -> None:
+ np.save(directory / "poses.npy", poses)
+ np.save(directory / "intrinsics.npy", intrinsics)
+
+
+def _trusted(directory: Path):
+ return resolve_trusted_action_directory(".", directory)
+
+
+def _trajectory(
+ poses: torch.Tensor,
+ intrinsics: torch.Tensor | None = None,
+) -> CameraTrajectory:
+ if intrinsics is None:
+ intrinsics = torch.tensor([[832.0 / 3.0, 160.0, 416.0, 240.0]]).repeat(poses.shape[0], 1)
+ return CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+
+@pytest.mark.parametrize("present_file", [None, "poses.npy"])
+def test_load_camera_trajectory_rejects_missing_files(tmp_path: Path, present_file: str | None) -> None:
+ if present_file == "poses.npy":
+ np.save(tmp_path / present_file, _identity_poses(1))
+
+ missing_file = "poses.npy" if present_file is None else "intrinsics.npy"
+ with pytest.raises(FileNotFoundError, match=missing_file):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+@pytest.mark.parametrize(
+ ("poses", "intrinsics", "message"),
+ [
+ (np.zeros((2, 3, 4)), np.zeros((2, 4)), "poses.npy"),
+ (_identity_poses(2), np.zeros((2, 3)), "intrinsics.npy"),
+ (_identity_poses(0), np.zeros((0, 4)), "1 and 4096 source frames"),
+ (_identity_poses(2), np.zeros((1, 4)), "same number of frames"),
+ ],
+)
+def test_load_camera_trajectory_rejects_invalid_shapes(
+ tmp_path: Path,
+ poses: np.ndarray,
+ intrinsics: np.ndarray,
+ message: str,
+) -> None:
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ with pytest.raises(ValueError, match=message):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+@pytest.mark.parametrize("array_name", ["poses", "intrinsics"])
+def test_load_camera_trajectory_rejects_non_finite_values(tmp_path: Path, array_name: str) -> None:
+ poses = _identity_poses(1)
+ intrinsics = np.ones((1, 4), dtype=np.float64)
+ if array_name == "poses":
+ poses[0, 0, 0] = np.nan
+ else:
+ intrinsics[0, 0] = np.inf
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ with pytest.raises(ValueError, match="finite"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_preserves_raw_c2w_poses(tmp_path: Path) -> None:
+ first_c2w = np.array(
+ [
+ [0.0, -1.0, 0.0, 2.0],
+ [1.0, 0.0, 0.0, 3.0],
+ [0.0, 0.0, 1.0, 4.0],
+ [0.0, 0.0, 0.0, 1.0],
+ ]
+ )
+ second_c2w = np.array(
+ [
+ [0.0, 0.0, 1.0, 5.0],
+ [0.0, 1.0, 0.0, 4.0],
+ [-1.0, 0.0, 0.0, 6.0],
+ [0.0, 0.0, 0.0, 1.0],
+ ]
+ )
+ poses = np.stack((first_c2w, second_c2w))
+ intrinsics = np.array([[100.0, 110.0, 120.0, 130.0]] * 2)
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ assert trajectory.poses.dtype == torch.float32
+ assert trajectory.intrinsics.dtype == torch.float32
+ torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
+ torch.testing.assert_close(trajectory.intrinsics, torch.from_numpy(intrinsics).float())
+
+
+def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_path: Path) -> None:
+ poses = _identity_poses(1).astype(object)
+ intrinsics = np.ones((1, 4), dtype=np.float32)
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ with pytest.raises(ValueError, match="poses.npy.*real numeric"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None:
+ _write_trajectory(
+ tmp_path,
+ _identity_poses(269),
+ np.arange(269 * 4, dtype=np.float32).reshape(269, 4),
+ )
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ assert trajectory.poses.shape == (117, 4, 4)
+ assert trajectory.intrinsics.shape == (117, 4)
+ np.testing.assert_array_equal(
+ trajectory.intrinsics.numpy(),
+ np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117],
+ )
+
+
+def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None:
+ _write_trajectory(
+ tmp_path,
+ _identity_poses(4097),
+ np.ones((4097, 4), dtype=np.float32),
+ )
+
+ with pytest.raises(ValueError, match="poses.npy.*bounded camera-array limit|4096 source frames"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_truncated_payload_before_loading(tmp_path: Path) -> None:
+ _write_trajectory(tmp_path, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ poses_path = tmp_path / "poses.npy"
+ poses_path.write_bytes(poses_path.read_bytes()[:-1])
+
+ with pytest.raises(ValueError, match="poses.npy byte size"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_oversized_npy_header(tmp_path: Path) -> None:
+ header_size = 20_000
+ (tmp_path / "poses.npy").write_bytes(
+ np.lib.format.magic(2, 0) + header_size.to_bytes(4, "little") + b" " * (header_size - 1) + b"\n"
+ )
+ np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
+
+ with pytest.raises(ValueError, match="poses.npy.*oversized NPY header"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_symlink_swap_after_resolution(tmp_path: Path) -> None:
+ trusted_root = tmp_path / "trusted"
+ action_directory = trusted_root / "action"
+ outside = tmp_path / "outside"
+ action_directory.mkdir(parents=True)
+ outside.mkdir()
+ _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ _write_trajectory(outside, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ resolved = resolve_trusted_action_directory("action", trusted_root)
+ action_directory.rename(trusted_root / "original-action")
+ action_directory.symlink_to(outside, target_is_directory=True)
+
+ with pytest.raises(ValueError, match="securely open.*action directory"):
+ load_camera_trajectory(resolved)
+
+
+def test_load_camera_trajectory_rejects_symlinked_camera_file(tmp_path: Path) -> None:
+ trusted_root = tmp_path / "trusted"
+ action_directory = trusted_root / "action"
+ outside = tmp_path / "outside"
+ action_directory.mkdir(parents=True)
+ outside.mkdir()
+ _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ np.save(outside / "poses.npy", _identity_poses(1))
+ resolved = resolve_trusted_action_directory("action", trusted_root)
+ (action_directory / "poses.npy").unlink()
+ (action_directory / "poses.npy").symlink_to(outside / "poses.npy")
+
+ with pytest.raises(ValueError, match="poses.npy.*opened securely"):
+ load_camera_trajectory(resolved)
+
+
+@pytest.mark.skipif(not hasattr(__import__("os"), "mkfifo"), reason="requires POSIX FIFO support")
+def test_load_camera_trajectory_rejects_fifo_without_blocking(tmp_path: Path) -> None:
+ import os
+
+ os.mkfifo(tmp_path / "poses.npy")
+ np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
+ script = """
+import importlib.util
+import sys
+
+spec = importlib.util.spec_from_file_location("_lingbot_camera_fifo_test", sys.argv[1])
+module = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)
+action = module.resolve_trusted_action_directory(".", sys.argv[2])
+try:
+ module.load_camera_trajectory(action)
+except ValueError as error:
+ print(error)
+else:
+ raise AssertionError("FIFO camera input was accepted")
+"""
+
+ completed = subprocess.run(
+ [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
+ capture_output=True,
+ text=True,
+ timeout=3,
+ check=False,
+ )
+
+ assert completed.returncode == 0, completed.stderr
+ assert "poses.npy must be a regular NPY file" in completed.stdout
+
+
+def test_load_camera_trajectory_decodes_the_validated_file_snapshot(tmp_path: Path, monkeypatch) -> None:
+ original_poses = _identity_poses(1)
+ replacement_poses = original_poses.copy()
+ replacement_poses[0, 0, 3] = 123.0
+ _write_trajectory(tmp_path, original_poses, np.ones((1, 4), dtype=np.float32))
+ poses_path = tmp_path / "poses.npy"
+ original_load = _CAMERA_MODULE.np.load
+ mutation_count = 0
+
+ def mutate_source_then_decode(source, *args, **kwargs):
+ nonlocal mutation_count
+ if mutation_count == 0:
+ mutation_count += 1
+ np.save(poses_path, replacement_poses)
+ return original_load(source, *args, **kwargs)
+
+ monkeypatch.setattr(_CAMERA_MODULE.np, "load", mutate_source_then_decode)
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ assert mutation_count == 1
+ torch.testing.assert_close(trajectory.poses, torch.from_numpy(original_poses).float())
+
+
+def test_load_camera_trajectory_ignores_non_camera_event_files(tmp_path: Path) -> None:
+ poses = _identity_poses(1)
+ intrinsics = np.ones((1, 4), dtype=np.float32)
+ _write_trajectory(tmp_path, poses, intrinsics)
+ np.save(tmp_path / "wasd_events.npy", np.array([{"key": "w"}], dtype=object))
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
+
+
+def test_interpolate_camera_trajectory_uses_linear_and_spherical_interpolation() -> None:
+ poses = torch.eye(4).repeat(2, 1, 1)
+ poses[1, :3, :3] = torch.tensor(
+ [
+ [0.0, -1.0, 0.0],
+ [1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ poses[1, :3, 3] = torch.tensor([2.0, 4.0, 6.0])
+ intrinsics = torch.tensor([[10.0, 20.0, 30.0, 40.0], [20.0, 40.0, 60.0, 80.0]])
+ trajectory = CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+ result = interpolate_camera_trajectory(trajectory, 3)
+
+ torch.testing.assert_close(result.poses[0], poses[0])
+ torch.testing.assert_close(result.poses[-1], poses[-1])
+ torch.testing.assert_close(result.poses[1, :3, 3], torch.tensor([1.0, 2.0, 3.0]))
+ expected_mid_rotation = torch.tensor(
+ [
+ [math.sqrt(0.5), -math.sqrt(0.5), 0.0],
+ [math.sqrt(0.5), math.sqrt(0.5), 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ torch.testing.assert_close(result.poses[1, :3, :3], expected_mid_rotation)
+ torch.testing.assert_close(result.intrinsics[1], torch.tensor([15.0, 30.0, 45.0, 60.0]))
+
+
+def test_build_plucker_embedding_has_forward_center_ray() -> None:
+ trajectory = _trajectory(torch.eye(4).unsqueeze(0))
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=3,
+ width=3,
+ target_height=3,
+ target_width=3,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ torch.testing.assert_close(result[0, :3, 1, 1], torch.zeros(3))
+ torch.testing.assert_close(result[0, 3:, 1, 1], torch.tensor([0.0, 0.0, 1.0]))
+
+
+def test_build_plucker_embedding_scales_reference_intrinsics() -> None:
+ trajectory = _trajectory(
+ torch.eye(4).unsqueeze(0),
+ torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=4,
+ width=8,
+ target_height=4,
+ target_width=8,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_corner_direction = torch.tensor([-0.875, -0.75, 1.0])
+ expected_corner_direction /= torch.linalg.vector_norm(expected_corner_direction)
+ torch.testing.assert_close(result[0, 3:, 0, 0], expected_corner_direction)
+
+
+def test_build_plucker_embedding_samples_target_pixel_centers() -> None:
+ trajectory = _trajectory(
+ torch.eye(4).unsqueeze(0),
+ torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=6,
+ width=10,
+ target_height=2,
+ target_width=5,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ # Target centers map through the source scale: (0.5 * 2, 0.5 * 3).
+ expected_direction = torch.tensor([-0.8, -0.5, 1.0])
+ expected_direction /= torch.linalg.vector_norm(expected_direction)
+ torch.testing.assert_close(result[0, 3:, 0, 0], expected_direction)
+
+
+def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_translations() -> None:
+ poses = torch.eye(4).repeat(3, 1, 1)
+ poses[0, :3, 3] = torch.tensor([5.0, 7.0, 11.0])
+ poses[1, :3, :3] = torch.tensor(
+ [
+ [0.0, -1.0, 0.0],
+ [1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ poses[1, :3, 3] = torch.tensor([5.0, 9.0, 11.0])
+ poses[2, :3, :3] = torch.tensor(
+ [
+ [-1.0, 0.0, 0.0],
+ [0.0, -1.0, 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ poses[2, :3, 3] = torch.tensor([3.0, 9.0, 11.0])
+ trajectory = _trajectory(
+ poses,
+ torch.tensor([[416.0, 480.0, 208.0, 240.0]]).repeat(3, 1),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=1,
+ width=2,
+ target_height=1,
+ target_width=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_direction = torch.tensor([0.0, 1.0, 1.0]) / math.sqrt(2.0)
+ torch.testing.assert_close(result[0, :3, 0, 1], torch.zeros(3))
+ torch.testing.assert_close(result[1, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
+ torch.testing.assert_close(result[2, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
+ torch.testing.assert_close(result[1, 3:, 0, 1], expected_direction)
+ torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
+
+
+def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
+ poses = torch.eye(4).repeat(3, 1, 1)
+ poses[:, :3, 3] = torch.tensor(
+ [
+ [10.0, -2.0, 5.0],
+ [13.0, 2.0, 5.0],
+ [19.0, 10.0, 5.0],
+ ]
+ )
+
+ result = build_plucker_embedding(
+ _trajectory(poses),
+ height=1,
+ width=1,
+ target_height=1,
+ target_width=1,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
+ torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
+
+
+def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:
+ poses = torch.eye(4).repeat(2, 1, 1)
+ poses[0, :3, 3] = torch.tensor([20.0, 4.0, -2.0])
+ poses[1, :3, 3] = torch.tensor([26.0, 12.0, -2.0])
+ raw = _trajectory(poses)
+
+ interpolated = interpolate_camera_trajectory(raw, 3)
+ result = build_plucker_embedding(
+ interpolated,
+ height=1,
+ width=1,
+ target_height=1,
+ target_width=1,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_delta = torch.tensor([0.6, 0.8, 0.0])
+ torch.testing.assert_close(result[0, :3, 0, 0], torch.zeros(3))
+ torch.testing.assert_close(result[1, :3, 0, 0], expected_delta)
+ torch.testing.assert_close(result[2, :3, 0, 0], expected_delta)
+
+
+def test_build_plucker_embedding_uses_half_pixel_centers_and_origin_direction_channel_order() -> None:
+ poses = torch.eye(4).repeat(2, 1, 1)
+ poses[1, :3, 3] = torch.tensor([2.0, 0.0, 0.0])
+ trajectory = _trajectory(
+ poses,
+ torch.tensor([[832.0, 480.0, 0.0, 0.0]]).repeat(2, 1),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=1,
+ width=2,
+ target_height=1,
+ target_width=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_direction = torch.tensor([0.25, 0.5, 1.0])
+ expected_direction /= torch.linalg.vector_norm(expected_direction)
+ torch.testing.assert_close(result[1, :3, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
+ torch.testing.assert_close(result[1, 3:, 0, 0], expected_direction)
+
+
+def test_build_plucker_embedding_honors_shape_dtype_and_device() -> None:
+ trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=12,
+ width=20,
+ target_height=3,
+ target_width=5,
+ device=torch.device("cpu"),
+ dtype=torch.float64,
+ )
+
+ assert result.shape == (2, 6, 3, 5)
+ assert result.dtype == torch.float64
+ assert result.device == torch.device("cpu")
+
+
+def test_build_plucker_embedding_is_deterministic() -> None:
+ trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
+ kwargs = {
+ "height": 12,
+ "width": 20,
+ "target_height": 3,
+ "target_width": 5,
+ "device": torch.device("cpu"),
+ "dtype": torch.float32,
+ }
+
+ first = build_plucker_embedding(trajectory, **kwargs)
+ second = build_plucker_embedding(trajectory, **kwargs)
+
+ torch.testing.assert_close(first, second, rtol=0.0, atol=0.0)
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py
new file mode 100644
index 000000000..e2e2a08b6
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py
@@ -0,0 +1,685 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import hashlib
+import inspect
+import json
+import math
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture"
+_SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture"
+
+
+@pytest.fixture(autouse=True)
+def _inference_context():
+ with torch.inference_mode():
+ yield
+
+
+def _tiny_model(
+ module,
+ *,
+ num_layers: int = 2,
+ num_frames_per_block: int = 1,
+ sliding_window_num_frames: int = 3,
+):
+ return module.CausalLingBotWorldTransformer3DModel(
+ patch_size=(1, 2, 2),
+ num_attention_heads=2,
+ attention_head_dim=2,
+ in_channels=36,
+ out_channels=2,
+ text_dim=6,
+ freq_dim=4,
+ ffn_dim=8,
+ num_layers=num_layers,
+ cross_attn_norm=True,
+ eps=1e-6,
+ rope_max_seq_len=16,
+ sink_size=1,
+ num_frames_per_block=num_frames_per_block,
+ sliding_window_num_frames=sliding_window_num_frames,
+ local_attn_size=-1,
+ )
+
+
+def _cache(module, model, *, max_tokens: int | None = None):
+ if max_tokens is None:
+ window_frames = (
+ model.config.local_attn_size
+ if model.config.local_attn_size != -1
+ else model.config.sliding_window_num_frames
+ )
+ max_tokens = window_frames * 2 * 2
+ return module.allocate_lingbot_cache(
+ batch_size=1,
+ num_layers=len(model.blocks),
+ max_tokens=max_tokens,
+ num_local_heads=2,
+ head_dim=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+
+def _checkpoint_parameter_specs(model) -> dict[str, tuple[torch.Size, torch.dtype]]:
+ specs: dict[str, tuple[torch.Size, torch.dtype]] = {}
+ for name, parameter in model.named_parameters():
+ marker = ".self_attn.qkv."
+ if marker not in name:
+ specs[name] = (parameter.shape, parameter.dtype)
+ continue
+ shard_shape = torch.Size((parameter.shape[0] // 3, *parameter.shape[1:]))
+ for projection in ("q", "k", "v"):
+ specs[name.replace(marker, f".self_attn.{projection}.")] = (shard_shape, parameter.dtype)
+ return specs
+
+
+def _checkpoint_weights_for_model(model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+ checkpoint_weights: dict[str, torch.Tensor] = {}
+ expected_parameters: dict[str, torch.Tensor] = {}
+ next_value = 1
+ for name, parameter in model.named_parameters():
+ marker = ".self_attn.qkv."
+ if marker not in name:
+ value = torch.full_like(parameter, next_value)
+ checkpoint_weights[name] = value
+ expected_parameters[name] = value
+ next_value += 1
+ continue
+ shards = []
+ shard_shape = (parameter.shape[0] // 3, *parameter.shape[1:])
+ for projection in ("q", "k", "v"):
+ value = torch.full(shard_shape, next_value, dtype=parameter.dtype, device=parameter.device)
+ checkpoint_weights[name.replace(marker, f".self_attn.{projection}.")] = value
+ shards.append(value)
+ next_value += 1
+ expected_parameters[name] = torch.cat(shards)
+ return checkpoint_weights, expected_parameters
+
+
+def _self_cache_snapshot(cache) -> list[tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]]:
+ return [attention_tests._cache_snapshot(layer_cache) for layer_cache in cache.self_attention]
+
+
+def _assert_self_cache_unchanged(cache, snapshot) -> None:
+ for layer_cache, layer_snapshot in zip(cache.self_attention, snapshot, strict=True):
+ attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot)
+
+
+def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None:
+ torch.manual_seed(7)
+ module = attention_tests._load_module()
+ model = _tiny_model(module).eval()
+ cache = _cache(module, model)
+ encoder_hidden_states = torch.randn(1, 3, 6)
+ camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+ cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
+ text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
+ outputs = []
+
+ assert model.patch_embedding.in_channels == 36
+ assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2
+
+ for start_frame in range(4):
+ hidden_states = torch.randn(1, 36, 1, 4, 4)
+ timestep = torch.tensor([float(start_frame + 1)])
+ snapshot = _self_cache_snapshot(cache)
+
+ transient_output = model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=start_frame,
+ update_cache=False,
+ )
+
+ assert transient_output.shape == (1, 2, 1, 4, 4)
+ _assert_self_cache_unchanged(cache, snapshot)
+ assert all(layer_cache is not None for layer_cache in cache.cross_attention)
+
+ committed_output = model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=start_frame,
+ update_cache=True,
+ )
+ torch.testing.assert_close(committed_output, transient_output)
+ outputs.append(committed_output)
+
+ expected_token_end = (start_frame + 1) * 2 * 2
+ assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention)
+
+ assert len(outputs) == 4
+ assert all(len(outputs) == 1 for outputs in cross_key_outputs)
+ assert len(text_projection_outputs) == 1
+
+ alternate_cache = _cache(module, model)
+ alternate_output = model(
+ torch.zeros(1, 36, 1, 4, 4),
+ torch.tensor([1.0]),
+ encoder_hidden_states,
+ torch.zeros_like(camera_hidden_states),
+ cache=alternate_cache,
+ start_frame=0,
+ update_cache=False,
+ )
+ camera_output = model(
+ torch.zeros(1, 36, 1, 4, 4),
+ torch.tensor([1.0]),
+ encoder_hidden_states,
+ torch.ones_like(camera_hidden_states),
+ cache=_cache(module, model),
+ start_frame=0,
+ update_cache=False,
+ )
+ assert not torch.equal(camera_output, alternate_output)
+
+
+def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=2).eval()
+ cache = _cache(module, model)
+ hidden_states = torch.randn(1, 36, 1, 4, 4)
+ timestep = torch.tensor([1.0])
+ encoder_hidden_states = torch.randn(1, 3, 6)
+ camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+
+ model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=0,
+ update_cache=False,
+ )
+ cache.cross_attention[1] = None
+ text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
+ cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
+
+ model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=0,
+ update_cache=False,
+ )
+
+ assert len(text_projection_outputs) == 1
+ assert [len(outputs) for outputs in cross_key_outputs] == [0, 1]
+ assert all(layer_cache is not None for layer_cache in cache.cross_attention)
+
+
+@pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks"))
+def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(
+ module,
+ num_layers=1,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ )
+ cache = _cache(module, model)
+ snapshot = _self_cache_snapshot(cache)
+ patch_inputs = attention_tests._record_inputs(model.patch_embedding)
+
+ with pytest.raises(ValueError, match="exactly 3 post-patch frames"):
+ model(
+ torch.randn(1, 36, frames, 4, 4),
+ torch.tensor([1.0]),
+ torch.randn(1, 3, 6),
+ torch.randn(1, 6 * 8 * 8, frames, 4, 4),
+ cache=cache,
+ start_frame=0,
+ update_cache=True,
+ )
+
+ assert patch_inputs == []
+ _assert_self_cache_unchanged(cache, snapshot)
+ assert cache.cross_attention == [None]
+
+
+def test_forward_accepts_exactly_one_configured_frame_block() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(
+ module,
+ num_layers=1,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ )
+
+ output = model(
+ torch.randn(1, 36, 3, 4, 4),
+ torch.tensor([1.0]),
+ torch.randn(1, 3, 6),
+ torch.randn(1, 6 * 8 * 8, 3, 4, 4),
+ cache=_cache(module, model),
+ start_frame=0,
+ update_cache=False,
+ )
+
+ assert output.shape == (1, 2, 3, 4, 4)
+
+
+def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(
+ module,
+ num_layers=2,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ )
+
+ cache = model.allocate_cache(
+ batch_size=2,
+ latent_height=8,
+ latent_width=12,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ assert len(cache.self_attention) == 2
+ assert len(cache.cross_attention) == 2
+ assert cache.cross_attention == [None, None]
+ assert all(layer.key.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
+ assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
+
+
+def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+ hidden_states = torch.zeros(1, 36, 2, 2, 4)
+ hidden_states[0, 0] = torch.tensor(
+ [
+ [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
+ [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]],
+ ]
+ )
+ with torch.no_grad():
+ model.patch_embedding.weight.zero_()
+ model.patch_embedding.bias.zero_()
+ model.patch_embedding.weight[0, 0, 0, 0, 0] = 1
+
+ tokens = model.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
+
+ torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0]))
+
+
+def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+ hidden_states = (
+ torch.stack(
+ (
+ torch.arange(0, 8),
+ torch.arange(10, 18),
+ torch.arange(20, 28),
+ torch.arange(30, 38),
+ )
+ )
+ .unsqueeze(0)
+ .float()
+ )
+ expected = torch.tensor(
+ [
+ [
+ [[0, 2, 10, 12], [4, 6, 14, 16]],
+ [[20, 22, 30, 32], [24, 26, 34, 36]],
+ ],
+ [
+ [[1, 3, 11, 13], [5, 7, 15, 17]],
+ [[21, 23, 31, 33], [25, 27, 35, 37]],
+ ],
+ ],
+ dtype=torch.float32,
+ ).unsqueeze(0)
+
+ output = model._unpatchify(
+ hidden_states,
+ batch_size=1,
+ frames=2,
+ height=1,
+ width=2,
+ )
+
+ torch.testing.assert_close(output, expected)
+
+
+def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+ model.head.norm = torch.nn.Identity()
+ with torch.no_grad():
+ model.head.modulation.zero_()
+ model.head.head.weight.zero_()
+ model.head.head.bias.zero_()
+ model.head.head.weight[:, 0] = 1
+ hidden_states = torch.ones(1, 4, model.dim)
+ timestep_embedding = torch.tensor([[[10.0, 0.0, 0.0, 0.0], [20.0, 0.0, 0.0, 0.0]]])
+ expected = torch.tensor([21.0, 21.0, 41.0, 41.0]).view(1, 4, 1).expand(1, 4, 8)
+
+ output = model.head(hidden_states, timestep_embedding)
+
+ torch.testing.assert_close(output, expected)
+
+
+def test_constructor_defaults_match_official_checkpoint_config() -> None:
+ module = attention_tests._load_module()
+ parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters
+ expected = {
+ "in_channels": 36,
+ "out_channels": 16,
+ "num_layers": 40,
+ "num_attention_heads": 40,
+ "attention_head_dim": 128,
+ "patch_size": (1, 2, 2),
+ "text_dim": 4096,
+ "ffn_dim": 13824,
+ "freq_dim": 256,
+ "rope_max_seq_len": 1024,
+ "eps": 1e-6,
+ "cross_attn_norm": True,
+ "sink_size": 9,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 18,
+ "local_attn_size": -1,
+ }
+
+ assert {name: parameters[name].default for name in expected} == expected
+
+
+def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module)
+
+ assert model.dtype == next(model.parameters()).dtype
+
+
+def test_transformer_declares_regional_compile_block() -> None:
+ module = attention_tests._load_module()
+
+ assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"]
+
+
+def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
+ module = attention_tests._load_module()
+ reduced_values: list[torch.Tensor] = []
+
+ monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+
+ def all_reduce(value: torch.Tensor) -> torch.Tensor:
+ reduced_values.append(value.detach().clone())
+ return value * 2
+
+ monkeypatch.setattr(module, "tensor_model_parallel_all_reduce", all_reduce)
+ norm = module._LingBotRMSNorm(hidden_size=2, eps=0.0)
+ value = torch.tensor([[3.0, 4.0]])
+
+ output = norm(value)
+
+ torch.testing.assert_close(output, value / math.sqrt(12.5))
+ assert len(reduced_values) == 1
+ torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]]))
+
+
+def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
+ module = attention_tests._load_module()
+ monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
+
+ with pytest.raises(ValueError, match="num_heads=2.*tp_size=3"):
+ module.LingBotSelfAttention(dim=4, num_heads=2)
+
+
+def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
+ module = attention_tests._load_module()
+
+ with pytest.raises(ValueError, match="qk_norm.*rms_norm_across_heads"):
+ module.CausalLingBotWorldTransformer3DModel(qk_norm="layer_norm")
+
+
+@pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"])
+def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None:
+ module = attention_tests._load_module()
+
+ with pytest.raises(ValueError, match=field):
+ module.CausalLingBotWorldTransformer3DModel(**{field: 4})
+
+
+def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
+ module = attention_tests._load_module()
+
+ with pytest.raises(RuntimeError, match="quant_config.*not supported"):
+ module.CausalLingBotWorldTransformer3DModel(quant_config=object())
+
+
+def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None:
+ module = attention_tests._load_module()
+
+ class ConfigProbe(module.CausalLingBotWorldTransformer3DModel):
+ def __init__(self, **kwargs) -> None:
+ self.received_kwargs = kwargs
+
+ config = {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "_diffusers_version": "0.35.0.dev0",
+ "patch_size": [1, 2, 2],
+ "num_attention_heads": 40,
+ "attention_head_dim": 128,
+ "in_channels": 36,
+ "out_channels": 16,
+ "text_dim": 4096,
+ "freq_dim": 256,
+ "ffn_dim": 13824,
+ "num_layers": 40,
+ "cross_attn_norm": True,
+ "eps": 1e-6,
+ "image_dim": None,
+ "added_kv_proj_dim": None,
+ "rope_max_seq_len": 1024,
+ "pos_embed_seq_len": None,
+ "qk_norm": "rms_norm_across_heads",
+ "sink_size": 9,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 18,
+ "local_attn_size": -1,
+ }
+
+ model = ConfigProbe.from_config(config, prefix="transformer")
+
+ assert model.received_kwargs["patch_size"] == (1, 2, 2)
+ assert model.received_kwargs["num_layers"] == 40
+ assert model.received_kwargs["prefix"] == "transformer"
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("in_channels", 35),
+ ("num_layers", 39),
+ ("num_attention_heads", 32),
+ ("attention_head_dim", 64),
+ ("ffn_dim", 10240),
+ ("num_frames_per_block", 4),
+ ("sliding_window_num_frames", 24),
+ ("sink_size", 0),
+ ],
+)
+def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None:
+ module = attention_tests._load_module()
+ config = {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "_diffusers_version": "0.35.0.dev0",
+ "patch_size": [1, 2, 2],
+ "num_attention_heads": 40,
+ "attention_head_dim": 128,
+ "in_channels": 36,
+ "out_channels": 16,
+ "text_dim": 4096,
+ "freq_dim": 256,
+ "ffn_dim": 13824,
+ "num_layers": 40,
+ "cross_attn_norm": True,
+ "eps": 1e-6,
+ "image_dim": None,
+ "added_kv_proj_dim": None,
+ "rope_max_seq_len": 1024,
+ "pos_embed_seq_len": None,
+ "qk_norm": "rms_norm_across_heads",
+ "sink_size": 9,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 18,
+ "local_attn_size": -1,
+ }
+ config[field] = value
+
+ with pytest.raises(ValueError, match=field):
+ module.CausalLingBotWorldTransformer3DModel.from_config(config)
+
+
+def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
+ module = attention_tests._load_module()
+ config = {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "architectures": ["CausalLingBotWorldTransformer3DModel"],
+ "future_diffusers_metadata": {"version": 1},
+ }
+
+ with torch.device("meta"):
+ model = module.CausalLingBotWorldTransformer3DModel.from_config(config)
+
+ assert model.config.in_channels == 36
+ assert model.config.num_layers == 40
+
+
+def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1)
+ checkpoint_weights, expected_parameters = _checkpoint_weights_for_model(model)
+ loader_calls: list[str] = []
+ qkv_weight = model.blocks[0].self_attn.qkv.weight
+
+ def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
+ loader_calls.append(shard_id)
+ shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
+ shard_size = param.shape[0] // 3
+ param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
+
+ qkv_weight.weight_loader = record_loader
+ loaded = model.load_weights(iter(checkpoint_weights.items()))
+
+ assert set(checkpoint_weights) <= loaded
+ assert set(dict(model.named_parameters())) <= loaded
+ assert loader_calls == ["q", "k", "v"]
+ for name, param in model.named_parameters():
+ torch.testing.assert_close(param, expected_parameters[name])
+
+ first_name = next(iter(checkpoint_weights))
+ partial = _tiny_model(module, num_layers=1)
+ partial_loaded = partial.load_weights([(first_name, checkpoint_weights[first_name])])
+ expected_partial = {first_name}
+ for projection in ("q", "k", "v"):
+ marker = f".self_attn.{projection}."
+ if marker in first_name:
+ expected_partial.add(first_name.replace(marker, ".self_attn.qkv."))
+ break
+ assert partial_loaded == expected_partial
+
+ with pytest.raises(KeyError, match="unexpected_model.weight"):
+ model.load_weights([("unexpected_model.weight", torch.ones(1))])
+
+
+def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1)
+ first_name, first_param = next(iter(model.named_parameters()))
+ loaded_weight = torch.full_like(first_param, 17)
+
+ def weights():
+ yield first_name, loaded_weight
+ torch.testing.assert_close(first_param, loaded_weight)
+
+ assert model.load_weights(weights()) == {first_name}
+
+
+def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None:
+ utils = pytest.importorskip("vllm.model_executor.models.utils")
+ module = attention_tests._load_module()
+ transformer = _tiny_model(module, num_layers=1)
+ checkpoint_weights, _ = _checkpoint_weights_for_model(transformer)
+
+ class Pipeline(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.transformer = transformer
+
+ pipeline = Pipeline()
+ loaded = utils.AutoWeightsLoader(pipeline).load_weights(
+ (f"transformer.{name}", value) for name, value in checkpoint_weights.items()
+ )
+
+ assert set(dict(pipeline.named_parameters())) <= loaded
+
+
+def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
+ module = attention_tests._load_module()
+ fixture = json.loads(_FIXTURE_PATH.read_text())
+ model = _tiny_model(module, num_layers=1)
+ parameter_names = set(_checkpoint_parameter_specs(model))
+
+ assert {name.split(".", 1)[0] for name in parameter_names} == set(fixture["top_level_modules"])
+ assert {name for name in parameter_names if not name.startswith("blocks.")} == set(fixture["non_block_parameters"])
+ assert {name.removeprefix("blocks.0.") for name in parameter_names if name.startswith("blocks.0.")} == set(
+ fixture["block_parameter_suffixes"]
+ )
+ expanded_names = set(fixture["non_block_parameters"])
+ expanded_names.update(
+ f"blocks.{layer}.{suffix}" for layer in range(40) for suffix in fixture["block_parameter_suffixes"]
+ )
+ canonical_names = "\n".join(sorted(expanded_names))
+
+ assert len(expanded_names) == fixture["index_parameter_count"] == 1421
+ assert hashlib.sha256(canonical_names.encode()).hexdigest() == fixture["canonical_name_set_sha256"]
+ assert fixture["canonical_name_to_shard_sha256"] == (
+ "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992"
+ )
+ assert list(fixture["shards"]) == [
+ f"diffusion_pytorch_model-{index:05}-of-00008.safetensors" for index in range(1, 9)
+ ]
+ assert sum(shard["parameter_count"] for shard in fixture["shards"].values()) == 1421
+ assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values())
+
+
+def test_official_default_shapes_match_public_safetensors_header_fixture() -> None:
+ module = attention_tests._load_module()
+ fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text())
+
+ with torch.device("meta"):
+ model = module.CausalLingBotWorldTransformer3DModel()
+ parameter_specs = _checkpoint_parameter_specs(model)
+ expected_dtype = {"F32": torch.float32}[fixture["dtype"]]
+
+ assert len(fixture["shard_representatives"]) == 8
+ assert set(fixture["shard_representatives"].values()) <= set(fixture["parameter_shapes"])
+
+ for name, expected_shape in fixture["parameter_shapes"].items():
+ shape, dtype = parameter_specs[name]
+ assert tuple(shape) == tuple(expected_shape), name
+ assert dtype == expected_dtype
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py
new file mode 100644
index 000000000..56c25fee3
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py
@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_ROOT = Path(__file__).parents[4]
+_PRODUCTION_FILES = (
+ _ROOT / "vllm_omni/diffusion/models/lingbot_world/camera.py",
+ _ROOT / "vllm_omni/diffusion/models/lingbot_world/transformer.py",
+ _ROOT / "vllm_omni/diffusion/models/lingbot_world/pipeline.py",
+)
+
+
+def test_all_lingbot_production_functions_have_complete_annotations() -> None:
+ missing: list[str] = []
+ for path in _PRODUCTION_FILES:
+ tree = ast.parse(path.read_text())
+ for node in ast.walk(tree):
+ if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+ continue
+ arguments = (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs)
+ for argument in arguments:
+ if argument.arg not in {"self", "cls"} and argument.annotation is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:{argument.arg}")
+ if node.args.vararg is not None and node.args.vararg.annotation is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:*{node.args.vararg.arg}")
+ if node.args.kwarg is not None and node.args.kwarg.annotation is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:**{node.args.kwarg.arg}")
+ if node.returns is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:return")
+
+ assert missing == []
diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py
new file mode 100644
index 000000000..1a08e8270
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py
@@ -0,0 +1,1765 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import gc
+import weakref
+from pathlib import Path
+from threading import Lock
+from types import SimpleNamespace
+
+import pytest
+import torch
+from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
+from PIL import Image
+from torch import nn
+
+import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline
+from tests.diffusion.models.wan2_2.conftest import noop_progress_bar
+from vllm_omni.diffusion.models.lingbot_world.actions import (
+ integrate_lingbot_camera_actions,
+)
+from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as _CameraTrajectory
+from vllm_omni.diffusion.models.lingbot_world.camera import (
+ build_plucker_embedding as _real_build_plucker_embedding,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_ROOT = Path(__file__).parents[4]
+_LINGBOT_INIT_PATH = _ROOT / "vllm_omni/diffusion/models/lingbot_world/__init__.py"
+_MODEL_INDEX_FIXTURE = Path(__file__).with_name("fixtures") / "lingbot_world_model_index.json.fixture"
+_SCHEDULER_FIXTURE = Path(__file__).with_name("fixtures") / "lingbot_world_scheduler_config.json.fixture"
+
+
+def _scheduler_config(**overrides):
+ values = {
+ "_class_name": "UniPCMultistepScheduler",
+ "num_train_timesteps": 1000,
+ "flow_shift": 5.0,
+ "prediction_type": "flow_prediction",
+ "predict_x0": True,
+ "use_flow_sigmas": True,
+ "use_dynamic_shifting": False,
+ "use_beta_sigmas": False,
+ "use_exponential_sigmas": False,
+ "use_karras_sigmas": False,
+ "final_sigmas_type": "zero",
+ "timestep_spacing": "linspace",
+ "solver_order": 2,
+ "solver_type": "bh2",
+ "lower_order_final": True,
+ "disable_corrector": [],
+ "time_shift_type": "exponential",
+ }
+ values.update(overrides)
+ return values
+
+
+class _AutoWeightsLoader:
+ def __init__(self, module):
+ self.module = module
+
+ def load_weights(self, weights):
+ loaded = self.module.transformer.load_weights(
+ (name.removeprefix("transformer."), value) for name, value in weights if name.startswith("transformer.")
+ )
+ return {f"transformer.{name}" for name in loaded}
+
+
+class _FakePretrained:
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ del args, kwargs
+ return cls()
+
+ def to(self, *args, **kwargs):
+ self.to_calls = getattr(self, "to_calls", [])
+ self.to_calls.append((args, kwargs))
+ return self
+
+
+class _FakeScheduler:
+ def __init__(self, **kwargs):
+ self.config = SimpleNamespace(**kwargs)
+
+
+class _FakeTransformerFactory:
+ @classmethod
+ def from_config(cls, config, *, quant_config=None, prefix=""):
+ cls.last_call = (config, quant_config, prefix)
+ return _RecordingTransformer()
+
+
+class _Cache:
+ def __init__(self, *, num_layers: int, max_tokens: int, num_local_heads: int, head_dim: int):
+ shape = (1, max_tokens, num_local_heads, head_dim)
+ self.self_attention = [
+ SimpleNamespace(key=torch.zeros(shape), value=torch.zeros(shape)) for _ in range(num_layers)
+ ]
+ self.cross_attention = [None] * num_layers
+
+
+class _RecordingTransformer(nn.Module):
+ def __init__(self, *, raise_on_call: int | None = None, dtype: torch.dtype = torch.float32):
+ super().__init__()
+ self.config = SimpleNamespace(
+ patch_size=(1, 2, 2),
+ in_channels=36,
+ out_channels=16,
+ text_dim=8,
+ num_layers=2,
+ num_attention_heads=2,
+ attention_head_dim=4,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ local_attn_size=-1,
+ sink_size=3,
+ )
+ self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()])
+ for block in self.blocks:
+ block.self_attn = SimpleNamespace(num_local_heads=2, head_dim=4)
+ self.calls: list[dict] = []
+ self.cache_allocations: list[dict] = []
+ self.raise_on_call = raise_on_call
+ self._dtype = dtype
+ self.loaded_weights: list[tuple[str, torch.Tensor]] = []
+
+ @property
+ def dtype(self) -> torch.dtype:
+ return self._dtype
+
+ def forward(self, **kwargs):
+ call = {
+ "hidden_states": kwargs["hidden_states"].detach().clone(),
+ "timestep": kwargs["timestep"].detach().clone(),
+ "encoder_hidden_states": kwargs["encoder_hidden_states"].detach().clone(),
+ "camera_hidden_states": kwargs["camera_hidden_states"].detach().clone(),
+ "cache_id": id(kwargs["cache"]),
+ "start_frame": kwargs["start_frame"],
+ "update_cache": kwargs["update_cache"],
+ }
+ self.calls.append(call)
+ if self.raise_on_call == len(self.calls):
+ raise RuntimeError("forced transformer failure")
+ return torch.ones_like(kwargs["hidden_states"][:, :16])
+
+ def allocate_cache(self, **kwargs):
+ self.cache_allocations.append(dict(kwargs))
+ patch_height, patch_width = self.config.patch_size[1:]
+ window_frames = (
+ self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames
+ )
+ max_tokens = window_frames * (kwargs["latent_height"] // patch_height) * (kwargs["latent_width"] // patch_width)
+ return _Cache(
+ num_layers=self.config.num_layers,
+ max_tokens=max_tokens,
+ num_local_heads=2,
+ head_dim=4,
+ )
+
+ def load_weights(self, weights):
+ self.loaded_weights = list(weights)
+ return {name for name, _ in self.loaded_weights}
+
+
+class _StubVAE(_FakePretrained):
+ dtype = torch.float32
+
+ def __init__(self):
+ self.config = SimpleNamespace(
+ z_dim=16,
+ scale_factor_temporal=4,
+ scale_factor_spatial=8,
+ latents_mean=[float(index) - 3.0 for index in range(16)],
+ latents_std=[1.0 + index / 4.0 for index in range(16)],
+ )
+ self.encode_inputs: list[torch.Tensor] = []
+ self.decode_inputs: list[torch.Tensor] = []
+ self.on_decode = None
+
+ def encode(self, video: torch.Tensor):
+ self.encode_inputs.append(video.detach().clone())
+ latent_frames = (video.shape[2] - 1) // 4 + 1
+ latents = torch.zeros(
+ video.shape[0],
+ 16,
+ latent_frames,
+ video.shape[-2] // 8,
+ video.shape[-1] // 8,
+ dtype=video.dtype,
+ device=video.device,
+ )
+ latents[:, :, 0] = 2.0
+ return SimpleNamespace(latents=latents)
+
+ def decode(self, latents: torch.Tensor, return_dict: bool = False):
+ del return_dict
+ if self.on_decode is not None:
+ self.on_decode()
+ self.decode_inputs.append(latents.detach().clone())
+ pixel_frames = (latents.shape[2] - 1) * 4 + 1
+ decoded = torch.zeros(
+ latents.shape[0],
+ 3,
+ pixel_frames,
+ latents.shape[-2] * 8,
+ latents.shape[-1] * 8,
+ dtype=latents.dtype,
+ device=latents.device,
+ )
+ return (decoded,)
+
+
+class _SamplingParams:
+ def __init__(
+ self,
+ *,
+ height: int | None = 16,
+ width: int | None = 16,
+ num_frames: int = 9,
+ num_inference_steps: int | None = 4,
+ num_outputs_per_prompt: int = 1,
+ seed: int | None = 17,
+ generator: torch.Generator | None = None,
+ output_type: str | None = "latent",
+ max_sequence_length: int | None = 512,
+ extra_args: dict | None = None,
+ include_action: bool = True,
+ ) -> None:
+ self.height = height
+ self.width = width
+ self.num_frames = num_frames
+ self.num_inference_steps = num_inference_steps
+ self.num_outputs_per_prompt = num_outputs_per_prompt
+ self.seed = seed
+ self.generator = (
+ generator
+ if generator is not None
+ else (torch.Generator(device="cpu").manual_seed(seed) if seed is not None else None)
+ )
+ self.output_type = output_type
+ self.max_sequence_length = max_sequence_length
+ self.extra_args = {"action_path": "."} if include_action else {}
+ self.extra_args["_lingbot_camera_trajectory"] = _CameraTrajectory(
+ poses=torch.eye(4).repeat(32, 1, 1),
+ intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(32, 1),
+ )
+ if extra_args is not None:
+ self.extra_args.update(extra_args)
+ self.latents = None
+
+
+class _RequestBatch:
+ def __init__(self, prompt, sampling_params: _SamplingParams, *, num_reqs: int = 1):
+ self._prompts = [prompt] * num_reqs
+ self._sampling = sampling_params
+ self.num_reqs = num_reqs
+
+ @property
+ def prompts(self):
+ return self._prompts
+
+ @property
+ def sampling_params(self):
+ return self._sampling
+
+ @property
+ def request_id(self):
+ return "legacy-lingbot-request"
+
+
+def _load_pipeline_module():
+ return lingbot_pipeline
+
+
+@pytest.fixture(autouse=True)
+def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Replace only heavyweight component construction, not the import graph."""
+
+ loader_state = SimpleNamespace(prefetch_calls=[])
+
+ def prefetch_subfolders(model, subfolders, *, local_files_only):
+ loader_state.prefetch_calls.append((model, tuple(subfolders), local_files_only))
+
+ def from_pretrained_with_prefetch(callable_, model, **kwargs):
+ kwargs.pop("prefetch_list", None)
+ return callable_(model, **kwargs)
+
+ def load_transformer_config(model, subfolder, local_files_only):
+ del model, subfolder, local_files_only
+ return {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "patch_size": [1, 2, 2],
+ "in_channels": 36,
+ "out_channels": 16,
+ "text_dim": 8,
+ "num_layers": 2,
+ "num_attention_heads": 2,
+ "attention_head_dim": 4,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 6,
+ "local_attn_size": -1,
+ }
+
+ def retrieve_latents(value, sample_mode="argmax"):
+ assert sample_mode == "argmax"
+ return value.latents
+
+ def load_json(model, filename, local_files_only):
+ del model, local_files_only
+ assert filename == "scheduler/scheduler_config.json"
+ return _scheduler_config()
+
+ trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(32, 1, 1),
+ intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(32, 1),
+ )
+
+ def load_camera_trajectory(action_path):
+ assert action_path
+ return trajectory
+
+ def interpolate_camera_trajectory(value, num_frames):
+ return SimpleNamespace(
+ poses=value.poses[:num_frames],
+ intrinsics=value.intrinsics[:num_frames],
+ )
+
+ def build_plucker_embedding(value, *, height, width, target_height, target_width, device, dtype):
+ del target_height, target_width
+ frames = value.poses.shape[0]
+ data = torch.arange(frames * 6 * height * width, device=device, dtype=torch.float32)
+ return data.reshape(frames, 6, height, width).to(dtype=dtype)
+
+ replacements = {
+ "AutoTokenizer": _FakePretrained,
+ "UMT5EncoderModel": _FakePretrained,
+ "AutoWeightsLoader": _AutoWeightsLoader,
+ "DistributedAutoencoderKLWan": _StubVAE,
+ "FlowUniPCMultistepScheduler": _FakeScheduler,
+ "CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory,
+ "get_local_device": lambda: torch.device("cpu"),
+ "set_forward_context_denoise_step_idx": lambda index: None,
+ "prefetch_subfolders": prefetch_subfolders,
+ "from_pretrained_with_prefetch": from_pretrained_with_prefetch,
+ "load_transformer_config": load_transformer_config,
+ "retrieve_latents": retrieve_latents,
+ "_load_json": load_json,
+ "load_camera_trajectory": load_camera_trajectory,
+ "interpolate_camera_trajectory": interpolate_camera_trajectory,
+ "build_plucker_embedding": build_plucker_embedding,
+ "randn_tensor": _diffusers_randn_tensor,
+ }
+ for name, value in replacements.items():
+ monkeypatch.setattr(lingbot_pipeline, name, value)
+ monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False)
+
+
+def _od_config(**overrides):
+ values = {
+ "model": "checkpoint",
+ "dtype": torch.float32,
+ "flow_shift": None,
+ "quantization_config": None,
+ "enable_diffusion_pipeline_profiler": False,
+ "model_config": {
+ "lingbot_action_root": str(_ROOT),
+ "ar_diffusion_height": 16,
+ "ar_diffusion_width": 16,
+ },
+ "enable_cpu_offload": False,
+ "enable_layerwise_offload": False,
+ "enforce_eager": True,
+ "parallel_config": SimpleNamespace(
+ pipeline_parallel_size=1,
+ sequence_parallel_size=1,
+ cfg_parallel_size=1,
+ vae_patch_parallel_size=1,
+ use_hsdp=False,
+ enable_expert_parallel=False,
+ ),
+ }
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
+def _prompt(*, action_path: str | None = None, images=None):
+ if images is None:
+ images = Image.new("RGB", (16, 16), color=(255, 128, 0))
+ prompt = {
+ "prompt": "move through the room",
+ "multi_modal_data": {"image": images},
+ "additional_information": {},
+ }
+ if action_path is not None:
+ prompt["additional_information"]["action_path"] = action_path
+ return prompt
+
+
+def _pipeline(module, *, transformer=None, od_config=None):
+ pipeline = module.LingBotWorldCausalDMDPipeline(od_config=od_config or _od_config())
+ if transformer is not None:
+ pipeline.transformer = transformer
+ pipeline.encode_prompt = lambda *args, **kwargs: torch.ones(1, 512, 8)
+ pipeline.progress_bar = noop_progress_bar
+ return pipeline
+
+
+def _preprocess_request(module, *, prompt=None, sampling=None, od_config=None):
+ request = SimpleNamespace(
+ prompt=_prompt() if prompt is None else prompt,
+ sampling_params=_SamplingParams() if sampling is None else sampling,
+ )
+ preprocess = module.get_lingbot_world_pre_process_func(od_config or _od_config())
+ return preprocess(request)
+
+
+@pytest.mark.parametrize("offload_field", ["enable_cpu_offload", "enable_layerwise_offload"])
+def test_pipeline_respects_loader_managed_component_placement(offload_field: str) -> None:
+ module = _load_pipeline_module()
+
+ pipeline = _pipeline(module, od_config=_od_config(**{offload_field: True}))
+
+ assert getattr(pipeline.text_encoder, "to_calls", []) == []
+ assert getattr(pipeline.vae, "to_calls", []) == []
+
+
+def test_ar_diffusion_capability_uses_fixed_tp_local_lingbot_geometry() -> None:
+ module = _load_pipeline_module()
+ module.get_tensor_model_parallel_world_size = lambda: 1
+ pipeline = _pipeline(module)
+
+ spec = pipeline.ar_diffusion_kv_cache_spec()
+
+ assert spec.num_layers == 2
+ assert spec.num_kv_heads == 2
+ assert spec.head_size == 4
+ assert spec.tokens_per_frame == 1
+ assert spec.frames_per_block == 3
+ assert spec.window_frames == 3
+ assert spec.sink_frames == 3
+ assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)]
+ assert spec.cross_attention_lengths == {"text": 512}
+ assert spec.model_owned_state_bytes_per_session == 9_600
+
+
+def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ action_root = tmp_path / "trusted-actions"
+ action_dir = action_root / "forward"
+ action_dir.mkdir(parents=True)
+ image_path = tmp_path / "first-frame.png"
+ Image.new("RGBA", (16, 16), color=(1, 2, 3, 4)).save(image_path)
+ sampling = _SamplingParams(extra_args={"action_path": "forward"})
+ sampling.extra_args.pop("_lingbot_camera_trajectory")
+ request = SimpleNamespace(
+ prompt={"prompt": "move", "multi_modal_data": {"image": str(image_path)}},
+ sampling_params=sampling,
+ )
+ trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(9, 1, 1),
+ intrinsics=torch.ones(9, 4),
+ )
+ load_calls = []
+ module.load_camera_trajectory = lambda action: load_calls.append(action) or trajectory
+
+ preprocess = module.get_lingbot_world_pre_process_func(
+ _od_config(model_config={"lingbot_action_root": str(action_root)})
+ )
+ result = preprocess(request)
+
+ assert result is request
+ assert isinstance(request.prompt["multi_modal_data"]["image"], Image.Image)
+ assert request.prompt["multi_modal_data"]["image"].mode == "RGB"
+ assert sampling.extra_args["action_path"] == "forward"
+ assert sampling.extra_args["_lingbot_camera_trajectory"] is trajectory
+ assert len(load_calls) == 1
+
+
+def test_preprocess_materializes_camera_from_typed_tick_without_action_path() -> None:
+ module = _load_pipeline_module()
+ poses = torch.eye(4).repeat(9, 1, 1)
+ intrinsics = torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(9, 1)
+ tick = ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="tick-request-0",
+ chunk_index=0,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_trajectory.v1",
+ data={
+ "poses": poses.tolist(),
+ "intrinsics": intrinsics.tolist(),
+ },
+ ),
+ ),
+ )
+ sampling = _SamplingParams(include_action=False)
+ sampling.extra_args.update(tick.to_extra_args())
+ request = SimpleNamespace(
+ request_id="tick-request-0",
+ prompt=_prompt(),
+ sampling_params=sampling,
+ )
+
+ result = module.get_lingbot_world_pre_process_func(_od_config())(request)
+
+ trajectory = result.sampling_params.extra_args["_lingbot_camera_trajectory"]
+ torch.testing.assert_close(trajectory.poses, poses)
+ torch.testing.assert_close(trajectory.intrinsics, intrinsics)
+
+
+def test_preprocess_materializes_chunk_sized_actions_from_typed_tick() -> None:
+ module = _load_pipeline_module()
+ tick = ARDiffusionTickRequest(
+ session_id="world-actions",
+ request_id="tick-request-0",
+ chunk_index=0,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_actions.v1",
+ data={
+ "mode": "frames",
+ "frames": [["w"], ["w", "j"], []],
+ },
+ ),
+ ),
+ )
+ sampling = _SamplingParams(include_action=False)
+ sampling.extra_args.update(tick.to_extra_args())
+ request = SimpleNamespace(
+ request_id="tick-request-0",
+ prompt=_prompt(),
+ sampling_params=sampling,
+ )
+
+ result = module.get_lingbot_world_pre_process_func(_od_config())(request)
+
+ assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None
+ assert result.sampling_params.extra_args["_lingbot_camera_actions"] == (
+ ("w",),
+ ("w", "j"),
+ (),
+ )
+
+
+def _request(*, sampling=None, prompt=None, num_reqs: int = 1):
+ return _RequestBatch(
+ _prompt() if prompt is None else prompt,
+ _SamplingParams() if sampling is None else sampling,
+ num_reqs=num_reqs,
+ )
+
+
+def _resolve_pipeline_through_real_registry(pipeline_module):
+ from vllm_omni.diffusion import registry as registry_module
+
+ resolved = registry_module.DiffusionModelRegistry._try_load_model_cls("LingBotWorldCausalDMDPipeline")
+ entry = registry_module._DIFFUSION_MODELS["LingBotWorldCausalDMDPipeline"]
+ cache_acceleration_disabled = "LingBotWorldCausalDMDPipeline" in registry_module._NO_CACHE_ACCELERATION
+ preprocess_name = registry_module._DIFFUSION_PRE_PROCESS_FUNCS["LingBotWorldCausalDMDPipeline"]
+ preprocess = registry_module.get_diffusion_pre_process_func(
+ SimpleNamespace(
+ model_class_name="LingBotWorldCausalDMDPipeline",
+ model_config={"lingbot_action_root": str(_ROOT)},
+ )
+ )
+ return resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess
+
+
+def test_component_discovery_uses_official_checkpoint_contract() -> None:
+ module = _load_pipeline_module()
+ pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config())
+
+ assert pipeline._dit_modules == ["transformer"]
+ assert pipeline._encoder_modules == ["text_encoder"]
+ assert pipeline._vae_modules == ["vae"]
+ assert pipeline.dummy_run_num_frames == 0
+ assert pipeline.weights_sources == [
+ module.DiffusersPipelineLoader.ComponentSource("checkpoint", "transformer", None, "transformer.", True)
+ ]
+ assert _FakeTransformerFactory.last_call[1:] == (None, "transformer")
+ assert pipeline.scheduler.config.shift == 5.0
+ assert pipeline.scheduler.config.num_train_timesteps == 1000
+ assert module._loader_state.prefetch_calls == [("checkpoint", ("tokenizer", "text_encoder", "vae"), False)]
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "feature"),
+ [
+ ("pipeline_parallel_size", 2, "pipeline parallelism"),
+ ("sequence_parallel_size", 2, "sequence parallelism"),
+ ("cfg_parallel_size", 2, "CFG parallelism"),
+ ("vae_patch_parallel_size", 2, "VAE parallelism"),
+ ("use_hsdp", True, "HSDP"),
+ ("enable_expert_parallel", True, "expert parallelism"),
+ ],
+)
+def test_unsupported_parallel_modes_fail_before_component_loading(field: str, value: object, feature: str) -> None:
+ module = _load_pipeline_module()
+ parallel_config = _od_config().parallel_config
+ setattr(parallel_config, field, value)
+
+ with pytest.raises(NotImplementedError, match=feature):
+ module.LingBotWorldCausalDMDPipeline(od_config=_od_config(parallel_config=parallel_config))
+
+ assert module._loader_state.prefetch_calls == []
+
+
+def test_unsupported_quantization_fails_before_component_loading() -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(NotImplementedError, match="quantization"):
+ module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=object()))
+
+ assert module._loader_state.prefetch_calls == []
+
+
+def test_official_scheduler_config_matches_fixed_dmd_contract() -> None:
+ import json
+
+ module = _load_pipeline_module()
+ scheduler_config = json.loads(_SCHEDULER_FIXTURE.read_text())
+ module._load_json = lambda *args, **kwargs: scheduler_config.copy()
+
+ pipeline = _pipeline(module)
+
+ assert pipeline.scheduler.config.num_train_timesteps == 1000
+ assert pipeline.scheduler.config.shift == 5.0
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("_class_name", "FlowMatchEulerDiscreteScheduler"),
+ ("num_train_timesteps", 999),
+ ("prediction_type", "epsilon"),
+ ("predict_x0", False),
+ ("use_flow_sigmas", False),
+ ("use_dynamic_shifting", True),
+ ("final_sigmas_type", "sigma_min"),
+ ],
+)
+def test_scheduler_config_rejects_semantic_drift(field: str, value: object) -> None:
+ module = _load_pipeline_module()
+ module._load_json = lambda *args, **kwargs: _scheduler_config(**{field: value})
+
+ with pytest.raises(ValueError, match=field):
+ _pipeline(module)
+
+
+def test_official_model_index_discovers_only_declared_components() -> None:
+ import json
+
+ module = _load_pipeline_module()
+ model_index = json.loads(_MODEL_INDEX_FIXTURE.read_text())
+ resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
+ module
+ )
+
+ assert model_index["_class_name"] == "LingBotWorldCausalDMDPipeline"
+ assert resolved is module.LingBotWorldCausalDMDPipeline
+ assert entry == ("lingbot_world", "pipeline", "LingBotWorldCausalDMDPipeline")
+ assert cache_acceleration_disabled
+ assert preprocess_name == "get_lingbot_world_pre_process_func"
+ assert callable(preprocess)
+ assert model_index["tokenizer"] == ["transformers", "T5TokenizerFast"]
+ assert model_index["text_encoder"] == ["transformers", "UMT5EncoderModel"]
+ assert model_index["vae"] == ["diffusers", "AutoencoderKLWan"]
+ assert model_index["scheduler"] == ["diffusers", "UniPCMultistepScheduler"]
+ assert model_index["transformer"] == ["diffusers", "CausalLingBotWorldTransformer3DModel"]
+ assert model_index["image_encoder"] == [None, None]
+ assert model_index["image_processor"] == [None, None]
+ assert model_index["transformer_2"] == [None, None]
+
+
+def test_postprocess_returns_latents_without_video_conversion() -> None:
+ module = _load_pipeline_module()
+ postprocess = module.get_lingbot_world_post_process_func(_od_config())
+ latents = torch.randn(1, 16, 3, 2, 2)
+
+ output = postprocess(latents, sampling_params=SimpleNamespace(output_type="latent"))
+
+ assert output is latents
+
+
+def test_postprocess_uses_the_standard_diffusion_output_envelope(monkeypatch) -> None:
+ import diffusers.video_processor as video_processor_module
+
+ module = _load_pipeline_module()
+ processed_video = object()
+ calls = []
+
+ class VideoProcessor:
+ def __init__(self, *, vae_scale_factor):
+ assert vae_scale_factor == 8
+
+ def postprocess_video(self, video, *, output_type):
+ calls.append((video, output_type))
+ return processed_video
+
+ monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+ postprocess = module.get_lingbot_world_post_process_func(_od_config())
+ video = torch.randn(1, 3, 9, 8, 8)
+
+ output = postprocess(video, output_type="np")
+
+ assert output == {"payload": {"video": processed_video}, "metadata": {}}
+ assert calls == [(video, "np")]
+
+
+def test_forward_propagates_pipeline_profiler_stage_durations() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._profiler_lock = Lock()
+ pipeline._stage_durations = {"_generate_block": 0.25, "vae.decode": 0.5}
+
+ output = pipeline(_request())
+
+ assert output.stage_durations == pipeline.stage_durations
+
+
+def test_pipeline_weight_loader_preserves_component_parameter_prefixes() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ weight = torch.randn(2, 2)
+
+ loaded = pipeline.load_weights([("transformer.blocks.0.weight", weight)])
+
+ assert loaded == {"transformer.blocks.0.weight"}
+ assert pipeline.transformer.loaded_weights == [("blocks.0.weight", weight)]
+
+
+def test_request_requires_runner_provided_generator() -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(seed=None, generator=None)
+
+ with pytest.raises(ValueError, match="runner-provided torch.Generator"):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+@pytest.mark.parametrize("unsupported", ["generator-list", "caller-latents"])
+def test_request_rejects_unsupported_sampling_state(unsupported: str) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams()
+ if unsupported == "generator-list":
+ sampling.generator = [sampling.generator]
+ message = "generator list"
+ else:
+ sampling.latents = torch.zeros(1, 16, 3, 2, 2)
+ message = "caller-provided latents"
+
+ with pytest.raises(ValueError, match=message):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer(dtype=torch.bfloat16)
+ pipeline = _pipeline(module, transformer=transformer)
+ requested_noise_dtypes: list[torch.dtype] = []
+
+ def randn(shape, *, generator, device, dtype):
+ del generator
+ assert device == pipeline.device
+ requested_noise_dtypes.append(dtype)
+ return torch.zeros(shape, device=device, dtype=dtype)
+
+ module.randn_tensor = randn
+ result = pipeline(_request())
+
+ assert requested_noise_dtypes == [torch.float32] * 4
+ assert result.output.dtype == torch.float32
+ assert all(call["hidden_states"].dtype == torch.bfloat16 for call in transformer.calls)
+
+
+def test_path_image_rejects_oversized_source_before_decode_or_convert(monkeypatch, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ source_path = tmp_path / "oversized-compressed.png"
+ Image.new("1", (4097, 4097)).save(source_path)
+ decode_calls: list[str] = []
+
+ def forbidden_convert(*args, **kwargs):
+ del args, kwargs
+ decode_calls.append("convert")
+ raise AssertionError("oversized source reached convert")
+
+ def forbidden_load(*args, **kwargs):
+ del args, kwargs
+ decode_calls.append("load")
+ raise AssertionError("oversized source reached load")
+
+ monkeypatch.setattr(module.PIL.Image.Image, "convert", forbidden_convert)
+ monkeypatch.setattr(module.PIL.Image.Image, "load", forbidden_load)
+
+ with pytest.raises(ValueError, match="source image.*4096.*4096"):
+ _preprocess_request(module, prompt=_prompt(images=str(source_path)))
+
+ assert decode_calls == []
+
+
+def test_normal_path_image_is_decoded_after_source_size_validation(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ source_path = tmp_path / "normal.png"
+ Image.new("RGBA", (32, 24), color=(1, 2, 3, 128)).save(source_path)
+
+ request = _preprocess_request(module, prompt=_prompt(images=str(source_path)))
+ parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params))
+
+ assert module._MAX_SOURCE_IMAGE_PIXELS == 4096 * 4096
+ assert isinstance(parsed.image, Image.Image)
+ assert parsed.image.mode == "RGB"
+ assert parsed.image.size == (32, 24)
+
+
+@pytest.mark.parametrize(
+ ("source_size", "expected_size"),
+ [
+ ((832, 480), (464, 832)),
+ ((480, 832), (832, 480)),
+ ((512, 512), (624, 624)),
+ ],
+)
+def test_default_resolution_matches_official_480p_aspect_derivation(source_size, expected_size) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(height=None, width=None)
+
+ parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(images=Image.new("RGB", source_size)), sampling))
+
+ assert (parsed.height, parsed.width) == expected_size
+ assert parsed.height % 16 == 0
+ assert parsed.width % 16 == 0
+ assert parsed.height * parsed.width <= 480 * 832
+
+
+def test_pil_and_tensor_images_share_official_bicubic_preprocess() -> None:
+ module = _load_pipeline_module()
+ height, width = 7, 11
+ source = torch.arange(3 * height * width, dtype=torch.uint8).reshape(3, height, width)
+ pil_image = Image.fromarray(source.permute(1, 2, 0).numpy(), mode="RGB")
+ sampling = _SamplingParams(height=16, width=32)
+
+ pil_pipeline = _pipeline(module)
+ tensor_pipeline = _pipeline(module)
+ pil_inputs = pil_pipeline._parse_request(_RequestBatch(_prompt(images=pil_image), sampling))
+ tensor_inputs = tensor_pipeline._parse_request(
+ _RequestBatch(_prompt(images=source.clone()), _SamplingParams(height=16, width=32))
+ )
+ pil_condition = pil_pipeline._prepare_image_tensor(pil_inputs.image, height=16, width=32)
+ tensor_condition = tensor_pipeline._prepare_image_tensor(tensor_inputs.image, height=16, width=32)
+ expected = torch.nn.functional.interpolate(
+ source.unsqueeze(0).float() / 255.0,
+ size=(16, 32),
+ mode="bicubic",
+ align_corners=False,
+ )
+ expected = expected * 2.0 - 1.0
+
+ torch.testing.assert_close(pil_condition, expected)
+ torch.testing.assert_close(tensor_condition, expected)
+
+
+def test_path_image_decode_error_is_sanitized(monkeypatch) -> None:
+ module = _load_pipeline_module()
+ unsafe_path = "/private/source/customer-secret.png"
+ monkeypatch.setattr(
+ module.PIL.Image,
+ "open",
+ lambda path: (_ for _ in ()).throw(module.PIL.Image.DecompressionBombError(f"unsafe {path}")),
+ )
+
+ with pytest.raises(ValueError, match="Unable to load multi_modal_data.image") as exc_info:
+ _preprocess_request(module, prompt=_prompt(images=unsafe_path))
+
+ assert unsafe_path not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+@pytest.mark.parametrize("source_size", [(4096, 4096), (4097, 4097)])
+def test_supplied_pil_image_obeys_documented_source_pixel_ceiling(source_size) -> None:
+ module = _load_pipeline_module()
+ source_image = Image.new("1", source_size)
+ close_calls: list[bool] = []
+ source_image.close = lambda: close_calls.append(True)
+
+ if source_size[0] * source_size[1] <= 4096 * 4096:
+ request = _preprocess_request(module, prompt=_prompt(images=source_image))
+ parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params))
+ assert parsed.image is not source_image
+ assert parsed.image.mode == "RGB"
+ assert close_calls == []
+ else:
+ with pytest.raises(ValueError, match="source image.*4096.*4096"):
+ _preprocess_request(module, prompt=_prompt(images=source_image))
+ assert close_calls == []
+
+
+def test_supplied_pil_decode_error_is_sanitized_without_closing_caller(monkeypatch) -> None:
+ module = _load_pipeline_module()
+ source_image = Image.new("RGB", (16, 16))
+ unsafe_detail = "/private/source/customer-secret.png"
+ close_calls: list[bool] = []
+ monkeypatch.setattr(source_image, "convert", lambda mode: (_ for _ in ()).throw(OSError(unsafe_detail)))
+ monkeypatch.setattr(source_image, "close", lambda: close_calls.append(True))
+
+ with pytest.raises(ValueError, match="Unable to load multi_modal_data.image") as exc_info:
+ _preprocess_request(module, prompt=_prompt(images=source_image))
+
+ assert unsafe_detail not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+ assert close_calls == []
+
+
+def test_action_path_resolves_from_sampling_extra_args(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ action_root = tmp_path / "trusted-actions"
+ contained_action = action_root / "actions"
+ contained_action.mkdir(parents=True)
+ sampling = _SamplingParams(extra_args={"action_path": "actions"})
+ prompt = _prompt()
+ original_prompt = prompt.copy()
+ trajectory = _CameraTrajectory(torch.eye(4).repeat(9, 1, 1), torch.ones(9, 4))
+ resolved_actions = []
+ module.load_camera_trajectory = lambda action: resolved_actions.append(action) or trajectory
+
+ request = _preprocess_request(
+ module,
+ prompt=prompt,
+ sampling=sampling,
+ od_config=_od_config(model_config={"lingbot_action_root": str(action_root)}),
+ )
+ parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params))
+
+ assert resolved_actions[0].root == action_root.resolve()
+ assert resolved_actions[0].relative == Path("actions")
+ assert parsed.camera_trajectory is trajectory
+ assert sampling.extra_args["action_path"] == "actions"
+ assert sampling.extra_args["_lingbot_camera_trajectory"] is trajectory
+ assert prompt == original_prompt
+
+
+@pytest.mark.parametrize("with_extra", [False, True], ids=["additional-only", "extra-and-additional"])
+def test_additional_information_action_path_is_not_an_input_source(with_extra: bool) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(
+ extra_args={"action_path": "."} if with_extra else None,
+ include_action=with_extra,
+ )
+
+ with pytest.raises(ValueError, match="additional_information.*action_path.*not supported"):
+ _preprocess_request(module, prompt=_prompt(action_path="legacy-actions"), sampling=sampling)
+
+
+def test_action_path_is_required_in_sampling_extra_args() -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(ValueError, match="sampling_params.extra_args.action_path"):
+ _preprocess_request(module, sampling=_SamplingParams(include_action=False))
+
+
+@pytest.mark.parametrize("escape_kind", ["traversal", "absolute", "symlink"])
+def test_online_action_path_rejects_escape_from_trusted_root(escape_kind: str, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ root = tmp_path / "trusted"
+ root.mkdir()
+ outside = tmp_path / "outside"
+ outside.mkdir()
+ if escape_kind == "traversal":
+ action_path = "../outside"
+ elif escape_kind == "absolute":
+ action_path = str(outside)
+ else:
+ (root / "escape-link").symlink_to(outside, target_is_directory=True)
+ action_path = "escape-link"
+ with pytest.raises(ValueError, match="trusted.*root|contained"):
+ _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": action_path}),
+ od_config=_od_config(model_config={"lingbot_action_root": str(root)}),
+ )
+
+
+def test_online_action_path_uses_environment_root_fallback(monkeypatch, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ root = tmp_path / "trusted"
+ action_dir = root / "forward"
+ action_dir.mkdir(parents=True)
+ monkeypatch.setenv("VLLM_OMNI_LINGBOT_ACTION_ROOT", str(root))
+ resolved_actions = []
+ module.load_camera_trajectory = lambda action: (
+ resolved_actions.append(action) or _CameraTrajectory(torch.eye(4).repeat(9, 1, 1), torch.ones(9, 4))
+ )
+
+ request = _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": "forward"}),
+ od_config=_od_config(model_config={}),
+ )
+
+ assert request.sampling_params.extra_args["_lingbot_camera_trajectory"] is not None
+ assert resolved_actions[0].root == root.resolve()
+ assert resolved_actions[0].relative == Path("forward")
+
+
+def test_online_action_path_error_suppresses_path_bearing_filesystem_cause(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ root = tmp_path / "trusted"
+ root.mkdir()
+
+ with pytest.raises(ValueError, match="trusted action root") as exc_info:
+ _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": "does-not-exist"}),
+ od_config=_od_config(model_config={"lingbot_action_root": str(root)}),
+ )
+
+ assert str(root) not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+@pytest.mark.parametrize("source", ["root", "candidate"])
+def test_online_action_path_unknown_user_is_sanitized(source: str, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ unknown_user_path = "~__vllm_omni_user_that_does_not_exist__/actions"
+ root = tmp_path / "trusted"
+ root.mkdir()
+ configured_root = unknown_user_path if source == "root" else str(root)
+ action_path = "actions" if source == "root" else unknown_user_path
+
+ with pytest.raises(ValueError, match="trusted action root") as exc_info:
+ _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": action_path}),
+ od_config=_od_config(model_config={"lingbot_action_root": configured_root}),
+ )
+
+ assert "__vllm_omni_user_that_does_not_exist__" not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+def test_action_path_requires_a_trusted_root_for_every_request() -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(ValueError, match="lingbot_action_root|VLLM_OMNI_LINGBOT_ACTION_ROOT"):
+ _preprocess_request(module, od_config=_od_config(model_config={}))
+
+
+@pytest.mark.parametrize(
+ ("request_batch", "message"),
+ [
+ (_request(num_reqs=2), "single prompt"),
+ (_request(prompt=_prompt(images=[Image.new("RGB", (16, 16))])), "image.*list"),
+ (_request(prompt=_prompt(images=[])), "image.*list"),
+ (
+ _request(prompt=_prompt(images=[Image.new("RGB", (16, 16)), Image.new("RGB", (16, 16))])),
+ "image.*list",
+ ),
+ (_request(sampling=_SamplingParams(num_outputs_per_prompt=2)), "num_outputs_per_prompt"),
+ (_request(sampling=_SamplingParams(num_inference_steps=5)), "num_inference_steps"),
+ (_request(sampling=_SamplingParams(height=15)), "height.*divisible"),
+ (_request(sampling=_SamplingParams(width=15)), "width.*divisible"),
+ (_request(sampling=_SamplingParams(height=None, width=16)), "height and width.*both"),
+ (_request(sampling=_SamplingParams(num_frames=13)), "num_frames.*three-frame"),
+ ],
+)
+def test_request_validation_rejects_unsupported_contracts(request_batch, message: str) -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(ValueError, match=message):
+ _pipeline(module)._parse_request(request_batch)
+
+
+def test_resource_limits_accept_exact_documented_boundaries() -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(height=480, width=832, num_frames=117, max_sequence_length=512)
+
+ parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+ assert (parsed.height, parsed.width, parsed.num_frames, parsed.max_sequence_length) == (480, 832, 117, 512)
+
+
+@pytest.mark.parametrize(
+ ("sampling", "message"),
+ [
+ (_SamplingParams(height=480, width=848), "pixel area|480.*832"),
+ (_SamplingParams(num_frames=129), "num_frames.*117"),
+ (_SamplingParams(max_sequence_length=511), "max_sequence_length.*512"),
+ (_SamplingParams(max_sequence_length=513), "max_sequence_length.*512"),
+ (_SamplingParams(max_sequence_length=512.0), "max_sequence_length.*512"),
+ ],
+)
+def test_resource_limits_reject_oversize_before_any_component_call(sampling, message: str) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ calls: list[str] = []
+ pipeline.encode_prompt = lambda *args, **kwargs: calls.append("text")
+ pipeline._prepare_condition = lambda *args, **kwargs: calls.append("vae")
+ pipeline._prepare_camera = lambda *args, **kwargs: calls.append("camera")
+ pipeline.transformer.allocate_cache = lambda *args, **kwargs: calls.append("cache")
+
+ with pytest.raises(ValueError, match=message):
+ pipeline(_RequestBatch(_prompt(), sampling))
+
+ assert calls == []
+
+
+def test_request_validation_rejects_insufficient_camera_frames() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ short_trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(8, 1, 1),
+ intrinsics=torch.ones(8, 4),
+ )
+ sampling = _SamplingParams(extra_args={"_lingbot_camera_trajectory": short_trajectory})
+
+ with pytest.raises(ValueError, match="camera.*frames.*num_frames"):
+ pipeline(_request(sampling=sampling))
+
+
+def test_camera_load_error_is_actionable_without_echoing_path_contents() -> None:
+ module = _load_pipeline_module()
+ module.load_camera_trajectory = lambda path: (_ for _ in ()).throw(FileNotFoundError(f"missing {path}/poses.npy"))
+
+ with pytest.raises(ValueError, match="camera trajectory.*action_path") as exc_info:
+ _preprocess_request(module)
+
+ assert str(_ROOT) not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ module.randn_tensor = lambda shape, **kwargs: torch.full(
+ shape,
+ -99.0,
+ device=kwargs["device"],
+ dtype=kwargs["dtype"],
+ )
+
+ result = pipeline(_request())
+
+ assert result.output.shape == (1, 16, 3, 2, 2)
+ first_input = transformer.calls[0]["hidden_states"]
+ assert first_input.shape == (1, 36, 3, 2, 2)
+ mean = torch.tensor(pipeline.vae.config.latents_mean).view(1, 16, 1, 1)
+ std = torch.tensor(pipeline.vae.config.latents_std).view(1, 16, 1, 1)
+ expected_first = (torch.full((1, 16, 2, 2), 2.0) - mean) / std
+ expected_future = (torch.zeros(1, 16, 2, 2, 2) - mean.unsqueeze(2)) / std.unsqueeze(2)
+ torch.testing.assert_close(first_input[:, :16], torch.full((1, 16, 3, 2, 2), -99.0))
+ torch.testing.assert_close(first_input[:, 16:20, 0], torch.ones(1, 4, 2, 2))
+ torch.testing.assert_close(first_input[:, 16:20, 1:], torch.zeros(1, 4, 2, 2, 2))
+ torch.testing.assert_close(first_input[:, 20:36, 0], expected_first)
+ torch.testing.assert_close(first_input[:, 20:36, 1:], expected_future)
+
+ raw_camera = module.build_plucker_embedding(
+ SimpleNamespace(poses=torch.empty(3, 4, 4)),
+ height=16,
+ width=16,
+ target_height=16,
+ target_width=16,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+ expected_camera = module._fold_camera_embedding(raw_camera)
+ reference_camera = torch.nn.functional.pixel_unshuffle(raw_camera, 8).permute(1, 0, 2, 3).unsqueeze(0)
+ torch.testing.assert_close(expected_camera, reference_camera)
+ torch.testing.assert_close(transformer.calls[0]["camera_hidden_states"], expected_camera)
+ assert expected_camera.shape == (1, 384, 3, 2, 2)
+
+
+def test_fixed_dmd_transition_and_cache_commit_trace() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+
+ result = pipeline(_request())
+
+ generator = torch.Generator(device="cpu").manual_seed(17)
+ current = torch.randn((1, 16, 3, 2, 2), generator=generator)
+ warped_schedule = (
+ (1000.0, 1.0),
+ (937.5, 0.9375),
+ (2500.0 / 3.0, 5.0 / 6.0),
+ (625.0, 0.625),
+ )
+ for index, (_, sigma) in enumerate(warped_schedule):
+ x0 = current - sigma
+ if index + 1 < len(module.LINGBOT_DMD_TIMESTEPS):
+ next_sigma = warped_schedule[index + 1][1]
+ noise = torch.randn(current.shape, generator=generator)
+ current = (1.0 - next_sigma) * x0 + next_sigma * noise
+ else:
+ current = x0
+ torch.testing.assert_close(result.output, current)
+
+ torch.testing.assert_close(
+ torch.cat([call["timestep"] for call in transformer.calls]),
+ torch.tensor([*(timestep for timestep, _ in warped_schedule), 0.0]),
+ )
+ assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True]
+ assert [call["start_frame"] for call in transformer.calls] == [0, 0, 0, 0, 0]
+ assert len({call["cache_id"] for call in transformer.calls}) == 1
+ torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], result.output)
+
+
+@pytest.mark.parametrize(
+ ("flow_shift", "expected_timesteps"),
+ [
+ (5.0, (1000.0, 937.5, 2500.0 / 3.0, 625.0)),
+ (2.5, (1000.0, 15000.0 / 17.0, 5000.0 / 7.0, 5000.0 / 11.0)),
+ ],
+)
+def test_request_flow_shift_warps_transformer_timesteps(flow_shift: float, expected_timesteps) -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+
+ pipeline(
+ _request(
+ sampling=_SamplingParams(extra_args={"action_path": ".", "flow_shift": flow_shift}),
+ )
+ )
+
+ torch.testing.assert_close(
+ torch.cat([call["timestep"] for call in transformer.calls[:4]]),
+ torch.tensor(expected_timesteps),
+ )
+
+
+def test_request_tiny_positive_flow_shift_keeps_schedule_finite() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+
+ pipeline(
+ _request(
+ sampling=_SamplingParams(extra_args={"action_path": ".", "flow_shift": 1e-20}),
+ )
+ )
+
+ timesteps = torch.cat([call["timestep"] for call in transformer.calls[:4]])
+ assert torch.isfinite(timesteps).all()
+ assert timesteps[0].item() == 1000.0
+
+
+def test_request_flow_shift_override_has_precedence_without_mutating_scheduler() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module, od_config=_od_config(flow_shift=3.0))
+
+ default_inputs = pipeline._parse_request(_request())
+ sampling = _SamplingParams(extra_args={"flow_shift": 2.0})
+ override_inputs = pipeline._parse_request(_RequestBatch(_prompt(), sampling))
+
+ assert default_inputs.flow_shift == 3.0
+ assert override_inputs.flow_shift == 2.0
+ assert sampling.extra_args["action_path"] == "."
+ assert sampling.extra_args["flow_shift"] == 2.0
+ assert pipeline.scheduler.config.shift == 3.0
+
+
+def test_checkpoint_flow_shift_default_and_engine_request_override_precedence() -> None:
+ module = _load_pipeline_module()
+ checkpoint_scheduler_config = _scheduler_config(flow_shift=6.25, shift=7.5)
+ module._load_json = lambda *args, **kwargs: checkpoint_scheduler_config.copy()
+
+ checkpoint_default = _pipeline(module)
+ engine_override = _pipeline(module, od_config=_od_config(flow_shift=3.5))
+ request_override = checkpoint_default._parse_request(
+ _RequestBatch(
+ _prompt(),
+ _SamplingParams(extra_args={"flow_shift": 2.25}),
+ )
+ )
+
+ assert checkpoint_default.scheduler.config.shift == 6.25
+ assert checkpoint_default._parse_request(_request()).flow_shift == 6.25
+ assert engine_override.scheduler.config.shift == 3.5
+ assert engine_override._parse_request(_request()).flow_shift == 3.5
+ assert request_override.flow_shift == 2.25
+
+
+@pytest.mark.parametrize(
+ ("scheduler_config", "expected"),
+ [
+ ({"flow_shift": 6.0, "shift": 7.0}, 6.0),
+ ({"shift": 7.0}, 7.0),
+ ({}, 5.0),
+ ],
+)
+def test_checkpoint_scheduler_shift_fallback_order(scheduler_config, expected: float) -> None:
+ module = _load_pipeline_module()
+ checkpoint_config = _scheduler_config()
+ checkpoint_config.pop("flow_shift")
+ checkpoint_config.update(scheduler_config)
+ module._load_json = lambda *args, **kwargs: checkpoint_config.copy()
+
+ pipeline = _pipeline(module)
+
+ assert pipeline.scheduler.config.shift == expected
+
+
+@pytest.mark.parametrize("flow_shift", [0, -1, float("nan"), float("inf"), "invalid"])
+def test_request_flow_shift_must_be_positive_and_finite(flow_shift) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(extra_args={"flow_shift": flow_shift})
+
+ with pytest.raises(ValueError, match="flow_shift.*positive.*finite"):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+@pytest.mark.parametrize(
+ "flow_shift",
+ [True, False, 10**10000],
+ ids=["true", "false", "overflowing-integer"],
+)
+def test_request_flow_shift_rejects_booleans_and_normalizes_integer_overflow(flow_shift) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(extra_args={"flow_shift": flow_shift})
+
+ with pytest.raises(ValueError, match="flow_shift.*positive.*finite"):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+def test_flow_shift_is_request_local_without_mutating_scheduler() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+
+ def generate(flow_shift=None):
+ extra_args = {} if flow_shift is None else {"flow_shift": flow_shift}
+ return pipeline(
+ _RequestBatch(
+ _prompt(),
+ _SamplingParams(num_frames=21, extra_args=extra_args),
+ )
+ ).output
+
+ default_before = generate()
+ shifted_two = generate(2.0)
+ shifted_seven = generate(7.0)
+ default_after = generate()
+
+ torch.testing.assert_close(default_before, default_after)
+ assert not torch.equal(default_before, shifted_two)
+ assert not torch.equal(shifted_two, shifted_seven)
+ assert pipeline.scheduler.config.shift == 5.0
+
+
+def test_encode_prompt_zeroes_padded_umt5_states_to_exactly_512_tokens() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ attention_mask = torch.zeros(1, 512, dtype=torch.long)
+ attention_mask[:, :3] = 1
+ pipeline.tokenizer = lambda *args, **kwargs: SimpleNamespace(
+ input_ids=torch.arange(512).view(1, 512),
+ attention_mask=attention_mask,
+ )
+ raw_states = torch.arange(512 * 8, dtype=torch.float32).view(1, 512, 8) + 1.0
+ pipeline.text_encoder = lambda input_ids, mask: SimpleNamespace(last_hidden_state=raw_states.clone())
+
+ encoded = module.LingBotWorldCausalDMDPipeline.encode_prompt(
+ pipeline,
+ "move",
+ max_sequence_length=512,
+ dtype=torch.float32,
+ )
+
+ assert encoded.shape == (1, 512, 8)
+ torch.testing.assert_close(encoded[:, :3], raw_states[:, :3])
+ torch.testing.assert_close(encoded[:, 3:], torch.zeros_like(encoded[:, 3:]))
+
+
+def test_multi_chunk_generation_uses_one_request_local_cache_and_decodes_accumulated_latents() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ allocations = []
+ original_allocate = transformer.allocate_cache
+
+ def allocate(**kwargs):
+ cache = original_allocate(**kwargs)
+ allocations.append((kwargs, weakref.ref(cache)))
+ return cache
+
+ transformer.allocate_cache = allocate
+ sampling = _SamplingParams(num_frames=21, output_type="np")
+
+ result = pipeline(_request(sampling=sampling))
+
+ assert result.output.shape == (1, 3, 21, 16, 16)
+ assert len(transformer.calls) == 10
+ assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
+ expected_timesteps = torch.tensor([1000.0, 937.5, 2500.0 / 3.0, 625.0, 0.0] * 2)
+ torch.testing.assert_close(
+ torch.cat([call["timestep"] for call in transformer.calls]),
+ expected_timesteps,
+ )
+ assert len(allocations) == 1
+ kwargs, cache_ref = allocations[0]
+ assert kwargs == {
+ "batch_size": 1,
+ "latent_height": 2,
+ "latent_width": 2,
+ "device": torch.device("cpu"),
+ "dtype": torch.float32,
+ }
+ assert not hasattr(pipeline, "cache")
+ assert not hasattr(pipeline, "transformer_cache")
+ assert pipeline.vae.decode_inputs[0].shape == (1, 16, 6, 2, 2)
+ normalized_latents = torch.cat(
+ (transformer.calls[4]["hidden_states"][:, :16], transformer.calls[9]["hidden_states"][:, :16]),
+ dim=2,
+ )
+ mean = torch.tensor(pipeline.vae.config.latents_mean).view(1, 16, 1, 1, 1)
+ std = torch.tensor(pipeline.vae.config.latents_std).view(1, 16, 1, 1, 1)
+ torch.testing.assert_close(pipeline.vae.decode_inputs[0], normalized_latents * std + mean)
+ transformer.calls.clear()
+ gc.collect()
+ assert cache_ref() is None
+
+
+def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room"):
+ return ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="legacy-lingbot-request",
+ chunk_index=chunk_index,
+ applied_event_ids=(chunk_index,),
+ prompt=prompt,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_trajectory.v1",
+ data={"poses": [], "intrinsics": []},
+ ),
+ ),
+ ).to_extra_args()
+
+
+def test_typed_ticks_generate_one_global_block_and_return_standard_metadata(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ cross_calls = []
+ generated = []
+
+ def ar_text_caches(prompt_embeds, *, invalidate):
+ del prompt_embeds
+ cross_calls.append(invalidate)
+ return [SimpleNamespace()]
+
+ def generate_block(**kwargs):
+ block = torch.randn(
+ (1, 16, 3, 2, 2),
+ generator=kwargs["generator"],
+ )
+ generated.append(
+ {
+ "start_frame": kwargs["start_frame"],
+ "condition": kwargs["condition"].clone(),
+ "block": block.clone(),
+ }
+ )
+ return block
+
+ monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches)
+ monkeypatch.setattr(pipeline, "_generate_block", generate_block)
+
+ first_sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=0))
+ first = pipeline(_request(sampling=first_sampling))
+ switched_prompt = _prompt()
+ switched_prompt["prompt"] = "enter the snowy valley"
+ second_sampling = _SamplingParams(
+ extra_args=_tick_extra_args(
+ chunk_index=1,
+ prompt=switched_prompt["prompt"],
+ )
+ )
+ second = pipeline(_request(sampling=second_sampling, prompt=switched_prompt))
+
+ assert generated[0]["start_frame"] == 0
+ assert generated[1]["start_frame"] == 3
+ assert torch.count_nonzero(generated[0]["condition"]) > 0
+ torch.testing.assert_close(
+ generated[1]["condition"][:, :4],
+ torch.zeros_like(generated[1]["condition"][:, :4]),
+ )
+ assert torch.count_nonzero(generated[1]["condition"][:, 4:]) > 0
+ assert cross_calls == [False, True]
+ assert first.output["payload"]["latents"].shape == (1, 16, 3, 2, 2)
+ assert second.output["metadata"]["ar_diffusion"] == {
+ "session_id": "world-1",
+ "request_id": "legacy-lingbot-request",
+ "chunk_index": 1,
+ "applied_event_ids": [1],
+ }
+ assert pipeline._ar_sessions["world-1"].next_chunk_index == 2
+ assert pipeline._ar_sessions["world-1"].prompt == "enter the snowy valley"
+ assert pipeline._ar_sessions["world-1"].camera_tail is not None
+ assert pipeline._ar_sessions["world-1"].camera_tail.poses.shape == (1, 4, 4)
+ expected_generator = torch.Generator(device="cpu").manual_seed(17)
+ expected_first = torch.randn((1, 16, 3, 2, 2), generator=expected_generator)
+ expected_second = torch.randn((1, 16, 3, 2, 2), generator=expected_generator)
+ torch.testing.assert_close(generated[0]["block"], expected_first)
+ torch.testing.assert_close(generated[1]["block"], expected_second)
+
+
+def test_typed_action_ticks_integrate_camera_across_chunks(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ monkeypatch.setattr(
+ pipeline,
+ "_ar_text_caches",
+ lambda *args, **kwargs: [SimpleNamespace()],
+ )
+ monkeypatch.setattr(
+ pipeline,
+ "_generate_block",
+ lambda **kwargs: torch.zeros_like(kwargs["condition"][:, :16]),
+ )
+
+ def action_sampling(chunk_index: int, frames: list[list[str]]) -> _SamplingParams:
+ tick = ARDiffusionTickRequest(
+ session_id="world-actions",
+ request_id="legacy-lingbot-request",
+ chunk_index=chunk_index,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_actions.v1",
+ data={"mode": "frames", "frames": frames},
+ ),
+ ),
+ )
+ sampling = _SamplingParams(extra_args=tick.to_extra_args())
+ sampling.extra_args["_lingbot_camera_trajectory"] = None
+ sampling.extra_args["_lingbot_camera_actions"] = tuple(tuple(frame) for frame in frames)
+ return sampling
+
+ pipeline(_request(sampling=action_sampling(0, [["w"], ["w"], ["w"]])))
+ first_tail = pipeline._ar_sessions["world-actions"].camera_tail.poses.clone()
+ pipeline(_request(sampling=action_sampling(1, [["d"], ["d"], ["d"]])))
+ state = pipeline._ar_sessions["world-actions"]
+
+ assert state.next_chunk_index == 2
+ assert state.camera_tail is not None
+ torch.testing.assert_close(state.camera_tail.poses[0, 2, 3], first_tail[0, 2, 3])
+ assert state.camera_tail.poses[0, 0, 3] > first_tail[0, 0, 3]
+
+
+def test_first_typed_yaw_action_uses_pre_action_identity_anchor(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ monkeypatch.setattr(
+ module,
+ "build_plucker_embedding",
+ _real_build_plucker_embedding,
+ )
+ action_trajectory, _ = integrate_lingbot_camera_actions(
+ [["j"], [], []],
+ width=16,
+ height=16,
+ )
+ neutral_trajectory, _ = integrate_lingbot_camera_actions(
+ [[], [], []],
+ width=16,
+ height=16,
+ )
+
+ def inputs(trajectory, actions):
+ return SimpleNamespace(
+ camera_trajectory=trajectory,
+ camera_actions=actions,
+ num_latent_frames=3,
+ num_frames=9,
+ height=16,
+ width=16,
+ )
+
+ action_embedding, tail = pipeline._prepare_camera(
+ inputs(action_trajectory, (("j",), (), ())),
+ dtype=torch.float32,
+ )
+ neutral_embedding, _ = pipeline._prepare_camera(
+ inputs(neutral_trajectory, ((), (), ())),
+ dtype=torch.float32,
+ )
+
+ identity_anchor = torch.eye(4, dtype=action_trajectory.poses.dtype).unsqueeze(0)
+ explicit_trajectory = _CameraTrajectory(
+ poses=torch.cat((identity_anchor, action_trajectory.poses), dim=0),
+ intrinsics=torch.cat(
+ (action_trajectory.intrinsics[:1], action_trajectory.intrinsics),
+ dim=0,
+ ),
+ )
+ explicit = _real_build_plucker_embedding(
+ explicit_trajectory,
+ height=16,
+ width=16,
+ target_height=16,
+ target_width=16,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )[1:]
+
+ assert not torch.equal(action_embedding, neutral_embedding)
+ torch.testing.assert_close(action_embedding, module._fold_camera_embedding(explicit))
+ torch.testing.assert_close(tail.poses, action_trajectory.poses[-1:])
+
+
+def test_typed_tick_rejects_non_contiguous_chunk_index(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ monkeypatch.setattr(
+ pipeline,
+ "_ar_text_caches",
+ lambda *args, **kwargs: [SimpleNamespace()],
+ )
+
+ sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=2))
+ with pytest.raises(ValueError, match="must be contiguous"):
+ pipeline(_request(sampling=sampling))
+
+
+def test_typed_tick_rejects_chunk_beyond_realtime_condition_horizon(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ pipeline._ar_sessions["world-1"] = module._LingBotARSessionState(
+ next_chunk_index=10,
+ )
+ monkeypatch.setattr(
+ pipeline,
+ "encode_prompt",
+ lambda *args, **kwargs: pytest.fail("the realtime horizon must be validated before prompt encoding"),
+ )
+
+ sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=10))
+ with pytest.raises(
+ ValueError,
+ match=r"at most 10 ticks.*chunk_index 0 through 9.*117 pixel frames",
+ ):
+ pipeline(_request(sampling=sampling))
+
+
+def test_request_cache_is_released_before_vae_decode() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ cache_refs: list[weakref.ReferenceType] = []
+ original_allocate = transformer.allocate_cache
+
+ def allocate(**kwargs):
+ cache = original_allocate(**kwargs)
+ cache_refs.append(weakref.ref(cache))
+ return cache
+
+ transformer.allocate_cache = allocate
+
+ def assert_cache_released() -> None:
+ gc.collect()
+ assert len(cache_refs) == 1
+ assert cache_refs[0]() is None
+
+ pipeline.vae.on_decode = assert_cache_released
+
+ result = pipeline(_request(sampling=_SamplingParams(num_frames=21, output_type="np")))
+
+ assert result.output.shape[2] == 21
+
+
+def test_117_frame_request_generates_ten_complete_latent_blocks() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(117, 1, 1),
+ intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(117, 1),
+ )
+ sampling = _SamplingParams(
+ num_frames=117,
+ output_type="latent",
+ extra_args={"_lingbot_camera_trajectory": trajectory},
+ )
+
+ result = pipeline(_request(sampling=sampling))
+
+ assert result.output.shape == (1, 16, 30, 2, 2)
+ assert len(transformer.calls) == 50
+ assert [call["start_frame"] for call in transformer.calls[::5]] == list(range(0, 30, 3))
+
+
+def test_request_cache_becomes_unreachable_after_transformer_error() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer(raise_on_call=2)
+ pipeline = _pipeline(module, transformer=transformer)
+ cache_refs = []
+ original_allocate = transformer.allocate_cache
+
+ def allocate(**kwargs):
+ cache = original_allocate(**kwargs)
+ cache_refs.append(weakref.ref(cache))
+ return cache
+
+ transformer.allocate_cache = allocate
+
+ with pytest.raises(RuntimeError, match="forced transformer failure") as exc_info:
+ pipeline(_request())
+
+ exc_info.value.__traceback__ = None
+ del exc_info
+ transformer.calls.clear()
+ gc.collect()
+ assert len(cache_refs) == 1
+ assert cache_refs[0]() is None
+ assert not hasattr(pipeline, "cache")
+ assert not hasattr(pipeline, "transformer_cache")
+
+
+def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
+ module = _load_pipeline_module()
+ resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
+ module
+ )
+
+ assert entry == ("lingbot_world", "pipeline", "LingBotWorldCausalDMDPipeline")
+ assert resolved is module.LingBotWorldCausalDMDPipeline
+ assert cache_acceleration_disabled
+ assert preprocess_name == "get_lingbot_world_pre_process_func"
+ request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams())
+ assert preprocess(request) is request
+ assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory)
+ assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline"
+ lingbot_init = _LINGBOT_INIT_PATH.read_text()
+ assert "from .pipeline import" in lingbot_init
+ assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init
+ assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init
diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2.py b/tests/e2e/offline_inference/test_lingbot_world_v2.py
new file mode 100644
index 000000000..39d9b9d95
--- /dev/null
+++ b/tests/e2e/offline_inference/test_lingbot_world_v2.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end smoke test for LingBot-World v2 generation."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunnerHandler
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = os.environ.get(
+ "VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH",
+ "robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
+)
+_IMAGE_PATH = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH")
+_ACTION_DIR_VALUE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR")
+_ACTION_DIR = Path(_ACTION_DIR_VALUE).expanduser().resolve() if _ACTION_DIR_VALUE else None
+_HAS_INPUT_ASSETS = _IMAGE_PATH is not None and _ACTION_DIR is not None
+
+_OMNI_RUNNER_PARAM = (
+ MODEL,
+ None,
+ {
+ "flow_shift": 5.0,
+ "parallel_config": DiffusionParallelConfig(tensor_parallel_size=1),
+ "model_config": {
+ "lingbot_action_root": str(_ACTION_DIR.parent) if _ACTION_DIR is not None else "",
+ },
+ },
+)
+
+pytestmark = [
+ pytest.mark.slow,
+ pytest.mark.diffusion,
+ pytest.mark.skipif(
+ not _HAS_INPUT_ASSETS,
+ reason=("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH and VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR are required"),
+ ),
+ pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
+]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None:
+ """Load the real checkpoint and generate one causal block on one GPU."""
+
+ assert _IMAGE_PATH is not None
+ assert _ACTION_DIR is not None
+ response = omni_runner_handler.send_diffusion_request(
+ {
+ "model": omni_runner_handler.runner.model_name,
+ "prompt": "The camera moves slowly forward through the scene.",
+ "multi_modal_data": {"image": str(Path(_IMAGE_PATH).expanduser().resolve())},
+ "sampling_params": OmniDiffusionSamplingParams(
+ height=480,
+ width=832,
+ num_frames=9,
+ num_inference_steps=4,
+ max_sequence_length=512,
+ seed=42,
+ fps=16,
+ extra_args={
+ "action_path": _ACTION_DIR.name,
+ "flow_shift": 5.0,
+ },
+ ),
+ }
+ )
+
+ assert response.success
+ assert response.images is not None
+ assert len(response.images) == 1
+ video = np.asarray(response.images[0])
+ assert video.shape == (1, 9, 480, 832, 3)
+ assert np.isfinite(video).all()
+ assert float(np.ptp(video)) > 1e-6
+ assert float(np.mean(np.abs(np.diff(video, axis=1)))) > 1e-6
diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py
index ccc6f69e8..803097dcf 100644
--- a/tests/entrypoints/test_utils.py
+++ b/tests/entrypoints/test_utils.py
@@ -417,8 +417,9 @@ class TestTryResolveOmniModelType:
class TestLoadAndResolveStageConfigs:
def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
- """Ensure that dtype survives default stage creation."""
- kwargs = {"dtype": torch.float32}
+ """Ensure that kwargs survive default stage creation."""
+ engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+ kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
@@ -431,6 +432,7 @@ class TestLoadAndResolveStageConfigs:
assert config_path is None
assert len(stage_configs) == 1
assert "dtype" in stage_configs[0]["engine_args"]
+ assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
deploy_path = tmp_path / "qwen3_multi.yaml"
diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py
index fd13fd482..423ccc7a8 100644
--- a/tests/model_tests/diffusion/test_alignment.py
+++ b/tests/model_tests/diffusion/test_alignment.py
@@ -64,6 +64,9 @@ EXCLUDED_MODELS = [
"HunyuanVideo15Pipeline",
"HunyuanVideo15ImageToVideoPipeline",
"LingBotVideoPipeline",
+ # Requires camera/action assets and persistent AR state; covered by its
+ # dedicated contract tests and real-checkpoint E2E smoke instead.
+ "LingBotWorldCausalDMDPipeline",
"MagiHumanPipeline",
"OmniVoicePipeline",
"OmniVoice",
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index b2668b0a5..693766332 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -70,6 +70,9 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
logger = logging.getLogger(__name__)
MAX_DREAMZERO_SESSIONS = 64
+# Shipped DreamZero geometry retains 24 Wan VAE causal-convolution cache
+# entries. This is the measured persistent CUDA upper bound per live session.
+DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION = 603 * 1024 * 1024
# The pipeline's per-session state is a bespoke ``DreamZeroState`` by default, or
# a ``DreamZeroStateAdapter`` view when the opt-in session manager is enabled.
@@ -149,6 +152,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
session_capacity=MAX_DREAMZERO_SESSIONS,
cross_attention=tuple(cross_attention),
max_scratch_tokens_per_branch=int(transformer.num_action_per_block + transformer.num_state_per_block),
+ model_owned_state_bytes_per_session=DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
)
@contextmanager
diff --git a/vllm_omni/diffusion/models/lingbot_world/__init__.py b/vllm_omni/diffusion/models/lingbot_world/__init__.py
new file mode 100644
index 000000000..a7fe05807
--- /dev/null
+++ b/vllm_omni/diffusion/models/lingbot_world/__init__.py
@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .actions import LingBotCameraControlReducer
+from .pipeline import (
+ LingBotWorldCausalDMDPipeline,
+ get_lingbot_world_post_process_func,
+ get_lingbot_world_pre_process_func,
+)
+from .transformer import CausalLingBotWorldTransformer3DModel
+
+__all__ = [
+ "CausalLingBotWorldTransformer3DModel",
+ "LingBotCameraControlReducer",
+ "LingBotWorldCausalDMDPipeline",
+ "get_lingbot_world_post_process_func",
+ "get_lingbot_world_pre_process_func",
+]
diff --git a/vllm_omni/diffusion/models/lingbot_world/actions.py b/vllm_omni/diffusion/models/lingbot_world/actions.py
new file mode 100644
index 000000000..e725c698a
--- /dev/null
+++ b/vllm_omni/diffusion/models/lingbot_world/actions.py
@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Realtime keyboard controls for LingBot World 2.0."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.models.lingbot_world.camera import (
+ CameraTrajectory,
+)
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionPreparedControls,
+ ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionControlInput,
+)
+
+LINGBOT_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
+LINGBOT_CAMERA_TRAJECTORY_SCHEMA = "lingbot.camera_trajectory.v1"
+_ACTION_ORDER = ("w", "a", "s", "d", "i", "j", "k", "l")
+_VALID_ACTIONS = frozenset(_ACTION_ORDER)
+_REFERENCE_HEIGHT = 480
+_REFERENCE_WIDTH = 832
+
+
+def _normalize_actions(value: object, *, field: str) -> tuple[str, ...]:
+ if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+ raise ValueError(f"{field} must be a sequence of LingBot action keys.")
+ actions: set[str] = set()
+ for item in value:
+ if not isinstance(item, str) or item.lower() not in _VALID_ACTIONS:
+ raise ValueError(f"{field} supports only W/A/S/D/I/J/K/L action keys.")
+ actions.add(item.lower())
+ return tuple(action for action in _ACTION_ORDER if action in actions)
+
+
+def _normalize_frames(value: object, *, field: str) -> tuple[tuple[str, ...], ...]:
+ if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+ raise ValueError(f"{field} must be a sequence of per-frame action lists.")
+ return tuple(_normalize_actions(actions, field=f"{field}[{index}]") for index, actions in enumerate(value))
+
+
+def parse_lingbot_camera_action_frames(
+ data: Mapping[str, Any],
+ *,
+ expected_frames: int,
+) -> tuple[tuple[str, ...], ...]:
+ """Validate the chunk-sized payload emitted by the session reducer."""
+
+ if data.get("mode") != "frames":
+ raise ValueError("lingbot.camera_actions.v1 model ticks require mode='frames'.")
+ frames = _normalize_frames(data.get("frames"), field="camera action frames")
+ if len(frames) != expected_frames:
+ raise ValueError(
+ "lingbot.camera_actions.v1 must contain exactly one action list per "
+ f"latent frame; expected {expected_frames}, got {len(frames)}."
+ )
+ return frames
+
+
+@dataclass(frozen=True)
+class _LingBotCameraReducerState:
+ mode: str | None = None
+ held_actions: tuple[str, ...] = ()
+ script_frames: tuple[tuple[str, ...], ...] = ()
+
+
+class LingBotCameraControlReducer:
+ """Sample live state/script controls into one LingBot AR block.
+
+ State mode preserves held keys across chunks and guarantees a one-frame
+ pulse when a press and release both arrive before the next chunk. Script
+ mode is a finite FIFO padded with neutral frames after exhaustion.
+ """
+
+ def __init__(self, *, frames_per_block: int = 3) -> None:
+ if isinstance(frames_per_block, bool) or not isinstance(frames_per_block, int) or frames_per_block <= 0:
+ raise ValueError("frames_per_block must be a positive integer.")
+ self._frames_per_block = frames_per_block
+ self._state = _LingBotCameraReducerState()
+
+ @staticmethod
+ def _state_transitions(data: Mapping[str, Any]) -> tuple[tuple[str, ...], ...]:
+ transitions = data.get("transitions")
+ if not isinstance(transitions, Sequence) or isinstance(transitions, (str, bytes)) or not transitions:
+ raise ValueError("LingBot state-mode camera actions require non-empty transitions.")
+ normalized: list[tuple[str, ...]] = []
+ previous_timestamp: int | None = None
+ for index, transition in enumerate(transitions):
+ if not isinstance(transition, Mapping):
+ raise ValueError("Each LingBot camera state transition must be a mapping.")
+ timestamp = transition.get("client_ts_ms")
+ if timestamp is not None:
+ if isinstance(timestamp, bool) or not isinstance(timestamp, int) or timestamp < 0:
+ raise ValueError("client_ts_ms must be a non-negative integer when provided.")
+ if previous_timestamp is not None and timestamp < previous_timestamp:
+ raise ValueError("client_ts_ms values must be monotonic within one event.")
+ previous_timestamp = timestamp
+ normalized.append(
+ _normalize_actions(
+ transition.get("actions"),
+ field=f"camera state transitions[{index}].actions",
+ )
+ )
+ return tuple(normalized)
+
+ @staticmethod
+ def _script(data: Mapping[str, Any]) -> tuple[tuple[str, ...], ...]:
+ return _normalize_frames(data.get("frames"), field="camera action script frames")
+
+ def prepare(
+ self,
+ *,
+ current_controls: Mapping[str, ARDiffusionControlInput],
+ events: Sequence[ARDiffusionSessionEvent],
+ chunk_index: int,
+ ) -> ARDiffusionPreparedControls:
+ del chunk_index
+ controls = {
+ track: control
+ for track, control in current_controls.items()
+ if not (track == "camera" and control.schema == LINGBOT_CAMERA_ACTION_SCHEMA)
+ }
+ state = self._state
+ pending_transitions: tuple[tuple[str, ...], ...] = ()
+
+ for event in events:
+ for control in event.controls:
+ if control.track != "camera" or control.schema != LINGBOT_CAMERA_ACTION_SCHEMA:
+ controls[control.track] = control
+ if control.track == "camera":
+ state = _LingBotCameraReducerState()
+ pending_transitions = ()
+ continue
+
+ mode = control.data.get("mode")
+ controls.pop("camera", None)
+ if mode == "script":
+ state = _LingBotCameraReducerState(
+ mode="script",
+ script_frames=self._script(control.data),
+ )
+ pending_transitions = ()
+ elif mode == "state":
+ transitions = self._state_transitions(control.data)
+ if state.mode != "state":
+ state = _LingBotCameraReducerState(mode="state")
+ else:
+ state = _LingBotCameraReducerState(
+ mode="state",
+ held_actions=state.held_actions,
+ )
+ pending_transitions += transitions
+ else:
+ raise ValueError("lingbot.camera_actions.v1 events require mode='state' or mode='script'.")
+
+ if state.mode == "script":
+ frames = state.script_frames[: self._frames_per_block]
+ frames += ((),) * (self._frames_per_block - len(frames))
+ state = _LingBotCameraReducerState(
+ mode="script",
+ script_frames=state.script_frames[self._frames_per_block :],
+ )
+ elif state.mode == "state":
+ final_actions = pending_transitions[-1] if pending_transitions else state.held_actions
+ pulse = next(
+ (actions for actions in reversed(pending_transitions) if actions),
+ None,
+ )
+ if pulse is not None and pulse != final_actions:
+ frames = (pulse,) + (final_actions,) * (self._frames_per_block - 1)
+ else:
+ frames = (final_actions,) * self._frames_per_block
+ state = _LingBotCameraReducerState(
+ mode="state",
+ held_actions=final_actions,
+ )
+ else:
+ frames = ()
+
+ if frames:
+ controls["camera"] = ARDiffusionControlInput(
+ track="camera",
+ schema=LINGBOT_CAMERA_ACTION_SCHEMA,
+ data={
+ "mode": "frames",
+ "frames": [list(actions) for actions in frames],
+ },
+ )
+ return ARDiffusionPreparedControls(
+ controls=tuple(controls[track] for track in sorted(controls)),
+ private_state=state,
+ )
+
+ def commit(self, prepared: ARDiffusionPreparedControls) -> None:
+ if not isinstance(prepared.private_state, _LingBotCameraReducerState):
+ raise TypeError("LingBot reducer prepared state has an unexpected type.")
+ self._state = prepared.private_state
+
+ def reset(self) -> None:
+ self._state = _LingBotCameraReducerState()
+
+
+def _rotation_matrix(axis: str, angle: float) -> np.ndarray:
+ cosine = math.cos(angle)
+ sine = math.sin(angle)
+ if axis == "x":
+ return np.array(
+ ((1.0, 0.0, 0.0), (0.0, cosine, -sine), (0.0, sine, cosine)),
+ dtype=np.float64,
+ )
+ return np.array(
+ ((cosine, 0.0, sine), (0.0, 1.0, 0.0), (-sine, 0.0, cosine)),
+ dtype=np.float64,
+ )
+
+
+def integrate_lingbot_camera_actions(
+ frames: Sequence[Sequence[str]],
+ *,
+ width: int,
+ height: int,
+ initial_pose: torch.Tensor | None = None,
+ initial_pitch: float = 0.0,
+) -> tuple[CameraTrajectory, float]:
+ """Convert latent-frame WASD/IJKL controls to cumulative C2W poses.
+
+ Calibration and motion constants intentionally match SGLang's LingBot
+ adapter: movement 0.05, pitch 4 degrees, yaw 6 degrees, pitch clamp 85
+ degrees, and target-resolution focal lengths of 500 pixels.
+ """
+
+ normalized_frames = tuple(
+ _normalize_actions(actions, field=f"camera action frames[{index}]") for index, actions in enumerate(frames)
+ )
+ if not normalized_frames:
+ raise ValueError("camera action frames must not be empty.")
+ if width <= 0 or height <= 0:
+ raise ValueError("camera action resolution must be positive.")
+
+ if initial_pose is None:
+ current_pose = np.eye(4, dtype=np.float64)
+ else:
+ if initial_pose.shape != (4, 4) or not torch.isfinite(initial_pose).all():
+ raise ValueError("initial_pose must be one finite 4x4 camera-to-world matrix.")
+ current_pose = initial_pose.detach().cpu().double().numpy().copy()
+ current_pitch = float(initial_pitch)
+ pitch_limit = math.radians(85.0)
+ poses: list[np.ndarray] = []
+
+ for actions in normalized_frames:
+ rotation = current_pose[:3, :3]
+ translation = current_pose[:3, 3]
+ pitch_delta = math.radians(4.0) * (("i" in actions) - ("k" in actions))
+ if not -pitch_limit <= current_pitch + pitch_delta <= pitch_limit:
+ pitch_delta = 0.0
+ else:
+ current_pitch += pitch_delta
+ yaw_delta = math.radians(6.0) * (("l" in actions) - ("j" in actions))
+ new_rotation = _rotation_matrix("y", yaw_delta) @ rotation @ _rotation_matrix("x", pitch_delta)
+
+ forward = np.array((new_rotation[0, 2], 0.0, new_rotation[2, 2]))
+ right = np.array((new_rotation[0, 0], 0.0, new_rotation[2, 0]))
+ forward_norm = np.linalg.norm(forward)
+ right_norm = np.linalg.norm(right)
+ if forward_norm > 0:
+ forward /= forward_norm + 1e-6
+ if right_norm > 0:
+ right /= right_norm + 1e-6
+
+ movement = np.zeros(3, dtype=np.float64)
+ movement += forward * 0.05 * (("w" in actions) - ("s" in actions))
+ movement += right * 0.05 * (("d" in actions) - ("a" in actions))
+ current_pose = np.eye(4, dtype=np.float64)
+ current_pose[:3, :3] = new_rotation
+ current_pose[:3, 3] = translation + movement
+ poses.append(current_pose)
+
+ # build_plucker_embedding expects intrinsics in the 832x480 reference
+ # coordinate system. These values become SGLang's [500, 500, W/2, H/2]
+ # after that function scales them to the requested resolution.
+ reference_intrinsics = (
+ 500.0 * _REFERENCE_WIDTH / width,
+ 500.0 * _REFERENCE_HEIGHT / height,
+ _REFERENCE_WIDTH / 2,
+ _REFERENCE_HEIGHT / 2,
+ )
+ trajectory = CameraTrajectory(
+ # Keep the cumulative controller pose in FP64 between chunks, matching
+ # SGLang's NumPy integration and avoiding long-session drift.
+ poses=torch.from_numpy(np.stack(poses)),
+ intrinsics=torch.tensor(reference_intrinsics, dtype=torch.float32).repeat(len(poses), 1),
+ )
+ return trajectory, current_pitch
diff --git a/vllm_omni/diffusion/models/lingbot_world/camera.py b/vllm_omni/diffusion/models/lingbot_world/camera.py
new file mode 100644
index 000000000..0e78e53e6
--- /dev/null
+++ b/vllm_omni/diffusion/models/lingbot_world/camera.py
@@ -0,0 +1,481 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Camera trajectory loading and geometry for LingBot World v2 conditioning."""
+
+from __future__ import annotations
+
+import io
+import math
+import os
+import stat
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import BinaryIO
+
+import numpy as np
+import torch
+
+_REFERENCE_HEIGHT = 480
+_REFERENCE_WIDTH = 832
+_MAX_ACTION_FRAMES = 117
+_MAX_SOURCE_ACTION_FRAMES = 4096
+_MAX_NPY_HEADER_BYTES = 10_000
+_DIRECTORY_OPEN_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
+_FILE_OPEN_FLAGS = os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK
+
+
+@dataclass(frozen=True)
+class CameraTrajectory:
+ """Camera-to-world poses and ``(fx, fy, cx, cy)`` intrinsics by frame."""
+
+ poses: torch.Tensor
+ intrinsics: torch.Tensor
+
+
+@dataclass(frozen=True)
+class TrustedActionDirectory:
+ """Canonical action location tied to the trusted root's filesystem identity."""
+
+ root: Path
+ relative: Path
+ root_device: int
+ root_inode: int
+
+
+def resolve_trusted_action_directory(
+ action_path: str | os.PathLike[str],
+ trusted_root: str | os.PathLike[str],
+) -> TrustedActionDirectory:
+ """Resolve an action directory beneath a trusted root without exposing paths."""
+
+ try:
+ root = Path(trusted_root).expanduser().resolve(strict=True)
+ except (OSError, RuntimeError):
+ raise ValueError("The configured LingBot trusted action root is unavailable.") from None
+ if not root.is_dir():
+ raise ValueError("The configured LingBot trusted action root must be a directory.")
+
+ try:
+ candidate = Path(action_path).expanduser()
+ if not candidate.is_absolute():
+ candidate = root / candidate
+ candidate = candidate.resolve(strict=True)
+ relative = candidate.relative_to(root)
+ except (OSError, RuntimeError, ValueError):
+ raise ValueError(
+ "sampling_params.extra_args.action_path must be contained by the trusted action root."
+ ) from None
+ if not candidate.is_dir():
+ raise ValueError("sampling_params.extra_args.action_path must identify a directory in the trusted action root.")
+
+ try:
+ root_stat = root.stat()
+ except OSError:
+ raise ValueError("The configured LingBot trusted action root is unavailable.") from None
+ return TrustedActionDirectory(
+ root=root,
+ relative=relative,
+ root_device=root_stat.st_dev,
+ root_inode=root_stat.st_ino,
+ )
+
+
+def _validate_trajectory(
+ trajectory: CameraTrajectory,
+ *,
+ file_suffix: str = "",
+) -> None:
+ poses = trajectory.poses
+ intrinsics = trajectory.intrinsics
+ if poses.ndim != 3 or poses.shape[1:] != (4, 4):
+ raise ValueError(f"poses{file_suffix} must have shape [frames, 4, 4], got {tuple(poses.shape)}")
+ if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
+ raise ValueError(
+ f"intrinsics{file_suffix} must have shape [frames, 4] ordered as fx, fy, cx, cy, "
+ f"got {tuple(intrinsics.shape)}"
+ )
+ if poses.shape[0] == 0:
+ raise ValueError("camera trajectory must contain at least one frame")
+ if poses.shape[0] != intrinsics.shape[0]:
+ raise ValueError("poses and intrinsics must contain the same number of frames")
+ if not torch.isfinite(poses).all() or not torch.isfinite(intrinsics).all():
+ raise ValueError("camera trajectory values must all be finite")
+
+
+@contextmanager
+def _open_action_directory(action_directory: TrustedActionDirectory) -> Iterator[int]:
+ root_fd: int | None = None
+ action_fd: int | None = None
+ try:
+ root_fd = os.open(action_directory.root, _DIRECTORY_OPEN_FLAGS)
+ root_stat = os.fstat(root_fd)
+ if (root_stat.st_dev, root_stat.st_ino) != (
+ action_directory.root_device,
+ action_directory.root_inode,
+ ):
+ raise ValueError("The configured LingBot trusted action root changed before camera files were opened.")
+
+ action_fd = os.dup(root_fd)
+ for component in action_directory.relative.parts:
+ next_fd = os.open(component, _DIRECTORY_OPEN_FLAGS, dir_fd=action_fd)
+ os.close(action_fd)
+ action_fd = next_fd
+ except OSError:
+ if action_fd is not None:
+ os.close(action_fd)
+ if root_fd is not None:
+ os.close(root_fd)
+ raise ValueError("Unable to securely open the LingBot action directory within the trusted root.") from None
+ except ValueError:
+ if action_fd is not None:
+ os.close(action_fd)
+ if root_fd is not None:
+ os.close(root_fd)
+ raise
+ assert action_fd is not None
+ assert root_fd is not None
+ try:
+ yield action_fd
+ finally:
+ os.close(action_fd)
+ os.close(root_fd)
+
+
+@contextmanager
+def _open_camera_file(action_fd: int, filename: str) -> Iterator[tuple[BinaryIO, os.stat_result]]:
+ file_fd = -1
+ try:
+ file_fd = os.open(filename, _FILE_OPEN_FLAGS, dir_fd=action_fd)
+ file_stat = os.fstat(file_fd)
+ if not stat.S_ISREG(file_stat.st_mode):
+ raise ValueError(f"{filename} must be a regular NPY file.")
+ with os.fdopen(file_fd, "rb", closefd=True) as file_handle:
+ file_fd = -1
+ yield file_handle, file_stat
+ except FileNotFoundError:
+ raise FileNotFoundError(f"camera trajectory file not found: {filename}") from None
+ except OSError:
+ raise ValueError(f"{filename} could not be opened securely from the trusted action directory.") from None
+ finally:
+ if file_fd >= 0:
+ os.close(file_fd)
+
+
+def _load_bounded_npy(
+ action_fd: int,
+ filename: str,
+ *,
+ trailing_shape: tuple[int, ...],
+) -> np.ndarray:
+ with _open_camera_file(action_fd, filename) as (file_handle, file_stat):
+ max_payload_bytes = _MAX_SOURCE_ACTION_FRAMES * math.prod(trailing_shape) * 8
+ max_file_bytes = _MAX_NPY_HEADER_BYTES + max_payload_bytes + 32
+ if file_stat.st_size > max_file_bytes:
+ raise ValueError(f"{filename} byte size exceeds the bounded camera-array limit.")
+ snapshot = file_handle.read(max_file_bytes + 1)
+ if len(snapshot) != file_stat.st_size:
+ raise ValueError(f"{filename} changed while its bounded contents were read.")
+ snapshot_file = io.BytesIO(snapshot)
+ try:
+ version = np.lib.format.read_magic(snapshot_file)
+ if version == (1, 0):
+ shape, _fortran_order, dtype = np.lib.format.read_array_header_1_0(snapshot_file)
+ elif version == (2, 0):
+ shape, _fortran_order, dtype = np.lib.format.read_array_header_2_0(snapshot_file)
+ else:
+ raise ValueError(f"{filename} must use NPY format version 1.0 or 2.0.")
+ except (EOFError, TypeError, ValueError):
+ raise ValueError(f"{filename} has an invalid or oversized NPY header.") from None
+
+ if len(shape) != len(trailing_shape) + 1 or tuple(shape[1:]) != trailing_shape:
+ expected = ", ".join(str(value) for value in trailing_shape)
+ raise ValueError(f"{filename} must have shape [frames, {expected}], got {shape}.")
+ frames = shape[0]
+ if isinstance(frames, bool) or not isinstance(frames, int) or not 1 <= frames <= _MAX_SOURCE_ACTION_FRAMES:
+ raise ValueError(f"{filename} must contain between 1 and {_MAX_SOURCE_ACTION_FRAMES} source frames.")
+ if dtype.kind not in "fiu" or dtype.hasobject or dtype.itemsize > 8:
+ raise ValueError(f"{filename} must contain real numeric values with at most 8 bytes per element.")
+
+ data_offset = snapshot_file.tell()
+ expected_bytes = math.prod(shape) * dtype.itemsize
+ if len(snapshot) != data_offset + expected_bytes:
+ raise ValueError(f"{filename} byte size does not match its NPY header.")
+ snapshot_file.seek(0)
+ try:
+ array = np.load(snapshot_file, allow_pickle=False)
+ except (EOFError, TypeError, ValueError):
+ raise ValueError(f"{filename} does not contain a valid numeric NPY array.") from None
+ # Official LingBot trajectories may be longer than one request. Only
+ # materialize the bounded prefix that any supported request can consume.
+ result: np.ndarray = np.asarray(array[:_MAX_ACTION_FRAMES], dtype=np.float32)
+ return result
+
+
+def load_camera_trajectory(action_directory: TrustedActionDirectory) -> CameraTrajectory:
+ """Load bounded numeric camera arrays through a trusted directory handle."""
+
+ if not isinstance(action_directory, TrustedActionDirectory):
+ raise TypeError("action_directory must be resolved against a trusted action root before loading.")
+ with _open_action_directory(action_directory) as action_fd:
+ poses_array = _load_bounded_npy(action_fd, "poses.npy", trailing_shape=(4, 4))
+ intrinsics_array = _load_bounded_npy(action_fd, "intrinsics.npy", trailing_shape=(4,))
+
+ trajectory = CameraTrajectory(
+ poses=torch.from_numpy(poses_array),
+ intrinsics=torch.from_numpy(intrinsics_array),
+ )
+ _validate_trajectory(trajectory, file_suffix=".npy")
+ return trajectory
+
+
+def _linear_interpolate(values: torch.Tensor, target_times: np.ndarray) -> torch.Tensor:
+ source_times = np.linspace(0.0, 1.0, values.shape[0])
+ values_array = values.detach().cpu().double().numpy()
+ columns = [
+ np.interp(target_times, source_times, values_array[:, column]) for column in range(values_array.shape[1])
+ ]
+ interpolated = torch.from_numpy(np.stack(columns, axis=1))
+ return interpolated.to(device=values.device, dtype=values.dtype)
+
+
+def _rotation_matrices_to_quaternions(matrices: np.ndarray) -> np.ndarray:
+ """Convert rotation matrices to normalized ``(w, x, y, z)`` quaternions."""
+
+ quaternions: np.ndarray = np.empty((matrices.shape[0], 4), dtype=np.float64)
+ for index, matrix in enumerate(matrices):
+ trace = float(np.trace(matrix))
+ if trace > 0.0:
+ scale = math.sqrt(trace + 1.0) * 2.0
+ quaternion = np.array(
+ [
+ 0.25 * scale,
+ (matrix[2, 1] - matrix[1, 2]) / scale,
+ (matrix[0, 2] - matrix[2, 0]) / scale,
+ (matrix[1, 0] - matrix[0, 1]) / scale,
+ ]
+ )
+ else:
+ axis = int(np.argmax(np.diag(matrix)))
+ if axis == 0:
+ scale = math.sqrt(max(0.0, 1.0 + matrix[0, 0] - matrix[1, 1] - matrix[2, 2])) * 2.0
+ quaternion = np.array(
+ [
+ (matrix[2, 1] - matrix[1, 2]) / scale,
+ 0.25 * scale,
+ (matrix[0, 1] + matrix[1, 0]) / scale,
+ (matrix[0, 2] + matrix[2, 0]) / scale,
+ ]
+ )
+ elif axis == 1:
+ scale = math.sqrt(max(0.0, 1.0 + matrix[1, 1] - matrix[0, 0] - matrix[2, 2])) * 2.0
+ quaternion = np.array(
+ [
+ (matrix[0, 2] - matrix[2, 0]) / scale,
+ (matrix[0, 1] + matrix[1, 0]) / scale,
+ 0.25 * scale,
+ (matrix[1, 2] + matrix[2, 1]) / scale,
+ ]
+ )
+ else:
+ scale = math.sqrt(max(0.0, 1.0 + matrix[2, 2] - matrix[0, 0] - matrix[1, 1])) * 2.0
+ quaternion = np.array(
+ [
+ (matrix[1, 0] - matrix[0, 1]) / scale,
+ (matrix[0, 2] + matrix[2, 0]) / scale,
+ (matrix[1, 2] + matrix[2, 1]) / scale,
+ 0.25 * scale,
+ ]
+ )
+ norm = np.linalg.norm(quaternion)
+ if not np.isfinite(norm) or norm <= np.finfo(np.float64).eps:
+ raise ValueError("camera rotation matrices must describe valid rotations")
+ quaternions[index] = quaternion / norm
+ return quaternions
+
+
+def _quaternions_to_rotation_matrices(quaternions: np.ndarray) -> np.ndarray:
+ quaternions = quaternions / np.linalg.norm(quaternions, axis=1, keepdims=True)
+ w, x, y, z = quaternions.T
+ matrices: np.ndarray = np.stack(
+ (
+ 1 - 2 * (y * y + z * z),
+ 2 * (x * y - z * w),
+ 2 * (x * z + y * w),
+ 2 * (x * y + z * w),
+ 1 - 2 * (x * x + z * z),
+ 2 * (y * z - x * w),
+ 2 * (x * z - y * w),
+ 2 * (y * z + x * w),
+ 1 - 2 * (x * x + y * y),
+ ),
+ axis=1,
+ ).reshape(-1, 3, 3)
+ return matrices
+
+
+def _slerp_rotations(matrices: np.ndarray, target_times: np.ndarray) -> np.ndarray:
+ quaternions = _rotation_matrices_to_quaternions(matrices)
+ interpolated = np.empty((target_times.shape[0], 4), dtype=np.float64)
+ for target_index, target_time in enumerate(target_times):
+ scaled_time = float(target_time) * (matrices.shape[0] - 1)
+ left = min(max(math.floor(scaled_time), 0), matrices.shape[0] - 2)
+ fraction = scaled_time - left
+ first = quaternions[left]
+ second = quaternions[left + 1]
+ cosine = float(np.dot(first, second))
+ if cosine < 0.0:
+ second = -second
+ cosine = -cosine
+ cosine = min(cosine, 1.0)
+ if cosine > 0.9995:
+ quaternion = first + fraction * (second - first)
+ interpolated[target_index] = quaternion / np.linalg.norm(quaternion)
+ continue
+ angle = math.acos(cosine)
+ denominator = math.sin(angle)
+ interpolated[target_index] = (
+ math.sin((1.0 - fraction) * angle) / denominator * first + math.sin(fraction * angle) / denominator * second
+ )
+ return _quaternions_to_rotation_matrices(interpolated)
+
+
+def interpolate_camera_trajectory(
+ trajectory: CameraTrajectory,
+ num_frames: int,
+) -> CameraTrajectory:
+ """Resample translations/intrinsics linearly and rotations with quaternion SLERP."""
+
+ _validate_trajectory(trajectory)
+ if num_frames <= 0:
+ raise ValueError(f"num_frames must be positive, got {num_frames}")
+
+ source_frames = trajectory.poses.shape[0]
+ if num_frames == source_frames:
+ return trajectory
+ if source_frames == 1:
+ return CameraTrajectory(
+ poses=trajectory.poses.repeat(num_frames, 1, 1),
+ intrinsics=trajectory.intrinsics.repeat(num_frames, 1),
+ )
+
+ # Rotations use quaternion SLERP; translations and intrinsics are linear.
+ target_times = np.linspace(0.0, 1.0, num_frames)
+ rotation_matrices = trajectory.poses[:, :3, :3].detach().cpu().double().numpy()
+ interpolated_rotations = _slerp_rotations(rotation_matrices, target_times)
+ translations = _linear_interpolate(trajectory.poses[:, :3, 3], target_times)
+ intrinsics = _linear_interpolate(trajectory.intrinsics, target_times)
+
+ poses = torch.eye(
+ 4,
+ device=trajectory.poses.device,
+ dtype=trajectory.poses.dtype,
+ ).repeat(num_frames, 1, 1)
+ poses[:, :3, :3] = torch.from_numpy(interpolated_rotations).to(
+ device=trajectory.poses.device,
+ dtype=trajectory.poses.dtype,
+ )
+ poses[:, :3, 3] = translations
+ return CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+
+def _invert_camera_poses(poses: torch.Tensor) -> torch.Tensor:
+ """Invert rigid transforms using R^-1 = R^T and t^-1 = -R^T t."""
+
+ rotations = poses[:, :3, :3]
+ translations = poses[:, :3, 3:]
+ inverse_rotations = rotations.transpose(-1, -2)
+ inverse = torch.eye(4, device=poses.device, dtype=poses.dtype).repeat(poses.shape[0], 1, 1)
+ inverse[:, :3, :3] = inverse_rotations
+ inverse[:, :3, 3:] = -torch.bmm(inverse_rotations, translations)
+ return inverse
+
+
+def _prepare_framewise_poses(raw_poses: torch.Tensor) -> torch.Tensor:
+ """Convert raw C2W poses to normalized first-relative framewise deltas."""
+
+ # Remove global placement, then convert the sequence to framewise deltas.
+ first_world_to_camera = _invert_camera_poses(raw_poses[:1])
+ relative_poses = torch.matmul(first_world_to_camera, raw_poses).clone()
+ relative_poses[0] = torch.eye(4, device=raw_poses.device, dtype=raw_poses.dtype)
+ if relative_poses.shape[0] > 1:
+ relative_poses[1:] = torch.bmm(
+ _invert_camera_poses(relative_poses[:-1]),
+ relative_poses[1:],
+ )
+
+ # Only relative direction/magnitude matters to the checkpoint; normalizing
+ # by the largest step removes the trajectory's arbitrary translation scale.
+ translations = relative_poses[:, :3, 3]
+ max_translation_norm = torch.linalg.vector_norm(translations, dim=-1).max()
+ if max_translation_norm > 0:
+ relative_poses[:, :3, 3] = translations / max_translation_norm
+ return relative_poses
+
+
+def build_plucker_embedding(
+ trajectory: CameraTrajectory,
+ *,
+ height: int,
+ width: int,
+ target_height: int,
+ target_width: int,
+ device: torch.device,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ """Build checkpoint-ordered ``(ray origin, ray direction)`` channels."""
+
+ _validate_trajectory(trajectory)
+ dimensions = {
+ "height": height,
+ "width": width,
+ "target_height": target_height,
+ "target_width": target_width,
+ }
+ for name, value in dimensions.items():
+ if value <= 0:
+ raise ValueError(f"{name} must be positive, got {value}")
+ if dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
+ raise ValueError(f"dtype must be floating point, got {dtype}")
+
+ # Compute geometry in at least FP32, then cast at the model boundary.
+ compute_dtype = torch.float32 if dtype in (torch.float16, torch.bfloat16) else dtype
+ poses = _prepare_framewise_poses(trajectory.poses.to(device=device, dtype=compute_dtype))
+ intrinsics = trajectory.intrinsics.to(device=device, dtype=compute_dtype).clone()
+ intrinsics[:, (0, 2)] *= width / _REFERENCE_WIDTH
+ intrinsics[:, (1, 3)] *= height / _REFERENCE_HEIGHT
+ if torch.any(intrinsics[:, :2] == 0):
+ raise ValueError("camera focal lengths fx and fy must be non-zero")
+
+ # Sample requested-pixel coordinates directly at the conditioning resolution.
+ x_coordinates = (torch.arange(target_width, device=device, dtype=compute_dtype) + 0.5) * (width / target_width)
+ y_coordinates = (torch.arange(target_height, device=device, dtype=compute_dtype) + 0.5) * (height / target_height)
+ grid_y, grid_x = torch.meshgrid(y_coordinates, x_coordinates, indexing="ij")
+ grid_x = grid_x.unsqueeze(0)
+ grid_y = grid_y.unsqueeze(0)
+
+ # Pinhole back-projection: pixel (x, y) becomes the camera-space direction
+ # ((x-cx)/fx, (y-cy)/fy, 1), which is normalized before rotation.
+ fx, fy, cx, cy = intrinsics.unbind(dim=1)
+ camera_x = (grid_x - cx[:, None, None]) / fx[:, None, None]
+ camera_y = (grid_y - cy[:, None, None]) / fy[:, None, None]
+ camera_directions = torch.stack(
+ (camera_x, camera_y, torch.ones_like(camera_x)),
+ dim=-1,
+ )
+ camera_directions = camera_directions / torch.linalg.vector_norm(camera_directions, dim=-1, keepdim=True)
+
+ # A ray is represented by its camera center (origin) and unit direction in
+ # the prepared relative coordinate frame. Channel order is checkpoint
+ # data, not an implementation preference: [origin_xyz, direction_xyz].
+ world_directions = torch.einsum(
+ "fij,fhwj->fhwi",
+ poses[:, :3, :3],
+ camera_directions,
+ )
+ world_directions = world_directions / torch.linalg.vector_norm(world_directions, dim=-1, keepdim=True)
+ ray_origins = poses[:, None, None, :3, 3].expand_as(world_directions)
+
+ embedding = torch.cat((ray_origins, world_directions), dim=-1)
+ return embedding.permute(0, 3, 1, 2).contiguous().to(dtype=dtype)
diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py
new file mode 100644
index 000000000..5f71979e8
--- /dev/null
+++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py
@@ -0,0 +1,1292 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Request-scoped LingBot-World v2 causal DMD pipeline."""
+
+from __future__ import annotations
+
+import math
+import os
+from collections.abc import Callable, Iterable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Any, ClassVar, cast
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import AutoTokenizer, UMT5EncoderModel
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
+from vllm_omni.diffusion.models.lingbot_world.actions import (
+ LINGBOT_CAMERA_ACTION_SCHEMA,
+ LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
+ integrate_lingbot_camera_actions,
+ parse_lingbot_camera_action_frames,
+)
+from vllm_omni.diffusion.models.lingbot_world.camera import (
+ CameraTrajectory,
+ build_plucker_embedding,
+ interpolate_camera_trajectory,
+ load_camera_trajectory,
+ resolve_trusted_action_directory,
+)
+from vllm_omni.diffusion.models.lingbot_world.transformer import (
+ CausalLingBotWorldTransformer3DModel,
+ LingBotAttentionCache,
+ LingBotTransformerCache,
+)
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
+from vllm_omni.diffusion.models.utils import _load_json
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import load_transformer_config, retrieve_latents
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.experimental.ar_diffusion.capability import (
+ ARDiffusionCrossAttentionKVSpec,
+ ARDiffusionKVBranchSpec,
+ ARDiffusionKVCacheSpec,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionTickRequest,
+)
+
+if TYPE_CHECKING:
+ from tqdm.std import tqdm as TqdmProgressBar
+
+ from vllm_omni.experimental.ar_diffusion.kv_cache.state import (
+ ARDiffusionKVState,
+ )
+
+LINGBOT_DMD_TIMESTEPS = (1000, 750, 500, 250)
+_CAMERA_SPATIAL_FOLD = 8
+_MAX_PIXEL_AREA = 480 * 832
+_MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096
+_MAX_RAW_FRAMES = 117
+_MAX_SEQUENCE_LENGTH = 512
+_ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT"
+_PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory"
+_PREPROCESSED_CAMERA_ACTIONS_KEY = "_lingbot_camera_actions"
+_SOURCE_IMAGE_ERROR = (
+ "Unable to load multi_modal_data.image; expected a decodable image within 4096 * 4096 source pixels."
+)
+
+
+@dataclass(frozen=True)
+class _LingBotRequestInputs:
+ """Validated model-specific request boundary."""
+
+ prompt: str
+ image: PIL.Image.Image | torch.Tensor
+ camera_trajectory: CameraTrajectory | None
+ camera_actions: tuple[tuple[str, ...], ...] | None
+ height: int
+ width: int
+ num_frames: int
+ num_latent_frames: int
+ output_type: str
+ max_sequence_length: int
+ flow_shift: float
+ generator: torch.Generator
+
+
+@dataclass
+class _LingBotARSessionState:
+ """Small model-owned state; attention tensors remain runner-owned."""
+
+ next_chunk_index: int = 0
+ prompt: str | None = None
+ generator_state: torch.Tensor | None = None
+ image_condition: torch.Tensor | None = None
+ camera_tail: CameraTrajectory | None = None
+ camera_pitch: float = 0.0
+
+
+def _positive_finite_flow_shift(value: Any) -> float:
+ if isinstance(value, bool):
+ raise ValueError("flow_shift must be a positive finite number.")
+ try:
+ flow_shift = float(value)
+ except (TypeError, ValueError, OverflowError) as exc:
+ raise ValueError("flow_shift must be a positive finite number.") from exc
+ if not math.isfinite(flow_shift) or flow_shift <= 0:
+ raise ValueError("flow_shift must be a positive finite number.")
+ return flow_shift
+
+
+def _build_shifted_flow_schedule(
+ *,
+ flow_shift: float,
+) -> tuple[tuple[float, float], ...]:
+ """Map checkpoint DMD labels to request-local warped timestep/sigma pairs."""
+
+ # The checkpoint labels index the unshifted [1, ..., 1 / N] training
+ # lattice. Apply the request's flow shift before using the value both as
+ # the transformer timestep and the sampling sigma.
+ base_sigmas = torch.tensor(LINGBOT_DMD_TIMESTEPS, dtype=torch.float64) / 1000
+ shifted_numerators = flow_shift * base_sigmas
+ shifted_sigmas = shifted_numerators / ((1.0 - base_sigmas) + shifted_numerators)
+ warped_timesteps = shifted_sigmas * 1000
+ return tuple(
+ (float(timestep), float(sigma))
+ for timestep, sigma in zip(warped_timesteps.tolist(), shifted_sigmas.tolist(), strict=True)
+ )
+
+
+def _validate_scheduler_config(config: dict[str, Any]) -> None:
+ contract = {
+ "_class_name": "UniPCMultistepScheduler",
+ "num_train_timesteps": 1000,
+ "prediction_type": "flow_prediction",
+ "predict_x0": True,
+ "use_flow_sigmas": True,
+ "use_dynamic_shifting": False,
+ "use_beta_sigmas": False,
+ "use_exponential_sigmas": False,
+ "use_karras_sigmas": False,
+ "final_sigmas_type": "zero",
+ "timestep_spacing": "linspace",
+ }
+ for name, expected in contract.items():
+ if config.get(name) != expected:
+ raise ValueError(f"LingBot scheduler config {name} must be {expected!r}, got {config.get(name)!r}.")
+
+
+def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None:
+ if getattr(od_config, "quantization_config", None) is not None:
+ raise NotImplementedError("LingBot World v1 does not support quantization.")
+ parallel_config = getattr(od_config, "parallel_config", None)
+ if parallel_config is None:
+ return
+ unsupported_sizes = {
+ "pipeline_parallel_size": "pipeline parallelism",
+ "sequence_parallel_size": "sequence parallelism",
+ "cfg_parallel_size": "CFG parallelism",
+ "vae_patch_parallel_size": "VAE parallelism",
+ }
+ for field, feature in unsupported_sizes.items():
+ size = getattr(parallel_config, field, 1) or 1
+ if size > 1:
+ raise NotImplementedError(f"LingBot World v1 does not support {feature} ({field}={size}).")
+ if getattr(parallel_config, "use_hsdp", False):
+ raise NotImplementedError("LingBot World v1 does not support HSDP.")
+ if getattr(parallel_config, "enable_expert_parallel", False):
+ raise NotImplementedError("LingBot World v1 does not support expert parallelism.")
+
+
+def _validate_source_image_size(image: PIL.Image.Image) -> None:
+ width, height = image.size
+ if (
+ isinstance(width, bool)
+ or not isinstance(width, int)
+ or width <= 0
+ or isinstance(height, bool)
+ or not isinstance(height, int)
+ or height <= 0
+ ):
+ raise ValueError("source image width and height must be positive integers.")
+ if width * height > _MAX_SOURCE_IMAGE_PIXELS:
+ raise ValueError("source image pixel count must not exceed 4096 * 4096.")
+
+
+def _decode_source_image(image: PIL.Image.Image) -> PIL.Image.Image:
+ _validate_source_image_size(image)
+ try:
+ return image.convert("RGB")
+ except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError):
+ raise ValueError(_SOURCE_IMAGE_ERROR) from None
+
+
+def _load_source_image(path: str | os.PathLike[str]) -> PIL.Image.Image:
+ try:
+ source_image = PIL.Image.open(path)
+ except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError):
+ raise ValueError(_SOURCE_IMAGE_ERROR) from None
+ try:
+ return _decode_source_image(source_image)
+ finally:
+ source_image.close()
+
+
+def get_lingbot_world_pre_process_func(
+ od_config: OmniDiffusionConfig,
+) -> Callable[[OmniDiffusionRequest], OmniDiffusionRequest]:
+ """Materialize request-local files once before dispatching GPU workers."""
+
+ model_config = getattr(od_config, "model_config", None) or {}
+ configured_action_root = model_config.get("lingbot_action_root") or os.environ.get(_ACTION_ROOT_ENV)
+
+ def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+ prompt = request.prompt
+ if not isinstance(prompt, dict):
+ raise ValueError("LingBot World requires a prompt mapping containing multi_modal_data.image.")
+ multi_modal_data = prompt.get("multi_modal_data") or {}
+ additional_information = prompt.get("additional_information") or {}
+ if not isinstance(multi_modal_data, dict):
+ raise ValueError("prompt.multi_modal_data must be a mapping containing image.")
+ if not isinstance(additional_information, dict):
+ raise ValueError("prompt.additional_information must be a mapping.")
+ if "action_path" in additional_information:
+ raise ValueError(
+ "prompt.additional_information.action_path is not supported; use "
+ "sampling_params.extra_args.action_path."
+ )
+
+ image = multi_modal_data.get("image")
+ if isinstance(image, list):
+ raise ValueError("LingBot World requires one image and does not accept an image list.")
+ if image is None:
+ raise ValueError("LingBot World requires exactly one image in multi_modal_data.image.")
+ if isinstance(image, (str, os.PathLike)):
+ image = _load_source_image(image)
+ elif isinstance(image, PIL.Image.Image):
+ image = _decode_source_image(image)
+ elif not isinstance(image, torch.Tensor):
+ raise ValueError("multi_modal_data.image must be a PIL image, tensor, or file path.")
+
+ extra_args = getattr(request.sampling_params, "extra_args", None) or {}
+ if not isinstance(extra_args, dict):
+ raise ValueError("sampling_params.extra_args must be a mapping.")
+ tick = ARDiffusionTickRequest.from_extra_args(extra_args)
+ if tick is not None:
+ camera_controls = [control for control in tick.controls if control.track == "camera"]
+ if len(camera_controls) != 1:
+ raise ValueError("LingBot realtime ticks require exactly one camera control.")
+ camera_control = camera_controls[0]
+ if camera_control.schema == LINGBOT_CAMERA_TRAJECTORY_SCHEMA:
+ camera_data = camera_control.data
+ try:
+ poses = torch.as_tensor(camera_data["poses"], dtype=torch.float32)
+ intrinsics = torch.as_tensor(
+ camera_data["intrinsics"],
+ dtype=torch.float32,
+ )
+ except (KeyError, TypeError, ValueError):
+ raise ValueError(
+ "lingbot.camera_trajectory.v1 requires numeric poses and intrinsics fields."
+ ) from None
+ if poses.ndim != 3 or poses.shape[1:] != (4, 4):
+ raise ValueError("camera control poses must have shape [frames, 4, 4].")
+ if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
+ raise ValueError("camera control intrinsics must have shape [frames, 4].")
+ if poses.shape[0] != intrinsics.shape[0] or poses.shape[0] == 0:
+ raise ValueError(
+ "camera control poses and intrinsics must contain the same positive number of frames."
+ )
+ if not torch.isfinite(poses).all() or not torch.isfinite(intrinsics).all():
+ raise ValueError("camera control values must all be finite.")
+ trajectory = CameraTrajectory(
+ poses=poses,
+ intrinsics=intrinsics,
+ )
+ camera_actions = None
+ elif camera_control.schema == LINGBOT_CAMERA_ACTION_SCHEMA:
+ trajectory = None
+ camera_actions = parse_lingbot_camera_action_frames(
+ camera_control.data,
+ expected_frames=3,
+ )
+ else:
+ raise ValueError(
+ "LingBot realtime camera controls require schema "
+ "'lingbot.camera_trajectory.v1' or 'lingbot.camera_actions.v1'."
+ )
+ else:
+ action_path = extra_args.get("action_path")
+ if not isinstance(action_path, (str, os.PathLike)) or not str(action_path):
+ raise ValueError("action_path is required in sampling_params.extra_args.action_path.")
+ if not configured_action_root:
+ raise ValueError(
+ "sampling_params.extra_args.action_path requires a trusted action root configured by "
+ f"model_config.lingbot_action_root or {_ACTION_ROOT_ENV}."
+ )
+ action_directory = resolve_trusted_action_directory(
+ action_path,
+ configured_action_root,
+ )
+ try:
+ trajectory = load_camera_trajectory(action_directory)
+ except OSError:
+ raise ValueError(
+ "Unable to load camera trajectory from action_path; expected poses.npy and intrinsics.npy."
+ ) from None
+ camera_actions = None
+
+ updated_prompt = dict(prompt)
+ updated_multi_modal_data = dict(multi_modal_data)
+ updated_multi_modal_data["image"] = image
+ updated_prompt["multi_modal_data"] = updated_multi_modal_data
+ request.prompt = updated_prompt
+ request.sampling_params.extra_args = {
+ **extra_args,
+ _PREPROCESSED_CAMERA_KEY: trajectory,
+ _PREPROCESSED_CAMERA_ACTIONS_KEY: camera_actions,
+ }
+ return request
+
+ return pre_process_func
+
+
+def _fold_camera_embedding(
+ camera_embedding: torch.Tensor,
+) -> torch.Tensor:
+ """Pixel-unshuffle ``[frames, 6, H, W]`` onto the Wan latent grid."""
+
+ folded = F.pixel_unshuffle(camera_embedding, _CAMERA_SPATIAL_FOLD)
+ return folded.permute(1, 0, 2, 3).unsqueeze(0).contiguous()
+
+
+def get_lingbot_world_post_process_func(od_config: OmniDiffusionConfig) -> Callable[..., Any]:
+ del od_config
+ from diffusers.video_processor import VideoProcessor
+
+ video_processor = VideoProcessor(vae_scale_factor=8)
+
+ def post_process_func(
+ video: torch.Tensor,
+ output_type: str = "np",
+ sampling_params: Any | None = None,
+ ) -> Any:
+ if isinstance(video, dict) and isinstance(video.get("payload"), dict):
+ return video
+ if sampling_params is not None:
+ output_type = getattr(sampling_params, "output_type", None) or output_type
+ if output_type == "latent":
+ return video
+ return {
+ "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)},
+ "metadata": {},
+ }
+
+ return post_process_func
+
+
+class LingBotWorldCausalDMDPipeline(
+ nn.Module,
+ SupportImageInput,
+ SupportsComponentDiscovery,
+ ProgressBarMixin,
+ DiffusionPipelineProfilerMixin,
+):
+ """LingBot-World v2 I2V generation with a request-local causal cache."""
+
+ _dit_modules: ClassVar[list[str]] = ["transformer"]
+ _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+ _vae_modules: ClassVar[list[str]] = ["vae"]
+ # Generic warmup cannot synthesize the required camera action directory.
+ dummy_run_num_frames: ClassVar[int] = 0
+ _AR_BRANCH = "main"
+ _AR_TEXT_CACHE = "text"
+
+ def __init__(
+ self,
+ *,
+ od_config: OmniDiffusionConfig,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ del prefix
+ _validate_parallel_config(od_config)
+ self.od_config = od_config
+ self.device = get_local_device()
+ dtype = getattr(od_config, "dtype", torch.bfloat16)
+ model = od_config.model
+ local_files_only = os.path.exists(model)
+ managed_component_placement = bool(
+ getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False)
+ )
+
+ # Standard components use from_pretrained; the custom transformer uses the loader.
+ self.weights_sources = [
+ DiffusersPipelineLoader.ComponentSource(
+ model_or_path=model,
+ subfolder="transformer",
+ revision=None,
+ prefix="transformer.",
+ fall_back_to_pt=True,
+ )
+ ]
+
+ subfolders = ["tokenizer", "text_encoder", "vae"]
+ prefetch_subfolders(model, subfolders, local_files_only=local_files_only)
+ self.tokenizer = from_pretrained_with_prefetch(
+ AutoTokenizer.from_pretrained,
+ model,
+ subfolder="tokenizer",
+ prefetch_list=subfolders,
+ local_files_only=local_files_only,
+ )
+ self.text_encoder = from_pretrained_with_prefetch(
+ UMT5EncoderModel.from_pretrained,
+ model,
+ subfolder="text_encoder",
+ prefetch_list=subfolders,
+ local_files_only=local_files_only,
+ torch_dtype=dtype,
+ )
+ if not managed_component_placement:
+ self.text_encoder = self.text_encoder.to(self.device)
+ self.vae = from_pretrained_with_prefetch(
+ DistributedAutoencoderKLWan.from_pretrained,
+ model,
+ subfolder="vae",
+ prefetch_list=subfolders,
+ local_files_only=local_files_only,
+ torch_dtype=dtype,
+ )
+ if not managed_component_placement:
+ self.vae = self.vae.to(self.device)
+
+ transformer_config = load_transformer_config(model, "transformer", local_files_only)
+ self.transformer = CausalLingBotWorldTransformer3DModel.from_config(
+ transformer_config,
+ quant_config=getattr(od_config, "quantization_config", None),
+ prefix="transformer",
+ )
+
+ scheduler_config = _load_json(model, "scheduler/scheduler_config.json", local_files_only)
+ _validate_scheduler_config(scheduler_config)
+ configured_shift = getattr(od_config, "flow_shift", None)
+ checkpoint_shift = scheduler_config.get("flow_shift", scheduler_config.get("shift", 5.0))
+ scheduler_shift = _positive_finite_flow_shift(
+ checkpoint_shift if configured_shift is None else configured_shift
+ )
+ scheduler_keys = {
+ "num_train_timesteps",
+ "solver_order",
+ "prediction_type",
+ "use_dynamic_shifting",
+ "thresholding",
+ "dynamic_thresholding_ratio",
+ "sample_max_value",
+ "predict_x0",
+ "solver_type",
+ "lower_order_final",
+ "disable_corrector",
+ "timestep_spacing",
+ "steps_offset",
+ "final_sigmas_type",
+ }
+ scheduler_kwargs = {name: value for name, value in scheduler_config.items() if name in scheduler_keys}
+ scheduler_kwargs["shift"] = scheduler_shift
+ self.scheduler = FlowUniPCMultistepScheduler(**scheduler_kwargs)
+
+ self.vae_scale_factor_temporal = int(getattr(self.vae.config, "scale_factor_temporal", 4))
+ self.vae_scale_factor_spatial = int(getattr(self.vae.config, "scale_factor_spatial", 8))
+ model_config = getattr(od_config, "model_config", None) or {}
+ self._ar_height = int(model_config.get("ar_diffusion_height", 480))
+ self._ar_width = int(model_config.get("ar_diffusion_width", 832))
+ self._ar_diffusion_kv_state: ARDiffusionKVState | None = None
+ self._ar_sessions: dict[str, _LingBotARSessionState] = {}
+ self.setup_diffusion_pipeline_profiler(
+ profiler_targets=[
+ "vae.encode",
+ "vae.decode",
+ "_generate_block",
+ "text_encoder.forward",
+ "tokenizer.forward",
+ ],
+ enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler,
+ )
+
+ def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec:
+ """Describe the fixed worker-local cache geometry for realtime ticks."""
+ patch_frames, patch_height, patch_width = self.transformer.config.patch_size
+ if patch_frames != 1:
+ raise ValueError("LingBot AR-Diffusion requires temporal patch size 1.")
+ spatial = self.vae_scale_factor_spatial
+ if self._ar_height % (spatial * patch_height) or self._ar_width % (spatial * patch_width):
+ raise ValueError("ar_diffusion_height/width must align with VAE and DiT patch sizes.")
+ latent_height = self._ar_height // spatial
+ latent_width = self._ar_width // spatial
+ tokens_per_frame = (latent_height // patch_height) * (latent_width // patch_width)
+ tp_size = get_tensor_model_parallel_world_size()
+ num_local_heads = int(self.transformer.config.num_attention_heads) // tp_size
+ total_window_frames = (
+ int(self.transformer.config.local_attn_size)
+ if int(self.transformer.config.local_attn_size) != -1
+ else int(self.transformer.config.sliding_window_num_frames)
+ )
+ sink_frames = int(self.transformer.config.sink_size)
+ recent_window_frames = total_window_frames - sink_frames
+ if recent_window_frames <= 0:
+ raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.")
+ horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1
+ condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels)
+ condition_bytes_per_session = (
+ condition_channels
+ * horizon_latent_frames
+ * latent_height
+ * latent_width
+ * torch.empty((), dtype=self.transformer.dtype).element_size()
+ )
+ return ARDiffusionKVCacheSpec(
+ num_layers=int(self.transformer.config.num_layers),
+ num_kv_heads=num_local_heads,
+ head_size=int(self.transformer.config.attention_head_dim),
+ tokens_per_frame=tokens_per_frame,
+ frames_per_block=int(self.transformer.config.num_frames_per_block),
+ window_frames=recent_window_frames,
+ sink_frames=sink_frames,
+ kv_branches=(ARDiffusionKVBranchSpec(self._AR_BRANCH, 0),),
+ session_capacity=2,
+ cross_attention=(
+ ARDiffusionCrossAttentionKVSpec(
+ self._AR_TEXT_CACHE,
+ _MAX_SEQUENCE_LENGTH,
+ ),
+ ),
+ model_owned_state_bytes_per_session=condition_bytes_per_session,
+ )
+
+ @contextmanager
+ def bind_ar_diffusion_state(
+ self,
+ session_id: str,
+ state: ARDiffusionKVState,
+ ) -> Iterator[None]:
+ if self._ar_diffusion_kv_state is not None:
+ raise RuntimeError("LingBot AR-Diffusion state is already bound.")
+ if state.session_id != session_id:
+ raise ValueError(f"LingBot bound session mismatch: {state.session_id!r} != {session_id!r}.")
+ self._ar_diffusion_kv_state = state
+ try:
+ yield
+ finally:
+ self._ar_diffusion_kv_state = None
+
+ def reset_ar_diffusion_session(self, session_id: str) -> None:
+ self._ar_sessions.pop(session_id, None)
+
+ def close_ar_diffusion_session(self, session_id: str) -> None:
+ self._ar_sessions.pop(session_id, None)
+
+ def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs:
+ if req.num_reqs != 1 or len(req.prompts) != 1:
+ raise ValueError("LingBot World supports a single prompt request, not request batching.")
+ sampling = req.sampling_params
+ if int(sampling.num_outputs_per_prompt or 1) != 1:
+ raise ValueError("LingBot World requires num_outputs_per_prompt=1.")
+ generator = getattr(sampling, "generator", None)
+ if isinstance(generator, list):
+ raise ValueError("LingBot World accepts one torch.Generator, not a generator list.")
+ if not isinstance(generator, torch.Generator):
+ raise ValueError("LingBot World requires the runner-provided torch.Generator.")
+ if getattr(sampling, "latents", None) is not None:
+ raise ValueError("LingBot World does not support caller-provided latents.")
+
+ prompt_value = req.prompts[0]
+ multi_modal_data: dict[str, Any]
+ if isinstance(prompt_value, str):
+ prompt = prompt_value
+ multi_modal_data = {}
+ elif isinstance(prompt_value, dict):
+ prompt = prompt_value.get("prompt") or ""
+ multi_modal_data = prompt_value.get("multi_modal_data") or {}
+ else:
+ raise ValueError("prompt must be a string or prompt mapping.")
+ if not isinstance(prompt, str) or not prompt.strip():
+ raise ValueError("prompt must contain non-empty text.")
+ if not isinstance(multi_modal_data, dict):
+ raise ValueError("prompt.multi_modal_data must be a mapping containing image.")
+
+ image = multi_modal_data.get("image")
+ if isinstance(image, list):
+ raise ValueError("LingBot World requires one image and does not accept an image list.")
+ if image is None:
+ raise ValueError("LingBot World requires exactly one image in multi_modal_data.image.")
+ if isinstance(image, (str, os.PathLike)):
+ raise ValueError("file-path images must be materialized by the LingBot pre-process function.")
+ if not isinstance(image, (PIL.Image.Image, torch.Tensor)):
+ raise ValueError("multi_modal_data.image must be a PIL image, tensor, or file path.")
+
+ extra_args = getattr(sampling, "extra_args", None) or {}
+ if not isinstance(extra_args, dict):
+ raise ValueError("sampling_params.extra_args must be a mapping.")
+ camera_trajectory = extra_args.get(_PREPROCESSED_CAMERA_KEY)
+ camera_actions = extra_args.get(_PREPROCESSED_CAMERA_ACTIONS_KEY)
+ if camera_trajectory is not None and not isinstance(camera_trajectory, CameraTrajectory):
+ raise ValueError("LingBot camera trajectory must be materialized by the pre-process function.")
+ if camera_actions is not None:
+ camera_actions = tuple(tuple(actions) for actions in camera_actions)
+ if (camera_trajectory is None) == (camera_actions is None):
+ raise ValueError("LingBot pre-processing must materialize exactly one camera input.")
+
+ request_flow_shift = (
+ extra_args["flow_shift"] if "flow_shift" in extra_args else getattr(self.scheduler.config, "shift", 5.0)
+ )
+ flow_shift = _positive_finite_flow_shift(request_flow_shift)
+
+ height = getattr(sampling, "height", None)
+ width = getattr(sampling, "width", None)
+ if isinstance(image, PIL.Image.Image):
+ image_width, image_height = image.size
+ elif image.ndim == 3 and image.shape[0] == 3:
+ image_height, image_width = image.shape[-2:]
+ elif image.ndim == 4 and image.shape[:2] == (1, 3):
+ image_height, image_width = image.shape[-2:]
+ else:
+ raise ValueError("tensor image must have shape [3, height, width] or [1, 3, height, width].")
+ if image_height <= 0 or image_width <= 0:
+ raise ValueError("source image width and height must be positive integers.")
+ if image_height * image_width > _MAX_SOURCE_IMAGE_PIXELS:
+ raise ValueError("source image pixel count must not exceed 4096 * 4096.")
+ if (height is None) != (width is None):
+ raise ValueError("height and width must either both be provided or both be omitted.")
+ patch_size = tuple(self.transformer.config.patch_size)
+ if len(patch_size) != 3 or patch_size[0] != 1:
+ raise RuntimeError(
+ "transformer.config.patch_size must be a three-dimensional tuple with temporal patch size 1."
+ )
+ height_divisor = self.vae_scale_factor_spatial * patch_size[1]
+ width_divisor = self.vae_scale_factor_spatial * patch_size[2]
+ if height is None or width is None:
+ aspect_ratio = image_height / image_width
+ latent_height = round(
+ math.sqrt(_MAX_PIXEL_AREA * aspect_ratio)
+ // self.vae_scale_factor_spatial
+ // patch_size[1]
+ * patch_size[1]
+ )
+ latent_width = round(
+ math.sqrt(_MAX_PIXEL_AREA / aspect_ratio)
+ // self.vae_scale_factor_spatial
+ // patch_size[2]
+ * patch_size[2]
+ )
+ height = latent_height * self.vae_scale_factor_spatial
+ width = latent_width * self.vae_scale_factor_spatial
+ if isinstance(height, bool) or not isinstance(height, int) or height <= 0:
+ raise ValueError(f"height must be a positive integer, got {height!r}.")
+ if isinstance(width, bool) or not isinstance(width, int) or width <= 0:
+ raise ValueError(f"width must be a positive integer, got {width!r}.")
+ if height % height_divisor:
+ raise ValueError(f"height must be divisible by {height_divisor}, got {height}.")
+ if width % width_divisor:
+ raise ValueError(f"width must be divisible by {width_divisor}, got {width}.")
+ if height * width > _MAX_PIXEL_AREA:
+ raise ValueError("height * width pixel area must not exceed 480 * 832.")
+
+ num_inference_steps = getattr(sampling, "num_inference_steps", None)
+ if num_inference_steps is None:
+ num_inference_steps = len(LINGBOT_DMD_TIMESTEPS)
+ if num_inference_steps != len(LINGBOT_DMD_TIMESTEPS):
+ raise ValueError(
+ "num_inference_steps must be 4 for LingBot World causal DMD "
+ f"timesteps {list(LINGBOT_DMD_TIMESTEPS)}, got {num_inference_steps}."
+ )
+
+ num_frames = getattr(sampling, "num_frames", None)
+ if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0:
+ raise ValueError(f"num_frames must be a positive integer, got {num_frames!r}.")
+ if num_frames > _MAX_RAW_FRAMES:
+ raise ValueError(f"num_frames must not exceed {_MAX_RAW_FRAMES}.")
+ temporal_factor = self.vae_scale_factor_temporal
+ if (num_frames - 1) % temporal_factor:
+ raise ValueError(
+ "num_frames must satisfy the causal Wan VAE geometry "
+ f"(num_frames - 1) divisible by {temporal_factor}, got {num_frames}."
+ )
+ num_latent_frames = (num_frames - 1) // temporal_factor + 1
+ block_frames = int(self.transformer.config.num_frames_per_block)
+ if num_latent_frames % block_frames:
+ raise ValueError(
+ "num_frames must map to a whole number of configured three-frame latent blocks; "
+ f"got num_frames={num_frames}, latent_frames={num_latent_frames}, block_frames={block_frames}."
+ )
+
+ max_sequence_length = getattr(sampling, "max_sequence_length", None)
+ if max_sequence_length is None:
+ max_sequence_length = _MAX_SEQUENCE_LENGTH
+ if isinstance(max_sequence_length, bool) or not isinstance(max_sequence_length, int):
+ raise ValueError(f"max_sequence_length must be exactly {_MAX_SEQUENCE_LENGTH}.")
+ if max_sequence_length != _MAX_SEQUENCE_LENGTH:
+ raise ValueError(f"max_sequence_length must be exactly {_MAX_SEQUENCE_LENGTH}.")
+
+ return _LingBotRequestInputs(
+ prompt=prompt.strip(),
+ image=image,
+ camera_trajectory=camera_trajectory,
+ camera_actions=camera_actions,
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ num_latent_frames=num_latent_frames,
+ output_type=getattr(sampling, "output_type", None) or "np",
+ max_sequence_length=max_sequence_length,
+ flow_shift=flow_shift,
+ generator=generator,
+ )
+
+ def _prepare_image_tensor(self, image: PIL.Image.Image | torch.Tensor, *, height: int, width: int) -> torch.Tensor:
+ if isinstance(image, PIL.Image.Image):
+ array = np.asarray(image, dtype=np.float32).copy()
+ image_tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) / 255.0
+ else:
+ image_tensor = image.detach()
+ if image_tensor.ndim == 3:
+ image_tensor = image_tensor.unsqueeze(0)
+ image_tensor = image_tensor.to(dtype=torch.float32)
+ if not torch.isfinite(image_tensor).all():
+ raise ValueError("tensor image values must all be finite.")
+ minimum = image_tensor.min().item()
+ maximum = image_tensor.max().item()
+ if minimum >= 0.0 and maximum > 255.0:
+ raise ValueError("tensor image values must be in [0, 1], [0, 255], or [-1, 1].")
+ if minimum < -1.0 or (minimum < 0.0 and maximum > 1.0):
+ raise ValueError("tensor image values must be in [0, 1], [0, 255], or [-1, 1].")
+ if maximum > 1.0:
+ image_tensor = image_tensor / 255.0
+ if image_tensor.min().item() >= 0.0:
+ image_tensor = image_tensor * 2.0 - 1.0
+ if image_tensor.shape[-2:] != (height, width):
+ image_tensor = F.interpolate(
+ image_tensor,
+ size=(height, width),
+ mode="bicubic",
+ align_corners=False,
+ )
+ return image_tensor.to(device=self.device, dtype=torch.float32)
+
+ def _vae_latent_stats(self, reference: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ shape = (1, -1, 1, 1, 1)
+ latent_mean = torch.as_tensor(
+ self.vae.config.latents_mean,
+ device=reference.device,
+ dtype=reference.dtype,
+ ).view(*shape)
+ latent_std = torch.as_tensor(
+ self.vae.config.latents_std,
+ device=reference.device,
+ dtype=reference.dtype,
+ ).view(*shape)
+ return latent_mean, latent_std
+
+ def _prepare_condition(self, inputs: _LingBotRequestInputs, *, dtype: torch.dtype) -> torch.Tensor:
+ """Encode the first frame as ``[mask4, image_latent16]``."""
+
+ image = self._prepare_image_tensor(inputs.image, height=inputs.height, width=inputs.width)
+ video_condition = image.new_zeros(1, 3, inputs.num_frames, inputs.height, inputs.width)
+ video_condition[:, :, 0] = image
+ latent_condition = retrieve_latents(
+ self.vae.encode(video_condition.to(dtype=self.vae.dtype)),
+ sample_mode="argmax",
+ )
+ if latent_condition.shape != (
+ 1,
+ self.transformer.config.out_channels,
+ inputs.num_latent_frames,
+ inputs.height // self.vae_scale_factor_spatial,
+ inputs.width // self.vae_scale_factor_spatial,
+ ):
+ raise RuntimeError(
+ f"vae.encode returned an incompatible image latent shape: got {tuple(latent_condition.shape)}."
+ )
+ latent_mean, latent_std = self._vae_latent_stats(latent_condition)
+ latent_condition = (latent_condition - latent_mean) / latent_std
+ temporal_mask = latent_condition.new_zeros(
+ 1,
+ self.vae_scale_factor_temporal,
+ inputs.num_latent_frames,
+ latent_condition.shape[-2],
+ latent_condition.shape[-1],
+ )
+ temporal_mask[:, :, 0] = 1
+ condition = torch.cat((temporal_mask, latent_condition), dim=1).to(dtype=dtype)
+ if condition.shape[1] != 20:
+ raise RuntimeError(
+ "LingBot image condition must contain 4 temporal-mask then 16 image-latent channels, "
+ f"got {condition.shape[1]}."
+ )
+ return condition
+
+ def _prepare_camera(
+ self,
+ inputs: _LingBotRequestInputs,
+ *,
+ dtype: torch.dtype,
+ previous: CameraTrajectory | None = None,
+ ) -> tuple[torch.Tensor, CameraTrajectory]:
+ """Convert raw camera frames to a latent-aligned ray tensor."""
+
+ trajectory = inputs.camera_trajectory
+ if trajectory is None:
+ raise RuntimeError("LingBot camera trajectory was not prepared before camera embedding.")
+ available_frames = int(trajectory.poses.shape[0])
+ if inputs.camera_actions is not None:
+ if available_frames != inputs.num_latent_frames:
+ raise ValueError(
+ "camera actions must produce exactly one pose per latent frame; "
+ f"got camera_frames={available_frames}, latent_frames={inputs.num_latent_frames}."
+ )
+ else:
+ if available_frames < inputs.num_frames:
+ raise ValueError(
+ "camera trajectory frames must be at least num_frames; "
+ f"got camera_frames={available_frames}, num_frames={inputs.num_frames}."
+ )
+ trajectory = CameraTrajectory(
+ poses=trajectory.poses[: inputs.num_frames],
+ intrinsics=trajectory.intrinsics[: inputs.num_frames],
+ )
+ trajectory = interpolate_camera_trajectory(trajectory, inputs.num_latent_frames)
+ embedding_trajectory = trajectory
+ drop_anchor = False
+ if previous is not None:
+ embedding_trajectory = CameraTrajectory(
+ poses=torch.cat((previous.poses, trajectory.poses), dim=0),
+ intrinsics=torch.cat(
+ (previous.intrinsics, trajectory.intrinsics),
+ dim=0,
+ ),
+ )
+ drop_anchor = True
+ elif inputs.camera_actions is not None:
+ # The action integrator returns post-action poses. Keep the first
+ # action visible to framewise-delta conditioning by prepending its
+ # known pre-action state (identity for a new realtime session).
+ identity = torch.eye(
+ 4,
+ device=trajectory.poses.device,
+ dtype=trajectory.poses.dtype,
+ ).unsqueeze(0)
+ embedding_trajectory = CameraTrajectory(
+ poses=torch.cat((identity, trajectory.poses), dim=0),
+ intrinsics=torch.cat(
+ (trajectory.intrinsics[:1], trajectory.intrinsics),
+ dim=0,
+ ),
+ )
+ drop_anchor = True
+ camera_embedding = build_plucker_embedding(
+ embedding_trajectory,
+ height=inputs.height,
+ width=inputs.width,
+ target_height=inputs.height,
+ target_width=inputs.width,
+ device=self.device,
+ dtype=dtype,
+ )
+ if drop_anchor:
+ camera_embedding = camera_embedding[1:]
+ tail = CameraTrajectory(
+ poses=trajectory.poses[-1:].clone(),
+ intrinsics=trajectory.intrinsics[-1:].clone(),
+ )
+ return _fold_camera_embedding(camera_embedding), tail
+
+ def _ar_text_caches(
+ self,
+ prompt_embeds: torch.Tensor,
+ *,
+ invalidate: bool,
+ ) -> list[LingBotAttentionCache]:
+ state = self._ar_diffusion_kv_state
+ if state is None:
+ raise RuntimeError("LingBot AR text cache requested without a bound state.")
+ if invalidate:
+ state.clear_cross_attention()
+ if not state.is_cross_attention_populated(
+ self._AR_BRANCH,
+ self._AR_TEXT_CACHE,
+ ):
+ projected_text = self.transformer.text_embedding(prompt_embeds)
+
+ def layer_kv() -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
+ for block in self.transformer.blocks:
+ cross_attention = block.cross_attn
+ key = cross_attention.norm_k(cross_attention.k(projected_text)).unflatten(
+ 2,
+ (
+ cross_attention.num_local_heads,
+ cross_attention.head_dim,
+ ),
+ )
+ value = cross_attention.v(projected_text).unflatten(
+ 2,
+ (
+ cross_attention.num_local_heads,
+ cross_attention.head_dim,
+ ),
+ )
+ yield key, value
+
+ state.populate_cross_attention(
+ self._AR_BRANCH,
+ self._AR_TEXT_CACHE,
+ layer_kv(),
+ )
+ return [
+ LingBotAttentionCache(
+ key=layer["k"],
+ value=layer["v"],
+ end=prompt_embeds.shape[1],
+ absolute_end=prompt_embeds.shape[1],
+ last_start=0,
+ )
+ for layer in state.get_cross_attention_kv(
+ self._AR_BRANCH,
+ self._AR_TEXT_CACHE,
+ )
+ ]
+
+ def _ar_transformer_cache(
+ self,
+ *,
+ condition: torch.Tensor,
+ cross_attention: list[LingBotAttentionCache],
+ commit_current: bool,
+ ) -> LingBotTransformerCache:
+ state = self._ar_diffusion_kv_state
+ if state is None:
+ raise RuntimeError("LingBot AR cache requested without a bound state.")
+ patch_frames, patch_height, patch_width = self.transformer.config.patch_size
+ frames = condition.shape[2] // patch_frames
+ height = condition.shape[3] // patch_height
+ width = condition.shape[4] // patch_width
+ seq_len = frames * height * width
+ return LingBotTransformerCache(
+ self_attention=state.get_kv_caches(
+ self._AR_BRANCH,
+ seq_len=seq_len,
+ commit_current=commit_current,
+ ),
+ cross_attention=cross_attention,
+ )
+
+ def _generate_block(
+ self,
+ *,
+ condition: torch.Tensor,
+ camera: torch.Tensor,
+ prompt_embeds: torch.Tensor,
+ cache: LingBotTransformerCache | None,
+ ar_cross_attention: list[LingBotAttentionCache] | None,
+ start_frame: int,
+ schedule: tuple[tuple[float, float], ...],
+ generator: torch.Generator,
+ progress_bar: TqdmProgressBar[Any],
+ ) -> torch.Tensor:
+ block_shape = (
+ 1,
+ self.transformer.config.out_channels,
+ condition.shape[2],
+ condition.shape[3],
+ condition.shape[4],
+ )
+ current_latents = randn_tensor(
+ block_shape,
+ generator=generator,
+ device=self.device,
+ dtype=torch.float32,
+ )
+ for step_index, (timestep_value, sigma) in enumerate(schedule):
+ if not self.od_config.enforce_eager:
+ torch.compiler.cudagraph_mark_step_begin()
+ set_forward_context_denoise_step_idx(step_index)
+ timestep = torch.full((1,), float(timestep_value), device=self.device, dtype=torch.float32)
+ # Checkpoint channel contract:
+ # [noise/x_t(16), temporal_mask(4), image_latent(16)] -> 36.
+ model_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1)
+ current_cache = (
+ self._ar_transformer_cache(
+ condition=condition,
+ cross_attention=cast(
+ list[LingBotAttentionCache],
+ ar_cross_attention,
+ ),
+ commit_current=False,
+ )
+ if ar_cross_attention is not None
+ else cast(LingBotTransformerCache, cache)
+ )
+ flow_prediction = self.transformer(
+ hidden_states=model_input,
+ timestep=timestep,
+ encoder_hidden_states=prompt_embeds,
+ camera_hidden_states=camera,
+ cache=current_cache,
+ start_frame=start_frame,
+ update_cache=False,
+ )
+ if flow_prediction.shape != current_latents.shape:
+ raise RuntimeError(
+ "transformer flow prediction shape must match the 16-channel noise latent, "
+ f"got {tuple(flow_prediction.shape)} and {tuple(current_latents.shape)}."
+ )
+ # The checkpoint's flow parameterization is inverted by
+ # x0 = x_t - sigma * flow. Intermediate steps re-noise that x0
+ # estimate at the next sigma; the final step keeps x0 as this
+ # block's generated latent.
+ x0 = current_latents - sigma * flow_prediction.float()
+ if step_index + 1 < len(schedule):
+ next_sigma = schedule[step_index + 1][1]
+ noise = randn_tensor(
+ current_latents.shape,
+ generator=generator,
+ device=self.device,
+ dtype=torch.float32,
+ )
+ current_latents = (1.0 - next_sigma) * x0 + next_sigma * noise
+ else:
+ current_latents = x0
+ progress_bar.update()
+
+ # Commit K/V only for the final clean block, never for noisy probes.
+ cache_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1)
+ commit_cache = (
+ self._ar_transformer_cache(
+ condition=condition,
+ cross_attention=cast(
+ list[LingBotAttentionCache],
+ ar_cross_attention,
+ ),
+ commit_current=True,
+ )
+ if ar_cross_attention is not None
+ else cast(LingBotTransformerCache, cache)
+ )
+ if not self.od_config.enforce_eager:
+ torch.compiler.cudagraph_mark_step_begin()
+ self.transformer(
+ hidden_states=cache_input,
+ timestep=torch.zeros(1, device=self.device, dtype=torch.float32),
+ encoder_hidden_states=prompt_embeds,
+ camera_hidden_states=camera,
+ cache=commit_cache,
+ start_frame=start_frame,
+ update_cache=True,
+ )
+ if ar_cross_attention is not None:
+ state = self._ar_diffusion_kv_state
+ if state is None:
+ raise RuntimeError("LingBot AR state disappeared before KV commit.")
+ state.commit_paged_context(self._AR_BRANCH)
+ return current_latents
+
+ def encode_prompt(
+ self,
+ prompt: str,
+ *,
+ max_sequence_length: int,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ text_inputs = self.tokenizer(
+ [" ".join(prompt.strip().split())],
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_attention_mask=True,
+ return_tensors="pt",
+ )
+ input_ids = text_inputs.input_ids.to(self.device)
+ attention_mask = text_inputs.attention_mask.to(self.device)
+ prompt_embeds = self.text_encoder(input_ids, attention_mask).last_hidden_state
+ prompt_embeds = prompt_embeds.to(device=self.device, dtype=dtype)
+ return prompt_embeds * attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype)
+
+ def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+ inputs = self._parse_request(req)
+ tick = ARDiffusionTickRequest.from_extra_args(req.sampling_params.extra_args)
+ if tick is not None and self._ar_diffusion_kv_state is None:
+ raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.")
+ if tick is None and self._ar_diffusion_kv_state is not None:
+ raise ValueError("LingBot ARDiffusionEngine requests must carry ar_diffusion_tick.")
+ session_state: _LingBotARSessionState | None = None
+ if tick is not None:
+ if tick.prompt is not None and " ".join(tick.prompt.split()) != " ".join(inputs.prompt.split()):
+ raise ValueError("ar_diffusion_tick.prompt must match the standard request prompt.")
+ block_frames = int(self.transformer.config.num_frames_per_block)
+ if inputs.num_latent_frames != block_frames:
+ raise ValueError("A LingBot AR-Diffusion request must generate exactly one three-latent-frame block.")
+ if (inputs.height, inputs.width) != (
+ self._ar_height,
+ self._ar_width,
+ ):
+ raise ValueError(
+ "LingBot AR-Diffusion request resolution must match the "
+ "fixed cache geometry "
+ f"{self._ar_height}x{self._ar_width}."
+ )
+ horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1
+ max_realtime_ticks = horizon_latent_frames // block_frames
+ if tick.chunk_index >= max_realtime_ticks:
+ raise ValueError(
+ "LingBot realtime generation currently supports at most "
+ f"{max_realtime_ticks} ticks per generation epoch "
+ f"(chunk_index 0 through {max_realtime_ticks - 1}) because "
+ f"the image-condition horizon is {_MAX_RAW_FRAMES} pixel "
+ "frames; reset or create a session to start a new world."
+ )
+ session_state = self._ar_sessions.setdefault(
+ tick.session_id,
+ _LingBotARSessionState(),
+ )
+ if tick.chunk_index != session_state.next_chunk_index:
+ raise ValueError(
+ "LingBot chunk_index must be contiguous within a session: "
+ f"expected {session_state.next_chunk_index}, got "
+ f"{tick.chunk_index}."
+ )
+ schedule = _build_shifted_flow_schedule(
+ flow_shift=inputs.flow_shift,
+ )
+ dtype = self.transformer.dtype
+ # Phase 1: turn all three user inputs into DiT-ready conditions.
+ prompt_embeds = self.encode_prompt(
+ inputs.prompt,
+ max_sequence_length=inputs.max_sequence_length,
+ dtype=dtype,
+ )
+ if tick is None:
+ condition = self._prepare_condition(inputs, dtype=dtype)
+ else:
+ assert session_state is not None
+ if session_state.image_condition is None:
+ horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1
+ session_state.image_condition = self._prepare_condition(
+ replace(
+ inputs,
+ num_frames=_MAX_RAW_FRAMES,
+ num_latent_frames=horizon_latent_frames,
+ ),
+ dtype=dtype,
+ )
+ block_frames = int(self.transformer.config.num_frames_per_block)
+ condition_start = tick.chunk_index * block_frames
+ condition_stop = condition_start + block_frames
+ if condition_stop > session_state.image_condition.shape[2]:
+ raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.")
+ condition = session_state.image_condition[
+ :,
+ :,
+ condition_start:condition_stop,
+ ]
+ camera_pitch: float | None = None
+ if inputs.camera_actions is not None:
+ assert session_state is not None
+ action_trajectory, camera_pitch = integrate_lingbot_camera_actions(
+ inputs.camera_actions,
+ width=inputs.width,
+ height=inputs.height,
+ initial_pose=(session_state.camera_tail.poses[-1] if session_state.camera_tail is not None else None),
+ initial_pitch=session_state.camera_pitch,
+ )
+ inputs = replace(inputs, camera_trajectory=action_trajectory)
+ camera, camera_tail = self._prepare_camera(
+ inputs,
+ dtype=dtype,
+ previous=session_state.camera_tail if session_state is not None else None,
+ )
+ if camera.shape[2:] != condition.shape[2:]:
+ raise RuntimeError(
+ "folded camera and image condition must share latent frame/height/width geometry; "
+ f"got camera={tuple(camera.shape)}, condition={tuple(condition.shape)}."
+ )
+
+ block_frames = int(self.transformer.config.num_frames_per_block)
+ cache: LingBotTransformerCache | None = None
+ ar_cross_attention: list[LingBotAttentionCache] | None = None
+ if tick is None:
+ # Direct/offline oracle keeps request-local contiguous K/V.
+ cache = self.transformer.allocate_cache(
+ batch_size=1,
+ latent_height=condition.shape[-2],
+ latent_width=condition.shape[-1],
+ device=self.device,
+ dtype=dtype,
+ )
+ else:
+ assert session_state is not None
+ prompt_changed = session_state.prompt is not None and session_state.prompt != inputs.prompt
+ if session_state.generator_state is not None:
+ inputs.generator.set_state(session_state.generator_state)
+ ar_cross_attention = self._ar_text_caches(
+ prompt_embeds,
+ invalidate=prompt_changed,
+ )
+ generated_blocks: list[torch.Tensor] = []
+ total_steps = (inputs.num_latent_frames // block_frames) * len(LINGBOT_DMD_TIMESTEPS)
+ first_start_frame = tick.chunk_index * block_frames if tick is not None else 0
+ with self.progress_bar(total=total_steps) as progress_bar:
+ # Phase 3: generate left-to-right in latent-frame blocks. The
+ # monotonically increasing ``start_frame`` is also the temporal
+ # position and cache offset used by the Transformer.
+ for local_start_frame in range(
+ 0,
+ inputs.num_latent_frames,
+ block_frames,
+ ):
+ start_frame = first_start_frame + local_start_frame
+ stop_frame = local_start_frame + block_frames
+ generated_blocks.append(
+ self._generate_block(
+ condition=condition[:, :, local_start_frame:stop_frame],
+ camera=camera[:, :, local_start_frame:stop_frame],
+ prompt_embeds=prompt_embeds,
+ cache=cache,
+ ar_cross_attention=ar_cross_attention,
+ start_frame=start_frame,
+ schedule=schedule,
+ generator=inputs.generator,
+ progress_bar=progress_bar,
+ )
+ )
+ generated_latents = torch.cat(generated_blocks, dim=2)
+ # The direct/offline cache is request-local and no longer needed once
+ # all latent blocks have been generated. Release it before VAE decode
+ # so the two large allocations do not overlap.
+ cache = None
+ if tick is not None:
+ assert session_state is not None
+ session_state.prompt = inputs.prompt
+ session_state.generator_state = inputs.generator.get_state()
+ session_state.camera_tail = camera_tail
+ if camera_pitch is not None:
+ session_state.camera_pitch = camera_pitch
+ session_state.next_chunk_index += 1
+
+ # Phase 4: either expose model-space latents or invert the checkpoint's
+ # latent normalization and decode to pixel-space video.
+ if tick is not None:
+ if inputs.output_type != "latent":
+ raise ValueError(
+ "LingBot realtime ticks currently require output_type='latent'; "
+ "stateful streaming VAE decode is a separate integration step."
+ )
+ output = {
+ "payload": {"latents": generated_latents},
+ "metadata": {"ar_diffusion": ARDiffusionChunkMetadata.from_tick(tick).to_dict()},
+ }
+ elif inputs.output_type == "latent":
+ output = generated_latents
+ else:
+ latent_mean, latent_std = self._vae_latent_stats(generated_latents)
+ vae_latents = (generated_latents * latent_std + latent_mean).to(dtype=self.vae.dtype)
+ output = self.vae.decode(vae_latents, return_dict=False)[0]
+ if output.shape[2] != inputs.num_frames:
+ raise RuntimeError(
+ "vae.decode returned an incompatible temporal geometry: "
+ f"expected {inputs.num_frames} frames, got {output.shape[2]}."
+ )
+ return DiffusionOutput(
+ output=output,
+ stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+ )
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ loader = AutoWeightsLoader(self)
+ return cast(set[str], loader.load_weights(weights))
diff --git a/vllm_omni/diffusion/models/lingbot_world/transformer.py b/vllm_omni/diffusion/models/lingbot_world/transformer.py
new file mode 100644
index 000000000..35ec18801
--- /dev/null
+++ b/vllm_omni/diffusion/models/lingbot_world/transformer.py
@@ -0,0 +1,1201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Checkpoint-compatible causal DiT for the LingBot World v2 model package."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any, Self, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.distributed import (
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+ tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.conv import Conv3dLayer
+from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.utils import set_weight_attrs
+
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.layers.norm import LayerNorm
+from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+ ARDiffusionPagedLayerContext,
+ ARDiffusionPagedLayerInputs,
+ ar_diffusion_paged_attention,
+ paged_write_attn,
+)
+
+
+@dataclass
+class LingBotAttentionCache:
+ """K/V storage plus its logical cursor metadata.
+
+ ``end`` is occupied storage, ``absolute_end`` is the next global token
+ position, and ``sink_end`` separates permanently retained prefix tokens
+ from the sliding local window. ``last_start`` rejects overlapping or
+ out-of-order causal chunks.
+ """
+
+ key: torch.Tensor
+ value: torch.Tensor
+ end: int = 0
+ absolute_end: int = 0
+ last_start: int | None = None
+ sink_end: int = 0
+
+
+@dataclass
+class LingBotTransformerCache:
+ """One request's per-layer video K/V and reusable text K/V."""
+
+ self_attention: list[LingBotAttentionCache | ARDiffusionPagedLayerContext | ARDiffusionPagedLayerInputs]
+ cross_attention: list[LingBotAttentionCache | None]
+
+
+def allocate_lingbot_cache(
+ *,
+ batch_size: int,
+ num_layers: int,
+ max_tokens: int,
+ num_local_heads: int,
+ head_dim: int,
+ device: torch.device,
+ dtype: torch.dtype,
+) -> LingBotTransformerCache:
+ # Cross-attention starts empty because its token count is known after text encoding.
+ shape = (batch_size, max_tokens, num_local_heads, head_dim)
+ self_attention = [
+ LingBotAttentionCache(
+ key=torch.zeros(shape, device=device, dtype=dtype),
+ value=torch.zeros(shape, device=device, dtype=dtype),
+ )
+ for _ in range(num_layers)
+ ]
+ return LingBotTransformerCache(
+ self_attention=self_attention,
+ cross_attention=[None for _ in range(num_layers)],
+ )
+
+
+class _LingBotRMSNorm(nn.Module):
+ """RMSNorm over a tensor-parallel sharded projection."""
+
+ def __init__(self, hidden_size: int, eps: float) -> None:
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ set_weight_attrs(self.weight, {"weight_loader": self.weight_loader})
+
+ def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+ if param.shape == loaded_weight.shape:
+ param.data.copy_(loaded_weight)
+ return
+
+ tp_size = get_tensor_model_parallel_world_size()
+ if loaded_weight.shape[0] % tp_size != 0:
+ raise ValueError(
+ f"Cannot shard RMSNorm weight of shape {tuple(loaded_weight.shape)} across tp_size={tp_size}."
+ )
+ shard_size = loaded_weight.shape[0] // tp_size
+ shard_start = get_tensor_model_parallel_rank() * shard_size
+ shard = loaded_weight.narrow(0, shard_start, shard_size)
+ if param.shape != shard.shape:
+ raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.")
+ param.data.copy_(shard)
+
+ def forward(self, value: torch.Tensor) -> torch.Tensor:
+ tp_size = get_tensor_model_parallel_world_size()
+ value_float = value.float()
+ sum_of_squares = value_float.pow(2).sum(dim=-1, keepdim=True)
+ element_count = value.shape[-1]
+ if tp_size > 1:
+ sum_of_squares = tensor_model_parallel_all_reduce(sum_of_squares)
+ element_count *= tp_size
+ rms = torch.sqrt(sum_of_squares / element_count + self.eps)
+ return (value_float / rms * self.weight.float()).to(value.dtype)
+
+
+def _projection_prefix(prefix: str, name: str) -> str:
+ return f"{prefix}.{name}" if prefix else name
+
+
+class LingBotSelfAttention(nn.Module):
+ """Block-causal self-attention over retained history and one full chunk."""
+
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ *,
+ eps: float = 1e-6,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ if dim % num_heads != 0:
+ raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
+
+ tp_size = get_tensor_model_parallel_world_size()
+ if num_heads % tp_size != 0:
+ raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
+
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.qkv = QKVParallelLinear(
+ hidden_size=dim,
+ head_size=self.head_dim,
+ total_num_heads=num_heads,
+ bias=True,
+ prefix=_projection_prefix(prefix, "qkv"),
+ )
+ self.num_local_heads = self.qkv.num_heads
+ self.tp_inner_dim = self.num_local_heads * self.head_dim
+ self.o = RowParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "o"),
+ )
+ self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)
+ self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)
+ self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
+ self.attn = Attention(
+ num_heads=self.num_local_heads,
+ head_size=self.head_dim,
+ num_kv_heads=self.num_local_heads,
+ softmax_scale=self.head_dim**-0.5,
+ causal=False,
+ role="self",
+ qkv_layout="BSND",
+ prefix=prefix,
+ skip_sequence_parallel=True,
+ )
+
+ def _update_cache(
+ self,
+ cache: LingBotAttentionCache,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ current_start: int,
+ *,
+ sink_tokens: int,
+ update_cache: bool,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ chunk_tokens = key.shape[1]
+ capacity = cache.key.shape[1]
+ next_sink_end = cache.sink_end
+
+ if cache.last_start is None:
+ if chunk_tokens > capacity:
+ raise ValueError(
+ f"Current chunk has {chunk_tokens} tokens but the cache can hold only {capacity}; "
+ "the full current chunk must remain visible."
+ )
+ next_key = key
+ next_value = value
+ next_sink_end = max(0, min(chunk_tokens, sink_tokens - current_start))
+ elif current_start == cache.last_start:
+ if current_start + chunk_tokens != cache.absolute_end:
+ raise ValueError("A repeated current_start must overwrite the same-size current chunk.")
+ if chunk_tokens > cache.end:
+ raise ValueError("The current chunk is no longer fully retained in the cache.")
+ prefix_end = cache.end - chunk_tokens
+ next_key = torch.cat((cache.key[:, :prefix_end], key), dim=1)
+ next_value = torch.cat((cache.value[:, :prefix_end], value), dim=1)
+ else:
+ if current_start < cache.last_start:
+ raise ValueError(f"current_start={current_start} precedes the latest chunk start {cache.last_start}.")
+ if current_start < cache.absolute_end:
+ raise ValueError(
+ f"current_start={current_start} overlaps cached tokens ending at {cache.absolute_end}."
+ )
+ if current_start > cache.absolute_end:
+ raise ValueError(
+ f"New chunks must be contiguous: current_start={current_start}, expected {cache.absolute_end}."
+ )
+
+ incoming_sink_tokens = max(0, min(chunk_tokens, sink_tokens - current_start))
+ old_sink_key = cache.key[:, : cache.sink_end]
+ old_sink_value = cache.value[:, : cache.sink_end]
+ new_sink_key = key[:, :incoming_sink_tokens]
+ new_sink_value = value[:, :incoming_sink_tokens]
+ next_sink_end = cache.sink_end + incoming_sink_tokens
+
+ old_local_key = cache.key[:, cache.sink_end : cache.end]
+ old_local_value = cache.value[:, cache.sink_end : cache.end]
+ new_local_key = key[:, incoming_sink_tokens:]
+ new_local_value = value[:, incoming_sink_tokens:]
+ local_capacity = capacity - next_sink_end
+ if new_local_key.shape[1] > local_capacity:
+ raise ValueError("The configured cache cannot retain all sink tokens and the full current chunk.")
+ retained_local_tokens = min(
+ old_local_key.shape[1],
+ local_capacity - new_local_key.shape[1],
+ )
+ if retained_local_tokens:
+ old_local_key = old_local_key[:, -retained_local_tokens:]
+ old_local_value = old_local_value[:, -retained_local_tokens:]
+ else:
+ old_local_key = old_local_key[:, :0]
+ old_local_value = old_local_value[:, :0]
+
+ next_key = torch.cat((old_sink_key, new_sink_key, old_local_key, new_local_key), dim=1)
+ next_value = torch.cat((old_sink_value, new_sink_value, old_local_value, new_local_value), dim=1)
+
+ next_end = next_key.shape[1]
+ if update_cache:
+ cache.key[:, :next_end].copy_(next_key)
+ cache.value[:, :next_end].copy_(next_value)
+ cache.end = next_end
+ cache.absolute_end = current_start + chunk_tokens
+ cache.last_start = current_start
+ cache.sink_end = next_sink_end
+ return next_key, next_value
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ *,
+ cache: LingBotAttentionCache | ARDiffusionPagedLayerInputs,
+ current_start: int,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ sink_tokens: int,
+ update_cache: bool = True,
+ ) -> torch.Tensor:
+ # Project logical [B, S, D] tokens into TP-local
+ # [B, S, num_local_heads, head_dim].
+ qkv, _ = self.qkv(hidden_states)
+ query, key, value = qkv.split((self.tp_inner_dim, self.tp_inner_dim, self.tp_inner_dim), dim=-1)
+ query = self.norm_q(query)
+ key = self.norm_k(key)
+
+ query = query.unflatten(2, (self.num_local_heads, self.head_dim))
+ key = key.unflatten(2, (self.num_local_heads, self.head_dim))
+ value = value.unflatten(2, (self.num_local_heads, self.head_dim))
+ if rotary_emb is not None:
+ cos, sin = rotary_emb
+ query = self.rotary_embedding(query, cos, sin)
+ key = self.rotary_embedding(key, cos, sin)
+
+ if isinstance(cache, ARDiffusionPagedLayerInputs):
+ if query.shape[0] != 1:
+ raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.")
+ output = paged_write_attn(
+ cache,
+ query[0],
+ key[0],
+ value[0],
+ None,
+ None,
+ self.head_dim**-0.5,
+ ).unsqueeze(0)
+ else:
+ # The direct path sees history + current K/V even when the caller
+ # asks not to mutate persistent cache state.
+ visible_key, visible_value = self._update_cache(
+ cache,
+ key,
+ value,
+ current_start,
+ sink_tokens=sink_tokens,
+ update_cache=update_cache,
+ )
+ if query.is_cuda and query.shape[0] == 1:
+ # Use the same block-table FlashAttention entry point as the
+ # realtime path so direct replay is a numerical oracle for
+ # paged execution, rather than a comparison between two
+ # different attention kernels.
+ block_size = key.shape[1]
+ key_cache = visible_key[0].unflatten(0, (-1, block_size))
+ value_cache = visible_value[0].unflatten(0, (-1, block_size))
+ block_count = key_cache.shape[0]
+ block_table = torch.arange(
+ block_count,
+ dtype=torch.int32,
+ device=query.device,
+ ).unsqueeze(0)
+ query_start_loc = torch.tensor(
+ [0, query.shape[1]],
+ dtype=torch.int32,
+ device=query.device,
+ )
+ seq_lens = torch.tensor(
+ [visible_key.shape[1]],
+ dtype=torch.int32,
+ device=query.device,
+ )
+ output = ar_diffusion_paged_attention(
+ query,
+ key_cache,
+ value_cache,
+ block_table=block_table,
+ query_start_loc=query_start_loc,
+ seq_lens=seq_lens,
+ max_query_len=query.shape[1],
+ max_seq_len=visible_key.shape[1],
+ softmax_scale=self.head_dim**-0.5,
+ )
+ else:
+ output = self.attn(query, visible_key, visible_value)
+ return self.o(output.flatten(2, 3))
+
+
+class LingBotCrossAttention(nn.Module):
+ """Cross-attention with caller-owned request-local encoder K/V reuse."""
+
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ *,
+ eps: float = 1e-6,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ if dim % num_heads != 0:
+ raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
+
+ tp_size = get_tensor_model_parallel_world_size()
+ if num_heads % tp_size != 0:
+ raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
+
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.num_local_heads = num_heads // tp_size
+ self.tp_inner_dim = self.num_local_heads * self.head_dim
+
+ self.q = ColumnParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "q"),
+ )
+ self.k = ColumnParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "k"),
+ )
+ self.v = ColumnParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "v"),
+ )
+ self.o = RowParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "o"),
+ )
+ self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)
+ self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)
+ self.attn = Attention(
+ num_heads=self.num_local_heads,
+ head_size=self.head_dim,
+ num_kv_heads=self.num_local_heads,
+ softmax_scale=self.head_dim**-0.5,
+ causal=False,
+ role="cross",
+ qkv_layout="BSND",
+ prefix=prefix,
+ skip_sequence_parallel=True,
+ disable_kv_quant=True,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor | None,
+ *,
+ cache: LingBotAttentionCache | None,
+ ) -> tuple[torch.Tensor, LingBotAttentionCache]:
+ query = self.norm_q(self.q(hidden_states))
+ query = query.unflatten(2, (self.num_local_heads, self.head_dim))
+
+ # Text K/V is constant within a request and is projected once per layer.
+ if cache is None:
+ if encoder_hidden_states is None:
+ raise ValueError("encoder_hidden_states are required when the cross-attention cache is empty.")
+ key = self.norm_k(self.k(encoder_hidden_states))
+ value = self.v(encoder_hidden_states)
+ key = key.unflatten(2, (self.num_local_heads, self.head_dim))
+ value = value.unflatten(2, (self.num_local_heads, self.head_dim))
+ cache = LingBotAttentionCache(
+ key=key,
+ value=value,
+ end=key.shape[1],
+ absolute_end=key.shape[1],
+ last_start=0,
+ )
+ else:
+ key = cache.key[:, : cache.end]
+ value = cache.value[:, : cache.end]
+
+ output = self.attn(query, key, value)
+ return self.o(output.flatten(2, 3)), cache
+
+
+class LingBotAttentionBlock(nn.Module):
+ """Checkpoint-compatible LingBot block with causal video attention."""
+
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int,
+ *,
+ ffn_dim: int | None = None,
+ cross_attn_norm: bool = True,
+ eps: float = 1e-6,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ ffn_dim = ffn_dim or dim * 4
+ self.dim = dim
+ self.norm1 = LayerNorm(dim, eps=eps, elementwise_affine=False)
+ self.self_attn = LingBotSelfAttention(
+ dim,
+ num_heads,
+ eps=eps,
+ prefix=_projection_prefix(prefix, "self_attn"),
+ )
+ self.cross_attn = LingBotCrossAttention(
+ dim,
+ num_heads,
+ eps=eps,
+ prefix=_projection_prefix(prefix, "cross_attn"),
+ )
+ self.norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False)
+ self.norm3 = LayerNorm(dim, eps=eps) if cross_attn_norm else nn.Identity()
+ self.ffn = nn.Sequential(
+ ColumnParallelLinear(
+ dim,
+ ffn_dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "ffn.0"),
+ ),
+ nn.GELU(approximate="tanh"),
+ RowParallelLinear(
+ ffn_dim,
+ dim,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "ffn.2"),
+ ),
+ )
+ self.modulation = nn.Parameter(torch.randn(1, 6, dim) / math.sqrt(dim))
+ self.cam_injector_layer1 = ColumnParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "cam_injector_layer1"),
+ )
+ self.cam_injector_layer2 = RowParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "cam_injector_layer2"),
+ )
+ self.cam_scale_layer = nn.Linear(dim, dim)
+ self.cam_shift_layer = nn.Linear(dim, dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor | None,
+ timestep_projection: torch.Tensor,
+ camera_hidden_states: torch.Tensor,
+ *,
+ self_cache: LingBotAttentionCache | ARDiffusionPagedLayerInputs,
+ cross_cache: LingBotAttentionCache | None,
+ current_start: int,
+ sink_tokens: int,
+ update_cache: bool,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor],
+ ) -> tuple[torch.Tensor, LingBotAttentionCache]:
+ batch_size, token_count, dim = hidden_states.shape
+ num_frames = timestep_projection.shape[1]
+ tokens_per_frame = token_count // num_frames
+ # Timestep, camera, and text remain separate conditioning paths.
+ modulation = self.modulation.unsqueeze(1) + timestep_projection.float()
+ shift_msa, scale_msa, gate_msa, shift_ffn, scale_ffn, gate_ffn = modulation.chunk(6, dim=2)
+
+ hidden_grid = hidden_states.unflatten(1, (num_frames, tokens_per_frame))
+ normalized = self.norm1(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame))
+ normalized = (normalized * (1 + scale_msa) + shift_msa).flatten(1, 2).to(hidden_states.dtype)
+ attention_output = self.self_attn(
+ normalized,
+ cache=self_cache,
+ current_start=current_start,
+ rotary_emb=rotary_emb,
+ sink_tokens=sink_tokens,
+ update_cache=update_cache,
+ )
+ hidden_grid = hidden_grid + attention_output.unflatten(1, (num_frames, tokens_per_frame)) * gate_msa
+ hidden_states = hidden_grid.flatten(1, 2).to(hidden_states.dtype)
+
+ camera_features = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(camera_hidden_states)))
+ camera_features = camera_features + camera_hidden_states
+ camera_scale = self.cam_scale_layer(camera_features)
+ camera_shift = self.cam_shift_layer(camera_features)
+ hidden_states = ((1 + camera_scale) * hidden_states + camera_shift).to(hidden_states.dtype)
+
+ attention_output, cross_cache = self.cross_attn(
+ self.norm3(hidden_states),
+ encoder_hidden_states,
+ cache=cross_cache,
+ )
+ hidden_states = hidden_states + attention_output
+
+ hidden_grid = hidden_states.unflatten(1, (num_frames, tokens_per_frame))
+ normalized = self.norm2(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame))
+ normalized = (normalized * (1 + scale_ffn) + shift_ffn).flatten(1, 2).to(hidden_states.dtype)
+ ffn_output = self.ffn(normalized).unflatten(1, (num_frames, tokens_per_frame))
+ hidden_states = (hidden_grid + ffn_output * gate_ffn).flatten(1, 2).to(hidden_states.dtype)
+ return hidden_states, cast(LingBotAttentionCache, cross_cache)
+
+
+class _LingBotCameraPatchEmbedding(nn.Module):
+ """Linear camera patchifier matching the checkpoint's direct parameter names."""
+
+ def __init__(
+ self,
+ in_channels: int,
+ dim: int,
+ patch_size: tuple[int, int, int],
+ ) -> None:
+ super().__init__()
+ self.patch_size = patch_size
+ self.in_features = in_channels * math.prod(patch_size)
+ self.weight = nn.Parameter(torch.empty(dim, self.in_features))
+ self.bias = nn.Parameter(torch.empty(dim))
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+ bound = 1 / math.sqrt(self.in_features)
+ nn.init.uniform_(self.bias, -bound, bound)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ batch_size, channels, frames, height, width = hidden_states.shape
+ patch_frames, patch_height, patch_width = self.patch_size
+ # Preserve Conv3d patch order while retaining checkpoint weight/bias names.
+ hidden_states = hidden_states.reshape(
+ batch_size,
+ channels,
+ frames // patch_frames,
+ patch_frames,
+ height // patch_height,
+ patch_height,
+ width // patch_width,
+ patch_width,
+ )
+ hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7)
+ hidden_states = hidden_states.reshape(batch_size, -1, self.in_features)
+ return F.linear(hidden_states, self.weight, self.bias)
+
+
+class _LingBotHead(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ out_channels: int,
+ patch_size: tuple[int, int, int],
+ eps: float,
+ ) -> None:
+ super().__init__()
+ self.norm = LayerNorm(dim, eps=eps, elementwise_affine=False)
+ self.head = nn.Linear(dim, out_channels * math.prod(patch_size))
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / math.sqrt(dim))
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep_embedding: torch.Tensor,
+ ) -> torch.Tensor:
+ num_frames = timestep_embedding.shape[1]
+ tokens_per_frame = hidden_states.shape[1] // num_frames
+ modulation = self.modulation.unsqueeze(1) + timestep_embedding.unsqueeze(2).float()
+ shift, scale = modulation.chunk(2, dim=2)
+ normalized = self.norm(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame))
+ normalized = (normalized * (1 + scale) + shift).flatten(1, 2).to(hidden_states.dtype)
+ return self.head(normalized)
+
+
+def _sinusoidal_embedding(dim: int, timestep: torch.Tensor) -> torch.Tensor:
+ if dim % 2:
+ raise ValueError(f"freq_dim must be even, got {dim}.")
+ half_dim = dim // 2
+ timestep = timestep.to(torch.float64)
+ frequencies = torch.pow(
+ 10000,
+ -torch.arange(half_dim, device=timestep.device, dtype=torch.float64) / half_dim,
+ )
+ phase = torch.outer(timestep, frequencies)
+ return torch.cat((phase.cos(), phase.sin()), dim=1)
+
+
+def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+ if dim == 0:
+ empty = torch.empty(max_seq_len, 0, dtype=torch.float32)
+ return empty, empty.clone()
+ if dim % 2:
+ raise ValueError(f"RoPE axis dimension must be even, got {dim}.")
+ frequencies = 1.0 / torch.pow(
+ 10000,
+ torch.arange(0, dim, 2, dtype=torch.float64) / dim,
+ )
+ phase = torch.outer(torch.arange(max_seq_len, dtype=torch.float64), frequencies)
+ return phase.cos().float(), phase.sin().float()
+
+
+class CausalLingBotWorldTransformer3DModel(nn.Module):
+ """Checkpoint-compatible causal LingBot World video transformer."""
+
+ _repeated_blocks = ["LingBotAttentionBlock"]
+ packed_modules_mapping = {"qkv": ["q", "k", "v"]}
+ _layerwise_offload_blocks_attrs = ["blocks"]
+
+ def __init__(
+ self,
+ patch_size: tuple[int, int, int] = (1, 2, 2),
+ num_attention_heads: int = 40,
+ attention_head_dim: int = 128,
+ in_channels: int = 36,
+ out_channels: int = 16,
+ text_dim: int = 4096,
+ freq_dim: int = 256,
+ ffn_dim: int = 13824,
+ num_layers: int = 40,
+ cross_attn_norm: bool = True,
+ eps: float = 1e-6,
+ image_dim: int | None = None,
+ added_kv_proj_dim: int | None = None,
+ rope_max_seq_len: int = 1024,
+ pos_embed_seq_len: int | None = None,
+ qk_norm: str = "rms_norm_across_heads",
+ sink_size: int = 9,
+ num_frames_per_block: int = 3,
+ sliding_window_num_frames: int = 18,
+ local_attn_size: int = -1,
+ quant_config: object | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ if len(patch_size) != 3 or any(size <= 0 for size in patch_size):
+ raise ValueError("patch_size must contain three positive integers.")
+ if num_attention_heads <= 0 or attention_head_dim <= 0:
+ raise ValueError("Attention head counts and dimensions must be positive.")
+ if attention_head_dim % 2:
+ raise ValueError("attention_head_dim must be even for RoPE.")
+ if num_layers <= 0 or ffn_dim <= 0 or freq_dim <= 0:
+ raise ValueError("num_layers, ffn_dim, and freq_dim must be positive.")
+ if sink_size < 0 or num_frames_per_block <= 0 or sliding_window_num_frames <= 0:
+ raise ValueError("Causal frame-window sizes must be positive, except sink_size which may be zero.")
+ if local_attn_size != -1 and local_attn_size <= 0:
+ raise ValueError("local_attn_size must be -1 or a positive frame count.")
+ cache_window_frames = local_attn_size if local_attn_size != -1 else sliding_window_num_frames
+ if cache_window_frames < sink_size + num_frames_per_block:
+ raise ValueError("The causal cache window must fit all sink frames and one complete current block.")
+ if qk_norm != "rms_norm_across_heads":
+ raise ValueError(
+ f"qk_norm must be 'rms_norm_across_heads' for the LingBot World v2 checkpoint, got {qk_norm!r}."
+ )
+ for field_name, field_value in (
+ ("image_dim", image_dim),
+ ("added_kv_proj_dim", added_kv_proj_dim),
+ ("pos_embed_seq_len", pos_embed_seq_len),
+ ):
+ if field_value is not None:
+ raise ValueError(f"{field_name} must be None because LingBot World v2 has no image embedding path.")
+ if quant_config is not None:
+ raise RuntimeError(
+ "quant_config is not supported by the LingBot World transformer; construct the unquantized model."
+ )
+
+ dim = num_attention_heads * attention_head_dim
+ self.dim = dim
+ self.config = SimpleNamespace(
+ patch_size=patch_size,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ text_dim=text_dim,
+ freq_dim=freq_dim,
+ ffn_dim=ffn_dim,
+ num_layers=num_layers,
+ cross_attn_norm=cross_attn_norm,
+ eps=eps,
+ image_dim=image_dim,
+ added_kv_proj_dim=added_kv_proj_dim,
+ rope_max_seq_len=rope_max_seq_len,
+ pos_embed_seq_len=pos_embed_seq_len,
+ qk_norm=qk_norm,
+ sink_size=sink_size,
+ num_frames_per_block=num_frames_per_block,
+ sliding_window_num_frames=sliding_window_num_frames,
+ local_attn_size=local_attn_size,
+ )
+
+ self.patch_embedding = Conv3dLayer(
+ in_channels=in_channels,
+ out_channels=dim,
+ kernel_size=patch_size,
+ stride=patch_size,
+ )
+ self.patch_embedding_wancamctrl = _LingBotCameraPatchEmbedding(6 * 8 * 8, dim, patch_size)
+ self.c2ws_hidden_states_layer1 = ColumnParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer1"),
+ )
+ self.c2ws_hidden_states_layer2 = RowParallelLinear(
+ dim,
+ dim,
+ bias=True,
+ input_is_parallel=True,
+ return_bias=False,
+ prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer2"),
+ )
+ self.text_embedding = nn.Sequential(
+ nn.Linear(text_dim, dim),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(dim, dim),
+ )
+ self.time_embedding = nn.Sequential(
+ nn.Linear(freq_dim, dim),
+ nn.SiLU(),
+ nn.Linear(dim, dim),
+ )
+ self.time_projection = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(dim, dim * 6),
+ )
+ self.blocks = nn.ModuleList(
+ [
+ LingBotAttentionBlock(
+ dim,
+ num_attention_heads,
+ ffn_dim=ffn_dim,
+ cross_attn_norm=cross_attn_norm,
+ eps=eps,
+ prefix=_projection_prefix(prefix, f"blocks.{index}"),
+ )
+ for index in range(num_layers)
+ ]
+ )
+ self.head = _LingBotHead(dim, out_channels, patch_size, eps)
+
+ temporal_dim = attention_head_dim - 4 * (attention_head_dim // 6)
+ height_dim = width_dim = 2 * (attention_head_dim // 6)
+ for axis, axis_dim in (("temporal", temporal_dim), ("height", height_dim), ("width", width_dim)):
+ cosine, sine = _rope_axis(rope_max_seq_len, axis_dim)
+ self.register_buffer(f"_rope_{axis}_cosine", cosine, persistent=False)
+ self.register_buffer(f"_rope_{axis}_sine", sine, persistent=False)
+
+ @property
+ def dtype(self) -> torch.dtype:
+ """Return the dtype used by the transformer parameters."""
+
+ return next(self.parameters()).dtype
+
+ @classmethod
+ def from_config(
+ cls,
+ config: dict[str, Any],
+ *,
+ quant_config: Any | None = None,
+ prefix: str = "",
+ ) -> Self:
+ checkpoint_contract = {
+ "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,
+ }
+ class_name = config.get("_class_name")
+ if class_name is not None and class_name != "CausalLingBotWorldTransformer3DModel":
+ raise ValueError(
+ "_class_name must be 'CausalLingBotWorldTransformer3DModel' for the LingBot World v2 checkpoint."
+ )
+ kwargs = {name: value for name, value in config.items() if name in checkpoint_contract}
+ if "patch_size" in kwargs:
+ kwargs["patch_size"] = tuple(kwargs["patch_size"])
+ for name, expected in checkpoint_contract.items():
+ actual = kwargs.get(name, expected)
+ if actual != expected:
+ raise ValueError(f"LingBot checkpoint config {name} must be {expected!r}, got {actual!r}.")
+ return cls(**kwargs, quant_config=quant_config, prefix=prefix)
+
+ def allocate_cache(
+ self,
+ *,
+ batch_size: int,
+ latent_height: int,
+ latent_width: int,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> LingBotTransformerCache:
+ """Allocate one caller-owned cache using this transformer's geometry."""
+
+ if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size <= 0:
+ raise ValueError("batch_size must be a positive integer for LingBot cache allocation.")
+ if (
+ isinstance(latent_height, bool)
+ or not isinstance(latent_height, int)
+ or latent_height <= 0
+ or isinstance(latent_width, bool)
+ or not isinstance(latent_width, int)
+ or latent_width <= 0
+ ):
+ raise ValueError("latent_height and latent_width must be positive integers.")
+
+ # Cache capacity is measured in post-patch tokens, not raw/latent
+ # pixels: window_frames * (latent_height/patch_h) * (latent_width/patch_w).
+ patch_frames, patch_height, patch_width = self.config.patch_size
+ if patch_frames != 1 or latent_height % patch_height or latent_width % patch_width:
+ raise ValueError("latent height/width must align with the configured LingBot patch size.")
+ post_patch_height = latent_height // patch_height
+ post_patch_width = latent_width // patch_width
+ window_frames = (
+ self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames
+ )
+ max_tokens = int(window_frames * post_patch_height * post_patch_width)
+ tp_size = get_tensor_model_parallel_world_size()
+ num_local_heads = self.config.num_attention_heads // tp_size
+ return allocate_lingbot_cache(
+ batch_size=batch_size,
+ num_layers=self.config.num_layers,
+ max_tokens=max_tokens,
+ num_local_heads=num_local_heads,
+ head_dim=self.config.attention_head_dim,
+ device=device,
+ dtype=dtype,
+ )
+
+ def _rotary_embedding(
+ self,
+ *,
+ frames: int,
+ height: int,
+ width: int,
+ start_frame: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ if start_frame + frames > self.config.rope_max_seq_len:
+ raise ValueError("Temporal RoPE positions exceed rope_max_seq_len.")
+ if height > self.config.rope_max_seq_len or width > self.config.rope_max_seq_len:
+ raise ValueError("Spatial RoPE positions exceed rope_max_seq_len.")
+
+ def expand_axis(table: torch.Tensor, axis: str) -> torch.Tensor:
+ if axis == "temporal":
+ return (
+ table[start_frame : start_frame + frames].view(frames, 1, 1, -1).expand(frames, height, width, -1)
+ )
+ if axis == "height":
+ return table[:height].view(1, height, 1, -1).expand(frames, height, width, -1)
+ return table[:width].view(1, 1, width, -1).expand(frames, height, width, -1)
+
+ cosine = torch.cat(
+ (
+ expand_axis(self._rope_temporal_cosine, "temporal"),
+ expand_axis(self._rope_height_cosine, "height"),
+ expand_axis(self._rope_width_cosine, "width"),
+ ),
+ dim=-1,
+ )
+ sine = torch.cat(
+ (
+ expand_axis(self._rope_temporal_sine, "temporal"),
+ expand_axis(self._rope_height_sine, "height"),
+ expand_axis(self._rope_width_sine, "width"),
+ ),
+ dim=-1,
+ )
+ return (
+ cosine.reshape(frames * height * width, -1).to(device=device, dtype=dtype),
+ sine.reshape(frames * height * width, -1).to(device=device, dtype=dtype),
+ )
+
+ def _validate_inputs(
+ self,
+ hidden_states: torch.Tensor,
+ timestep: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ camera_hidden_states: torch.Tensor,
+ cache: LingBotTransformerCache,
+ start_frame: int,
+ ) -> tuple[int, int, int, int, int]:
+ if hidden_states.ndim != 5:
+ raise ValueError("hidden_states must be rank 5 [batch, channels, frames, height, width].")
+ if hidden_states.shape[1] != self.config.in_channels:
+ raise ValueError(f"hidden_states must have {self.config.in_channels} channels.")
+ if not isinstance(start_frame, int) or isinstance(start_frame, bool) or start_frame < 0:
+ raise ValueError("start_frame must be a non-negative integer latent-frame offset.")
+ patch_frames, patch_height, patch_width = self.config.patch_size
+ batch_size, _, frames, height, width = hidden_states.shape
+ if frames % patch_frames or height % patch_height or width % patch_width:
+ raise ValueError("hidden_states dimensions must be divisible by patch_size.")
+ post_patch_frames = frames // patch_frames
+ if post_patch_frames != self.config.num_frames_per_block:
+ raise ValueError(
+ f"Each LingBot forward must contain exactly {self.config.num_frames_per_block} post-patch frames, "
+ f"got {post_patch_frames}."
+ )
+ if start_frame % patch_frames:
+ raise ValueError("start_frame must align to the temporal patch size.")
+ if encoder_hidden_states.ndim != 3:
+ raise ValueError("encoder_hidden_states must be rank 3 [batch, tokens, text_dim].")
+ if encoder_hidden_states.shape[0] != batch_size or encoder_hidden_states.shape[2] != self.config.text_dim:
+ raise ValueError("encoder_hidden_states batch/width do not match the LingBot config.")
+ if encoder_hidden_states.shape[1] == 0:
+ raise ValueError("encoder_hidden_states must contain at least one token.")
+ expected_camera_shape = (batch_size, 6 * 8 * 8, frames, height, width)
+ if camera_hidden_states.shape != expected_camera_shape:
+ raise ValueError(
+ "camera_hidden_states must use folded [batch, 6*8*8, frames, height, width] layout "
+ f"matching the video; expected {expected_camera_shape}, got {tuple(camera_hidden_states.shape)}."
+ )
+ if hidden_states.device != encoder_hidden_states.device or hidden_states.device != camera_hidden_states.device:
+ raise ValueError("Video, text, and camera tensors must be on the same device.")
+ if hidden_states.dtype != encoder_hidden_states.dtype or hidden_states.dtype != camera_hidden_states.dtype:
+ raise ValueError("Video, text, and camera tensors must use the same dtype.")
+ if timestep.device != hidden_states.device:
+ raise ValueError("timestep must be on the same device as hidden_states.")
+ if len(cache.self_attention) != self.config.num_layers or len(cache.cross_attention) != self.config.num_layers:
+ raise ValueError("LingBotTransformerCache layer counts must match num_layers.")
+ return batch_size, frames, height, width, patch_frames
+
+ def _timestep_embeddings(
+ self,
+ timestep: torch.Tensor,
+ *,
+ batch_size: int,
+ frames: int,
+ dtype: torch.dtype,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ if timestep.ndim == 0:
+ if batch_size != 1:
+ raise ValueError("A scalar timestep is only valid for batch_size=1.")
+ timestep = timestep.reshape(1)
+ if timestep.ndim == 1:
+ if timestep.shape[0] != batch_size:
+ raise ValueError("A rank-1 timestep must contain one value per batch item.")
+ timestep = timestep.unsqueeze(1).expand(batch_size, frames)
+ elif timestep.ndim == 2:
+ if timestep.shape != (batch_size, frames):
+ raise ValueError(f"A rank-2 timestep must have shape {(batch_size, frames)}.")
+ else:
+ raise ValueError("timestep must be scalar, rank 1 [batch], or rank 2 [batch, frames].")
+
+ frequency_embedding = _sinusoidal_embedding(self.config.freq_dim, timestep.reshape(-1)).to(dtype=dtype)
+ timestep_embedding = self.time_embedding(frequency_embedding).unflatten(0, (batch_size, frames))
+ timestep_projection = self.time_projection(timestep_embedding).unflatten(2, (6, self.dim))
+ return timestep_embedding, timestep_projection
+
+ def _unpatchify(
+ self,
+ hidden_states: torch.Tensor,
+ *,
+ batch_size: int,
+ frames: int,
+ height: int,
+ width: int,
+ ) -> torch.Tensor:
+ patch_frames, patch_height, patch_width = self.config.patch_size
+ hidden_states = hidden_states.reshape(
+ batch_size,
+ frames,
+ height,
+ width,
+ patch_frames,
+ patch_height,
+ patch_width,
+ self.config.out_channels,
+ )
+ hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+ return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ timestep: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ camera_hidden_states: torch.Tensor,
+ *,
+ cache: LingBotTransformerCache,
+ start_frame: int,
+ update_cache: bool,
+ ) -> torch.Tensor:
+ if not isinstance(update_cache, bool):
+ raise ValueError("update_cache must be a boolean.")
+ batch_size, frames, height, width, patch_frames = self._validate_inputs(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache,
+ start_frame,
+ )
+ patch_height, patch_width = self.config.patch_size[1:]
+ patched_frames = frames // patch_frames
+ patched_height = height // patch_height
+ patched_width = width // patch_width
+ tokens_per_frame = patched_height * patched_width
+ patched_start_frame = start_frame // patch_frames
+ current_start = patched_start_frame * tokens_per_frame
+ sink_tokens = self.config.sink_size * tokens_per_frame
+ # Phase 1: create absolute 3D positions for this causal block. The
+ # spatial coordinates restart for every frame; temporal coordinates
+ # begin at ``start_frame`` so cached blocks retain their true ordering.
+ rotary_emb = self._rotary_embedding(
+ frames=patched_frames,
+ height=patched_height,
+ width=patched_width,
+ start_frame=patched_start_frame,
+ dtype=hidden_states.dtype,
+ device=hidden_states.device,
+ )
+ # Phase 2: independently patchify video and camera grids to identical
+ # token layouts, then project timestep and text conditions.
+ hidden_states = self.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
+ camera_hidden_states = self.patch_embedding_wancamctrl(camera_hidden_states)
+ camera_hidden_states = camera_hidden_states + self.c2ws_hidden_states_layer2(
+ F.silu(self.c2ws_hidden_states_layer1(camera_hidden_states))
+ )
+ if hidden_states.shape != camera_hidden_states.shape:
+ raise ValueError("Patched video and camera token shapes must match.")
+
+ timestep_embedding, timestep_projection = self._timestep_embeddings(
+ timestep,
+ batch_size=batch_size,
+ frames=patched_frames,
+ dtype=hidden_states.dtype,
+ )
+ projected_text = (
+ self.text_embedding(encoder_hidden_states)
+ if any(value is None for value in cache.cross_attention)
+ else None
+ )
+ if cache.self_attention and isinstance(cache.self_attention[0], ARDiffusionPagedLayerContext):
+ if batch_size != 1:
+ raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.")
+ forward_context = cache.self_attention[0].forward_ctx
+ expected_seq_len = patched_frames * tokens_per_frame
+ if forward_context.seq_len != expected_seq_len:
+ raise RuntimeError(
+ "LingBot AR-Diffusion paged context token count does not "
+ f"match this block: {forward_context.seq_len} != "
+ f"{expected_seq_len}."
+ )
+ forward_context.prepare(
+ device=hidden_states.device,
+ action_len=0,
+ query_len=hidden_states.shape[1],
+ )
+ cache.self_attention = [layer_context.to_layer_inputs() for layer_context in cache.self_attention]
+ # Phase 3: each layer receives its own cache entry. Text K/V is passed
+ # only when absent; the returned cache is stored for subsequent DMD
+ # steps and causal blocks in this request.
+ for index, block in enumerate(self.blocks):
+ hidden_states, cross_cache = block(
+ hidden_states,
+ projected_text if cache.cross_attention[index] is None else None,
+ timestep_projection,
+ camera_hidden_states,
+ self_cache=cache.self_attention[index],
+ cross_cache=cache.cross_attention[index],
+ current_start=current_start,
+ sink_tokens=sink_tokens,
+ update_cache=update_cache,
+ rotary_emb=rotary_emb,
+ )
+ cache.cross_attention[index] = cross_cache
+
+ # Phase 4: map tokens to per-patch 16-channel flow values and restore
+ # [B, C, F, H, W] for the Pipeline's sampler update.
+ hidden_states = self.head(hidden_states, timestep_embedding)
+ return self._unpatchify(
+ hidden_states,
+ batch_size=batch_size,
+ frames=patched_frames,
+ height=patched_height,
+ width=patched_width,
+ )
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ """Load exact checkpoint names, delegating TP sharding to parameters."""
+
+ params = dict(self.named_parameters())
+ loaded: set[str] = set()
+ qkv_mapping = (("q", "q"), ("k", "k"), ("v", "v"))
+ for checkpoint_name, loaded_weight in weights:
+ name = checkpoint_name
+ shard_id = None
+ for projection_name, projection_shard in qkv_mapping:
+ marker = f".self_attn.{projection_name}."
+ if marker in name:
+ name = name.replace(marker, ".self_attn.qkv.")
+ shard_id = projection_shard
+ break
+ if name not in params:
+ raise KeyError(f"Unexpected LingBot model weight name: {checkpoint_name}")
+ param = params[name]
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
+ if shard_id is None:
+ weight_loader(param, loaded_weight)
+ else:
+ if weight_loader is default_weight_loader:
+ raise RuntimeError(f"Fused LingBot QKV parameter {name} has no stacked weight loader.")
+ weight_loader(param, loaded_weight, shard_id)
+ # DiffusersPipelineLoader compares this return value with the
+ # model parameter namespace. Packed Q/K/V checkpoint tensors
+ # therefore need to report both their source names and the fused
+ # model parameter they populated, matching Wan's loader contract.
+ loaded.add(checkpoint_name)
+ loaded.add(name)
+ return loaded
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 9bac50323..2c4bffb4c 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -136,6 +136,11 @@ _DIFFUSION_MODELS = {
"pipeline_wan2_2_i2v",
"WanI2VDMD2Pipeline",
),
+ "LingBotWorldCausalDMDPipeline": (
+ "lingbot_world",
+ "pipeline",
+ "LingBotWorldCausalDMDPipeline",
+ ),
"LongCatImagePipeline": (
"longcat_image",
"pipeline_longcat_image",
@@ -353,6 +358,7 @@ _NO_CACHE_ACCELERATION = {
# Pipelines that do not support cache acceleration (cache_dit / tea_cache).
"NextStep11Pipeline",
"AudioXPipeline",
+ "LingBotWorldCausalDMDPipeline",
}
@@ -547,6 +553,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"WanS2VPipeline": "get_wan22_s2v_post_process_func",
"WanT2VDMD2Pipeline": "get_wan22_post_process_func",
"WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
+ "LingBotWorldCausalDMDPipeline": "get_lingbot_world_post_process_func",
"LongCatImagePipeline": "get_longcat_image_post_process_func",
"LongCatVideoAvatarPipeline": "get_longcat_video_avatar_post_process_func",
"BagelPipeline": "get_bagel_post_process_func",
@@ -607,6 +614,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
"WanS2VPipeline": "get_wan22_s2v_pre_process_func",
"WanT2VDMD2Pipeline": "get_wan22_pre_process_func",
"WanI2VDMD2Pipeline": "get_wan22_i2v_pre_process_func",
+ "LingBotWorldCausalDMDPipeline": "get_lingbot_world_pre_process_func",
"OmniGen2Pipeline": "get_omnigen2_pre_process_func",
"HeliosPipeline": "get_helios_pre_process_func",
"HeliosPyramidPipeline": "get_helios_pre_process_func",
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index fea78bcf9..436fcca9a 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -444,6 +444,25 @@ class DiffusionWorker:
else:
profiler.stop()
+ def _run_ar_diffusion_session_lifecycle(self, method: str, session_id: str) -> bool:
+ """Delegate an optional AR session lifecycle call to the model runner."""
+ if not isinstance(session_id, str) or not session_id.strip():
+ raise ValueError("session_id must be a non-empty string.")
+ assert self.model_runner is not None, "Model runner not initialized"
+ lifecycle_method = getattr(self.model_runner, method, None)
+ if not callable(lifecycle_method):
+ return False
+ lifecycle_method(session_id)
+ return True
+
+ def reset_ar_diffusion_session(self, session_id: str) -> bool:
+ """Reset runner-owned AR state through the collective RPC boundary."""
+ return self._run_ar_diffusion_session_lifecycle("reset_session", session_id)
+
+ def close_ar_diffusion_session(self, session_id: str) -> bool:
+ """Close runner-owned AR state through the collective RPC boundary."""
+ return self._run_ar_diffusion_session_lifecycle("close_session", session_id)
+
def execute_model(
self,
req: OmniDiffusionRequest | list[OmniDiffusionRequest],
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index a613ebc80..c457b7052 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -982,6 +982,9 @@ class AsyncOmniEngine:
stage_engine_args = {
"max_num_seqs": kwargs.get("max_num_seqs") or 1,
"parallel_config": parallel_config,
+ # Default-stage construction bypasses the structured projection.
+ # Runner selection remains owned by the selected engine/platform.
+ "engine_backend": kwargs.get("engine_backend", "default"),
"model_class_name": kwargs.get("model_class_name", None),
"task_type": kwargs.get("task_type", None),
"model_config": kwargs.get("model_config", None),
diff --git a/vllm_omni/experimental/ar_diffusion/__init__.py b/vllm_omni/experimental/ar_diffusion/__init__.py
index b87ac1ccc..bd0f2eac2 100644
--- a/vllm_omni/experimental/ar_diffusion/__init__.py
+++ b/vllm_omni/experimental/ar_diffusion/__init__.py
@@ -13,12 +13,32 @@ from vllm_omni.experimental.ar_diffusion.capability import (
SupportsARDiffusionPipeline,
SupportsARDiffusionWarmup,
)
+from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionSession,
+ ARDiffusionSessionManager,
+ ARDiffusionWorkerLifecycle,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ AR_DIFFUSION_TICK_KEY,
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
__all__ = [
"ARDiffusionKVBranchSpec",
"ARDiffusionCrossAttentionKVSpec",
"ARDiffusionEngine",
+ "ARDiffusionOmniTickConsumer",
+ "ARDiffusionSession",
+ "ARDiffusionSessionManager",
+ "ARDiffusionWorkerLifecycle",
+ "AR_DIFFUSION_TICK_KEY",
+ "ARDiffusionChunkMetadata",
+ "ARDiffusionControlInput",
+ "ARDiffusionTickRequest",
"ARDiffusionKVCacheSpec",
"SupportsARDiffusionPipeline",
"SupportsARDiffusionWarmup",
diff --git a/vllm_omni/experimental/ar_diffusion/capability.py b/vllm_omni/experimental/ar_diffusion/capability.py
index b30d10c1e..64eb62cf0 100644
--- a/vllm_omni/experimental/ar_diffusion/capability.py
+++ b/vllm_omni/experimental/ar_diffusion/capability.py
@@ -58,6 +58,9 @@ class ARDiffusionKVCacheSpec:
``max_scratch_tokens_per_branch`` is the maximum non-video KV (for example,
action/state registers) that must coexist with an uncommitted video block.
+ ``model_owned_state_bytes_per_session`` is a conservative upper bound for
+ persistent per-session CUDA state outside runner-owned self/cross-attention
+ KV and scratch pools.
``session_capacity`` is the pipeline's upper bound; a runner may retain
fewer sessions when its current memory and execution policy is stricter.
@@ -79,6 +82,7 @@ class ARDiffusionKVCacheSpec:
cross_attention: tuple[ARDiffusionCrossAttentionKVSpec, ...] = ()
max_model_len: int = 1 << 20
max_scratch_tokens_per_branch: int = 0
+ model_owned_state_bytes_per_session: int = 0
def __post_init__(self) -> None:
positive_fields = {
@@ -101,6 +105,11 @@ class ARDiffusionKVCacheSpec:
"AR-Diffusion max_scratch_tokens_per_branch must be non-negative, "
f"got {self.max_scratch_tokens_per_branch}"
)
+ if self.model_owned_state_bytes_per_session < 0:
+ raise ValueError(
+ "AR-Diffusion model_owned_state_bytes_per_session must be non-negative, "
+ f"got {self.model_owned_state_bytes_per_session}"
+ )
if not self.kv_branches:
raise ValueError("AR-Diffusion requires at least one KV branch")
kv_branch_names = [kv_branch.name for kv_branch in self.kv_branches]
diff --git a/vllm_omni/experimental/ar_diffusion/consumer.py b/vllm_omni/experimental/ar_diffusion/consumer.py
new file mode 100644
index 000000000..a819e0b31
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/consumer.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AsyncOmni consumer for typed realtime AR-Diffusion ticks."""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
+from typing import Any, Protocol
+
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionTickRequest,
+)
+from vllm_omni.inputs.data import (
+ OmniDiffusionSamplingParams,
+ OmniPromptType,
+ OmniSamplingParams,
+)
+from vllm_omni.outputs import OmniRequestOutput
+
+AR_DIFFUSION_OUTPUT_METADATA_KEY = "ar_diffusion"
+
+
+class ARDiffusionGenerateClient(Protocol):
+ """Internal subset of AsyncOmni used to execute one tick."""
+
+ def generate(
+ self,
+ prompt: OmniPromptType,
+ *,
+ request_id: str,
+ sampling_params_list: Sequence[OmniSamplingParams],
+ output_modalities: list[str] | None = None,
+ ) -> AsyncIterator[OmniRequestOutput]: ...
+
+
+def _metadata_mapping(output: OmniRequestOutput) -> Mapping[str, Any]:
+ multimodal_output = output.multimodal_output
+ if not isinstance(multimodal_output, Mapping):
+ raise ValueError("AR-Diffusion output must contain a multimodal output mapping.")
+ metadata = multimodal_output.get("metadata")
+ if not isinstance(metadata, Mapping):
+ raise ValueError("AR-Diffusion output must contain multimodal_output.metadata.")
+ value = metadata.get(AR_DIFFUSION_OUTPUT_METADATA_KEY)
+ if not isinstance(value, Mapping):
+ raise ValueError(f"AR-Diffusion output metadata must contain {AR_DIFFUSION_OUTPUT_METADATA_KEY!r}.")
+ return value
+
+
+def _parse_chunk_metadata(output: OmniRequestOutput) -> ARDiffusionChunkMetadata:
+ value = _metadata_mapping(output)
+ session_id = value.get("session_id")
+ request_id = value.get("request_id")
+ chunk_index = value.get("chunk_index")
+ applied_event_ids = value.get("applied_event_ids", ())
+ if not isinstance(session_id, str) or not session_id.strip():
+ raise ValueError("AR-Diffusion output metadata session_id must be a non-empty string.")
+ if not isinstance(request_id, str) or not request_id.strip():
+ raise ValueError("AR-Diffusion output metadata request_id must be a non-empty string.")
+ if isinstance(chunk_index, bool) or not isinstance(chunk_index, int) or chunk_index < 0:
+ raise ValueError("AR-Diffusion output metadata chunk_index must be a non-negative integer.")
+ if not isinstance(applied_event_ids, Sequence) or isinstance(applied_event_ids, (str, bytes, bytearray)):
+ raise ValueError("AR-Diffusion output metadata applied_event_ids must be a sequence.")
+ event_ids = tuple(applied_event_ids)
+ if any(isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0 for event_id in event_ids):
+ raise ValueError("AR-Diffusion output metadata applied_event_ids must contain non-negative integers.")
+ if event_ids != tuple(sorted(set(event_ids))):
+ raise ValueError("AR-Diffusion output metadata applied_event_ids must be unique and strictly increasing.")
+ return ARDiffusionChunkMetadata(
+ session_id=session_id,
+ request_id=request_id,
+ chunk_index=chunk_index,
+ applied_event_ids=event_ids,
+ )
+
+
+class ARDiffusionOmniTickConsumer:
+ """Converts a typed tick into one AsyncOmni diffusion request.
+
+ ``prompt_provider`` packages session initialization data such as an initial
+ image into the normal Omni prompt. Control ``track/schema/data`` remains in
+ the typed tick and is not interpreted by this class.
+ """
+
+ def __init__(
+ self,
+ engine_client: ARDiffusionGenerateClient,
+ *,
+ prompt_provider: Callable[[ARDiffusionTickRequest], OmniPromptType],
+ sampling_params_list: Sequence[OmniSamplingParams],
+ diffusion_stage_id: int,
+ output_modalities: Sequence[str] | None = None,
+ ) -> None:
+ templates = tuple(sampling_params_list)
+ if not templates:
+ raise ValueError("sampling_params_list must not be empty.")
+ if (
+ isinstance(diffusion_stage_id, bool)
+ or not isinstance(diffusion_stage_id, int)
+ or not 0 <= diffusion_stage_id < len(templates)
+ ):
+ raise ValueError("diffusion_stage_id must index sampling_params_list.")
+ if not isinstance(templates[diffusion_stage_id], OmniDiffusionSamplingParams):
+ raise TypeError("The selected diffusion stage must use OmniDiffusionSamplingParams.")
+ stage_pools = getattr(getattr(engine_client, "engine", None), "stage_pools", None)
+ if stage_pools is None or not 0 <= diffusion_stage_id < len(stage_pools):
+ raise ValueError(
+ "AR-Diffusion realtime execution requires an initialized AsyncOmni "
+ "engine with a matching diffusion stage pool."
+ )
+ replica_count = getattr(stage_pools[diffusion_stage_id], "num_replicas", None)
+ if isinstance(replica_count, bool) or not isinstance(replica_count, int) or replica_count < 1:
+ raise ValueError("The selected AR-Diffusion stage must expose a positive num_replicas.")
+ if replica_count != 1:
+ raise ValueError(
+ "AR-Diffusion realtime sessions currently require exactly one "
+ "replica for the selected diffusion stage; session-affine "
+ f"routing is not implemented (got num_replicas={replica_count})."
+ )
+ self._engine_client = engine_client
+ self._prompt_provider = prompt_provider
+ self._sampling_params_templates = templates
+ self._diffusion_stage_id = diffusion_stage_id
+ self._output_modalities = list(output_modalities) if output_modalities is not None else None
+
+ def _sampling_params_for_tick(
+ self,
+ tick: ARDiffusionTickRequest,
+ ) -> list[OmniSamplingParams]:
+ params = [copy.copy(template) for template in self._sampling_params_templates]
+ diffusion_params = params[self._diffusion_stage_id]
+ assert isinstance(diffusion_params, OmniDiffusionSamplingParams)
+ diffusion_params.extra_args = {
+ **(diffusion_params.extra_args or {}),
+ **tick.to_extra_args(),
+ }
+ return params
+
+ async def execute_tick(self, tick: ARDiffusionTickRequest) -> OmniRequestOutput:
+ prompt = self._prompt_provider(tick)
+ final_output: OmniRequestOutput | None = None
+ async for output in self._engine_client.generate(
+ prompt,
+ request_id=tick.request_id,
+ sampling_params_list=self._sampling_params_for_tick(tick),
+ output_modalities=self._output_modalities,
+ ):
+ if output.stage_id != self._diffusion_stage_id:
+ continue
+ if output.finished:
+ final_output = output
+
+ if final_output is None:
+ raise RuntimeError(
+ f"AR-Diffusion tick {tick.request_id!r} completed without a final "
+ f"output from stage {self._diffusion_stage_id}."
+ )
+ if final_output.error:
+ raise RuntimeError(f"AR-Diffusion tick {tick.request_id!r} failed: {final_output.error}")
+ if final_output.request_id != tick.request_id:
+ raise ValueError("AR-Diffusion output request_id does not match the submitted tick request_id.")
+ # Parse here so malformed/missing standard metadata fails at the engine
+ # boundary, before the session attempts to commit accepted events.
+ self.chunk_metadata(final_output)
+ return final_output
+
+ def chunk_metadata(self, output: OmniRequestOutput) -> ARDiffusionChunkMetadata:
+ return _parse_chunk_metadata(output)
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
index b407d24a3..94e6edf0d 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
@@ -32,7 +32,6 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
ARDiffusionPagedLayerInputs,
ar_diffusion_paged_attention,
paged_write_attn,
- set_current_paged_kv_cache,
)
__all__ = [
@@ -53,5 +52,4 @@ __all__ = [
"paged_write_attn",
"pool_write_chunk",
"resident_block_ids",
- "set_current_paged_kv_cache",
]
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
index 07be2546e..ca91a8bb6 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
@@ -23,7 +23,8 @@ class ARDiffusionKVConfig:
sink_chunks: int = 0
# Boundary reset vs. sliding replacement.
reset_at_boundary: bool = False
- # Fraction of free device memory budgeted for the AR-Diffusion KV pool.
+ # Fraction of free device memory used to admit additional resident
+ # sessions. One session is admitted whenever it fits actual free memory.
gpu_memory_fraction: float = 0.1
# When CUDA graph / torch.compile is on (not enforce_eager), pre-capture the
# DiT graphs for every window-fill shape at load time via a synthetic rollout,
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
index 9ef3cca0c..da61afe4d 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
@@ -174,6 +174,7 @@ class ARDiffusionKVCache:
device: torch.device | None = None,
frames_per_block: int = 1,
max_scratch_tokens_per_branch: int = 0,
+ model_owned_state_bytes_per_session: int = 0,
) -> None:
if not config.enable:
raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig")
@@ -197,13 +198,17 @@ class ARDiffusionKVCache:
self.config = config
self.kv_branches = kv_branches
- self.session_capacity = session_capacity
+ self.requested_session_capacity = session_capacity
self._kv_branch_local_indices = {kv_branch.name: kv_branch.local_index for kv_branch in kv_branches}
self.num_local_kv_branches = max(local_indices) + 1
if frames_per_block <= 0:
raise ValueError(f"frames_per_block must be positive, got {frames_per_block}")
if max_scratch_tokens_per_branch < 0:
raise ValueError(f"max_scratch_tokens_per_branch must be non-negative, got {max_scratch_tokens_per_branch}")
+ if model_owned_state_bytes_per_session < 0:
+ raise ValueError(
+ f"model_owned_state_bytes_per_session must be non-negative, got {model_owned_state_bytes_per_session}"
+ )
self.frames_per_block = int(frames_per_block)
self.block_size = block_size
self.num_layers = num_layers
@@ -217,28 +222,11 @@ class ARDiffusionKVCache:
self.device = device or torch.device("cpu")
self._allocate_tensors = device is not None
self._adapters: dict[str, ARDiffusionRequestAdapter] = {}
-
- # Named cross-attention K/V is allocated lazily per session and released
- # with that session. Reserve its worst-case capacity when sizing the
- # paged self-attention pool so the two stores share one memory budget.
self._cross_sessions: dict[
str,
dict[str, dict[str, tuple[list[torch.Tensor], list[torch.Tensor]]]],
] = {}
- def _cross_pool_bytes(length: int) -> int:
- return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers)
-
- cross_bytes_per_session = sum(_cross_pool_bytes(length) for length in self.cross_attention_lengths.values())
- cross_total_bytes = cross_bytes_per_session * session_capacity
- if cross_total_bytes:
- _log.info(
- "AR-Diffusion cross-attn reservation: %.1f MiB/session × %d sessions = %.1f MiB",
- cross_bytes_per_session / (1024 * 1024),
- session_capacity,
- cross_total_bytes / (1024 * 1024),
- )
-
self.spec = ChunkWindowSpec(
block_size=block_size,
num_kv_heads=num_kv_heads,
@@ -250,41 +238,7 @@ class ARDiffusionKVCache:
sink_chunks=config.sink_chunks,
reset_at_boundary=config.reset_at_boundary,
)
- # Each pool block spans all layers' K/V, so size against the per-layer
- # page size times the layer count.
- # Size the self-attn pool against memory reserved for the maximum number
- # of lazily allocated cross-attention sessions.
- num_blocks = compute_num_blocks(
- max(0, available_bytes - cross_total_bytes),
- config.gpu_memory_fraction,
- self.spec.page_size_bytes * num_layers,
- )
- # Floor: one forward needs the resident window plus the in-flight chunk
- # (frames_per_block frame-blocks) for every KV branch THIS rank runs,
- # with a little eviction-transient headroom. The memory-fraction
- # heuristic can under-size this once block_size grows — e.g. frame-granular
- # paging at the true frame_seqlen makes each block larger and the pool
- # fewer-blocks — so guarantee the minimum the rollout cannot run without,
- # otherwise allocate_chunk hits an exhausted pool mid-forward.
- resident_blocks = config.sink_chunks + config.window_chunks
- min_blocks = self.num_local_kv_branches * (resident_blocks + self.frames_per_block) + 2
- if num_blocks < min_blocks:
- _log.warning(
- "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
- "floor (%d local KV branch(es) x (sink_chunks=%d + window_chunks=%d "
- "+ frames_per_block=%d) + 2 headroom)",
- num_blocks,
- min_blocks,
- self.num_local_kv_branches,
- config.sink_chunks,
- config.window_chunks,
- self.frames_per_block,
- )
- num_blocks = min_blocks
- layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
- self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
- self.managed_num_blocks = num_blocks
- self.num_blocks = num_blocks
+
# Scratch blocks are outside KVCacheManager ownership. A non-committing
# forward needs one block per current frame plus space for any
# model-declared action/state tokens that coexist with video KV.
@@ -297,6 +251,101 @@ class ARDiffusionKVCache:
scratch_per_kv_branch = max(minimum_scratch_blocks, override_blocks)
self.scratch_blocks_per_kv_branch = scratch_per_kv_branch
self.scratch_num_blocks = self.num_local_kv_branches * scratch_per_kv_branch
+
+ # The self-attention pool, scratch pool, and lazily materialized
+ # cross-attention caches share one hard memory budget. Select the largest
+ # feasible resident-session count rather than raising a block-count floor
+ # past that budget. All resident sessions need a complete sink + window;
+ # only one request can be in flight, so frames_per_block is counted once.
+ page_size_bytes = self.spec.page_size_bytes * num_layers
+ self.available_memory_bytes = available_bytes
+ self.configured_memory_budget_bytes = int(available_bytes * config.gpu_memory_fraction)
+ # Reuse the public helper's validation for the memory fraction/page size.
+ compute_num_blocks(available_bytes, config.gpu_memory_fraction, page_size_bytes)
+ self.scratch_reserved_bytes = self.scratch_num_blocks * page_size_bytes
+ self.model_owned_state_bytes_per_session = model_owned_state_bytes_per_session
+
+ def _cross_pool_bytes(length: int) -> int:
+ return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers)
+
+ self.cross_attention_bytes_per_session = sum(
+ _cross_pool_bytes(length) for length in self.cross_attention_lengths.values()
+ )
+
+ def _required_managed_blocks(capacity: int) -> int:
+ resident_per_session = config.sink_chunks + config.window_chunks
+ return self.num_local_kv_branches * (capacity * resident_per_session + self.frames_per_block) + 2
+
+ def _required_bytes(capacity: int) -> int:
+ return (
+ self.scratch_reserved_bytes
+ + capacity * self.cross_attention_bytes_per_session
+ + capacity * self.model_owned_state_bytes_per_session
+ + _required_managed_blocks(capacity) * page_size_bytes
+ )
+
+ one_session_bytes = _required_bytes(1)
+ if one_session_bytes > available_bytes:
+ raise ValueError(
+ "AR-Diffusion available device memory cannot fit one session: "
+ f"available={available_bytes} bytes, required={one_session_bytes} bytes "
+ "(managed self-attention + cross-attention + scratch + model-owned state)."
+ )
+ self.memory_budget_bytes = max(self.configured_memory_budget_bytes, one_session_bytes)
+ if self.memory_budget_bytes > self.configured_memory_budget_bytes:
+ _log.warning(
+ "AR-Diffusion raised the configured memory budget from %d to %d bytes "
+ "to admit one session that fits actual free device memory",
+ self.configured_memory_budget_bytes,
+ self.memory_budget_bytes,
+ )
+
+ effective_capacity = 0
+ required_managed_blocks = 0
+ for candidate in range(session_capacity, 0, -1):
+ candidate_managed_blocks = _required_managed_blocks(candidate)
+ if _required_bytes(candidate) <= self.memory_budget_bytes:
+ effective_capacity = candidate
+ required_managed_blocks = candidate_managed_blocks
+ break
+ assert effective_capacity > 0
+
+ self.session_capacity = effective_capacity
+ self.cross_attention_reserved_bytes = self.cross_attention_bytes_per_session * effective_capacity
+ self.model_owned_state_reserved_bytes = self.model_owned_state_bytes_per_session * effective_capacity
+ self_attn_budget_bytes = (
+ self.memory_budget_bytes
+ - self.scratch_reserved_bytes
+ - self.cross_attention_reserved_bytes
+ - self.model_owned_state_reserved_bytes
+ )
+ num_blocks = self_attn_budget_bytes // page_size_bytes
+ assert num_blocks >= required_managed_blocks
+ if effective_capacity < session_capacity:
+ _log.warning(
+ "AR-Diffusion resident session capacity reduced from %d to %d by the KV memory budget",
+ session_capacity,
+ effective_capacity,
+ )
+ if self.cross_attention_reserved_bytes:
+ _log.info(
+ "AR-Diffusion cross-attn reservation: %.1f MiB/session × %d sessions = %.1f MiB",
+ self.cross_attention_bytes_per_session / (1024 * 1024),
+ effective_capacity,
+ self.cross_attention_reserved_bytes / (1024 * 1024),
+ )
+ if self.model_owned_state_reserved_bytes:
+ _log.info(
+ "AR-Diffusion model-owned state reservation: %.1f MiB/session × %d sessions = %.1f MiB",
+ self.model_owned_state_bytes_per_session / (1024 * 1024),
+ effective_capacity,
+ self.model_owned_state_reserved_bytes / (1024 * 1024),
+ )
+
+ layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
+ self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
+ self.managed_num_blocks = num_blocks
+ self.num_blocks = num_blocks
self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks
self.null_block_id = self.manager.block_pool.null_block.block_id
@@ -561,6 +610,11 @@ class ARDiffusionKVCache:
"""
_log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
adapter.on_chunk_committed()
+ self.manager.remove_skipped_blocks(
+ adapter.request_id,
+ adapter.num_computed_tokens,
+ num_prompt_tokens=adapter.num_prompt_tokens,
+ )
_log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
# -- pool-backed K/V access --------------------------------------------
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
index f16ecc9e9..b2e19a95f 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
@@ -10,19 +10,6 @@ import torch
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
-# Set by ARDiffusionPagedForwardContext.prepare() before each KV branch forward and
-# read by the fused write+attend custom op below. The pools are process-lifetime
-# allocations owned by one ARDiffusionKVCache per worker process (CFG-parallel /
-# TP ranks are separate processes), mirroring vLLM's forward-context pattern for
-# keeping multi-GiB cache tensors out of the compiled graph's inputs.
-_CURRENT_PAGED_KV_CACHE: Any = None
-
-
-def set_current_paged_kv_cache(kv_cache: Any) -> None:
- global _CURRENT_PAGED_KV_CACHE
- _CURRENT_PAGED_KV_CACHE = kv_cache
-
-
_LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
@@ -50,6 +37,9 @@ class ARDiffusionPagedLayerInputs(NamedTuple):
"""
layer_idx: torch.Tensor
+ key_pool: torch.Tensor
+ value_pool: torch.Tensor
+ block_size: int
seq_len: int
video_slots: torch.Tensor
action_slots: torch.Tensor
@@ -228,7 +218,6 @@ class ARDiffusionPagedForwardContext:
) = self.build_block_table(action_len=action_len, query_len=query_len, device=device)
if self.action_slot_mapping is None:
self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
- set_current_paged_kv_cache(self.kv_cache)
self._prepared = True
def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
@@ -236,6 +225,9 @@ class ARDiffusionPagedForwardContext:
raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()")
return ARDiffusionPagedLayerInputs(
layer_idx=_layer_idx_tensor(layer_idx),
+ key_pool=self.kv_cache._k_pools[layer_idx],
+ value_pool=self.kv_cache._v_pools[layer_idx],
+ block_size=int(self.kv_cache.block_size),
seq_len=int(self.seq_len),
video_slots=self.current_video_slot_mapping,
action_slots=self.action_slot_mapping,
@@ -476,17 +468,18 @@ def ar_diffusion_paged_attention(
# One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats
# it as a single graph node (no eager island, no graph breaks), and the K/V slot
# writes happen inside the op so write→read ordering with the block-table kernel
-# is internal — no hidden-mutation ordering hazards against a separate reader op
-# and no multi-GiB pool tensors as graph inputs (the flat and paged pool views
-# alias the same storage, which AOTAutograd handles poorly as inputs). Pattern
-# follows sage_attn3.py in-repo and vLLM's own unified attention ops.
+# is internal. The flat pools are explicit mutable inputs: Inductor/CUDA Graph
+# must track their storage lifetime instead of observing an undeclared mutation
+# through a process-global registry.
def _paged_write_attn_impl(
query: torch.Tensor,
k_curr: torch.Tensor,
v_curr: torch.Tensor,
k_act: torch.Tensor | None,
v_act: torch.Tensor | None,
- layer_idx: torch.Tensor,
+ key_pool: torch.Tensor,
+ value_pool: torch.Tensor,
+ block_size: int,
video_slots: torch.Tensor,
action_slots: torch.Tensor,
block_table: torch.Tensor,
@@ -496,21 +489,17 @@ def _paged_write_attn_impl(
max_seq_len: int,
softmax_scale: float,
) -> torch.Tensor:
- kv = _CURRENT_PAGED_KV_CACHE
- if kv is None:
- raise RuntimeError("ar_diffusion_paged_write_attn called before prepare() set the KV pool registry")
- layer_idx = int(layer_idx)
- k_pool = kv._k_pools[layer_idx]
- v_pool = kv._v_pools[layer_idx]
- k_pool[video_slots] = k_curr.to(k_pool.dtype)
- v_pool[video_slots] = v_curr.to(v_pool.dtype)
+ key_pool[video_slots] = k_curr.to(key_pool.dtype)
+ value_pool[video_slots] = v_curr.to(value_pool.dtype)
if k_act is not None and v_act is not None and k_act.shape[0] > 0:
- k_pool[action_slots] = k_act.to(k_pool.dtype)
- v_pool[action_slots] = v_act.to(v_pool.dtype)
+ key_pool[action_slots] = k_act.to(key_pool.dtype)
+ value_pool[action_slots] = v_act.to(value_pool.dtype)
+ key_cache = key_pool.unflatten(0, (-1, block_size))
+ value_cache = value_pool.unflatten(0, (-1, block_size))
return ar_diffusion_paged_attention(
query,
- kv.key_cache(layer_idx),
- kv.value_cache(layer_idx),
+ key_cache,
+ value_cache,
block_table=block_table,
query_start_loc=query_start_loc,
seq_lens=seq_lens,
@@ -525,14 +514,19 @@ def _paged_write_attn_impl(
# the module from sys.modules (same as sage_attn3.py).
if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
- @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=())
+ @torch.library.custom_op(
+ "vllm_omni::ar_diffusion_paged_write_attn",
+ mutates_args=("key_pool", "value_pool"),
+ )
def _paged_write_attn_op(
query: torch.Tensor,
k_curr: torch.Tensor,
v_curr: torch.Tensor,
k_act: torch.Tensor | None,
v_act: torch.Tensor | None,
- layer_idx: torch.Tensor,
+ key_pool: torch.Tensor,
+ value_pool: torch.Tensor,
+ block_size: int,
video_slots: torch.Tensor,
action_slots: torch.Tensor,
block_table: torch.Tensor,
@@ -548,7 +542,9 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
v_curr,
k_act,
v_act,
- layer_idx,
+ key_pool,
+ value_pool,
+ block_size,
video_slots,
action_slots,
block_table,
@@ -566,7 +562,9 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
v_curr,
k_act,
v_act,
- layer_idx,
+ key_pool,
+ value_pool,
+ block_size,
video_slots,
action_slots,
block_table,
@@ -589,7 +587,9 @@ def paged_write_attn(
v_curr,
k_act,
v_act,
- inputs.layer_idx,
+ inputs.key_pool,
+ inputs.value_pool,
+ inputs.block_size,
inputs.video_slots,
inputs.action_slots,
inputs.block_table,
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
index d114212bb..ec14a9177 100644
--- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -130,7 +130,7 @@ class ARDiffusionKVState:
if ctx.commit_current and ctx._allocated_video:
n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size
for _ in range(n_chunks):
- ctx.adapter.on_chunk_committed()
+ self.kv_cache.commit_chunk(ctx.adapter)
self._committed[kv_branch] += ctx.seq_len
_log.debug(
"AR-Diffusion COMMIT [%s] new_tokens=%d chunks=%d resident=%d/%d",
@@ -173,6 +173,17 @@ class ARDiffusionKVState:
for i in range(self.num_layers)
]
+ def clear_cross_attention(self) -> None:
+ """Invalidate all named cross-attention caches without dropping self-KV.
+
+ Prompt changes may preserve the autoregressive world history while
+ requiring fresh text projections. This operation deliberately leaves
+ branch adapters and their resident self-attention blocks untouched.
+ """
+ if self._closed:
+ raise RuntimeError(f"AR-Diffusion session {self.session_id!r} is closed")
+ self.kv_cache.release_cross_attention(self.session_id)
+
def close(self) -> None:
"""Release all self- and cross-attention storage owned by this session."""
if self._closed:
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
index 31a23eaf7..410c45fb1 100644
--- a/vllm_omni/experimental/ar_diffusion/runner.py
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -21,6 +21,7 @@ from vllm_omni.experimental.ar_diffusion.capability import (
from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
logger = init_logger(__name__)
@@ -51,7 +52,6 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
"""
_WARMUP_SID = "__ardiffusion_warmup__"
- _MAX_RESIDENT_SESSIONS = 1
def __init__(self, vllm_config: object, od_config: OmniDiffusionConfig, device: torch.device) -> None:
super().__init__(vllm_config, od_config, device)
@@ -123,7 +123,6 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
self.ar_diffusion_kv_config = config
self._ar_diffusion_capability = capability
self._ar_diffusion_kv_cache_spec = spec
- self._session_capacity = min(spec.session_capacity, self._MAX_RESIDENT_SESSIONS)
self.kv_cache = ARDiffusionKVCache(
config,
num_layers=spec.num_layers,
@@ -134,12 +133,14 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
max_model_len=spec.max_model_len,
available_bytes=self._available_memory_bytes() if available_bytes is None else available_bytes,
kv_branches=spec.kv_branches,
- session_capacity=self._session_capacity,
+ session_capacity=spec.session_capacity,
cross_attention_lengths=spec.cross_attention_lengths,
frames_per_block=spec.frames_per_block,
max_scratch_tokens_per_branch=spec.max_scratch_tokens_per_branch,
+ model_owned_state_bytes_per_session=spec.model_owned_state_bytes_per_session,
device=self.device,
)
+ self._session_capacity = self.kv_cache.session_capacity
logger.info(
"AR-Diffusion KV cache: blocks=%d layers=%d local_kv_heads=%d head_size=%d "
"tokens/frame=%d frames/block=%d window=%d sink=%d kv_branches=%s cross=%s "
@@ -230,9 +231,14 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
return state
@staticmethod
- def _request_session(req: OmniDiffusionRequest) -> tuple[str, dict]:
+ def _request_session(
+ req: OmniDiffusionRequest,
+ ) -> tuple[str, dict, ARDiffusionTickRequest | None]:
extra_args = req.sampling_params.extra_args or {}
- return str(extra_args.get("session_id") or "default"), extra_args
+ tick = ARDiffusionTickRequest.from_extra_args(extra_args)
+ if tick is not None:
+ return tick.session_id, extra_args, tick
+ return str(extra_args.get("session_id") or "default"), extra_args, None
def execute_model(
self,
@@ -245,8 +251,10 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
if capability is None:
raise RuntimeError("AR-Diffusion capability missing after KV cache initialization")
- session_id, extra_args = self._request_session(req)
- if extra_args.get("reset", False):
+ session_id, extra_args, tick = self._request_session(req)
+ reset = tick.reset if tick is not None else bool(extra_args.get("reset", False))
+ close_session = tick.close_session if tick is not None else bool(extra_args.get("close_session", False))
+ if reset:
self.reset_session(session_id)
state = self._get_or_create_session(session_id)
started = time.perf_counter()
@@ -268,7 +276,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
)
raise
self._perf_e2e_times.append(time.perf_counter() - started)
- if extra_args.get("close_session", False):
+ if close_session:
self.close_session(session_id)
return output
diff --git a/vllm_omni/experimental/ar_diffusion/session.py b/vllm_omni/experimental/ar_diffusion/session.py
new file mode 100644
index 000000000..6bb4b27d9
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/session.py
@@ -0,0 +1,578 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Internal realtime session control plane for AR-Diffusion.
+
+This module owns transport-level identity, event ordering, backpressure, and
+chunk-boundary snapshots. Model-defined control payloads remain opaque here.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from collections import OrderedDict
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Generic, Protocol, TypeVar
+
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+
+class ARDiffusionSessionError(RuntimeError):
+ """Base error for realtime AR-Diffusion session operations."""
+
+
+class ARDiffusionSessionStateError(ARDiffusionSessionError):
+ """Raised when an operation is invalid for the current session state."""
+
+
+class ARDiffusionEventQueueFullError(ARDiffusionSessionError):
+ """Raised when accepting an event would exceed the bounded queue."""
+
+
+class ARDiffusionStaleEventError(ARDiffusionSessionError):
+ """Raised when an event is older than the current monotonic boundary."""
+
+
+class ARDiffusionChunkMetadataError(ARDiffusionSessionError):
+ """Raised when a tick consumer returns metadata for a different chunk."""
+
+
+class ARDiffusionEventAcceptanceStatus(str, Enum):
+ ACCEPTED = "accepted"
+ DUPLICATE = "duplicate"
+
+
+class ARDiffusionSessionStatus(str, Enum):
+ ACTIVE = "active"
+ RESETTING = "resetting"
+ FAILED = "failed"
+ CLOSING = "closing"
+ CLEANUP_FAILED = "cleanup_failed"
+ CLOSED = "closed"
+
+
+def _copy_control(control: ARDiffusionControlInput) -> ARDiffusionControlInput:
+ return ARDiffusionControlInput(
+ track=control.track,
+ schema=control.schema,
+ data=control.data,
+ )
+
+
+@dataclass(frozen=True)
+class ARDiffusionSessionEvent:
+ """One atomically accepted prompt/control update.
+
+ ``prompt=None`` means that this event does not update the prompt. Controls
+ are keyed by ``track`` when the chunk snapshot is assembled; their schema
+ and data are otherwise opaque to this layer.
+ """
+
+ event_id: int
+ prompt: str | None = None
+ controls: tuple[ARDiffusionControlInput, ...] = ()
+
+ def __post_init__(self) -> None:
+ if isinstance(self.event_id, bool) or not isinstance(self.event_id, int) or self.event_id < 0:
+ raise ValueError("event_id must be a non-negative integer.")
+ if self.prompt is not None and not isinstance(self.prompt, str):
+ raise ValueError("prompt must be a string or None.")
+ controls = tuple(self.controls)
+ if any(not isinstance(control, ARDiffusionControlInput) for control in controls):
+ raise ValueError("controls must contain ARDiffusionControlInput values.")
+ tracks = [control.track for control in controls]
+ if len(tracks) != len(set(tracks)):
+ raise ValueError("A composite event may update each control track at most once.")
+ if self.prompt is None and not controls:
+ raise ValueError("An event must update the prompt, one or more controls, or both.")
+ object.__setattr__(self, "controls", controls)
+
+
+@dataclass(frozen=True)
+class ARDiffusionEventAcceptance:
+ """Immediate queue acceptance result; this does not mean applied."""
+
+ session_id: str
+ event_id: int
+ status: ARDiffusionEventAcceptanceStatus
+ pending_event_count: int
+
+
+TickOutputT = TypeVar("TickOutputT")
+
+
+class ARDiffusionTickConsumer(Protocol[TickOutputT]):
+ """Consumes one immutable chunk-boundary tick."""
+
+ async def execute_tick(self, tick: ARDiffusionTickRequest) -> TickOutputT: ...
+
+ def chunk_metadata(self, output: TickOutputT) -> ARDiffusionChunkMetadata: ...
+
+
+@dataclass(frozen=True)
+class ARDiffusionPreparedControls:
+ """One reducer-owned, transactional control snapshot.
+
+ ``private_state`` is opaque to the session layer and becomes durable only
+ after the corresponding model tick succeeds.
+ """
+
+ controls: tuple[ARDiffusionControlInput, ...]
+ private_state: Any = None
+
+
+class ARDiffusionControlReducer(Protocol):
+ """Optional model adapter for stateful controls sampled per AR block."""
+
+ def prepare(
+ self,
+ *,
+ current_controls: Mapping[str, ARDiffusionControlInput],
+ events: Sequence[ARDiffusionSessionEvent],
+ chunk_index: int,
+ ) -> ARDiffusionPreparedControls: ...
+
+ def commit(self, prepared: ARDiffusionPreparedControls) -> None: ...
+
+ def reset(self) -> None: ...
+
+
+class ARDiffusionSessionLifecycle(Protocol):
+ """Worker-facing session lifecycle operations."""
+
+ async def reset_session(self, session_id: str) -> None: ...
+
+ async def close_session(self, session_id: str) -> None: ...
+
+
+class ARDiffusionCollectiveRPCClient(Protocol):
+ """Subset of the engine control RPC surface used by session lifecycle."""
+
+ async def collective_rpc(
+ self,
+ method: str,
+ timeout: float | None = None,
+ args: tuple[Any, ...] = (),
+ kwargs: dict[str, Any] | None = None,
+ stage_ids: list[int] | None = None,
+ ) -> list[Any]: ...
+
+
+def _collect_rpc_support_results(value: Any, *, errors: list[str], supported: list[bool]) -> None:
+ if isinstance(value, bool):
+ supported.append(value)
+ return
+ if isinstance(value, (list, tuple)):
+ for item in value:
+ _collect_rpc_support_results(item, errors=errors, supported=supported)
+ return
+ if isinstance(value, Mapping):
+ if value.get("supported") is False:
+ errors.append(str(value.get("error") or "target stage does not support this lifecycle RPC"))
+ return
+ if "result" in value:
+ _collect_rpc_support_results(value["result"], errors=errors, supported=supported)
+
+
+class ARDiffusionWorkerLifecycle:
+ """Routes reset/close through the existing orchestrator collective RPC."""
+
+ def __init__(
+ self,
+ rpc_client: ARDiffusionCollectiveRPCClient,
+ *,
+ stage_ids: Sequence[int],
+ timeout: float | None = None,
+ ) -> None:
+ normalized_stage_ids = tuple(stage_ids)
+ if not normalized_stage_ids:
+ raise ValueError("stage_ids must select at least one AR-Diffusion stage.")
+ if any(isinstance(stage_id, bool) or not isinstance(stage_id, int) or stage_id < 0 for stage_id in stage_ids):
+ raise ValueError("stage_ids must contain non-negative integers.")
+ self._rpc_client = rpc_client
+ self._stage_ids = normalized_stage_ids
+ self._timeout = timeout
+
+ async def _invoke(self, method: str, session_id: str) -> None:
+ results = await self._rpc_client.collective_rpc(
+ method=method,
+ timeout=self._timeout,
+ args=(session_id,),
+ stage_ids=list(self._stage_ids),
+ )
+ errors: list[str] = []
+ supported: list[bool] = []
+ _collect_rpc_support_results(results, errors=errors, supported=supported)
+ if errors:
+ raise ARDiffusionSessionError(f"{method} failed: {'; '.join(errors)}")
+ if not supported or not all(supported):
+ raise ARDiffusionSessionError(f"{method} is not supported by every selected worker.")
+
+ async def reset_session(self, session_id: str) -> None:
+ await self._invoke("reset_ar_diffusion_session", session_id)
+
+ async def close_session(self, session_id: str) -> None:
+ await self._invoke("close_ar_diffusion_session", session_id)
+
+
+def _default_request_id(session_id: str, chunk_index: int) -> str:
+ return f"ar-{session_id}-{chunk_index}-{uuid.uuid4().hex}"
+
+
+class ARDiffusionSession(Generic[TickOutputT]):
+ """Transactional event queue and chunk snapshot for one realtime world."""
+
+ def __init__(
+ self,
+ session_id: str,
+ *,
+ tick_consumer: ARDiffusionTickConsumer[TickOutputT],
+ lifecycle: ARDiffusionSessionLifecycle,
+ max_pending_events: int,
+ request_id_factory: Callable[[str, int], str] | None = None,
+ control_reducer: ARDiffusionControlReducer | None = None,
+ ) -> None:
+ if not isinstance(session_id, str) or not session_id.strip():
+ raise ValueError("session_id must be a non-empty string.")
+ if isinstance(max_pending_events, bool) or not isinstance(max_pending_events, int) or max_pending_events <= 0:
+ raise ValueError("max_pending_events must be a positive integer.")
+ self.session_id = session_id
+ self._tick_consumer = tick_consumer
+ self._lifecycle = lifecycle
+ self._max_pending_events = max_pending_events
+ self._request_id_factory = request_id_factory or _default_request_id
+ self._control_reducer = control_reducer
+
+ self._status = ARDiffusionSessionStatus.ACTIVE
+ self._pending_events: dict[int, ARDiffusionSessionEvent] = {}
+ self._recent_event_ids: OrderedDict[int, None] = OrderedDict()
+ self._event_floor = -1
+ self._inflight_event_floor: int | None = None
+ self._chunk_index = 0
+ self._prompt: str | None = None
+ self._controls: dict[str, ARDiffusionControlInput] = {}
+
+ self._state_lock = asyncio.Lock()
+ self._tick_lock = asyncio.Lock()
+
+ @property
+ def status(self) -> ARDiffusionSessionStatus:
+ return self._status
+
+ @property
+ def chunk_index(self) -> int:
+ return self._chunk_index
+
+ @property
+ def pending_event_count(self) -> int:
+ return len(self._pending_events)
+
+ def _require_active(self) -> None:
+ if self._status is not ARDiffusionSessionStatus.ACTIVE:
+ raise ARDiffusionSessionStateError(f"AR-Diffusion session {self.session_id!r} is {self._status.value}.")
+
+ def _require_resettable(self) -> None:
+ if self._status not in {
+ ARDiffusionSessionStatus.ACTIVE,
+ ARDiffusionSessionStatus.FAILED,
+ }:
+ raise ARDiffusionSessionStateError(f"AR-Diffusion session {self.session_id!r} is {self._status.value}.")
+
+ async def accept_event(self, event: ARDiffusionSessionEvent) -> ARDiffusionEventAcceptance:
+ """Atomically enqueue an event or return an idempotent duplicate result."""
+ if not isinstance(event, ARDiffusionSessionEvent):
+ raise TypeError("event must be an ARDiffusionSessionEvent.")
+ async with self._state_lock:
+ self._require_active()
+ if event.event_id in self._pending_events:
+ return ARDiffusionEventAcceptance(
+ session_id=self.session_id,
+ event_id=event.event_id,
+ status=ARDiffusionEventAcceptanceStatus.DUPLICATE,
+ pending_event_count=len(self._pending_events),
+ )
+ if event.event_id in self._recent_event_ids:
+ return ARDiffusionEventAcceptance(
+ session_id=self.session_id,
+ event_id=event.event_id,
+ status=ARDiffusionEventAcceptanceStatus.DUPLICATE,
+ pending_event_count=len(self._pending_events),
+ )
+
+ event_floor = self._event_floor
+ if self._inflight_event_floor is not None:
+ event_floor = max(event_floor, self._inflight_event_floor)
+ if event.event_id <= event_floor:
+ raise ARDiffusionStaleEventError(
+ f"event_id {event.event_id} is stale; the session event floor is {event_floor}."
+ )
+ if len(self._pending_events) >= self._max_pending_events:
+ raise ARDiffusionEventQueueFullError(
+ f"AR-Diffusion session {self.session_id!r} event queue is full "
+ f"({self._max_pending_events} pending)."
+ )
+
+ stored_event = ARDiffusionSessionEvent(
+ event_id=event.event_id,
+ prompt=event.prompt,
+ controls=tuple(_copy_control(control) for control in event.controls),
+ )
+ self._pending_events[event.event_id] = stored_event
+ self._recent_event_ids[event.event_id] = None
+ while len(self._recent_event_ids) > self._max_pending_events:
+ self._recent_event_ids.popitem(last=False)
+ return ARDiffusionEventAcceptance(
+ session_id=self.session_id,
+ event_id=event.event_id,
+ status=ARDiffusionEventAcceptanceStatus.ACCEPTED,
+ pending_event_count=len(self._pending_events),
+ )
+
+ def _build_tick_locked(
+ self,
+ ) -> tuple[
+ ARDiffusionTickRequest,
+ str | None,
+ dict[str, ARDiffusionControlInput],
+ ARDiffusionPreparedControls | None,
+ ]:
+ event_ids = tuple(sorted(self._pending_events))
+ events = tuple(self._pending_events[event_id] for event_id in event_ids)
+ prompt = self._prompt
+ for event in events:
+ if event.prompt is not None:
+ prompt = event.prompt
+
+ prepared_controls: ARDiffusionPreparedControls | None = None
+ if self._control_reducer is None:
+ controls = {track: _copy_control(control) for track, control in self._controls.items()}
+ for event in events:
+ for control in event.controls:
+ controls[control.track] = _copy_control(control)
+ else:
+ prepared_controls = self._control_reducer.prepare(
+ current_controls={track: _copy_control(control) for track, control in self._controls.items()},
+ events=events,
+ chunk_index=self._chunk_index,
+ )
+ if not isinstance(prepared_controls, ARDiffusionPreparedControls):
+ raise TypeError("control reducer prepare() must return ARDiffusionPreparedControls.")
+ controls = {}
+ for control in prepared_controls.controls:
+ if not isinstance(control, ARDiffusionControlInput):
+ raise TypeError("prepared controls must contain ARDiffusionControlInput values.")
+ if control.track in controls:
+ raise ValueError("prepared controls must contain at most one value per track.")
+ controls[control.track] = _copy_control(control)
+
+ chunk_index = self._chunk_index
+ request_id = self._request_id_factory(self.session_id, chunk_index)
+ tick = ARDiffusionTickRequest(
+ session_id=self.session_id,
+ request_id=request_id,
+ chunk_index=chunk_index,
+ applied_event_ids=event_ids,
+ prompt=prompt,
+ controls=tuple(_copy_control(controls[track]) for track in sorted(controls)),
+ )
+ self._inflight_event_floor = event_ids[-1] if event_ids else None
+ return tick, prompt, controls, prepared_controls
+
+ async def _commit_tick(
+ self,
+ tick: ARDiffusionTickRequest,
+ prompt: str | None,
+ controls: dict[str, ARDiffusionControlInput],
+ prepared_controls: ARDiffusionPreparedControls | None,
+ ) -> None:
+ async with self._state_lock:
+ # Reducer state and the session snapshot form one logical commit.
+ # Run the only fallible operation before mutating session fields so
+ # a reducer failure cannot acknowledge events or advance the chunk.
+ if prepared_controls is not None:
+ assert self._control_reducer is not None
+ self._control_reducer.commit(prepared_controls)
+ for event_id in tick.applied_event_ids:
+ self._pending_events.pop(event_id, None)
+ if tick.applied_event_ids:
+ self._event_floor = tick.applied_event_ids[-1]
+ self._prompt = prompt
+ self._controls = controls
+ self._chunk_index += 1
+ self._inflight_event_floor = None
+
+ async def _fail_tick(self, tick_error: BaseException) -> None:
+ async with self._state_lock:
+ self._inflight_event_floor = None
+ self._status = ARDiffusionSessionStatus.FAILED
+ try:
+ # A model forward may already have released its state. This
+ # idempotent close also covers successful forwards whose metadata
+ # or local reducer commit could not be accepted.
+ await self._lifecycle.close_session(self.session_id)
+ except Exception as cleanup_error:
+ raise ARDiffusionSessionError(
+ "AR-Diffusion tick failed "
+ f"({type(tick_error).__name__}: {tick_error}) and worker "
+ f"session cleanup also failed: {cleanup_error}"
+ ) from tick_error
+
+ async def next_chunk(self) -> TickOutputT:
+ """Execute and atomically commit one chunk-boundary snapshot."""
+ async with self._tick_lock:
+ async with self._state_lock:
+ self._require_active()
+ tick, prompt, controls, prepared_controls = self._build_tick_locked()
+
+ expected_metadata = ARDiffusionChunkMetadata.from_tick(tick)
+ try:
+ output = await self._tick_consumer.execute_tick(tick)
+ metadata = self._tick_consumer.chunk_metadata(output)
+ if not isinstance(metadata, ARDiffusionChunkMetadata) or metadata != expected_metadata:
+ raise ARDiffusionChunkMetadataError(
+ "Tick consumer metadata does not match the submitted AR-Diffusion tick."
+ )
+ except BaseException as tick_error:
+ await self._fail_tick(tick_error)
+ raise
+
+ commit_task = asyncio.create_task(
+ self._commit_tick(
+ tick,
+ prompt,
+ controls,
+ prepared_controls,
+ )
+ )
+ try:
+ await asyncio.shield(commit_task)
+ except asyncio.CancelledError:
+ # Once a consumer reports success, the local chunk/event state
+ # must commit even if the transport task is being cancelled.
+ try:
+ await commit_task
+ except BaseException as commit_error:
+ await self._fail_tick(commit_error)
+ raise
+ raise
+ except BaseException as commit_error:
+ await self._fail_tick(commit_error)
+ raise
+ return output
+
+ async def reset(self) -> None:
+ """Reset worker state and discard queued/current transport state."""
+ async with self._tick_lock:
+ async with self._state_lock:
+ self._require_resettable()
+ self._status = ARDiffusionSessionStatus.RESETTING
+ try:
+ await self._lifecycle.reset_session(self.session_id)
+ except BaseException:
+ async with self._state_lock:
+ self._status = ARDiffusionSessionStatus.FAILED
+ raise
+ async with self._state_lock:
+ if self._pending_events:
+ self._event_floor = max(self._event_floor, max(self._pending_events))
+ self._pending_events.clear()
+ self._prompt = None
+ self._controls.clear()
+ if self._control_reducer is not None:
+ self._control_reducer.reset()
+ self._inflight_event_floor = None
+ self._chunk_index = 0
+ self._status = ARDiffusionSessionStatus.ACTIVE
+
+ async def close(self) -> None:
+ """Close worker state, retaining a tombstone until cleanup succeeds."""
+ async with self._tick_lock:
+ async with self._state_lock:
+ if self._status is ARDiffusionSessionStatus.CLOSED:
+ return
+ self._status = ARDiffusionSessionStatus.CLOSING
+ try:
+ await self._lifecycle.close_session(self.session_id)
+ except BaseException:
+ async with self._state_lock:
+ self._status = ARDiffusionSessionStatus.CLEANUP_FAILED
+ raise
+ else:
+ async with self._state_lock:
+ self._pending_events.clear()
+ self._prompt = None
+ self._controls.clear()
+ self._inflight_event_floor = None
+ self._status = ARDiffusionSessionStatus.CLOSED
+
+ async def disconnect(self) -> None:
+ """Disconnect uses the same owning-worker cleanup as explicit close."""
+ await self.close()
+
+
+class ARDiffusionSessionManager(Generic[TickOutputT]):
+ """Creates and owns transport sessions sharing one tick/lifecycle backend."""
+
+ def __init__(
+ self,
+ *,
+ tick_consumer: ARDiffusionTickConsumer[TickOutputT],
+ lifecycle: ARDiffusionSessionLifecycle,
+ max_pending_events: int,
+ request_id_factory: Callable[[str, int], str] | None = None,
+ session_id_factory: Callable[[], str] | None = None,
+ control_reducer_factory: Callable[[], ARDiffusionControlReducer] | None = None,
+ ) -> None:
+ self._tick_consumer = tick_consumer
+ self._lifecycle = lifecycle
+ self._max_pending_events = max_pending_events
+ self._request_id_factory = request_id_factory
+ self._session_id_factory = session_id_factory or (lambda: uuid.uuid4().hex)
+ self._control_reducer_factory = control_reducer_factory
+ self._sessions: dict[str, ARDiffusionSession[TickOutputT]] = {}
+ self._lock = asyncio.Lock()
+
+ async def create_session(self, session_id: str | None = None) -> ARDiffusionSession[TickOutputT]:
+ resolved_session_id = session_id or self._session_id_factory()
+ async with self._lock:
+ if resolved_session_id in self._sessions:
+ raise ARDiffusionSessionStateError(f"AR-Diffusion session {resolved_session_id!r} already exists.")
+ session = ARDiffusionSession(
+ resolved_session_id,
+ tick_consumer=self._tick_consumer,
+ lifecycle=self._lifecycle,
+ max_pending_events=self._max_pending_events,
+ request_id_factory=self._request_id_factory,
+ control_reducer=(
+ self._control_reducer_factory() if self._control_reducer_factory is not None else None
+ ),
+ )
+ self._sessions[resolved_session_id] = session
+ return session
+
+ async def get_session(self, session_id: str) -> ARDiffusionSession[TickOutputT]:
+ async with self._lock:
+ try:
+ return self._sessions[session_id]
+ except KeyError as exc:
+ raise ARDiffusionSessionStateError(f"AR-Diffusion session {session_id!r} does not exist.") from exc
+
+ async def reset_session(self, session_id: str) -> None:
+ session = await self.get_session(session_id)
+ await session.reset()
+
+ async def close_session(self, session_id: str) -> None:
+ session = await self.get_session(session_id)
+ await session.close()
+ async with self._lock:
+ if self._sessions.get(session_id) is session:
+ self._sessions.pop(session_id)
+
+ async def disconnect(self, session_id: str) -> None:
+ """Release the owning worker route when its transport disconnects."""
+ await self.close_session(session_id)
diff --git a/vllm_omni/experimental/ar_diffusion/tick_protocol.py b/vllm_omni/experimental/ar_diffusion/tick_protocol.py
new file mode 100644
index 000000000..aa8315118
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/tick_protocol.py
@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Typed request boundary shared by AR-Diffusion models and serving APIs.
+
+The transport layer owns event ordering and snapshots. Model adapters consume
+the resulting immutable tick without learning about WebSocket messages.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Any
+
+AR_DIFFUSION_TICK_KEY = "ar_diffusion_tick"
+
+
+def _non_empty_string(value: object, *, field: str) -> str:
+ if not isinstance(value, str) or not value.strip():
+ raise ValueError(f"{field} must be a non-empty string.")
+ return value
+
+
+def _non_negative_int(value: object, *, field: str) -> int:
+ if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+ raise ValueError(f"{field} must be a non-negative integer.")
+ return value
+
+
+def _freeze_transport_value(value: Any, *, path: str) -> Any:
+ """Copy JSON-like transport data into recursively immutable containers."""
+ if value is None or isinstance(value, (str, bool, int, float)):
+ return value
+ if isinstance(value, Mapping):
+ frozen: dict[str, Any] = {}
+ for key, item in value.items():
+ if not isinstance(key, str):
+ raise ValueError(f"{path} mapping keys must be strings.")
+ frozen[key] = _freeze_transport_value(item, path=f"{path}.{key}")
+ return MappingProxyType(frozen)
+ if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
+ return tuple(_freeze_transport_value(item, path=f"{path}[]") for item in value)
+ raise ValueError(
+ f"{path} must contain only transport-safe mappings, sequences, and scalar values; got {type(value).__name__}."
+ )
+
+
+def _thaw_transport_value(value: Any) -> Any:
+ """Return a detached JSON-serializable copy of an immutable snapshot."""
+ if isinstance(value, Mapping):
+ return {key: _thaw_transport_value(item) for key, item in value.items()}
+ if isinstance(value, tuple):
+ return [_thaw_transport_value(item) for item in value]
+ return value
+
+
+@dataclass(frozen=True)
+class ARDiffusionControlInput:
+ """One model-defined control track in a chunk-boundary snapshot."""
+
+ track: str
+ schema: str
+ data: Mapping[str, Any]
+
+ def __post_init__(self) -> None:
+ _non_empty_string(self.track, field="control.track")
+ _non_empty_string(self.schema, field="control.schema")
+ if not isinstance(self.data, Mapping):
+ raise ValueError("control.data must be a mapping.")
+ object.__setattr__(
+ self,
+ "data",
+ _freeze_transport_value(self.data, path=f"control[{self.track!r}].data"),
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ return {
+ "track": self.track,
+ "schema": self.schema,
+ "data": _thaw_transport_value(self.data),
+ }
+
+ @classmethod
+ def from_dict(cls, value: object) -> ARDiffusionControlInput:
+ if not isinstance(value, Mapping):
+ raise ValueError("Each control must be a mapping.")
+ return cls(
+ track=value.get("track"), # type: ignore[arg-type]
+ schema=value.get("schema"), # type: ignore[arg-type]
+ data=value.get("data"), # type: ignore[arg-type]
+ )
+
+
+@dataclass(frozen=True)
+class ARDiffusionTickRequest:
+ """A single atomic AR-block computation.
+
+ ``request_id`` identifies this computation, ``session_id`` identifies the
+ persistent rollout, and ``applied_event_ids`` records which accepted
+ interaction events are represented by this immutable snapshot.
+ """
+
+ session_id: str
+ request_id: str
+ chunk_index: int
+ applied_event_ids: tuple[int, ...] = ()
+ prompt: str | None = None
+ controls: tuple[ARDiffusionControlInput, ...] = ()
+ reset: bool = False
+ close_session: bool = False
+
+ def __post_init__(self) -> None:
+ _non_empty_string(self.session_id, field="session_id")
+ _non_empty_string(self.request_id, field="request_id")
+ _non_negative_int(self.chunk_index, field="chunk_index")
+ if self.prompt is not None and not isinstance(self.prompt, str):
+ raise ValueError("prompt must be a string or None.")
+ if not isinstance(self.reset, bool):
+ raise ValueError("reset must be a boolean.")
+ if not isinstance(self.close_session, bool):
+ raise ValueError("close_session must be a boolean.")
+
+ event_ids = tuple(
+ _non_negative_int(event_id, field="applied_event_ids item") for event_id in self.applied_event_ids
+ )
+ if event_ids != tuple(sorted(set(event_ids))):
+ raise ValueError("applied_event_ids must be unique and strictly increasing.")
+ controls = tuple(self.controls)
+ if any(not isinstance(control, ARDiffusionControlInput) for control in controls):
+ raise ValueError("controls must contain ARDiffusionControlInput values.")
+ object.__setattr__(self, "applied_event_ids", event_ids)
+ object.__setattr__(self, "controls", controls)
+
+ def to_extra_args(self) -> dict[str, Any]:
+ """Serialize the tick under one namespaced ``extra_args`` key."""
+ return {
+ AR_DIFFUSION_TICK_KEY: {
+ "session_id": self.session_id,
+ "request_id": self.request_id,
+ "chunk_index": self.chunk_index,
+ "applied_event_ids": list(self.applied_event_ids),
+ "prompt": self.prompt,
+ "controls": [control.to_dict() for control in self.controls],
+ "reset": self.reset,
+ "close_session": self.close_session,
+ }
+ }
+
+ @classmethod
+ def from_extra_args(
+ cls,
+ extra_args: Mapping[str, Any] | None,
+ ) -> ARDiffusionTickRequest | None:
+ """Parse a typed tick, returning ``None`` for legacy AR requests."""
+ if extra_args is None or AR_DIFFUSION_TICK_KEY not in extra_args:
+ return None
+ value = extra_args[AR_DIFFUSION_TICK_KEY]
+ if not isinstance(value, Mapping):
+ raise ValueError(f"{AR_DIFFUSION_TICK_KEY} must be a mapping.")
+ raw_event_ids = value.get("applied_event_ids", ())
+ raw_controls = value.get("controls", ())
+ if not isinstance(raw_event_ids, Sequence) or isinstance(raw_event_ids, (str, bytes)):
+ raise ValueError("applied_event_ids must be a sequence.")
+ if not isinstance(raw_controls, Sequence) or isinstance(raw_controls, (str, bytes)):
+ raise ValueError("controls must be a sequence.")
+ return cls(
+ session_id=value.get("session_id"), # type: ignore[arg-type]
+ request_id=value.get("request_id"), # type: ignore[arg-type]
+ chunk_index=value.get("chunk_index"), # type: ignore[arg-type]
+ applied_event_ids=tuple(raw_event_ids),
+ prompt=value.get("prompt"), # type: ignore[arg-type]
+ controls=tuple(ARDiffusionControlInput.from_dict(control) for control in raw_controls),
+ reset=value.get("reset", False), # type: ignore[arg-type]
+ close_session=value.get("close_session", False), # type: ignore[arg-type]
+ )
+
+
+@dataclass(frozen=True)
+class ARDiffusionChunkMetadata:
+ """Identity metadata returned with one generated AR block."""
+
+ session_id: str
+ request_id: str
+ chunk_index: int
+ applied_event_ids: tuple[int, ...] = ()
+
+ @classmethod
+ def from_tick(cls, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+ return cls(
+ session_id=tick.session_id,
+ request_id=tick.request_id,
+ chunk_index=tick.chunk_index,
+ applied_event_ids=tick.applied_event_ids,
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ return {
+ "session_id": self.session_id,
+ "request_id": self.request_id,
+ "chunk_index": self.chunk_index,
+ "applied_event_ids": list(self.applied_event_ids),
+ }