Add model selection

Fixes the sequence parallel sharding on t, now SP shards on t*h*w

.gitignore

@@ -40,6 +40,9 @@ dist/
 *.egg
 eggs/
 .eggs/
+node_modules/
+.vite
+vite.config.js.timestamp-*.mjs
 
 # MkDocs documentation
 site/

docs/attention/index.md

@@ -6,6 +6,8 @@ FastVideo provides highly optimized custom attention kernels to accelerate video
 
 * **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks.
 * **[Sliding Tile Attention (STA)](sta/index.md)**: Optimized attention for window-based video generation.
+* **Backend development guide**: See the developer guide at
+  [Attention Backend Development](../contributing/attention_backend.md).
 
 ## General Build Instructions
 

docs/contributing/attention_backend.md

@@ -0,0 +1,176 @@
+# Attention Backend Development
+
+This guide is for contributors adding a new attention backend (kernel or
+implementation) to FastVideo. If you just want to use existing kernels or build
+fastvideo-kernel, see [Attention overview](../attention/index.md).
+
+## When you need this guide
+
+Use this guide when you are:
+
+- Adding a new attention kernel or algorithm.
+- Wiring an existing kernel into FastVideo's attention selection.
+- Extending attention support to a new platform.
+
+## 0) Choose a backend name and scope
+
+Pick a backend name in `UPPER_SNAKE_CASE` and decide where it should run.
+Example: `MY_NEW_ATTN`.
+
+You will use this name in:
+
+- `AttentionBackendEnum` (global list of backends).
+- `get_name()` in your backend class (match the enum name).
+- Platform selectors (CUDA/ROCm/MPS/NPU) to return your backend class.
+
+## 1) Add enum + platform selection
+
+1) Add your backend to `fastvideo/platforms/interface.py`:
+
+```python
+class AttentionBackendEnum(enum.Enum):
+    ...
+    MY_NEW_ATTN = enum.auto()
+```
+
+1) Register it in platform selection (example: CUDA). Update
+`fastvideo/platforms/cuda.py` inside `get_attn_backend_cls`:
+
+```python
+elif selected_backend == AttentionBackendEnum.MY_NEW_ATTN:
+    try:
+        from fastvideo.attention.backends.my_new_attn import MyNewAttnBackend
+        return "fastvideo.attention.backends.my_new_attn.MyNewAttnBackend"
+    except ImportError as e:
+        logger.error("Failed to import MY_NEW_ATTN backend: %s", str(e))
+        raise
+```
+
+If you want support on other platforms, add a similar branch in
+`fastvideo/platforms/rocm.py`, `fastvideo/platforms/mps.py`, or `fastvideo/platforms/npu.py`.
+
+## 2) Implement the backend
+
+Create `fastvideo/attention/backends/my_new_attn.py` and implement the required
+classes.
+
+Minimal skeleton (no custom metadata):
+
+```python
+import torch
+from dataclasses import dataclass
+from fastvideo.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+
+class MyNewAttnBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "MY_NEW_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["MyNewAttnImpl"]:
+        return MyNewAttnImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return MyNewAttnMetadataBuilder
+
+
+@dataclass
+class MyNewAttnMetadata(AttentionMetadata):
+    current_timestep: int
+
+
+class MyNewAttnMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(self, current_timestep: int, **kwargs):
+        return MyNewAttnMetadata(current_timestep=current_timestep)
+
+
+class MyNewAttnImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.softmax_scale = softmax_scale
+        self.causal = causal
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: MyNewAttnMetadata,
+    ) -> torch.Tensor:
+        # Implement attention
+        return torch.nn.functional.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            is_causal=self.causal,
+            scale=self.softmax_scale,
+        ).transpose(1, 2)
+```
+
+Optional:
+
+- Implement `preprocess_qkv` / `postprocess_output` if your kernel needs tiling
+  or reshaping.
+- Use `fastvideo.forward_context.get_forward_context()` if you need dynamic
+  per-step data (e.g., window sizes).
+- Set `accept_output_buffer = True` if your backend writes into a provided
+  output buffer.
+
+## 3) Wire into attention layers
+
+Backends are used by `LocalAttention` and `DistributedAttention`. These layers
+accept a `supported_attention_backends` tuple. If your backend should be
+eligible, update the call sites that construct these layers (search for
+`supported_attention_backends=`).
+
+## 4) Add compiled kernels (optional)
+
+If you have a custom CUDA kernel:
+
+1) Add sources in `fastvideo-kernel/csrc/attention/`.
+2) Register bindings in `fastvideo-kernel/csrc/common_extension.cpp`.
+3) Add to `fastvideo-kernel/CMakeLists.txt` (and any feature flags).
+4) Expose in `fastvideo-kernel/python/fastvideo_kernel/ops.py`.
+5) Export in `fastvideo-kernel/python/fastvideo_kernel/__init__.py`.
+
+Keep a Python/Triton fallback so the backend runs even when the kernel is not
+available.
+
+## 5) Testing and debugging
+
+- Add a small parity test or microbenchmark comparing to SDPA.
+- Force your backend with the env var:
+  `FASTVIDEO_ATTENTION_BACKEND=MY_NEW_ATTN`.
+- Check logs from `fastvideo/attention/selector.py` to confirm selection.
+
+## Checklist
+
+- [ ] Added enum entry in `fastvideo/platforms/interface.py`.
+- [ ] Implemented backend in `fastvideo/attention/backends/`.
+- [ ] Registered selection in platform(s).
+- [ ] Updated layer call sites to include the backend where appropriate.
+- [ ] Added tests and documentation.

docs/contributing/coding_agents.md

@@ -0,0 +1,480 @@
+# FastVideo + Coding Agents
+
+Coding agents are now strong at navigating large codebases and iterating fast
+with parity tests and examples. This guide shows how to use them to add new
+model pipelines and ship PRs in a production-grade video diffusion framework.
+
+FastVideo is a great project to contribute to, with production-grade
+infrastructure, active collaborations (including NVIDIA), and a pipeline design
+and inference architecture that has been forked by [SGLang’s
+multimodal generation stack](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen).
+
+Goal: run the new pipeline with a minimal script like
+`examples/inference/basic/basic.py`. In production, FastVideo can download
+models automatically via `HF_HOME`; for development, use local directories so
+agents can run scripts and tests deterministically. We standardize local paths
+as:
+
+- `official_weights/<model_name>/` for official checkpoints
+- `converted_weights/<model_name>/` if conversion is required
+
+## Tips when prompting the agent
+
+When prompting the agent, include:
+
+- This guide and the [FastVideo design overview](../design/overview.md).
+- Exact file paths to edit.
+- A closest reference example file in FastVideo.
+- Expected behavior and acceptance criteria.
+- Repro steps (command, inputs, logs).
+- Constraints (performance, memory, compatibility).
+- Local paths (e.g., `official_weights/<model_name>/` or
+  `converted_weights/<model_name>/`) for parity tests.
+
+## FastVideo structure at a glance
+
+Before diving in, scan these references:
+
+- [Contributing overview](overview.md) for environment/setup context.
+- [FastVideo design overview](../design/overview.md) for pipeline architecture, configs, and HF layout.
+
+FastVideo maps a Diffusers-style repo into a pipeline like:
+
+- `fastvideo/models/*`: model implementations (DiT, VAE, encoders, upsamplers).
+- `fastvideo/configs/models/*`: arch configs and `param_names_mapping` for
+  weight name translation.
+- `fastvideo/configs/pipelines/*`: pipeline wiring (component classes + names).
+- `fastvideo/configs/sample/*`: default runtime sampling parameters.
+- `fastvideo/pipelines/basic/*`: end-to-end pipeline logic built from stages.
+- `model_index.json`: the HF repo entrypoint that maps component names to
+  classes and weight files.
+- Component loading happens in `VideoGenerator.from_pretrained`, which reads
+  `model_index.json`, resolves configs, and loads weights.
+
+Minimal usage example (based on `examples/inference/basic/basic.py`):
+
+```python
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
+
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # or official_weights/<model_name>/
+generator = VideoGenerator.from_pretrained(model_id, num_gpus=1)
+
+sampling = SamplingParam.from_pretrained(model_id)
+sampling.num_frames = 45
+video = generator.generate_video(
+    "A vibrant city street at sunset.",
+    sampling_param=sampling,
+    output_path="video_samples",
+    save_video=True,
+)
+```
+
+## Some questions to ask yourself before starting
+
+Answering these upfront clarifies the work and speeds up implementation.
+
+### Is the model already supported by SGLang's multimodal generation stack?
+If yes, you can port many components from SGLang. It is a FastVideo fork, so
+interfaces line up, but you still need to swap layers/modules to match
+FastVideo's architecture and attention stack.
+
+If not, implement the model directly in FastVideo.
+
+### Is there an official implementation of the model you are adding?
+
+If yes, use it as the numerical reference. For example, LTX‑2 has an official
+implementation here: https://github.com/Lightricks/LTX-2. Prefer official code
+even if Diffusers also has one.
+
+### Is there a HuggingFace repo for the model you are adding? Is it in Diffusers format?
+
+If yes, load it directly in FastVideo after setting tensor mapping rules in the
+config. Otherwise, convert the weights to Diffusers format. See [Weights and
+Diffusers format](../design/overview.md#weights-and-diffusers-format) for details.
+
+### Do I have official weights + local paths ready?
+
+Standardize local paths as:
+
+- `official_weights/<model_name>/` for official checkpoints
+- `converted_weights/<model_name>/` if conversion is required (can be created later)
+
+### What pipeline components are required for the model you are adding?
+
+Usually you need a transformer (DiT), VAE, text encoder, and tokenizer. Some
+models add extra components.
+
+### What tasks does the model support?
+
+Usually a video diffusion model supports text‑to‑video (T2V),
+image‑to‑video (I2V), and video‑to‑video (V2V). Some add extra tasks (two‑stage
+generation, keyframe interpolation), which require extra components.
+
+It's usually easiest to start with a T2V pipeline and add the other tasks later.
+
+You can refer to the [Pipeline system](../design/overview.md#pipeline-system)
+section for more details.
+
+### Am I able to generate videos with the official implementation?
+
+These videos and prompts are your reference. Once the FastVideo pipeline works,
+compare outputs to the official implementation. Due to seeding and other
+factors, outputs may not match exactly, but they should be comparable.
+
+## Workflow: adding a full pipeline
+
+This is an example workflow for adding a full model pipeline (model + configs +
+examples + tests). This guide is in active development; feedback is welcome.
+
+!!! note
+    If you get stuck, refer to existing models/pipelines in FastVideo or ask in Slack.
+
+### 0) Fetch official model's code and weights
+
+Purpose:
+
+- Keep official checkpoints and source code local so conversion, parity tests,
+  and reference runs are reproducible.
+- Clone the official repo so you can use it as a numerical reference.
+
+Action:
+
+- Download official weights into `official_weights/<model_name>/`
+  (Diffusers format or not).
+- Clone the official repo under the project root (e.g., `FastVideo/LTX-2/`).
+- If a Diffusers-format HF repo already exists, you can skip manual weight
+  handling and download it directly with
+  `scripts/huggingface/download_hf.py`.
+
+!!! note
+    This step is best done manually because large downloads can time out.
+    Example:
+    ```bash
+    python scripts/huggingface/download_hf.py \
+      --repo_id Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+      --local_dir official_weights/Wan2.1-T2V-1.3B-Diffusers \
+      --repo_type model
+    ```
+  
+### 1) Implement the model + config mapping
+
+Purpose:
+
+- Model weights are a dictionary of named tensors (`state_dict`). If the names
+  don’t line up with FastVideo’s module names, weights won’t load correctly (or
+  will silently load into the wrong layer).
+- Official checkpoints often use different prefixes or module layouts than
+  FastVideo, so we translate names via the mapping (during load or conversion).
+- Mapping aligns three things:
+  1. the official implementation’s module names,
+  2. the checkpoint `state_dict` keys,
+  3. FastVideo’s model classes and layer naming conventions.
+- If names don’t align, weights won’t load; implement the FastVideo model and
+  define mapping rules first.
+
+Action:
+
+- Implement the FastVideo model + config mapping.
+  - Add/extend the model in `fastvideo/models/...` and config in
+    `fastvideo/configs/models/...` (including `param_names_mapping`).
+  - Reuse existing FastVideo layers/modules where possible.
+  - Use FastVideo’s attention layers:
+    - `DistributedAttention` only for full‑sequence self‑attention in the DiT.
+    - `LocalAttention` for cross‑attention and other attention layers.
+  - See the “Configuration System” and “Weights and Diffusers format” sections
+    in `docs/design/overview.md` for how these pieces connect.
+  - If you are using an agent, ask it to implement the model, config mapping,
+    and a parity test together so you can validate numerics immediately.
+
+!!! note
+    After the first component is aligned and parity‑tested, open a **DRAFT PR**
+    on FastVideo so the rest of the pipeline work can build on top of it.
+
+!!! note
+    If a Diffusers-format HF repo already exists and loads correctly, you can
+    skip conversion entirely (no conversion script needed) and just download it
+    with `scripts/huggingface/download_hf.py`. Otherwise, you may need a
+    conversion script + a `converted_weights/<model>/` staging directory.
+
+Example (key renaming via arch config mapping, Wan2.1‑style):
+
+```python
+# Official model (simplified) in the upstream repo.
+class OfficialWanTransformer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_embedding = torch.nn.Conv3d(16, 1536, kernel_size=2, padding=0)
+
+    def forward(self, x):
+        return self.patch_embedding(x)
+
+# FastVideo model (simplified) in fastvideo/models/dits/wanvideo.py
+class PatchEmbed(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = torch.nn.Conv3d(16, 1536, kernel_size=2, padding=0)
+
+    def forward(self, x):
+        return self.proj(x)
+
+class WanTransformer3DModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_embedding = PatchEmbed()
+
+    def forward(self, x):
+        return self.patch_embedding(x)
+
+# Mapping defined in a config (simplified; see the real mapping in
+# fastvideo/configs/models/dits/wanvideo.py)
+param_names_mapping = {
+    r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
+    r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2",
+}
+
+def apply_regex_map(state_dict, mapping):
+    # Pseudocode: apply regex substitutions in order
+    ...
+
+# Official checkpoint keys (example)
+official = {
+    "patch_embedding.weight": ...,
+    "blocks.0.attn1.to_q.weight": ...,
+}
+
+# Apply mapping so keys match FastVideo modules
+converted = apply_regex_map(official, param_names_mapping)
+
+```
+
+Optional helper (print a few checkpoint keys quickly):
+
+```bash
+python - <<'PY'
+import safetensors.torch as st
+keys = list(st.load_file("official_weights/<model>/transformer/diffusion_pytorch_model.safetensors").keys())
+print(keys[:20])
+PY
+```
+
+Example agent prompt (task request):
+
+```
+Please add the Wan2.1 T2V 1.3B Diffusers pipeline to FastVideo:
+- Add a FastVideo native Wan2.1 DiT implementation + config mapping.
+- Make sure to use the existing FastVideo layers and attention modules where possible.
+- Add a parity test that loads the official model alongside the FastVideo model and compares outputs numerically with fixed seeds and inputs.
+
+Paths:
+  - Official repo: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+  - Local download: official_weights/Wan2.1-T2V-1.3B-Diffusers
+Mapping steps:
+  - Load the official DiT weights from
+    official_weights/Wan2.1-T2V-1.3B-Diffusers/transformer/diffusion_pytorch_model.safetensors.
+  - Instantiate the FastVideo DiT (`WanTransformer3DModel`) and compare
+    its `state_dict().keys()` to the official keys.
+  - Update `param_names_mapping` in
+    fastvideo/configs/models/dits/wanvideo.py to resolve missing/unexpected keys.
+  - Use `load_state_dict(strict=False)` during iteration to surface mismatches.
+```
+
+External examples of the same pattern:
+- SGLang uses prefix-based routing in its weight loader to map checkpoint keys
+  into internal submodules (e.g., stripping a top-level prefix before delegating).
+- vLLM includes model-specific renamers for certain checkpoints that adjust
+  key prefixes so weights match its internal naming.
+
+### 2) Test numerical alignment with the official implementation
+
+Purpose:
+
+- Verify that the FastVideo component is numerically aligned with the official
+  implementation.
+
+Action:
+
+- Add or reuse a numerical parity test that loads the official model and the
+  FastVideo model and compares outputs.
+- See examples in `tests/local_tests/` (e.g., `tests/local_tests/upsamplers/`)
+  and the commands in `tests/local_tests/README.md`.
+- If there are discrepancies, add opt‑in logging to both models and compare
+  activation summaries (layer output sums, per‑stage logs).
+- First align the loaded weights (validate `param_names_mapping`).
+- Then align forward outputs using fixed seeds and inputs.
+  - Start with `atol=1e-4, rtol=1e-4` in `assert_close`.
+  - Keep dtype consistent (bf16 if available; otherwise fp32).
+  - If attention parity is unstable, align backends (e.g.,
+    `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`).
+
+### 3) Repeat the process for each component
+
+If the model requires additional components, repeat Steps 1–2 for each one.
+For example, implement the VAE in `fastvideo/models/vaes/` and its config in
+`fastvideo/configs/models/vaes/`, then add parity coverage for it.
+
+### 4) Add a pipeline config + sample defaults
+
+Purpose:
+
+- `fastvideo/configs/pipelines/` describes pipeline wiring and model module
+  names.
+- `fastvideo/configs/sample/` defines default runtime parameters.
+
+Action:
+
+- Add a new pipeline config + sampling params.
+- Register them in `fastvideo/configs/pipelines/registry.py` and
+  `fastvideo/configs/sample/registry.py`.
+
+### 5) Wire pipeline stages
+
+Purpose:
+
+- `fastvideo/pipelines/basic/<pipeline>/` contains the actual pipeline logic.
+- `fastvideo/pipelines/stages/` holds reusable, testable stages.
+
+Action:
+
+- Build the pipeline using stages; keep new stages isolated and documented.
+- Prefer opt‑in flags for expensive or optional steps.
+
+### 6) Add pipeline‑level tests
+
+Purpose:
+
+- Ensure the end‑to‑end pipeline works and stays aligned as pieces evolve.
+
+Action:
+
+- Add a pipeline parity test under `tests/local_tests/pipelines/`.
+- See the [Testing Guide](testing.md) for test conventions.
+
+### 7) Add user‑facing examples
+
+Purpose:
+
+- `examples/inference/basic/` is the entry point for simple, runnable scripts.
+
+Action:
+
+- Provide a minimal “hello world” example plus advanced variations.
+- Use fixed seeds and stable prompts.
+- Run the example locally to confirm end‑to‑end behavior.
+
+### 8) Add SSIM tests for CI checks
+
+Purpose:
+
+- Ensure visual similarity stays within expected bounds for regression testing.
+- SSIM tests act as a higher‑level guardrail beyond unit/parity tests.
+
+Action:
+
+- Add SSIM tests under `fastvideo/tests/ssim/` and include reference videos
+  (see the structure in the Testing Guide).
+- Use stable prompts/seeds and document any GPU‑specific requirements.
+- Follow the [Testing Guide](testing.md) for reference video placement and
+  execution details.
+
+### 9) Document it
+
+Purpose:
+
+- `docs/` is where users find the new pipeline usage and limitations.
+
+Action:
+
+- Add a short doc page or update an existing one.
+- Mention any caveats (memory, speed, constraints).
+
+## Common pitfalls when porting models
+
+- **Attention backend mismatch**: parity can fail if the official model uses a
+  different attention backend (e.g., SDPA vs custom). Align backends before
+  debugging deeper issues.
+- **Patchifier shape mistakes**: wrong patchification or reshape lengths can
+  silently corrupt outputs. Validate patch shapes early.
+- **Mask handling**: attention masks must match the official behavior (padding,
+  causal masks, and broadcast shapes).
+- **Scheduler / sigma schedule mismatch**: even small differences in schedules
+  or timestep shapes can cause noticeable drift.
+
+## Diffusers vs manual conversion
+
+If a model already ships in Diffusers format (with a proper `model_index.json`),
+prefer downloading it directly and loading it via FastVideo. In that case:
+
+- You usually **do not need** a conversion script.
+- You still need a correct `param_names_mapping` if the internal module names
+  differ from FastVideo’s implementation.
+
+If the model does **not** have a Diffusers-format repo:
+
+- You will need a conversion script to rewrite `state_dict` keys into FastVideo
+  naming and stage the result (e.g., under `converted_weights/<model>/`).
+- You may still use the official repo for reference parity and debugging.
+
+In both cases, parity testing is required to validate correctness.
+
+If you want to publish a Diffusers‑style repo after conversion, use
+`scripts/checkpoint_conversion/create_hf_repo.py` to assemble a HuggingFace‑ready
+directory before uploading.
+
+## FAQ
+
+**Q: Why do we implement the FastVideo model before conversion?**  
+A: You can’t define the key‑mapping rules until the FastVideo module names are
+known. The implementation determines the target `state_dict` schema.
+
+**Q: Do we always need a conversion script?**  
+A: No. If a Diffusers‑format repo exists and loads correctly, download it and
+skip conversion.
+
+**Q: How do I figure out `param_names_mapping` quickly?**  
+A: Load the official weights, instantiate the FastVideo model, and diff
+`state_dict().keys()` on both sides. Add regex rules until missing/unexpected
+keys are resolved. Agents can help you with this.
+
+**Q: What if parity fails even after mapping?**  
+A: Align attention backends, sigma schedules, and timestep shapes first. Then
+add opt‑in activation logging to locate the first divergent layer.
+
+## Case study: LTX‑2 port (from PLAN.md)
+
+The LTX‑2 port in `PLAN.md` shows the real sequence of steps and backtracking
+that happened during integration. Use it as a reference for how parity work
+actually unfolds:
+
+- Ported components first (transformer, VAE, audio, text encoder).
+- Added parity tests per component; used SDPA for reference parity.
+- Added debug logging to compare per‑block activations and isolate divergence.
+- Fixed cross‑attention reshape and patch grid bounds issues after logging.
+- Aligned sigma schedule and masking behavior to match the official pipeline.
+
+Recommendation:
+
+- Keep raw step‑by‑step logs in your own local `PLAN.md` for large ports.
+
+## Worked example: Wan2.1 T2V 1.3B pipeline
+
+The Wan2.1 T2V 1.3B Diffusers pipeline is a good “standard” example for
+FastVideo integration.
+
+1. Verify model config + mapping.
+   - DiT mapping: `fastvideo/configs/models/dits/wanvideo.py`
+   - VAE: `fastvideo/models/vaes/wanvae.py`
+   - Text encoder: `fastvideo/models/encoders/t5.py`
+
+2. Parity test the core components.
+   - Example tests: `fastvideo/tests/transformers/test_wanvideo.py`,
+     `fastvideo/tests/vaes/test_wan_vae.py`,
+     `fastvideo/tests/encoders/test_t5_encoder.py`
+
+3. Pipeline wiring.
+   - Pipeline: `fastvideo/pipelines/basic/wan/wan_pipeline.py`
+   - Pipeline config: `fastvideo/configs/pipelines/wan.py`
+   - Sampling defaults: `fastvideo/configs/sample/wan.py`
+
+4. Minimal example.
+   - Script: `examples/inference/basic/basic.py`

docs/contributing/developer_env/runpod.md

@@ -1,16 +1,55 @@
 
 # 📦 Developing FastVideo on RunPod
 
-You can easily use the FastVideo Docker image as a custom container on [RunPod](https://www.runpod.io) for development or experimentation.
+You can easily use the FastVideo Pod Template on [RunPod](https://www.runpod.io) for development or experimentation.
 
 ## Creating a new pod
 
-Choose a GPU that supports CUDA 12.8
-
-Pick 1 or 2 L40S GPU(s)
+- Make sure you are using the correct RunPod account.
+![RunPod Account Selection](../../assets/images/runpod_account.png)
 
+- Use "Additional Filters" to select CUDA 12.8.
 ![RunPod CUDA selection](../../assets/images/runpod_cuda.png)
 
+- Click "Deploy" and Pick a single A40 or RTX 4090 GPU.
+![RunPod GPU Selection](../../assets/images/runpod_deploy.png)
+
+- Select the "FastVideo" or "fastvideo-dev" Pod Template.
+![RunPod Pod Template Selection](../../assets/images/runpod_create.png)
+
+- Set the Pod name to "<name>-<FastVideo>-<date>".
+
+- Finally, once the pod is deployed (will take a few minutes as the image is being pulled), you can SSH into it using "SSH exposed over TCP". You'll need to use the matching private ssh key you provided.
+![RunPod SSH](../../assets/images/runpod_ssh.png)
+
+## Working with the pod
+
+After SSH'ing into your pod, you'll find the correct `uv` environment already activated and you should be in /FastVideo/ directory. Make sure to use /FastVideo/ for all your work.
+
+To pull in the latest changes from the GitHub repo:
+
+```bash
+cd /FastVideo
+git pull
+```
+
+Run your development workflows as usual:
+
+```bash
+# Run linters
+pre-commit run --all-files
+
+# Run tests
+pytest tests/
+```
+
+Make sure to push your changes back to the GitHub repo as nothing will be saved to the pod when it is terminated.
+
+After you are done with your work, you can terminate the pod by clicking the "Terminate" and "Delete" buttons. Remember if the pod is not completely deleted, Runpod will keep charging you for it.
+
+## Extra Information:
+If you need to customize the pod template this section has some useful information. For the most part you can leave the defaults of the FastVideo Pod Template.
+
 When creating your pod template, use this image:
 
 ```
@@ -24,30 +63,3 @@ bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-serve
 ```
 
 ![RunPod template configuration](../../assets/images/runpod_template.png)
-
-After deploying, the pod will take a few minutes to pull the image and start the SSH service.
-
-![RunPod ssh](../../assets/images/runpod_ssh.png)
-
-## Working with the pod
-
-After SSH'ing into your pod, you'll find the `fastvideo-dev` Conda environment already activated.
-
-To pull in the latest changes from the GitHub repo:
-
-```bash
-cd /FastVideo
-git pull
-```
-
-`If you have a persistent volume and want to keep your code changes, you can move /FastVideo to /workspace/FastVideo, or simply clone the repository there.`
-
-Run your development workflows as usual:
-
-```bash
-# Run linters
-pre-commit run --all-files
-
-# Run tests
-pytest tests/
-```

docs/contributing/overview.md

@@ -1,76 +1,84 @@
 
 # 🛠️ Contributing to FastVideo
 
-Thank you for your interest in contributing to FastVideo. We want to make the process as smooth for you as possible and this is a guide to help get you started!
+Thank you for your interest in contributing to FastVideo. We want the process
+to be smooth and beginner‑friendly, whether you are adding a new pipeline,
+improving performance, or fixing a bug.
 
-Our community is open to everyone and welcomes any contributions no matter how large or small.
+## Quick prerequisites
 
-# Developer Environment:
-Do make sure you have CUDA 12.4 installed and supported. FastVideo currently only supports Linux and CUDA GPUs, but we hope to support other platforms in the future.
+- **OS**: Linux is the primary development target (WSL can work).
+- **GPU**: NVIDIA GPU recommended for inference and training workflows.
+- **CUDA**: Use a recent CUDA 12.x toolchain (see the installation guide for
+  the current recommendation).
 
-We recommend using a fresh Python 3.10 Conda environment to develop FastVideo:
+For a full install checklist, see `docs/getting_started/installation/gpu.md`.
+
+## Local development (Conda + editable install)
 
 Install Miniconda:
 
-```
+```bash
 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
 bash Miniconda3-latest-Linux-x86_64.sh
 source ~/.bashrc
 ```
 
-Create and activate a Conda environment for FastVideo:
+Create and activate a Conda environment:
 
-```
+```bash
 conda create -n fastvideo python=3.12 -y
 conda activate fastvideo
 ```
 
 Install `uv` (optional, but recommended):
 
-From instructions on [uv](https://astral.sh/uv/):
-
-```
+```bash
 curl -LsSf https://astral.sh/uv/install.sh | sh
-# or 
+# or
 wget -qO- https://astral.sh/uv/install.sh | sh
 ```
 
-Clone the FastVideo repository and go to the FastVideo directory:
+Clone the repo:
 
-```
+```bash
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
-
 ```
 
-Now you can install FastVideo and setup git hooks for running linting. By using `pre-commit`, the linters will run and have to pass before you'll be able to make a commit.
+Install FastVideo in editable mode and set up hooks:
 
 ```bash
 uv pip install -e .[dev]
 
-# Can also install flash-attn (optional)
-uv pip install flash-attn --no-build-isolation 
+# Optional: FlashAttention (builds native kernels)
+uv pip install flash-attn --no-build-isolation
 
-# Linting, formatting and static type checking
+# Linting, formatting, static typing
 pre-commit install --hook-type pre-commit --hook-type commit-msg
-
-# You can manually run pre-commit with
 pre-commit run --all-files
 
 # Unit tests
 pytest tests/
 ```
 
-If you are on a Hopper GPU, you should also install [FA3](https://github.com/Dao-AILab/flash-attention) for much better performance:
+If you are on a Hopper GPU, installing FlashAttention 3 can improve
+performance (see `docs/inference/optimizations.md`).
 
-```
-git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention/hopper
+## Docker development (optional)
 
-# make sure you have ninja installed
-uv pip install ninja
-
-python setup.py install
-```
+If you prefer a containerized environment, use the dev image documented in
+`docs/contributing/developer_env/docker.md`.
 
 ## Testing
 
-Please refer to the [Testing Guide](testing.md) for more information on how to add and run tests in FastVideo.
+See the [Testing Guide](testing.md) for how to add and run tests in FastVideo.
+
+## Attention backend development
+
+If you are adding a new attention kernel or backend, follow
+[Attention Backend Development](attention_backend.md).
+
+## Contributing with coding agents
+
+For a step‑by‑step workflow on adding pipelines or components with coding
+agents, see `docs/contributing/coding_agents.md`.

docs/design/overview.md

@@ -1,424 +1,178 @@
-# 🔍 FastVideo Overview
+# FastVideo Architecture Overview
 
-This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/` codebase.
+This document summarizes how FastVideo is structured and how a Diffusers-style
+model repo maps into a runnable pipeline. It is intended for contributors who
+need the high-level layout and key entrypoints, not every internal detail.
 
-## Table of Contents - Directory Structure and Files
+## FastVideo structure at a glance
 
-- [`fastvideo/pipelines/`](#pipeline-system) - Core diffusion pipeline components
-- [`fastvideo/models/`](#model-components) - Model implementations
-  - [`dits/`](#transformer-models) - Transformer-based diffusion models
-  - [`vaes/`](#vae-variational-auto-encoder) - Variational autoencoders
-  - [`encoders/`](#text-and-image-encoders) - Text and image encoders
-  - [`schedulers/`](#schedulers) - Diffusion schedulers
-- [`fastvideo/attention/`](#optimized-attention) - Optimized attention implementations
-- [`fastvideo/distributed/`](#distributed-processing) - Distributed computing utilities
-- [`fastvideo/layers/`](#tensor-parallelism) - Custom neural network layers
-- [`fastvideo/platforms/`](#platforms) - Hardware platform abstractions
-- [`fastvideo/worker/`](#executor-and-worker-system) - Multi-GPU process management
-- [`fastvideo/fastvideo_args.py`](#fastvideoargs) - Argument handling
-- [`fastvideo/forward_context.py`](#forward-context-management) - Forward pass context management
-- `fastvideo/utils.py` - Utility functions
-- [`fastvideo/logger.py`](#logger) - Logging infrastructure
+FastVideo maps a Diffusers-style repo into a pipeline like this:
 
-## Core Architecture
+- `fastvideo/models/*`: model implementations (DiT, VAE, encoders, upsamplers).
+- `fastvideo/configs/models/*`: arch configs and `param_names_mapping` for
+  weight name translation.
+- `fastvideo/configs/pipelines/*`: pipeline wiring (component classes + names).
+- `fastvideo/configs/sample/*`: default runtime sampling parameters.
+- `fastvideo/pipelines/basic/*`: end-to-end pipelines.
+- `fastvideo/pipelines/stages/*`: reusable pipeline stages.
+- `fastvideo/models/loader/*`: component loaders for Diffusers-style repos.
+- `model_index.json`: HF repo entrypoint mapping component names to classes.
 
-FastVideo separates model components from execution logic with these principles:
+Flow:
+`model_index.json` -> component loaders -> model modules -> pipeline stages ->
+sampling params.
 
-- **Component Isolation**: Models (encoders, VAEs, transformers) are isolated from execution (pipelines, stages, distributed processing)
-- **Modular Design**: Components can be independently replaced
-- **Distributed Execution**: Supports various parallelism strategies (Tensor, Sequence)
-- **Custom Attention Backends**: Components can support and use different Attention implementations
-- **Pipeline Abstraction**: Consistent interface across diffusion models
-
-## FastVideoArgs
-
-The `FastVideoArgs` class in `fastvideo/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.
-
-Key features include:
-
-- **Command-line Interface**: Automatic conversion between CLI arguments and dataclass fields
-- **Configuration Groups**: Organized by functional areas (model loading, video params, optimization settings)
-- **Context Management**: Global access to current settings via `get_current_fastvideo_args()`
-- **Parameter Validation**: Ensures valid combinations of settings
-
-Common configuration areas:
-
-- **Model paths and loading options**: `model_path`, `trust_remote_code`, `revision`
-- **Distributed execution settings**: `num_gpus`, `tp_size`, `sp_size`
-- **Video generation parameters**: `height`, `width`, `num_frames`, `num_inference_steps`
-- **Precision settings**: Control computation precision for different components
-
-Example usage:
+Minimal usage (from `examples/inference/basic/basic.py`):
 
 ```python
-# Load arguments from command line
-fastvideo_args = prepare_fastvideo_args(sys.argv[1:])
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
 
-# Access parameters
-model = load_model(fastvideo_args.model_path)
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # or official_weights/<model_name>/
+generator = VideoGenerator.from_pretrained(model_id, num_gpus=1)
 
-# Set as global context
-with set_current_fastvideo_args(fastvideo_args):
-    # Code that requires access to these arguments
-    result = generate_video()
-```
-
-## Pipeline System
-
-### `ComposedPipelineBase`
-
-This foundational class provides:
-
-- **Model Loading**: Automatically loads components from HuggingFace-Diffusers-compatible model directories
-- **Stage Management**: Creates and orchestrates processing stages
-- **Data Flow Coordination**: Ensures proper state flow between stages
-
-```python
-class MyCustomPipeline(ComposedPipelineBase):
-    _required_config_modules = [
-        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
-    ]
-    
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        # Pipeline-specific initialization
-        pass
-        
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        self.add_stage("input_validation_stage", InputValidationStage())
-        self.add_stage("text_encoding_stage", CLIPTextEncodingStage(
-            text_encoder=self.get_module("text_encoder"),
-            tokenizer=self.get_module("tokenizer")
-        ))
-        # Additional stages...
-```
-
-### Pipeline Stages
-
-Each stage handles a specific diffusion process component:
-
-- **Input Validation**: Parameter verification
-- **Text Encoding**: CLIP, LLaMA, or T5-based encoding
-- **Image Encoding**: Image input processing
-- **Timestep & Latent Preparation**: Setup for diffusion
-- **Denoising**: Core diffusion loop
-- **Decoding**: Latent-to-pixel conversion
-
-Each stage implements a standard interface:
-
-```python
-def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-    # Process batch and update state
-    return batch
-```
-
-![Pipeline execution and data flow](../assets/images/pipeline.png)
-
-### ForwardBatch
-
-Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
-
-- **Input Data**: Prompts, images, generation parameters
-- **Intermediate State**: Embeddings, latents, timesteps, accumulated during stage execution
-- **Output Storage**: Generated results and metadata
-- **Configuration**: Sampling parameters, precision settings
-
-This structure facilitates clear state transitions between stages.
-
-## Model Components
-
-The `fastvideo/models/` directory contains implementations of the core neural network models used in video diffusion:
-
-### Transformer Models
-
-Transformer networks perform the actual denoising during diffusion:
-
-- **Location**: `fastvideo/models/dits/`
-- **Examples**:
-  - `WanTransformer3DModel`
-  - `HunyuanVideoTransformer3DModel`
-
-Features include:
-
-- Text/image conditioning
-- Standardized interface for model-specific optimizations
-
-```python
-def forward(
-    self, 
-    latents,                    # [B, T, C, H, W]
-    encoder_hidden_states,      # Text embeddings
-    timestep,                   # Current diffusion timestep
-    encoder_hidden_states_image=None,  # Optional image embeddings
-    **kwargs
-):
-    # Perform denoising computation
-    return noise_pred  # Predicted noise residual
-```
-
-### VAE (Variational Auto-Encoder)
-
-VAEs handle conversion between pixel space and latent space:
-
-- **Location**: `fastvideo/models/vaes/`
-- **Examples**:
-  - `AutoencoderKLWan`
-  - `AutoencoderKLHunyuanVideo`
-
-These models compress image/video data to a more efficient latent representation (typically 4x-8x smaller in each dimension).
-
-FastVideo's VAE implementations include:
-
-- Efficient video batch processing
-- Memory optimization
-- Optional tiling for large frames
-- Distributed weight support
-
-### Text and Image Encoders
-
-Encoders process conditioning inputs into embeddings:
-
-- **Location**: `fastvideo/models/encoders/`
-- **Text Encoders**:
-  - `CLIPTextModel`
-  - `LlamaModel`
-  - `UMT5EncoderModel`
-- **Image Encoders**:
-  - `CLIPVisionModel`
-
-FastVideo implements optimizations such as:
-
-- Vocab parallelism for distributed processing
-- Caching for common prompts
-- Precision-tuned computation
-
-### Schedulers
-
-Schedulers manage the diffusion sampling process:
-
-- **Location**: `fastvideo/models/schedulers/`
-- **Examples**:
-  - `UniPCMultistepScheduler`
-  - `FlowMatchEulerDiscreteScheduler`
-
-These components control:
-
-- Diffusion timestep sequences
-- Noise prediction to latent update conversions
-- Quality/speed trade-offs
-
-```python
-def step(
-    self, 
-    model_output: torch.Tensor,
-    timestep: torch.LongTensor,
-    sample: torch.Tensor,
-    **kwargs
-) -> torch.Tensor:
-    # Process model output and update latents
-    # Return updated latents
-    return prev_sample
-```
-
-This diagram shows how models are discovered, validated, and loaded across entrypoints, executors, pipelines, and model loaders.
-
-![Model loading flow](../assets/images/load_models.png)
-
-## Optimized Attention
-
-The `fastvideo/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
-
-### Attention Backends
-
-Multiple implementations with automatic selection:
-
-- **FLASH_ATTN**: Optimized for supporting hardware
-- **TORCH_SDPA**: Built-in PyTorch scaled dot-product attention
-- **SLIDING_TILE_ATTN**: For very long sequences
-
-```python
-# Configure available attention backends for this layer
-self.attn = LocalAttention(
-    num_heads=num_heads,
-    head_size=head_dim,
-    causal=False,
-    supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
-)
-
-# Override via environment variable
-# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-```
-
-![Attention backend selector design](../assets/images/attention_backend.png)
-
-### Attention Patterns
-
-Supports various patterns with memory optimization techniques:
-
-- **Cross/Self/Temporal/Global-Local Attention**
-- Chunking, progressive computation, optimized masking
-
-## Distributed Processing
-
-The `fastvideo/distributed/` directory contains implementations for distributed model execution:
-
-### Tensor Parallelism
-
-Tensor parallelism splits model weights across devices:
-
-- **Implementation**: Through `RowParallelLinear` and `ColumnParallelLinear` layers
-- **Use cases**: Will be used by encoder models as their sequence lengths are shorter and enables efficient sharding.
-
-```python
-# Tensor-parallel layers in a transformer block
-from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear
-
-# Split along output dimension
-self.qkv_proj = ColumnParallelLinear(
-    input_size=hidden_size,
-    output_size=3 * hidden_size,
-    bias=True,
-    gather_output=False
-)
-
-# Split along input dimension
-self.out_proj = RowParallelLinear(
-    input_size=hidden_size,
-    output_size=hidden_size,
-    bias=True,
-    input_is_parallel=True
+sampling = SamplingParam.from_pretrained(model_id)
+sampling.num_frames = 45
+video = generator.generate_video(
+    "A vibrant city street at sunset.",
+    sampling_param=sampling,
+    output_path="video_samples",
+    save_video=True,
 )
 ```
 
-### Sequence Parallelism
+## Configuration system
 
-Sequence parallelism splits sequences across devices:
+FastVideo uses typed configs to keep model definitions, pipeline wiring, and
+runtime parameters consistent:
 
-- **Implementation**: Through `DistributedAttention` and sequence splitting
-- **Use cases**: Long video sequences or high-resolution processing. Used by DiT models.
+- `fastvideo/configs/models/`: architecture definitions, layer shapes, and
+  `param_names_mapping` rules for key renaming.
+- `fastvideo/configs/pipelines/`: pipeline wiring and required components.
+- `fastvideo/configs/sample/`: default sampling parameters (steps, frames,
+  guidance scale, resolution, fps).
+- `fastvideo/configs/registry.py`: pipeline registry.
+- `fastvideo/configs/sample/registry.py`: sampling param registry.
 
-```python
-# Distributed attention for long sequences
-from fastvideo.attention import DistributedAttention
+`FastVideoArgs` (in `fastvideo/fastvideo_args.py`) provides runtime settings and
+is passed into pipeline construction and stages.
 
-self.attn = DistributedAttention(
-    num_heads=num_heads,
-    head_size=head_dim,
-    causal=False,
-    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN)
-)
+## Weights and Diffusers format
+
+FastVideo follows the HuggingFace Diffusers repo layout. This keeps loaders
+compatible with HF repos and makes it easy to add new components.
+
+Typical Diffusers repo:
+
+```
+<model-repo>/
+  model_index.json
+  scheduler/
+    scheduler_config.json
+  transformer/               # or unet/ for image models
+    config.json
+    diffusion_pytorch_model.safetensors
+  vae/
+    config.json
+    diffusion_pytorch_model.safetensors
+  text_encoder/
+    config.json
+    model.safetensors
+  tokenizer/
+    tokenizer_config.json
+    tokenizer.json
 ```
 
-### Communication Primitives
+Key points:
 
-Efficient distributed operations via AllGather, AllReduce, and synchronization mechanisms.
+- `model_index.json` is the root map that tells FastVideo which components to
+  load and which classes implement them.
+- Each component lives in its own folder with a `config.json` and weights.
+- Weights are usually in `diffusion_pytorch_model.safetensors`.
 
-Efficient communication primitives minimize distributed overhead:
+Note on tensor names:
 
-- **Sequence-Parallel AllGather**: Collects sequence chunks
-- **Tensor-Parallel AllReduce**: Combines partial results
-- **Distributed Synchronization**: Coordinates execution
+Official checkpoints often use different `state_dict` names than FastVideo's
+module layout. We translate tensor names via the DiT arch config mapping
+(`param_names_mapping` under `fastvideo/configs/models/dits/`). This is similar
+in spirit to name-translation layers used in systems like vLLM and SGLang.
 
-## Forward Context Management
+Example HF repo (Wan 2.1 T2V 1.3B Diffusers):
 
-### ForwardContext
-
-Defined in `fastvideo/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.
-
-- **Attention Metadata**: Configuration for optimized attention kernels (`attn_metadata`)
-- **Profiling Data**: Potential hooks for performance metrics collection
-
-This context-based approach enables:
-
-- Dynamic optimization based on execution state (e.g., attention backend selection)
-- Step-specific customizations within model components
-
-Usage example:
-
-```python
-with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
-    # During this forward pass, components can access context
-    # through get_forward_context()
-    output = model(inputs)
+```
+https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/tree/main
 ```
 
-## Executor and Worker System
+Example `model_index.json` from that repo:
 
-The `fastvideo/worker/` directory contains the distributed execution framework:
-
-### Executor Abstraction
-
-FastVideo implements a flexible execution model for distributed processing:
-
-- **Executor Base Class**: An abstract base class defining the interface for all executors
-- **MultiProcExecutor**: Primary implementation that spawns and manages worker processes
-- **GPU Workers**: Handle actual model execution on individual GPUs
-
-The MultiProcExecutor implementation:
-
-1. Spawns worker processes for each GPU
-2. Establishes communication channels via pipes
-3. Coordinates distributed operations across workers
-4. Handles graceful startup and shutdown of the process group
-
-Each GPU worker:
-
-1. Initializes the distributed environment
-2. Builds the pipeline for the specified model
-3. Executes requested operations on its assigned GPU
-4. Manages local resources and communicates results back to the executor
-
-This design allows FastVideo to efficiently utilize multiple GPUs while providing a simple, unified interface for model execution.
-
-## Platforms
-
-The `fastvideo/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
-
-### Platform Abstraction
-
-FastVideo's platform abstraction layer enables:
-
-- **Hardware Detection**: Automatic detection of available hardware
-- **Backend Selection**: Appropriate selection of compute kernels
-- **Memory Management**: Efficient utilization of hardware-specific memory features
-
-The primary components include:
-
-- **Platform Interface**: Defines the common API for all platform implementations
-- **CUDA Platform**: Optimized implementation for NVIDIA GPUs
-- **Backend Enum**: Used throughout the codebase for feature selection
-
-Usage example:
-
-```python
-from fastvideo.platforms import current_platform, _Backend
-
-# Check hardware capabilities
-if current_platform.supports_backend(_Backend.FLASH_ATTN):
-    # Use FlashAttention implementation
-else:
-    # Fall back to standard implementation
+```json
+{
+  "_class_name": "WanPipeline",
+  "_diffusers_version": "0.33.0.dev0",
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "UMT5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "T5TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "WanTransformer3DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan"
+  ]
+}
 ```
 
-The platform system is designed to be extensible for future hardware targets.
+How this maps to FastVideo:
 
-## Logger
+- `WanPipeline` -> `fastvideo/pipelines/basic/wan/wan_pipeline.py`
+- `WanTransformer3DModel` -> `fastvideo/models/dits/wanvideo.py`
+- `AutoencoderKLWan` -> `fastvideo/models/vaes/wanvae.py`
+- `UMT5EncoderModel` -> `fastvideo/models/encoders/t5.py`
+- `T5TokenizerFast` -> loaded via HF in `fastvideo/models/loader/`
+- `UniPCMultistepScheduler` -> loaded via Diffusers scheduler utilities
+- Pipeline defaults -> `fastvideo/configs/pipelines/wan.py`
+- Sampling defaults -> `fastvideo/configs/sample/wan.py`
 
-See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)
+## Pipeline system
 
-*TODO*: (help wanted) Add an environment variable that disables process-aware logging.
+- `fastvideo/pipelines/basic/*` contains end-to-end pipelines for each model
+  family.
+- `fastvideo/pipelines/stages/*` contains reusable, testable stages.
+- Pipelines subclass `ComposedPipelineBase` and declare required components via
+  `_required_config_modules`.
+- `ForwardBatch` (in `fastvideo/pipelines/pipeline_batch_info.py`) carries
+  prompts, latents, timesteps, and intermediate state across stages.
 
-## Contributing to FastVideo
+## Model components
 
-If you're a new contributor, here are some common areas to explore:
+- DiT models: `fastvideo/models/dits/`
+- VAEs: `fastvideo/models/vaes/`
+- Text/image encoders: `fastvideo/models/encoders/`
+- Schedulers: `fastvideo/models/schedulers/`
+- Upsamplers: `fastvideo/models/upsamplers/`
+- Optional audio models: `fastvideo/models/audio/`
 
-1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/models/`
-2. **Optimizing performance**: Look at attention implementations or memory management
-3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/pipelines/`
-4. **Hardware support**: Extend the `platforms` module for new hardware targets
+## Attention and distributed execution
 
-When adding code, follow these practices:
+- Attention backends live in `fastvideo/attention/` and can be selected via
+  `FASTVIDEO_ATTENTION_BACKEND`.
+- `LocalAttention` is used for cross-attention and most attention layers.
+- `DistributedAttention` is used for full-sequence self-attention in the DiT.
+- Tensor-parallel layers live in `fastvideo/layers/`.
+- Sequence/tensor parallel utilities live in `fastvideo/distributed/`.
 
-- Use type hints for better code readability
-- Add appropriate docstrings
-- Maintain the separation between model components and execution logic
-- Follow existing patterns for distributed processing
+## Related docs
+
+- [Contributing overview](../contributing/overview.md)
+- [Coding agents workflow](../contributing/coding_agents.md)
+- [Testing guide](../contributing/testing.md)

docs/inference/offloading.md

@@ -0,0 +1,140 @@
+# Offloading
+
+This page describes how to use offloading techniques for inference to reduce GPU memory usage while maintaining acceptable performance.
+
+## Default Behavior
+
+```python
+dit_cpu_offload: bool = True
+use_fsdp_inference: bool = False
+dit_layerwise_offload: bool = True
+text_encoder_cpu_offload: bool = True
+image_encoder_cpu_offload: bool = True
+vae_cpu_offload: bool = True
+pin_cpu_memory: bool = True
+```
+
+## Behavior Explanation
+
+!!! note
+    For CLI usage, replace underscores (`_`) with hyphens (`-`).
+
+### `use_fsdp_inference`
+
+Enables [FSDP](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html) for inference. The model weights are sharded across multiple GPUs to reduce memory usage per GPU, and weights are broadcast to all GPUs layer by layer during inference.
+
+#### Performance Impact
+
+FSDP inference introduces negligible performance overhead due to weight prefetching. Performance overhead may be visible when GPU interconnect is slow (e.g., multiple consumer-level GPUs connected by slow PCIe without GPU P2P support).
+
+#### Usage Recommendation
+
+We recommend enabling this option when multiple GPUs are available.
+
+### `dit_cpu_offload`
+
+Enables CPU offloading for FSDP inference. When enabled, the model weights are offloaded to CPU memory, and the weight of each layer is moved to GPU memory only when that layer is being computed.
+
+#### Performance Impact
+
+The PyTorch FSDP implementation does not overlap computation and data transfer perfectly for inference, so enabling this option will harm performance.
+
+#### Usage Recommendation
+
+This option only takes effect when FSDP is enabled. For single GPU usage, we recommend using `dit_layerwise_offload` instead.
+
+### `dit_layerwise_offload`
+
+This option is similar to `dit_cpu_offload`, but with two key differences:
+
+1. It overlaps computation and PCIe data transfer
+2. It only works for single GPU inference
+
+#### Performance Impact
+
+This option introduces negligible performance overhead.
+
+#### Usage Recommendation
+
+We recommend enabling this option for single GPU usage. This option is not compatible with FSDP.
+
+### `text_encoder_cpu_offload`
+
+When enabled, the text encoder model weights are offloaded to CPU memory, and text encoding is computed on CPU.
+
+#### Performance Impact
+
+This option significantly slows down text encoding computation, but text encoding is usually not the bottleneck.
+
+#### Usage Recommendation
+
+We recommend enabling this option only when OOM happens.
+
+### `image_encoder_cpu_offload` and `vae_cpu_offload`
+
+When enabled, the weights are stored in CPU memory and moved to GPU memory when the corresponding module is being computed. After computation, the weights are moved back to CPU memory.
+
+#### Performance Impact
+
+These options introduce performance overhead due to PCIe data transfer.
+
+#### Usage Recommendation
+
+We recommend enabling these options when OOM happens.
+
+## General Recommendations
+
+### Single GPU Inference
+
+We recommend enabling `dit_layerwise_offload`. If OOM happens, also enable `image_encoder_cpu_offload` and `vae_cpu_offload`. If OOM still happens, consider enabling `text_encoder_cpu_offload`.
+
+### Multi-GPU Inference
+
+We recommend enabling `use_fsdp_inference` and disabling both `dit_layerwise_offload` and `dit_cpu_offload`. If OOM happens, consider enabling `text_encoder_cpu_offload`, `image_encoder_cpu_offload`, and `vae_cpu_offload`. If OOM still happens, consider enabling `dit_cpu_offload`.
+
+## Examples
+
+### Single GPU with Layerwise Offloading
+
+```python
+from fastvideo import VideoGenerator
+
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=1,
+    # Recommended for single GPU
+    dit_layerwise_offload=True,
+    # Enable if OOM happens
+    vae_cpu_offload=True,
+    image_encoder_cpu_offload=True,
+    text_encoder_cpu_offload=True,
+    # Speeds up CPU-GPU transfer
+    pin_cpu_memory=True,
+)
+
+prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers."
+video = generator.generate_video(prompt, output_path="output/", save_video=True)
+```
+
+### Multi-GPU with FSDP
+
+```python
+from fastvideo import VideoGenerator
+
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=2,
+    # Recommended for multi-GPU
+    use_fsdp_inference=True,
+    dit_layerwise_offload=False,
+    dit_cpu_offload=False,
+    # Enable if OOM happens
+    vae_cpu_offload=True,
+    image_encoder_cpu_offload=True,
+    text_encoder_cpu_offload=True,
+    pin_cpu_memory=True,
+)
+
+prompt = "A majestic lion strides across the golden savanna."
+video = generator.generate_video(prompt, output_path="output/", save_video=True)
+```

docs/inference/optimizations.md

@@ -103,7 +103,7 @@ python setup.py install  # or pip install -e .
 
 **`SAGE_ATTN_THREE`**
 
-[SageAttention 3](https://huggingface.co/jt-zhang/SageAttention3) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.
+[SageAttention 3](https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.
 
 #### Hardware Requirements
 
@@ -113,11 +113,7 @@ python setup.py install  # or pip install -e .
 
 Note that Sage Attention 3 requires `python>=3.13`, `torch>=2.8.0`, `CUDA >=12.8`. If you are using `uv` and using `torch==2.8.0` make sure that `sentencepiece==0.2.1` in the pyproject.toml file.
 
-To use Sage Attention 3 in FastVideo, first get access to the SageAttention3 code, then move `sageattn/` and `setup.py` to the directory `fastvideo/attention/backends`, then install from using:
-
-```bash
-python setup.py install
-```
+To use Sage Attention 3 in FastVideo, follow the `README.md` in the linked repository to install the package from source.
 
 ## Teacache
 

examples/inference/basic/basic_matrixgame.py

@@ -1,6 +1,5 @@
 from fastvideo import VideoGenerator
-from fastvideo.configs.pipelines.wan import MatrixGameI2V480PConfig
-from fastvideo.models.dits.matrix_game.utils import create_action_presets
+from fastvideo.models.dits.matrixgame.utils import create_action_presets
 
 import torch
 

examples/inference/basic/basic_matrixgame_streaming.py

@@ -1,5 +1,5 @@
 from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator
-from fastvideo.models.dits.matrix_game.utils import get_current_action_async, expand_action_to_frames
+from fastvideo.models.dits.matrixgame.utils import get_current_action_async, expand_action_to_frames
 
 import torch
 import asyncio

examples/inference/gradio/local/gradio_local_demo_matrixgame.py

@@ -10,7 +10,7 @@ from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import HTMLResponse, FileResponse
 
 from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator
-from fastvideo.models.dits.matrix_game.utils import expand_action_to_frames
+from fastvideo.models.dits.matrixgame.utils import expand_action_to_frames
 
 
 VARIANT_CONFIG = {

examples/training/finetune/MatrixGame2.0/finetune_i2v.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="FastVideo/Matrix-Game-2.0-Foundation-Diffusers"
+DATA_DIR="footsies-dataset/preprocessed/combined_parquet_dataset"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "matrixgame_finetune"
+  --output_dir "checkpoints/matrixgame_finetune"
+  --max_train_steps 1500
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 4
+  --num_latent_t 20
+  --num_height 352
+  --num_width 640
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 2
+  --tp_size 1
+  --hsdp_replicate_dim 4
+  --hsdp_shard_dim 2
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 100
+  --training_state_checkpointing_steps 100
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/matrixgame_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Matrix-Game-2.0-Foundation-Diffusers"
+DATA_MERGE_PATH="footsies-dataset/merge.txt"
+OUTPUT_DIR="footsies-dataset/preprocessed/"
+
+# export CUDA_VISIBLE_DEVICES=0
+export MASTER_ADDR=localhost
+export MASTER_PORT=29500
+export RANK=0
+export WORLD_SIZE=1
+
+python fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 4 \
+    --seed 42 \
+    --max_height 352 \
+    --max_width 640 \
+    --num_frames 77 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --samples_per_file 4 \
+    --train_fps 25 \
+    --flush_frequency 4 \
+    --preprocess_task matrixgame

fastvideo/attention/backends/sage_attn3.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
-from fastvideo.attention.backends.sageattn.api import sageattn_blackwell
+from sageattn3 import sageattn3_blackwell
 
 from fastvideo.attention.backends.abstract import (AttentionBackend,
                                                    AttentionImpl,
@@ -67,6 +67,6 @@ class SageAttention3Impl(AttentionImpl):
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
-        output = sageattn_blackwell(query, key, value, is_causal=self.causal)
+        output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
         output = output.transpose(1, 2)
         return output

fastvideo/configs/models/dits/__init__.py

@@ -7,9 +7,11 @@ from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
+from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
     "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
-    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig"
+    "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig",
+    "WanGameVideoConfig"
 ]

fastvideo/configs/models/dits/matrixgame.py

@@ -8,8 +8,8 @@ class MatrixGameWanVideoArchConfig(WanVideoArchConfig):
     # because MatrixGame checkpoints already have patch_embedding.proj format
     param_names_mapping: dict = field(
         default_factory=lambda: {
-            # Removed: r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1"
-            # because checkpoint already has correct format
+            r"^patch_embedding\.(?!proj\.)(.*)$":
+            r"patch_embedding.proj.\1",
             r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
             r"condition_embedder.text_embedder.fc_in.\1",
             r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":

fastvideo/configs/models/dits/wangamevideo.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class WanGameVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embedding\.(.*)$":
+            r"patch_embedding.proj.\1",
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
+            r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
+            r"condition_embedder.text_embedder.fc_out.\1",
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
+            r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
+            r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$":
+            r"condition_embedder.time_modulation.linear.\1",
+            r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_in.\1",
+            r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_out.\1",
+            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+            r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+            r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+            r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
+            r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
+            r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
+            r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
+            r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
+            r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
+            r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm2\.(.*)$":
+            r"blocks.\1.self_attn_residual_norm.norm.\2",
+        })
+
+    # Reverse mapping for saving checkpoints: custom -> hf
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    # Some LoRA adapters use the original official layer names instead of hf layer names,
+    # so apply this before the param_names_mapping
+    lora_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
+            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
+            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
+            r"blocks.\1.attn1.to_out.0.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
+            r"blocks.\1.attn2.to_out.0.\2",
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+        })
+
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len = 512
+    num_attention_heads: int = 40
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    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
+    qk_norm: str = "rms_norm_across_heads"
+    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
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
+
+    # Wan MoE
+    boundary_ratio: float | None = None
+
+    # Causal Wan
+    local_attn_size: int = 6  # Window size for temporal local attention (-1 indicates global attention)
+    sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
+    num_frames_per_block: int = 3
+    sliding_window_num_frames: int = 21
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class WanGameVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=WanGameVideoArchConfig)
+
+    prefix: str = "WanGame"
+
+
+@dataclass
+class WanLingBotVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=WanGameVideoArchConfig)
+
+    prefix: str = "WanLingBot"

fastvideo/configs/pipelines/registry.py

@@ -23,7 +23,7 @@ from fastvideo.configs.pipelines.wan import (
     Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
     WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
     SelfForcingWanT2V480PConfig, WANV2VConfig, SelfForcingWan2_2_T2V480PConfig,
-    MatrixGameI2V480PConfig)
+    MatrixGameI2V480PConfig, WanGameI2V480PConfig)
 # isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
@@ -42,6 +42,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/HY-WorldPlay-Bidirectional-Diffusers": HYWorldConfig,
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
     "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
+    "weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers": WanGameI2V480PConfig,
     "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": WANV2VConfig,
     "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
     "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
@@ -93,6 +94,8 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "hyworld" in id.lower(),
     "matrixgame":
     lambda id: "matrix-game" in id.lower() or "matrixgame" in id.lower(),
+    "wangameactionimagetovideo":
+    lambda id: "wangameactionimagetovideo" in id.lower(),
     "wanpipeline":
     lambda id: "wanpipeline" in id.lower(),
     "wanimagetovideo":
@@ -124,6 +127,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
     "matrixgame": MatrixGameI2V480PConfig,
+    "wangameactionimagetovideo": WanGameI2V480PConfig,
     "hunyuan15":
     Hunyuan15T2V480PConfig,  # Base Hunyuan15 config as fallback for any Hunyuan15 variant
     "hyworld":

fastvideo/configs/pipelines/wan.py

@@ -7,6 +7,7 @@ import torch
 from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.configs.models.dits import WanVideoConfig
 from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig
+from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig
 from fastvideo.configs.models.encoders import (BaseEncoderOutput,
                                                CLIPVisionConfig, T5Config,
                                                WAN2_1ControlCLIPVisionConfig)
@@ -196,6 +197,12 @@ class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config):
 # =============================================
 # ============= Matrix Game ===================
 # =============================================
+@dataclass
+class MatrixGameBaseI2V480PConfig(WanI2V480PConfig):
+    dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig)
+    flow_shift: float | None = 5.0
+
+
 @dataclass
 class MatrixGameI2V480PConfig(WanI2V480PConfig):
     dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig)
@@ -211,3 +218,23 @@ class MatrixGameI2V480PConfig(WanI2V480PConfig):
     context_noise: int = 0
     num_frames_per_block: int = 3
     # sliding_window_num_frames: int = 15
+
+
+# =============================================
+# ============= WanGame ======================
+# =============================================
+@dataclass
+class WanGameI2V480PConfig(WanI2V480PConfig):
+    """Configuration for WanGame image-to-video pipeline."""
+    dit_config: DiTConfig = field(default_factory=WanGameVideoConfig)
+
+    image_encoder_config: EncoderConfig = field(
+        default_factory=WAN2_1ControlCLIPVisionConfig)
+
+    is_causal: bool = True
+    flow_shift: float | None = 5.0
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 666, 333])
+    warp_denoising_step: bool = True
+    context_noise: int = 0
+    num_frames_per_block: int = 3

fastvideo/configs/sample/registry.py

@@ -92,6 +92,9 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "KyleShao/Cosmos-Predict2.5-2B-Diffusers":
     Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
 
+    # WanGame models
+    "weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers": MatrixGame2_SamplingParam,
+
     # MatrixGame2.0 models
     "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam,
     "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam,
@@ -130,6 +133,8 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
     lambda id: "wandmdpipeline" in id.lower(),
     "wancausaldmdpipeline":
     lambda id: "wancausaldmdpipeline" in id.lower(),
+    "wangame":
+    lambda id: "wangame" in id.lower(),
     "matrixgame":
     lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(),
     "turbodiffusion":
@@ -157,6 +162,7 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
     "wandmdpipeline": FastWanT2V480P_SamplingParam,
     "wancausaldmdpipeline": SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
     "stepvideo": StepVideoT2VSamplingParam,
+    "wangame": MatrixGame2_SamplingParam,
     "matrixgame": MatrixGame2_SamplingParam,
     "turbodiffusion":
     TurboDiffusionT2V_1_3B_SamplingParam,  # Default to T2V for fallback

fastvideo/dataset/dataloader/schema.py

@@ -116,3 +116,44 @@ pyarrow_schema_text_only = pa.schema([
     # --- Metadata ---
     pa.field("caption", pa.string()),
 ])
+
+
+pyarrow_schema_matrixgame = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Image/Video VAE latents ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("vae_latent_bytes", pa.binary()),
+    # e.g., [C, T, H, W] or [C, H, W]
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    # e.g., 'float32'
+    pa.field("vae_latent_dtype", pa.string()),
+    #I2V
+    pa.field("clip_feature_bytes", pa.binary()),
+    pa.field("clip_feature_shape", pa.list_(pa.int64())),
+    pa.field("clip_feature_dtype", pa.string()),
+    pa.field("first_frame_latent_bytes", pa.binary()),
+    pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
+    pa.field("first_frame_latent_dtype", pa.string()),
+    # --- Action ---
+    pa.field("mouse_cond_bytes", pa.binary()),
+    pa.field("mouse_cond_shape", pa.list_(pa.int64())),  # [T, 2]
+    pa.field("mouse_cond_dtype", pa.string()),
+    pa.field("keyboard_cond_bytes", pa.binary()),
+    pa.field("keyboard_cond_shape", pa.list_(pa.int64())),  # [T, 4]
+    pa.field("keyboard_cond_dtype", pa.string()),
+    # I2V Validation
+    pa.field("pil_image_bytes", pa.binary()),
+    pa.field("pil_image_shape", pa.list_(pa.int64())),
+    pa.field("pil_image_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    # -- Video-specific (can be null/default for images) ---
+    # Number of frames processed (e.g., 1 for image, N for video)
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])

fastvideo/dataset/preprocessing_datasets.py

@@ -40,6 +40,7 @@ class PreprocessBatch:
     num_frames: int | None = None
     sample_frame_index: list[int] | None = None
     sample_num_frames: int | None = None
+    action_path: str | None = None
 
     # Processed data
     pixel_values: torch.Tensor | None = None
@@ -365,9 +366,11 @@ class ImageTransformStage(DatasetStage):
             image = self.transform_topcrop(image)
         elif self.transform is not None:
             image = self.transform(image)
+            image = image.float() / 127.5 - 1.0
+        else:
+            image = image.float() / 127.5 - 1.0
 
         image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-        image = image.float() / 127.5 - 1.0
         batch.pixel_values = image
         return batch
 
@@ -470,8 +473,11 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
 
         # Initialize tokenizer
         tokenizer_path = os.path.join(args.model_path, "tokenizer")
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
-                                                  cache_dir=args.cache_dir)
+        if os.path.exists(tokenizer_path):
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
+                                                      cache_dir=args.cache_dir)
+        else:
+            tokenizer = None
 
         # Initialize processing stages
         self._init_stages(args, transform, transform_topcrop, tokenizer)
@@ -495,11 +501,14 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         self.video_transform_stage = VideoTransformStage(transform)
         self.image_transform_stage = ImageTransformStage(
             transform, transform_topcrop)
-        self.text_encoding_stage = TextEncodingStage(
-            tokenizer=tokenizer,
-            text_max_length=args.text_max_length,
-            cfg_rate=args.training_cfg_rate,
-            seed=self.seed)
+        if tokenizer is not None:
+            self.text_encoding_stage = TextEncodingStage(
+                tokenizer=tokenizer,
+                text_max_length=args.text_max_length,
+                cfg_rate=args.training_cfg_rate,
+                seed=self.seed)
+        else:
+            self.text_encoding_stage = None
 
     def _load_raw_data(self) -> list[dict]:
         """Load raw data from JSON files."""
@@ -521,6 +530,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         # Update paths with folder prefix
         for item in data_items:
             item["path"] = opj(folder, item["path"])
+            if "action_path" in item and item["action_path"]:
+                item["action_path"] = opj(folder, item["action_path"])
 
         return data_items
 
@@ -542,7 +553,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                                     cap=item["cap"],
                                     resolution=item.get("resolution"),
                                     fps=item.get("fps"),
-                                    duration=item.get("duration"))
+                                    duration=item.get("duration"),
+                                    action_path=item.get("action_path"))
 
             # Apply filtering stages
             if not self._apply_filter_stages(batch, filter_counts):
@@ -604,20 +616,27 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
         # Apply transformation stages
         batch = self.video_transform_stage.process(batch)
         batch = self.image_transform_stage.process(batch)
-        batch = self.text_encoding_stage.process(batch)
+        if self.text_encoding_stage is not None:
+            batch = self.text_encoding_stage.process(batch)
 
         # Build result dictionary
         result = {
             "pixel_values": batch.pixel_values,
-            "text": batch.text,
-            "input_ids": batch.input_ids,
-            "cond_mask": batch.cond_mask,
             "path": batch.path,
         }
 
+        if batch.text is not None:
+            result["text"] = batch.text
+            result["input_ids"] = batch.input_ids
+            result["cond_mask"] = batch.cond_mask
+
         # Add video-specific fields
         if batch.is_video:
             result.update({"fps": batch.fps, "duration": batch.duration})
+        
+        # Add action_path
+        if batch.action_path:
+            result["action_path"] = batch.action_path
 
         return result
 

fastvideo/entrypoints/streaming_generator.py

@@ -171,7 +171,7 @@ class StreamingVideoGenerator(VideoGenerator):
 
     def step(
             self, keyboard_cond: torch.Tensor,
-            mouse_cond: torch.Tensor) -> tuple[list[np.ndarray], Future | None]:
+            mouse_cond: torch.Tensor) -> tuple[list[np.ndarray], Future | None, dict | None]:
         if self.batch is None:
             raise RuntimeError("Call reset() before step()")
 
@@ -195,7 +195,10 @@ class StreamingVideoGenerator(VideoGenerator):
                 # Returns Future for block file, or None if no block_dir
                 block_future = self.writer.add_frames(frames)
 
-        return frames, block_future
+        # Extract stage timings if available
+        stage_timings = getattr(output_batch, 'stage_timings', None)
+
+        return frames, block_future, stage_timings
 
     async def step_async(
             self, keyboard_cond: torch.Tensor,

fastvideo/models/dits/matrixgame/action_module.py

@@ -257,6 +257,17 @@ class ActionModule(nn.Module):
         '''
         assert use_rope_keyboard
 
+        target_device = x.device
+        target_dtype = x.dtype
+        if mouse_condition is not None:
+            mouse_condition = mouse_condition.to(device=target_device,
+                                                 dtype=target_dtype)
+        if keyboard_condition is not None:
+            keyboard_condition = keyboard_condition.to(
+                device=target_device, dtype=target_dtype)
+        else:
+            return x
+
         B, N_frames, C = keyboard_condition.shape
         assert tt*th*tw == x.shape[1]
         assert ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0
@@ -272,7 +283,9 @@ class ActionModule(nn.Module):
         # Defined freqs_cis early so it's available for both mouse and keyboard
         freqs_cis = (self._freqs_cos, self._freqs_sin)
 
-        assert (N_feats == tt and ((is_causal and kv_cache_mouse is None) or not is_causal)) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block and is_causal)
+        if is_causal:
+            assert (N_feats == tt and kv_cache_mouse is None) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block)
+        # For non-causal (training), we trust that the caller provides correctly shaped inputs
         
         if self.enable_mouse and mouse_condition is not None:
             hidden_states = rearrange(x, "B (T S) C -> (B S) T C", T=tt, S=th*tw) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670
@@ -289,13 +302,15 @@ class ActionModule(nn.Module):
                 mouse_condition = mouse_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] 
                 group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)]
             else:
-                group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(N_feats)]
+                local_num_frames = tt
+                group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(local_num_frames)]
                 
             group_mouse = torch.stack(group_mouse, dim = 1)
+            actual_num_frames = group_mouse.shape[1]  # Use actual stacked frame count
 
             S = th * tw 
-            group_mouse = group_mouse.unsqueeze(-1).expand(B, num_frame_per_block, pad_t, C, S)
-            group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, num_frame_per_block, pad_t * C) 
+            group_mouse = group_mouse.unsqueeze(-1).expand(B, actual_num_frames, pad_t, C, S)
+            group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, actual_num_frames, pad_t * C) 
 
             group_mouse = torch.cat([hidden_states, group_mouse], dim = -1)
             group_mouse = self.mouse_mlp(group_mouse)
@@ -314,7 +329,7 @@ class ActionModule(nn.Module):
             ## TODO: adding cache here
             if is_causal:
                 if kv_cache_mouse is None:
-                    assert q.shape[0] ==  k.shape[0] and q.shape[0] % 880 == 0 # == 880, f"{q.shape[0]},{k.shape[0]}"
+                    assert q.shape[0] ==  k.shape[0] and q.shape[0] % S == 0 # == 880, f"{q.shape[0]},{k.shape[0]}"
                     padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
                     padded_q = torch.cat(
                         [q,
@@ -395,7 +410,8 @@ class ActionModule(nn.Module):
                 group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)]
             else:
                 keyboard_condition = self.keyboard_embed(keyboard_condition)
-                group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(N_feats)]
+                local_num_frames = tt
+                group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(local_num_frames)]
             group_keyboard = torch.stack(group_keyboard, dim = 1) # B F RW C
             group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0],group_keyboard.shape[1],-1))
             # apply cross attn
@@ -415,7 +431,7 @@ class ActionModule(nn.Module):
             q = self.key_attn_q_norm(q).to(v)
             k = self.key_attn_k_norm(k).to(v)
             S = th * tw
-            assert S == 880
+            # assert S == 880
             # position embed 
             if use_rope_keyboard: 
                 B, TS, H, D = q.shape
@@ -424,13 +440,13 @@ class ActionModule(nn.Module):
                 q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame)
 
                 k1, k2, k3, k4 = k.shape
-                k = k.expand(S, k2, k3, k4)
-                v = v.expand(S, k2, k3, k4)
+                k = k.repeat_interleave(S, dim=0)  
+                v = v.repeat_interleave(S, dim=0)
 
 
                 if is_causal:
                     if kv_cache_keyboard is None:
-                        assert q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0 
+                        assert q.shape[0] == k.shape[0] and q.shape[0] % S == 0 
 
                         padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
                         padded_q = torch.cat(
@@ -480,14 +496,14 @@ class ActionModule(nn.Module):
                         else:
                             local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item()
                             local_start_index = local_end_index - num_new_tokens
-                        assert k.shape[0] == 880 # BS == 1 or the cache should not be saved/ load method should be modified
-                        kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[:1]
-                        kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[:1]
+                        assert k.shape[0] % S == 0  # BS >= 1
+                        kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[::S]
+                        kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[::S]
 
                         attn = self.keyboard_attn_layer(
                             q,
-                            kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1),
-                            kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1),
+                            kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat_interleave(S, dim=0),
+                            kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat_interleave(S, dim=0),
                         )
 
                         kv_cache_keyboard["global_end_index"].fill_(current_end)

fastvideo/models/dits/matrixgame/causal_model.py

@@ -14,7 +14,7 @@ from torch.nn.attention.flex_attention import BlockMask
 from fastvideo.forward_context import get_forward_context
 
 flex_attention = torch.compile(
-    flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
+    flex_attention, dynamic=True, mode="max-autotune-no-cudagraphs")
 
 from fastvideo.attention import LocalAttention
 from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig
@@ -465,6 +465,7 @@ _DEFAULT_MATRIXGAME_CONFIG = MatrixGameWanVideoConfig()
 
 class CausalMatrixGameWanModel(BaseDiT):
     supports_action_input = True
+    _concatenates_image_latent = True  # Model handles image_latent internally via forward_context
 
     _fsdp_shard_conditions = _DEFAULT_MATRIXGAME_CONFIG._fsdp_shard_conditions
     _compile_conditions = _DEFAULT_MATRIXGAME_CONFIG._compile_conditions
@@ -1096,4 +1097,4 @@ class CausalMatrixGameWanModel(BaseDiT):
                     if block.action_model.proj_keyboard.bias is not None:
                         nn.init.zeros_(block.action_model.proj_keyboard.bias)
                 except AttributeError:
-                    pass
+                    pass

fastvideo/models/dits/matrixgame/model.py

@@ -61,14 +61,17 @@ class MatrixGameTimeImageEmbedding(nn.Module):
         timestep_proj = self.time_modulation(temb)
 
         # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states
-        # and return None for the text embedding part
         
         if encoder_hidden_states_image is not None:
             assert self.image_embedder is not None
             encoder_hidden_states_image = self.image_embedder(
                 encoder_hidden_states_image)
 
-        return temb, timestep_proj, None, encoder_hidden_states_image
+        encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]),
+                                            device=temb.device,
+                                            dtype=temb.dtype)
+
+        return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
 
 
 class MatrixGameCrossAttention(WanSelfAttention):
@@ -313,6 +316,7 @@ class MatrixGameWanModel(BaseDiT):
         inner_dim = config.num_attention_heads * config.attention_head_dim
         self.hidden_size = config.hidden_size
         self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.out_channels
         self.patch_size = config.patch_size
 
         # 1. Patch & position embedding
@@ -395,7 +399,8 @@ class MatrixGameWanModel(BaseDiT):
             self.num_attention_heads,
             rope_dim_list,
             dtype=torch.float32 if current_platform.is_mps() else torch.float64,
-            rope_theta=10000)
+            rope_theta=10000,
+            do_sp_sharding=True)
         freqs_cos = freqs_cos.to(hidden_states.device)
         freqs_sin = freqs_sin.to(hidden_states.device)
         freqs_cis = (freqs_cos.float(),

fastvideo/models/dits/matrixgame/utils.py

@@ -157,8 +157,8 @@ def load_initial_image(image_path: str = None) -> Image.Image:
     return Image.new("RGB", (640, 352), (128, 128, 128))
 
 def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None):
-    if keyboard_dim not in (2, 4, 7):
-        raise ValueError(f"keyboard_dim must be 2, 4, or 7, got {keyboard_dim}")
+    if keyboard_dim not in (2, 4, 6, 7):
+        raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}")
     if num_frames % 4 != 1:
         raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.")
 
@@ -181,6 +181,11 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No
         actions_double_action = []
         actions_single_camera = ["camera_l", "camera_r"]
         keyboard_idx = {"forward": 0, "back": 1}
+    elif keyboard_dim == 6:
+        actions_single_action = ["forward", "back", "left", "right"]
+        actions_double_action = ["forward_left", "forward_right"]
+        actions_single_camera = ["camera_l", "camera_r"]
+        keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3, "t1": 4, "t2": 5}
     else:  # keyboard_dim == 7
         # Temple Run model: still, w, s, left, right, a, d (no mouse)
         actions_single_action = ["forward", "back", "left", "right"]

fastvideo/models/dits/wangame/__init__.py

@@ -0,0 +1,12 @@
+from .model import WanGameActionTransformer3DModel
+from .causal_model import (CausalWanGameTransformer3DModel,
+                           CausalWanTransformer3DModel)
+from .hyworld_action_module import WanGameActionTimeImageEmbedding, WanGameActionSelfAttention
+
+__all__ = [
+    "WanGameActionTransformer3DModel",
+    "CausalWanTransformer3DModel",
+    "CausalWanGameTransformer3DModel",
+    "WanGameActionTimeImageEmbedding",
+    "WanGameActionSelfAttention",
+]

fastvideo/models/dits/wangame/causal_model.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any
+
+import torch
+
+from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig
+from fastvideo.models.dits.matrixgame.causal_model import CausalMatrixGameWanModel
+
+_DEFAULT_WANGAME_CAUSAL_CONFIG = WanGameVideoConfig()
+
+
+class CausalWanTransformer3DModel(CausalMatrixGameWanModel):
+    supports_action_input = False
+
+    _fsdp_shard_conditions = _DEFAULT_WANGAME_CAUSAL_CONFIG._fsdp_shard_conditions
+    _compile_conditions = _DEFAULT_WANGAME_CAUSAL_CONFIG._compile_conditions
+    _supported_attention_backends = _DEFAULT_WANGAME_CAUSAL_CONFIG._supported_attention_backends
+    param_names_mapping = _DEFAULT_WANGAME_CAUSAL_CONFIG.param_names_mapping
+    reverse_param_names_mapping = _DEFAULT_WANGAME_CAUSAL_CONFIG.reverse_param_names_mapping
+    lora_param_names_mapping = _DEFAULT_WANGAME_CAUSAL_CONFIG.lora_param_names_mapping
+
+    def __init__(self,
+                 config: WanGameVideoConfig,
+                 hf_config: dict[str, Any],
+                 **kwargs) -> None:
+        super().__init__(config=config, hf_config=hf_config, **kwargs)
+
+    def _normalize_action_inputs(
+        self,
+        mouse_cond: torch.Tensor | None,
+        keyboard_cond: torch.Tensor | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        # drop
+        if len(getattr(self, "action_config", {})) == 0:
+            return None, None
+        return mouse_cond, keyboard_cond
+
+    def _forward_inference(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        mouse_cond: torch.Tensor | None = None,
+        keyboard_cond: torch.Tensor | None = None,
+        kv_cache: dict | None = None,
+        kv_cache_mouse: dict | None = None,
+        kv_cache_keyboard: dict | None = None,
+        crossattn_cache: dict | None = None,
+        current_start: int = 0,
+        cache_start: int = 0,
+        start_frame: int = 0,
+        **kwargs,
+    ) -> torch.Tensor:
+        mouse_cond, keyboard_cond = self._normalize_action_inputs(
+            mouse_cond, keyboard_cond)
+        return super()._forward_inference(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            encoder_hidden_states_image=encoder_hidden_states_image,
+            mouse_cond=mouse_cond,
+            keyboard_cond=keyboard_cond,
+            kv_cache=kv_cache,
+            kv_cache_mouse=kv_cache_mouse,
+            kv_cache_keyboard=kv_cache_keyboard,
+            crossattn_cache=crossattn_cache,
+            current_start=current_start,
+            cache_start=cache_start,
+            start_frame=start_frame,
+            **kwargs,
+        )
+
+    def _forward_train(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        mouse_cond: torch.Tensor | None = None,
+        keyboard_cond: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        mouse_cond, keyboard_cond = self._normalize_action_inputs(
+            mouse_cond, keyboard_cond)
+        return super()._forward_train(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            encoder_hidden_states_image=encoder_hidden_states_image,
+            mouse_cond=mouse_cond,
+            keyboard_cond=keyboard_cond,
+            **kwargs,
+        )
+
+
+class CausalWanGameTransformer3DModel(CausalWanTransformer3DModel):
+    pass

fastvideo/models/dits/wangame/hyworld_action_module.py

@@ -0,0 +1,289 @@
+import math
+
+import torch
+import torch.nn as nn
+
+from fastvideo.layers.visual_embedding import TimestepEmbedder, ModulateProjection, timestep_embedding
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.attention import DistributedAttention
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.dits.wanvideo import WanImageEmbedding
+
+from fastvideo.models.dits.hyworld.camera_rope import prope_qkv
+from fastvideo.layers.rotary_embedding import _apply_rotary_emb
+from fastvideo.layers.mlp import MLP
+
+class WanGameActionTimeImageEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        time_freq_dim: int,
+        image_embed_dim: int | None = None,
+    ):
+        super().__init__()
+
+        self.time_freq_dim = time_freq_dim
+        self.time_embedder = TimestepEmbedder(
+            dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
+        self.time_modulation = ModulateProjection(dim,
+                                                  factor=6,
+                                                  act_layer="silu")
+        
+        self.image_embedder = None
+        if image_embed_dim is not None:
+            self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
+
+        self.action_embedder = MLP(
+            time_freq_dim,
+            dim,
+            dim,
+            bias=True,
+            act_type="silu"
+        )
+        # Initialize fc_in with kaiming_uniform (same as nn.Linear default)
+        nn.init.kaiming_uniform_(self.action_embedder.fc_in.weight, a=math.sqrt(5))
+        # Initialize fc_out with zeros for residual-like behavior
+        nn.init.zeros_(self.action_embedder.fc_out.weight)
+        if self.action_embedder.fc_out.bias is not None:
+            nn.init.zeros_(self.action_embedder.fc_out.bias)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        action: torch.Tensor,
+        encoder_hidden_states: torch.Tensor, # Kept for interface compatibility
+        encoder_hidden_states_image: torch.Tensor | None = None,
+        timestep_seq_len: int | None = None,
+    ):
+        """
+        Args:
+            timestep: [B] diffusion timesteps (one per batch sample)
+            action: [B, T] action labels (one per frame per batch sample)
+        
+        Returns:
+            temb: [B*T, dim] combined timestep + action embedding
+            timestep_proj: [B*T, 6*dim] modulation projection
+        """
+        # timestep: [B] -> temb: [B, dim]
+        temb = self.time_embedder(timestep, timestep_seq_len)
+        
+        # Handle action embedding for batch > 1
+        # action shape: [B, T] where B=batch_size, T=num_frames
+        batch_size = action.shape[0]
+        num_frames = action.shape[1]
+        
+        # Compute action embeddings: [B, T] -> [B*T] -> [B*T, dim]
+        action_flat = action.flatten()  # [B*T]
+        action_emb = timestep_embedding(action_flat, self.time_freq_dim)
+        action_embedder_dtype = next(iter(self.action_embedder.parameters())).dtype
+        if (
+            action_emb.dtype != action_embedder_dtype
+            and action_embedder_dtype != torch.int8
+        ):
+            action_emb = action_emb.to(action_embedder_dtype)
+        action_emb = self.action_embedder(action_emb).type_as(temb)  # [B*T, dim]
+        
+        # Expand temb to match action_emb: [B, dim] -> [B, T, dim] -> [B*T, dim]
+        # Each batch's temb is repeated for all its frames
+        temb_expanded = temb.unsqueeze(1).expand(-1, num_frames, -1)  # [B, T, dim]
+        temb_expanded = temb_expanded.reshape(batch_size * num_frames, -1)  # [B*T, dim]
+        
+        # Add action embedding to expanded temb
+        temb = temb_expanded + action_emb  # [B*T, dim]
+
+        timestep_proj = self.time_modulation(temb)  # [B*T, 6*dim]
+
+        # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states
+        
+        if encoder_hidden_states_image is not None:
+            assert self.image_embedder is not None
+            encoder_hidden_states_image = self.image_embedder(
+                encoder_hidden_states_image)
+
+        encoder_hidden_states = torch.zeros((batch_size, 0, temb.shape[-1]),
+                                            device=temb.device,
+                                            dtype=temb.dtype)
+
+        return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image
+
+class WanGameActionSelfAttention(nn.Module):
+    """
+    Self-attention module with support for:
+    - Standard RoPE-based attention
+    - Camera PRoPE-based attention (when viewmats and Ks are provided)
+    - KV caching for autoregressive generation
+    """
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 local_attn_size: int = -1,
+                 sink_size: int = 0,
+                 qk_norm=True,
+                 eps=1e-6) -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560
+
+        # Scaled dot product attention (using DistributedAttention for SP support)
+        self.attn = DistributedAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
+                                          AttentionBackendEnum.TORCH_SDPA))
+
+    def forward(self,
+                q: torch.Tensor,
+                k: torch.Tensor,
+                v: torch.Tensor,
+                freqs_cis: tuple[torch.Tensor, torch.Tensor],
+                kv_cache: dict | None = None,
+                current_start: int = 0,
+                cache_start: int | None = None,
+                viewmats: torch.Tensor | None = None,
+                Ks: torch.Tensor | None = None,
+                is_cache: bool = False,
+                attention_mask: torch.Tensor | None = None):
+        """
+        Forward pass with camera PRoPE attention combining standard RoPE and projective positional encoding.
+        
+        Args:
+            q, k, v: Query, key, value tensors [B, L, num_heads, head_dim]
+            freqs_cis: RoPE frequency cos/sin tensors
+            kv_cache: KV cache dict (may have None values for training)
+            current_start: Current position for KV cache
+            cache_start: Cache start position
+            viewmats: Camera view matrices for PRoPE [B, cameras, 4, 4]
+            Ks: Camera intrinsics for PRoPE [B, cameras, 3, 3]
+            is_cache: Whether to store to KV cache (for inference)
+            attention_mask: Attention mask [B, L] (1 = attend, 0 = mask)
+        """
+        if cache_start is None:
+            cache_start = current_start
+
+        # Apply RoPE manually
+        cos, sin = freqs_cis
+        query_rope = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
+        key_rope = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)
+        value_rope = v
+
+        # # DEBUG: Check camera matrices
+        # if self.training and torch.distributed.get_rank() == 0:
+        #     vm_info = f"viewmats={viewmats.shape if viewmats is not None else None}"
+        #     ks_info = f"Ks={Ks.shape if Ks is not None else None}"
+        #     vm_nonzero = (viewmats != 0).sum().item() if viewmats is not None else 0
+        #     ks_nonzero = (Ks != 0).sum().item() if Ks is not None else 0
+        #     print(f"[DEBUG] PRoPE input: {vm_info} nonzero={vm_nonzero}, {ks_info} nonzero={ks_nonzero}", flush=True)
+        
+        # Get PRoPE transformed q, k, v
+        query_prope, key_prope, value_prope, apply_fn_o = prope_qkv(
+            q.transpose(1, 2),  # [B, num_heads, L, head_dim]
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+            viewmats=viewmats,
+            Ks=Ks,
+            patches_x=40,  # hardcoded for now
+            patches_y=22,
+        )
+        # PRoPE returns [B, num_heads, L, head_dim], convert to [B, L, num_heads, head_dim]
+        query_prope = query_prope.transpose(1, 2)
+        key_prope = key_prope.transpose(1, 2)
+        value_prope = value_prope.transpose(1, 2)
+        
+        # # DEBUG: Check prope_qkv output
+        # if self.training and torch.distributed.get_rank() == 0:
+        #     q_nz = (query_prope != 0).sum().item()
+        #     k_nz = (key_prope != 0).sum().item()
+        #     v_nz = (value_prope != 0).sum().item()
+        #     print(f"[DEBUG] prope_qkv output: q_nonzero={q_nz}, k_nonzero={k_nz}, v_nonzero={v_nz}", flush=True)
+
+        # KV cache handling
+        if kv_cache is not None:
+            cache_key = kv_cache.get("k", None)
+            cache_value = kv_cache.get("v", None)
+
+            if cache_value is not None and not is_cache:
+                cache_key_rope, cache_key_prope = cache_key.chunk(2, dim=-1)
+                cache_value_rope, cache_value_prope = cache_value.chunk(2, dim=-1)
+
+                key_rope = torch.cat([cache_key_rope, key_rope], dim=1)
+                value_rope = torch.cat([cache_value_rope, value_rope], dim=1)
+                key_prope = torch.cat([cache_key_prope, key_prope], dim=1)
+                value_prope = torch.cat([cache_value_prope, value_prope], dim=1)
+
+            if is_cache:
+                # Store to cache (update input dict directly)
+                kv_cache["k"] = torch.cat([key_rope, key_prope], dim=-1)
+                kv_cache["v"] = torch.cat([value_rope, value_prope], dim=-1)
+
+        # Concatenate rope and prope paths (matching original)
+        query_all = torch.cat([query_rope, query_prope], dim=0)
+        key_all = torch.cat([key_rope, key_prope], dim=0)
+        value_all = torch.cat([value_rope, value_prope], dim=0)
+
+        # Check if Q and KV have different sequence lengths (KV cache mode)
+        # In this case, use LocalAttention (supports different Q/KV lengths)
+        if query_all.shape[1] != key_all.shape[1]:
+            # KV cache mode: Q has new tokens only, KV has cached + new tokens
+            # Use LocalAttention which supports different Q/KV lengths
+            # LocalAttention will use the appropriate backend (SageAttn, FlashAttn, etc.)
+            if not hasattr(self, '_kv_cache_attn'):
+                from fastvideo.attention import LocalAttention
+                self._kv_cache_attn = LocalAttention(
+                    num_heads=self.num_heads,
+                    head_size=self.head_dim,
+                    causal=False,
+                    supported_attention_backends=(AttentionBackendEnum.SAGE_ATTN,
+                                                  AttentionBackendEnum.FLASH_ATTN,
+                                                  AttentionBackendEnum.TORCH_SDPA)
+                )
+            hidden_states_all = self._kv_cache_attn(query_all, key_all, value_all)
+        else:
+            # Same sequence length: use DistributedAttention (supports SP)
+            # Create default attention mask if not provided
+            # NOTE: query_all has shape [2*B, L, ...] (rope+prope concatenated), so mask needs 2*B
+            if attention_mask is None:
+                batch_size, seq_len = q.shape[0], q.shape[1]
+                attention_mask = torch.ones(batch_size * 2, seq_len, device=q.device, dtype=q.dtype)
+        
+            from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder
+            attn_metadata_builder = SDPAMetadataBuilder
+            if q.dtype != torch.float32:
+                try:
+                    from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+                    attn_metadata_builder = FlashAttnMetadataBuilder
+                except ImportError:
+                    pass  # fall back to SDPA
+            attn_metadata = attn_metadata_builder().build(
+                current_timestep=0,
+                attn_mask=attention_mask,
+            )
+            with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
+                hidden_states_all, _ = self.attn(query_all, key_all, value_all, attention_mask=attention_mask)
+
+        hidden_states_rope, hidden_states_prope = hidden_states_all.chunk(2, dim=0)
+        
+        # # DEBUG: Check attention output and apply_fn_o
+        # if self.training and torch.distributed.get_rank() == 0:
+        #     attn_all_nz = (hidden_states_all != 0).sum().item()
+        #     rope_nz = (hidden_states_rope != 0).sum().item()
+        #     prope_before = (hidden_states_prope != 0).sum().item()
+        #     print(f"[DEBUG] attn output: all_nonzero={attn_all_nz}, rope_nonzero={rope_nz}, prope_before_apply={prope_before}", flush=True)
+        
+        hidden_states_prope = apply_fn_o(hidden_states_prope.transpose(1, 2)).transpose(1, 2)
+        
+        # # DEBUG: Check after apply_fn_o
+        # if self.training and torch.distributed.get_rank() == 0:
+        #     prope_after = (hidden_states_prope != 0).sum().item()
+        #     print(f"[DEBUG] prope_after_apply_fn_o={prope_after}", flush=True)
+
+        return hidden_states_rope, hidden_states_prope

fastvideo/models/dits/wangame/model.py

@@ -0,0 +1,427 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig
+from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift,
+                                        RMSNorm, ScaleResidual,
+                                        ScaleResidualLayerNormScaleShift)
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.mlp import MLP
+from fastvideo.layers.rotary_embedding import (_apply_rotary_emb,
+                                               get_rotary_pos_embed)
+from fastvideo.layers.visual_embedding import PatchEmbed
+from fastvideo.logger import init_logger
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.dits.wanvideo import WanI2VCrossAttention
+from fastvideo.platforms import AttentionBackendEnum, current_platform
+
+# Import ActionModule
+from fastvideo.models.dits.wangame.hyworld_action_module import WanGameActionTimeImageEmbedding, WanGameActionSelfAttention
+
+logger = init_logger(__name__)
+
+
+class WanGameCrossAttention(WanI2VCrossAttention):
+    def forward(self, x, context, context_lens=None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        context_img = context
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
+        k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view(
+            b, -1, n, d)
+        v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
+        img_x = self.attn(q, k_img, v_img)
+
+        # output
+        x = img_x.flatten(2)
+        x, _ = self.to_out(x)
+        return x
+
+class WanGameActionTransformerBlock(nn.Module):
+    """
+    Transformer block for WAN Action model with support for:
+    - Self-attention with RoPE and camera PRoPE
+    - Cross-attention with text/image context
+    - Feed-forward network with AdaLN modulation
+    """
+
+    def __init__(self,
+                 dim: int,
+                 ffn_dim: int,
+                 num_heads: int,
+                 local_attn_size: int = -1,
+                 sink_size: int = 0,
+                 qk_norm: str = "rms_norm_across_heads",
+                 cross_attn_norm: bool = False,
+                 eps: float = 1e-6,
+                 added_kv_proj_dim: int | None = None,
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+                 prefix: str = ""):
+        super().__init__()
+
+        # 1. Self-attention
+        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True)
+        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        
+        self.attn1 = WanGameActionSelfAttention(
+            dim,
+            num_heads,
+            local_attn_size=local_attn_size,
+            sink_size=sink_size,
+            qk_norm=qk_norm,
+            eps=eps)
+        
+        self.hidden_dim = dim
+        self.num_attention_heads = num_heads
+        self.local_attn_size = local_attn_size
+        dim_head = dim // num_heads
+
+        if qk_norm == "rms_norm":
+            self.norm_q = RMSNorm(dim_head, eps=eps)
+            self.norm_k = RMSNorm(dim_head, eps=eps)
+        elif qk_norm == "rms_norm_across_heads":
+            self.norm_q = RMSNorm(dim, eps=eps)
+            self.norm_k = RMSNorm(dim, eps=eps)
+        else:
+            raise ValueError(f"QK Norm type {qk_norm} not supported")
+
+        assert cross_attn_norm is True
+        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim,
+            norm_type="layer",
+            eps=eps,
+            elementwise_affine=True,
+            compute_dtype=torch.float32)
+
+        # 2. Cross-attention (I2V only for now)
+        self.attn2 = WanGameCrossAttention(dim,
+                                          num_heads,
+                                          qk_norm=qk_norm,
+                                          eps=eps)
+        # norm3 for FFN input 
+        self.norm3 = LayerNormScaleShift(dim, norm_type="layer", eps=eps,
+                                         elementwise_affine=False)
+
+        # 3. Feed-forward
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.mlp_residual = ScaleResidual()
+
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+        # PRoPE output projection (initialized via add_discrete_action_parameters on the model)
+        self.to_out_prope = ReplicatedLinear(dim, dim, bias=True)
+        nn.init.zeros_(self.to_out_prope.weight)
+        if self.to_out_prope.bias is not None:
+            nn.init.zeros_(self.to_out_prope.bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        kv_cache: dict | None = None,
+        crossattn_cache: dict | None = None,
+        current_start: int = 0,
+        cache_start: int | None = None,
+        viewmats: torch.Tensor | None = None,
+        Ks: torch.Tensor | None = None,
+        is_cache: bool = False,
+    ) -> torch.Tensor:
+        if hidden_states.dim() == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        num_frames = temb.shape[1]
+        frame_seqlen = hidden_states.shape[1] // num_frames
+        bs, seq_length, _ = hidden_states.shape
+        orig_dtype = hidden_states.dtype
+
+        # Cast temb to float32 for scale/shift computation
+        e = self.scale_shift_table + temb.float()
+        assert e.shape == (bs, num_frames, 6, self.hidden_dim)
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2)
+
+        # 1. Self-attention
+        norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+            (1 + scale_msa) + shift_msa).to(orig_dtype).flatten(1, 2)
+        
+        query, _ = self.to_q(norm_hidden_states)
+        key, _ = self.to_k(norm_hidden_states)
+        value, _ = self.to_v(norm_hidden_states)
+
+        if self.norm_q is not None:
+            query = self.norm_q.forward_native(query)
+        if self.norm_k is not None:
+            key = self.norm_k.forward_native(key)
+
+        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+
+        # Self-attention with camera PRoPE
+        attn_output_rope, attn_output_prope = self.attn1(
+            query, key, value, freqs_cis,
+            kv_cache, current_start, cache_start, viewmats, Ks,
+            is_cache=is_cache
+        )
+        # Combine rope and prope outputs
+        attn_output_rope = attn_output_rope.flatten(2)
+        attn_output_rope, _ = self.to_out(attn_output_rope)
+        attn_output_prope = attn_output_prope.flatten(2)
+        
+        # # DEBUG: Check if prope input is zero
+        # if self.training and torch.distributed.get_rank() == 0:
+        #     prope_nonzero = (attn_output_prope != 0).sum().item()
+        #     prope_total = attn_output_prope.numel()
+        #     if prope_nonzero == 0:
+        #         print(f"[DEBUG] to_out_prope INPUT is ALL ZEROS! shape={attn_output_prope.shape}", flush=True)
+        
+        attn_output_prope, _ = self.to_out_prope(attn_output_prope)
+        attn_output = attn_output_rope.squeeze(1) + attn_output_prope.squeeze(1)
+
+        # Self-attention residual + norm in float32
+        null_shift = null_scale = torch.zeros(1, device=hidden_states.device, dtype=torch.float32)
+        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
+            hidden_states.float(), attn_output.float(), gate_msa, null_shift, null_scale)
+        hidden_states = hidden_states.type_as(attn_output)
+        norm_hidden_states = norm_hidden_states.type_as(attn_output)
+
+        # 2. Cross-attention
+        attn_output = self.attn2(norm_hidden_states.to(orig_dtype),
+                                 context=encoder_hidden_states,
+                                 context_lens=None)
+        # Cross-attention residual in bfloat16
+        hidden_states = hidden_states + attn_output
+        
+        # norm3 for FFN input in float32
+        norm_hidden_states = self.norm3(
+            hidden_states.float(), c_shift_msa, c_scale_msa
+        ).type_as(hidden_states)
+
+        # 3. Feed-forward
+        ff_output = self.ffn(norm_hidden_states.to(orig_dtype))
+        hidden_states = self.mlp_residual(hidden_states.float(), ff_output.float(), c_gate_msa)
+        hidden_states = hidden_states.to(orig_dtype)  # Cast back to original dtype
+
+        return hidden_states
+
+class WanGameActionTransformer3DModel(BaseDiT):
+    """
+    WAN Action Transformer 3D Model for video generation with action conditioning.
+
+    Extends the base WAN video model with:
+    - Action embedding support for controllable generation
+    - camera PRoPE attention for 3D-aware generation
+    - KV caching for autoregressive inference
+    """
+    _uses_discrete_action = True  # Uses single action tensor [B, T] instead of keyboard/mouse KV caches
+    _kv_cache_head_dim_multiplier = 2  # PRoPE stores cat([rope, prope], dim=-1) in KV cache
+
+    _fsdp_shard_conditions = WanGameVideoConfig()._fsdp_shard_conditions
+    _compile_conditions = WanGameVideoConfig()._compile_conditions
+    _supported_attention_backends = WanGameVideoConfig()._supported_attention_backends
+    param_names_mapping = WanGameVideoConfig().param_names_mapping
+    reverse_param_names_mapping = WanGameVideoConfig().reverse_param_names_mapping
+    lora_param_names_mapping = WanGameVideoConfig().lora_param_names_mapping
+
+    def __init__(self, config: WanGameVideoConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        inner_dim = config.num_attention_heads * config.attention_head_dim
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_dim = config.attention_head_dim
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_channels_latents = config.num_channels_latents
+        self.patch_size = config.patch_size
+        self.local_attn_size = config.local_attn_size
+        self.inner_dim = inner_dim
+
+        # 1. Patch & position embedding
+        self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
+                                          embed_dim=inner_dim,
+                                          patch_size=config.patch_size,
+                                          flatten=False)
+
+        # 2. Condition embeddings (with action support)
+        self.condition_embedder = WanGameActionTimeImageEmbedding(
+            dim=inner_dim,
+            time_freq_dim=config.freq_dim,
+            image_embed_dim=config.image_dim,
+        )
+
+        # 3. Transformer blocks
+        self.blocks = nn.ModuleList([
+            WanGameActionTransformerBlock(
+                inner_dim,
+                config.ffn_dim,
+                config.num_attention_heads,
+                config.local_attn_size,
+                config.sink_size,
+                config.qk_norm,
+                config.cross_attn_norm,
+                config.eps,
+                config.added_kv_proj_dim,
+                supported_attention_backends=self._supported_attention_backends,
+                prefix=f"{config.prefix}.blocks.{i}")
+            for i in range(config.num_layers)
+        ])
+
+        # 4. Output norm & projection
+        self.norm_out = LayerNormScaleShift(inner_dim,
+                                            norm_type="layer",
+                                            eps=config.eps,
+                                            elementwise_affine=False,
+                                            dtype=torch.float32)
+        self.proj_out = nn.Linear(
+            inner_dim, config.out_channels * math.prod(config.patch_size))
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+        self.gradient_checkpointing = False
+
+        # Causal-specific
+        self.num_frame_per_block = config.arch_config.num_frames_per_block
+        assert self.num_frame_per_block <= 3
+
+        self.__post_init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor],
+        guidance=None,
+        action: torch.Tensor | None = None,
+        viewmats: torch.Tensor | None = None,
+        Ks: torch.Tensor | None = None,
+        kv_cache: list[dict] | None = None,
+        crossattn_cache: list[dict] | None = None,
+        current_start: int = 0,
+        cache_start: int = 0,
+        start_frame: int = 0,
+        is_cache: bool = False,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward pass for both training and inference with KV caching.
+        
+        Args:
+            hidden_states: Video latents [B, C, T, H, W]
+            encoder_hidden_states: Text embeddings [B, L, D]
+            timestep: Timestep tensor
+            encoder_hidden_states_image: Optional image embeddings
+            action: Action tensor [B, T] for per-frame conditioning
+            viewmats: Camera view matrices for PRoPE [B, T, 4, 4]
+            Ks: Camera intrinsics for PRoPE [B, T, 3, 3]
+            kv_cache: KV cache for autoregressive inference (list of dicts per layer)
+            crossattn_cache: Cross-attention cache for inference
+            current_start: Current position for KV cache
+            cache_start: Cache start position
+            start_frame: RoPE offset for new frames in autoregressive mode
+            is_cache: If True, populate KV cache and return early (cache-only mode)
+        """
+        orig_dtype = hidden_states.dtype
+        # if not isinstance(encoder_hidden_states, torch.Tensor):
+        #     encoder_hidden_states = encoder_hidden_states[0]
+        if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0:
+            encoder_hidden_states_image = encoder_hidden_states_image[0]
+        # else:
+        #     encoder_hidden_states_image = None
+
+        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Get rotary embeddings
+        d = self.hidden_size // self.num_attention_heads
+        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
+        freqs_cos, freqs_sin = get_rotary_pos_embed(
+            (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width),
+            self.hidden_size,
+            self.num_attention_heads,
+            rope_dim_list,
+            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            rope_theta=10000,
+            start_frame=start_frame
+        )
+        freqs_cos = freqs_cos.to(hidden_states.device)
+        freqs_sin = freqs_sin.to(hidden_states.device)
+        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        if timestep.dim() == 2:
+            # condition_embedder expects [B] timestep, not [B*T].
+            # All frames share the same timestep, so take first column.
+            timestep = timestep[:, 0]
+
+        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
+            timestep, action, encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image)
+        
+        # condition_embedder returns:
+        # - temb: [B*T, dim] where T = post_patch_num_frames
+        # - timestep_proj: [B*T, 6*dim]
+        # Reshape to [B, T, 6, dim] for transformer blocks
+        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size))  # [B*T, 6, dim]
+        timestep_proj = timestep_proj.view(batch_size, post_patch_num_frames, 6, self.hidden_size)  # [B, T, 6, dim]
+
+        encoder_hidden_states = encoder_hidden_states_image
+
+        # Transformer blocks
+        for block_idx, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis,
+                    kv_cache[block_idx] if kv_cache else None,
+                    crossattn_cache[block_idx] if crossattn_cache else None,
+                    current_start, cache_start,
+                    viewmats, Ks, is_cache)
+            else:
+                hidden_states = block(
+                    hidden_states, encoder_hidden_states, timestep_proj, freqs_cis,
+                    kv_cache[block_idx] if kv_cache else None,
+                    crossattn_cache[block_idx] if crossattn_cache else None,
+                    current_start, cache_start,
+                    viewmats, Ks, is_cache)
+
+        # If cache-only mode, return early
+        if is_cache:
+            return kv_cache
+
+        # Output norm, projection & unpatchify
+        # temb is [B*T, dim], reshape to [B, T, 1, dim]
+        temb = temb.view(batch_size, post_patch_num_frames, -1).unsqueeze(2)  # [B, T, 1, dim]
+        
+        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2)
+        hidden_states = self.norm_out(hidden_states, shift, scale)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
+                                              post_patch_height,
+                                              post_patch_width, p_t, p_h, p_w,
+                                              -1)
+        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return output

fastvideo/models/dits/wanvideo.py

@@ -76,7 +76,7 @@ class WanTimeTextImageEmbedding(nn.Module):
                                  dim,
                                  dim,
                                  bias=True,
-                                 act_type="gelu_pytorch_tanh")
+                                 act_type="gelu_pytorch_tanh") if text_embed_dim > 0 else None
 
         self.image_embedder = None
         if image_embed_dim is not None:
@@ -92,7 +92,10 @@ class WanTimeTextImageEmbedding(nn.Module):
         temb = self.time_embedder(timestep, timestep_seq_len)
         timestep_proj = self.time_modulation(temb)
 
-        encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        if self.text_embedder is not None:
+            encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        else:
+            encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), device=temb.device, dtype=temb.dtype)
         if encoder_hidden_states_image is not None:
             assert self.image_embedder is not None
             encoder_hidden_states_image = self.image_embedder(
@@ -179,7 +182,10 @@ class WanT2VCrossAttention(WanSelfAttention):
             v = self.to_v(context)[0].view(b, -1, n, d)
 
         # compute attention
-        x = self.attn(q, k, v)
+        if k.size(1) > 0:
+            x = self.attn(q, k, v)
+        else:
+            x = torch.zeros_like(q)
 
         # output
         x = x.flatten(2)
@@ -227,7 +233,10 @@ class WanI2VCrossAttention(WanSelfAttention):
         v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
         img_x = self.attn(q, k_img, v_img)
         # compute attention
-        x = self.attn(q, k, v)
+        if k.size(1) > 0:
+            x = self.attn(q, k, v)
+        else:
+            x = torch.zeros_like(q)
 
         # output
         x = x.flatten(2)
@@ -633,7 +642,7 @@ class WanTransformer3DModel(CachableDiT):
         enable_teacache = forward_batch is not None and forward_batch.enable_teacache
 
         orig_dtype = hidden_states.dtype
-        if not isinstance(encoder_hidden_states, torch.Tensor):
+        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
         if isinstance(encoder_hidden_states_image,
                       list) and len(encoder_hidden_states_image) > 0:
@@ -705,8 +714,11 @@ class WanTransformer3DModel(CachableDiT):
             timestep_proj = timestep_proj.unflatten(1, (6, -1))
 
         if encoder_hidden_states_image is not None:
-            encoder_hidden_states = torch.concat(
-                [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+            if encoder_hidden_states is not None:
+                encoder_hidden_states = torch.concat(
+                    [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+            else:
+                encoder_hidden_states = encoder_hidden_states_image
 
         if current_platform.is_mps() or current_platform.is_npu():
             encoder_hidden_states = encoder_hidden_states.to(orig_dtype)

fastvideo/models/registry.py

@@ -42,8 +42,10 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
-    "MatrixGameWanModel": ("dits", "matrix_game", "MatrixGameWanModel"),
-    "CausalMatrixGameWanModel": ("dits", "matrix_game", "CausalMatrixGameWanModel"),
+    "MatrixGameWanModel": ("dits", "matrixgame", "MatrixGameWanModel"),
+    "CausalMatrixGameWanModel": ("dits", "matrixgame", "CausalMatrixGameWanModel"),
+    "WanGameActionTransformer3DModel": ("dits", "wangame", "WanGameActionTransformer3DModel"),
+    "CausalWanGameTransformer3DModel": ("dits", "wangame", "CausalWanGameTransformer3DModel"),
 }
 
 _TEXT_ENCODER_MODELS = {

fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py

@@ -96,28 +96,47 @@ class MatrixGameCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
         self._vae_cache = None
 
     def streaming_step(self, keyboard_action, mouse_action) -> ForwardBatch:
+        import time
         denoiser = self._stage_name_mapping["denoising_stage"]
         ctx = denoiser._streaming_ctx
         assert ctx is not None, "streaming_ctx must be set"
 
         start_idx = ctx.start_index
+
+        # Time DiT forward pass
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
         batch = denoiser.streaming_step(keyboard_action, mouse_action)
+        torch.cuda.synchronize()
+        dit_time_ms = (time.perf_counter() - t0) * 1000
+
         end_idx = ctx.start_index
 
         # Decode only the new generated block
+        vae_time_ms = 0.0
         if end_idx > start_idx:
             current_latents = batch.latents[:, :, start_idx:end_idx, :, :]
             args = ctx.fastvideo_args
             decoder = self._stage_name_mapping["decoding_stage"]
+
+            # Time VAE decode
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
             decoded_frames, self._vae_cache = decoder.streaming_decode(
                 current_latents,
                 args,
                 cache=self._vae_cache,
                 is_first_chunk=(start_idx == 0))
+            torch.cuda.synchronize()
+            vae_time_ms = (time.perf_counter() - t0) * 1000
+
             batch.output = decoded_frames
         else:
             batch.output = None
 
+        # Store timings in batch for caller to access
+        batch.stage_timings = {"dit_ms": dit_time_ms, "vae_ms": vae_time_ms}
+
         return batch
 
     def streaming_clear(self) -> None:

fastvideo/pipelines/basic/wan/wangame_causal_dmd_pipeline.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""WanGame causal DMD pipeline implementation."""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+import torch
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import ComposedPipelineBase, ForwardBatch, LoRAPipeline
+
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        MatrixGameCausalDenoisingStage)
+from fastvideo.pipelines.stages.image_encoding import (
+    MatrixGameImageEncodingStage, MatrixGameImageVAEEncodingStage)
+
+logger = init_logger(__name__)
+
+
+class WanGameCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
+    _required_config_modules = [
+        "vae", "transformer", "scheduler", "image_encoder", "image_processor"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        if (self.get_module("text_encoder", None) is not None
+                and self.get_module("tokenizer", None) is not None):
+            self.add_stage(stage_name="prompt_encoding_stage",
+                           stage=TextEncodingStage(
+                               text_encoders=[self.get_module("text_encoder")],
+                               tokenizers=[self.get_module("tokenizer")],
+                           ))
+
+        if (self.get_module("image_encoder", None) is not None
+                and self.get_module("image_processor", None) is not None):
+            self.add_stage(
+                stage_name="image_encoding_stage",
+                stage=MatrixGameImageEncodingStage(
+                    image_encoder=self.get_module("image_encoder"),
+                    image_processor=self.get_module("image_processor"),
+                ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+
+        self.add_stage(
+            stage_name="image_latent_preparation_stage",
+            stage=MatrixGameImageVAEEncodingStage(vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=MatrixGameCausalDenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
+                           scheduler=self.get_module("scheduler"),
+                           pipeline=self,
+                           vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+        logger.info("WanGameCausalDMDPipeline initialized")
+
+    @torch.no_grad()
+    def streaming_reset(self, batch: ForwardBatch,
+                        fastvideo_args: FastVideoArgs):
+        if not self.post_init_called:
+            self.post_init()
+
+        stages_to_run = [
+            "input_validation_stage", "prompt_encoding_stage",
+            "image_encoding_stage", "conditioning_stage",
+            "latent_preparation_stage", "image_latent_preparation_stage"
+        ]
+
+        for stage_name in stages_to_run:
+            if stage_name in self._stage_name_mapping:
+                batch = self._stage_name_mapping[stage_name].forward(
+                    batch, fastvideo_args)
+
+        denoiser = self._stage_name_mapping["denoising_stage"]
+        denoiser.streaming_reset(batch, fastvideo_args)
+        self._vae_cache = None
+
+    def streaming_step(self, keyboard_action, mouse_action) -> ForwardBatch:
+        denoiser = self._stage_name_mapping["denoising_stage"]
+        ctx = denoiser._streaming_ctx
+        assert ctx is not None, "streaming_ctx must be set"
+
+        start_idx = ctx.start_index
+        batch = denoiser.streaming_step(keyboard_action, mouse_action)
+        end_idx = ctx.start_index
+
+        if end_idx > start_idx:
+            current_latents = batch.latents[:, :, start_idx:end_idx, :, :]
+            args = ctx.fastvideo_args
+            decoder = self._stage_name_mapping["decoding_stage"]
+            decoded_frames, self._vae_cache = decoder.streaming_decode(
+                current_latents,
+                args,
+                cache=self._vae_cache,
+                is_first_chunk=(start_idx == 0))
+            batch.output = decoded_frames
+        else:
+            batch.output = None
+
+        return batch
+
+    def streaming_clear(self) -> None:
+        denoiser = self._stage_name_mapping.get("denoising_stage")
+        if denoiser is not None and hasattr(denoiser, "streaming_clear"):
+            denoiser.streaming_clear()
+        self._vae_cache = None
+
+EntryClass = [WanGameCausalDMDPipeline]

fastvideo/pipelines/basic/wan/wangame_i2v_pipeline.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Wan video diffusion pipeline implementation.
+
+This module contains an implementation of the Wan video diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+
+# isort: off
+from fastvideo.pipelines.stages import (
+    ImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage,
+    ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage,
+    TimestepPreparationStage)
+# isort: on
+from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
+    FlowUniPCMultistepScheduler)
+
+logger = init_logger(__name__)
+
+
+class WanGameActionImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
+
+    _required_config_modules = [
+        "vae", "transformer", "scheduler", \
+        "image_encoder", "image_processor"
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(
+            stage_name="image_encoding_stage",
+            stage=ImageEncodingStage(
+                image_encoder=self.get_module("image_encoder"),
+                image_processor=self.get_module("image_processor"),
+            ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer")))
+
+        self.add_stage(stage_name="image_latent_preparation_stage",
+                       stage=ImageVAEEncodingStage(vae=self.get_module("vae")))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+class WanLingBotImageToVideoPipeline(WanGameActionImageToVideoPipeline):
+    pass
+
+
+EntryClass = [WanGameActionImageToVideoPipeline, WanLingBotImageToVideoPipeline]

fastvideo/pipelines/pipeline_registry.py

@@ -33,6 +33,8 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "Cosmos2_5Pipeline": "cosmos",
     "MatrixGamePipeline": "matrixgame",
     "MatrixGameCausalDMDPipeline": "matrixgame",
+    "WanGameActionImageToVideoPipeline": "wan",
+    "WanGameCausalDMDPipeline": "wan",
     "LongCatPipeline": "longcat",
     "LongCatImageToVideoPipeline": "longcat",
     "LongCatVideoContinuationPipeline": "longcat",

fastvideo/pipelines/preprocess/matrixgame/matrixgame_preprocess_pipeline.py

@@ -0,0 +1,303 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
+    BasePreprocessPipeline)
+from fastvideo.pipelines.stages import ImageEncodingStage
+
+
+class PreprocessPipeline_MatrixGame(BasePreprocessPipeline):
+    """I2V preprocessing pipeline implementation."""
+
+    _required_config_modules = ["vae", "image_encoder", "image_processor"]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        self.add_stage(stage_name="image_encoding_stage",
+                       stage=ImageEncodingStage(
+                           image_encoder=self.get_module("image_encoder"),
+                           image_processor=self.get_module("image_processor"),
+                       ))
+
+    def get_pyarrow_schema(self):
+        """Return the PyArrow schema for I2V pipeline."""
+        return pyarrow_schema_matrixgame
+
+    def get_extra_features(self, valid_data: dict[str, Any],
+                           fastvideo_args: FastVideoArgs) -> dict[str, Any]:
+
+        # TODO(will): move these to cpu at some point
+        self.get_module("image_encoder").to(get_local_torch_device())
+        self.get_module("vae").to(get_local_torch_device())
+
+        features = {}
+        """Get CLIP features from the first frame of each video."""
+        first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(
+            0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
+        _, _, num_frames, height, width = valid_data["pixel_values"].shape
+        # latent_height = height // self.get_module(
+        #     "vae").spatial_compression_ratio
+        # latent_width = width // self.get_module("vae").spatial_compression_ratio
+
+        processed_images = []
+        # Frame has values between -1 and 1
+        for frame in first_frame:
+            frame = (frame + 1) * 127.5
+            frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
+            processed_img = self.get_module("image_processor")(
+                images=frame_pil, return_tensors="pt")
+            processed_images.append(processed_img)
+
+        # Get CLIP features
+        pixel_values = torch.cat(
+            [img['pixel_values'] for img in processed_images],
+            dim=0).to(get_local_torch_device())
+        with torch.no_grad():
+            image_inputs = {'pixel_values': pixel_values}
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                clip_features = self.get_module("image_encoder")(**image_inputs)
+            clip_features = clip_features.last_hidden_state
+
+        features["clip_feature"] = clip_features
+        """Get VAE features from the first frame of each video"""
+        video_conditions = []
+        for frame in first_frame:
+            processed_img = frame.to(device="cpu", dtype=torch.float32)
+            processed_img = processed_img.unsqueeze(0).permute(0, 3, 1,
+                                                               2).unsqueeze(2)
+            # (B, H, W, C) -> (B, C, 1, H, W)
+            video_condition = torch.cat([
+                processed_img,
+                processed_img.new_zeros(processed_img.shape[0],
+                                        processed_img.shape[1], num_frames - 1,
+                                        height, width)
+            ],
+                                        dim=2)
+            video_condition = video_condition.to(
+                device=get_local_torch_device(), dtype=torch.float32)
+            video_conditions.append(video_condition)
+
+        video_conditions = torch.cat(video_conditions, dim=0)
+
+        with torch.autocast(device_type="cuda",
+                            dtype=torch.float32,
+                            enabled=True):
+            encoder_outputs = self.get_module("vae").encode(video_conditions)
+
+        latent_condition = encoder_outputs.mean
+        if (hasattr(self.get_module("vae"), "shift_factor")
+                and self.get_module("vae").shift_factor is not None):
+            if isinstance(self.get_module("vae").shift_factor, torch.Tensor):
+                latent_condition -= self.get_module("vae").shift_factor.to(
+                    latent_condition.device, latent_condition.dtype)
+            else:
+                latent_condition -= self.get_module("vae").shift_factor
+
+        if isinstance(self.get_module("vae").scaling_factor, torch.Tensor):
+            latent_condition = latent_condition * self.get_module(
+                "vae").scaling_factor.to(latent_condition.device,
+                                         latent_condition.dtype)
+        else:
+            latent_condition = latent_condition * self.get_module(
+                "vae").scaling_factor
+
+        # mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
+        #                            latent_width)
+        # mask_lat_size[:, :, list(range(1, num_frames))] = 0
+        # first_frame_mask = mask_lat_size[:, :, 0:1]
+        # first_frame_mask = torch.repeat_interleave(
+        #     first_frame_mask,
+        #     dim=2,
+        #     repeats=self.get_module("vae").temporal_compression_ratio)
+        # mask_lat_size = torch.concat(
+        #     [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+        # mask_lat_size = mask_lat_size.view(
+        #     batch_size, -1,
+        #     self.get_module("vae").temporal_compression_ratio, latent_height,
+        #     latent_width)
+        # mask_lat_size = mask_lat_size.transpose(1, 2)
+        # mask_lat_size = mask_lat_size.to(latent_condition.device)
+
+        # image_latent = torch.concat([mask_lat_size, latent_condition], dim=1)
+
+        features["first_frame_latent"] = latent_condition
+
+        if "action_path" in valid_data and valid_data["action_path"]:
+            keyboard_cond_list = []
+            mouse_cond_list = []
+            num_bits = 6
+            for action_path in valid_data["action_path"]:
+                if action_path:
+                    action_data = np.load(action_path, allow_pickle=True)
+                    if isinstance(
+                            action_data,
+                            np.ndarray) and action_data.dtype == np.dtype('O'):
+                        action_dict = action_data.item()
+                        if "keyboard" in action_dict:
+                            keyboard_raw = action_dict["keyboard"]
+                            # Convert 1D bit-flag values to 2D multi-hot encoding
+                            if isinstance(keyboard_raw, np.ndarray):
+                                if keyboard_raw.ndim == 1:
+                                    # [T] -> [T, num_bits]
+                                    T = len(keyboard_raw)
+                                    multi_hot = np.zeros((T, num_bits),
+                                                         dtype=np.float32)
+                                    action_values = keyboard_raw.astype(int)
+                                    for bit_idx in range(num_bits):
+                                        target_idx = (
+                                            2 -
+                                            (bit_idx % 3)) + 3 * (bit_idx // 3)
+                                        if target_idx < num_bits:
+                                            multi_hot[:, target_idx] = (
+                                                (action_values >> bit_idx)
+                                                & 1).astype(np.float32)
+                                    keyboard_cond_list.append(multi_hot)
+                                else:
+                                    # If already 2D, pad to num_bits if necessary
+                                    k_data = keyboard_raw.astype(np.float32)
+                                    if k_data.ndim == 2 and k_data.shape[
+                                            -1] < num_bits:
+                                        padding = np.zeros(
+                                            (k_data.shape[0],
+                                             num_bits - k_data.shape[-1]),
+                                            dtype=np.float32)
+                                        k_data = np.concatenate(
+                                            [k_data, padding], axis=-1)
+                                    keyboard_cond_list.append(k_data)
+                            else:
+                                keyboard_cond_list.append(keyboard_raw)
+                        if "mouse" in action_dict:
+                            mouse_cond_list.append(action_dict["mouse"])
+                    else:
+                        if isinstance(action_data,
+                                      np.ndarray) and action_data.ndim == 1:
+                            T = len(action_data)
+                            multi_hot = np.zeros((T, num_bits),
+                                                 dtype=np.float32)
+                            action_values = action_data.astype(int)
+                            for bit_idx in range(num_bits):
+                                target_idx = (
+                                    2 - (bit_idx % 3)) + 3 * (bit_idx // 3)
+                                if target_idx < num_bits:
+                                    multi_hot[:, target_idx] = (
+                                        (action_values >> bit_idx) & 1).astype(
+                                            np.float32)
+                            keyboard_cond_list.append(multi_hot)
+                        else:
+                            # If already 2D, pad to num_bits if necessary
+                            k_data = action_data.astype(np.float32)
+                            if k_data.ndim == 2 and k_data.shape[-1] < num_bits:
+                                padding = np.zeros(
+                                    (k_data.shape[0],
+                                     num_bits - k_data.shape[-1]),
+                                    dtype=np.float32)
+                                k_data = np.concatenate([k_data, padding],
+                                                        axis=-1)
+                            keyboard_cond_list.append(k_data)
+            if keyboard_cond_list:
+                features["keyboard_cond"] = keyboard_cond_list
+            if mouse_cond_list:
+                features["mouse_cond"] = mouse_cond_list
+
+        return features
+
+    def create_record(
+            self,
+            video_name: str,
+            vae_latent: np.ndarray,
+            text_embedding: np.ndarray,
+            valid_data: dict[str, Any],
+            idx: int,
+            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Create a record for the Parquet dataset with CLIP features."""
+        record = super().create_record(video_name=video_name,
+                                       vae_latent=vae_latent,
+                                       text_embedding=text_embedding,
+                                       valid_data=valid_data,
+                                       idx=idx,
+                                       extra_features=extra_features)
+
+        if extra_features and "clip_feature" in extra_features:
+            clip_feature = extra_features["clip_feature"]
+            record.update({
+                "clip_feature_bytes": clip_feature.tobytes(),
+                "clip_feature_shape": list(clip_feature.shape),
+                "clip_feature_dtype": str(clip_feature.dtype),
+            })
+        else:
+            record.update({
+                "clip_feature_bytes": b"",
+                "clip_feature_shape": [],
+                "clip_feature_dtype": "",
+            })
+
+        if extra_features and "first_frame_latent" in extra_features:
+            first_frame_latent = extra_features["first_frame_latent"]
+            record.update({
+                "first_frame_latent_bytes":
+                first_frame_latent.tobytes(),
+                "first_frame_latent_shape":
+                list(first_frame_latent.shape),
+                "first_frame_latent_dtype":
+                str(first_frame_latent.dtype),
+            })
+        else:
+            record.update({
+                "first_frame_latent_bytes": b"",
+                "first_frame_latent_shape": [],
+                "first_frame_latent_dtype": "",
+            })
+
+        if extra_features and "pil_image" in extra_features:
+            pil_image = extra_features["pil_image"]
+            record.update({
+                "pil_image_bytes": pil_image.tobytes(),
+                "pil_image_shape": list(pil_image.shape),
+                "pil_image_dtype": str(pil_image.dtype),
+            })
+        else:
+            record.update({
+                "pil_image_bytes": b"",
+                "pil_image_shape": [],
+                "pil_image_dtype": "",
+            })
+
+        if extra_features and "keyboard_cond" in extra_features:
+            keyboard_cond = extra_features["keyboard_cond"]
+            record.update({
+                "keyboard_cond_bytes": keyboard_cond.tobytes(),
+                "keyboard_cond_shape": list(keyboard_cond.shape),
+                "keyboard_cond_dtype": str(keyboard_cond.dtype),
+            })
+        else:
+            record.update({
+                "keyboard_cond_bytes": b"",
+                "keyboard_cond_shape": [],
+                "keyboard_cond_dtype": "",
+            })
+
+        if extra_features and "mouse_cond" in extra_features:
+            mouse_cond = extra_features["mouse_cond"]
+            record.update({
+                "mouse_cond_bytes": mouse_cond.tobytes(),
+                "mouse_cond_shape": list(mouse_cond.shape),
+                "mouse_cond_dtype": str(mouse_cond.dtype),
+            })
+        else:
+            record.update({
+                "mouse_cond_bytes": b"",
+                "mouse_cond_shape": [],
+                "mouse_cond_dtype": "",
+            })
+
+        return record
+
+
+EntryClass = PreprocessPipeline_MatrixGame

fastvideo/pipelines/preprocess/preprocess_pipeline_base.py

@@ -320,12 +320,18 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                     "pixel_values":
                     torch.stack(
                         [data["pixel_values"][i] for i in valid_indices]),
-                    "text": [data["text"][i] for i in valid_indices],
+                    "text": [data["text"][i] for i in valid_indices]
+                    if "text" in data else ["" for _ in valid_indices],
                     "path": [data["path"][i] for i in valid_indices],
                     "fps": [data["fps"][i] for i in valid_indices],
                     "duration": [data["duration"][i] for i in valid_indices],
                 }
 
+                if "action_path" in data:
+                    valid_data["action_path"] = [
+                        data["action_path"][i] for i in valid_indices
+                    ]
+
                 # VAE
                 with torch.autocast("cuda", dtype=torch.float32):
                     latents = self.get_module("vae").encode(
@@ -343,28 +349,35 @@ class BasePreprocessPipeline(ComposedPipelineBase):
                     prompt_embeds=[],
                     prompt_attention_mask=[],
                 )
-                assert hasattr(self, "prompt_encoding_stage")
-                result_batch = self.prompt_encoding_stage(batch, fastvideo_args)
-                prompt_embeds, prompt_attention_mask = result_batch.prompt_embeds[
-                    0], result_batch.prompt_attention_mask[0]
-                assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]
 
-                # Get sequence lengths from attention masks (number of 1s)
-                seq_lens = prompt_attention_mask.sum(dim=1)
+                if hasattr(self, "prompt_encoding_stage"):
+                    result_batch = self.prompt_encoding_stage(
+                        batch, fastvideo_args)
+                    prompt_embeds, prompt_attention_mask = result_batch.prompt_embeds[
+                        0], result_batch.prompt_attention_mask[0]
+                    assert prompt_embeds.shape[
+                        0] == prompt_attention_mask.shape[0]
 
-                non_padded_embeds = []
-                non_padded_masks = []
+                    # Get sequence lengths from attention masks (number of 1s)
+                    seq_lens = prompt_attention_mask.sum(dim=1)
 
-                # Process each item in the batch
-                for i in range(prompt_embeds.size(0)):
-                    seq_len = seq_lens[i].item()
-                    # Slice the embeddings and masks to keep only non-padding parts
-                    non_padded_embeds.append(prompt_embeds[i, :seq_len])
-                    non_padded_masks.append(prompt_attention_mask[i, :seq_len])
+                    non_padded_embeds = []
+                    non_padded_masks = []
 
-                # Update the tensors with non-padded versions
-                prompt_embeds = non_padded_embeds
-                prompt_attention_mask = non_padded_masks
+                    # Process each item in the batch
+                    for i in range(prompt_embeds.size(0)):
+                        seq_len = seq_lens[i].item()
+                        # Slice the embeddings and masks to keep only non-padding parts
+                        non_padded_embeds.append(prompt_embeds[i, :seq_len])
+                        non_padded_masks.append(
+                            prompt_attention_mask[i, :seq_len])
+
+                    # Update the tensors with non-padded versions
+                    prompt_embeds = non_padded_embeds
+                    prompt_attention_mask = non_padded_masks
+                else:
+                    bs = len(valid_indices)
+                    prompt_embeds = [torch.zeros(0) for _ in range(bs)]
 
             # Prepare batch data for Parquet dataset
             batch_data = []

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -16,6 +16,8 @@ from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (
     PreprocessPipeline_T2V)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_text import (
     PreprocessPipeline_Text)
+from fastvideo.pipelines.preprocess.matrixgame.matrixgame_preprocess_pipeline import (
+    PreprocessPipeline_MatrixGame)
 from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -60,9 +62,12 @@ def main(args) -> None:
         assert args.flow_shift is not None, "flow_shift is required for ode_trajectory"
         fastvideo_args.pipeline_config.flow_shift = args.flow_shift
         PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
+    elif args.preprocess_task == "matrixgame":
+        PreprocessPipeline = PreprocessPipeline_MatrixGame
     else:
-        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. "
-                         f"Valid options: t2v, i2v, ode_trajectory, text_only")
+        raise ValueError(
+            f"Invalid preprocess task: {args.preprocess_task}. "
+            f"Valid options: t2v, i2v, ode_trajectory, text_only, matrixgame")
 
     logger.info("Preprocess task: %s using %s", args.preprocess_task,
                 PreprocessPipeline.__name__)
@@ -106,11 +111,12 @@ if __name__ == "__main__":
     parser.add_argument("--group_frame", action="store_true")  # TODO
     parser.add_argument("--group_resolution", action="store_true")  # TODO
     parser.add_argument("--flow_shift", type=float, default=None)
-    parser.add_argument("--preprocess_task",
-                        type=str,
-                        default="t2v",
-                        choices=["t2v", "i2v", "text_only", "ode_trajectory"],
-                        help="Type of preprocessing task to run")
+    parser.add_argument(
+        "--preprocess_task",
+        type=str,
+        default="t2v",
+        choices=["t2v", "i2v", "text_only", "ode_trajectory", "matrixgame"],
+        help="Type of preprocessing task to run")
     parser.add_argument("--train_fps", type=int, default=30)
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--text_max_length", type=int, default=256)

fastvideo/pipelines/stages/image_encoding.py

@@ -845,7 +845,18 @@ class MatrixGameImageVAEEncodingStage(ImageVAEEncodingStage):
             num_frames = batch.num_frames if isinstance(
                 batch.num_frames, int) else batch.num_frames[0]
 
+        def _gpu_mem(label=""):
+            a = torch.cuda.memory_allocated() / 1024**3
+            r = torch.cuda.memory_reserved() / 1024**3
+            logger.info("VAE stage [%s]: alloc=%.2fGiB, reserved=%.2fGiB", label, a, r)
+
+        _gpu_mem("before vae.to(GPU)")
         self.vae = self.vae.to(get_local_torch_device())
+        _gpu_mem("after vae.to(GPU)")
+        logger.info("VAE use_feature_cache=%s, use_tiling=%s, use_temporal_tiling=%s",
+                     getattr(self.vae, 'use_feature_cache', 'N/A'),
+                     getattr(self.vae, 'use_tiling', 'N/A'),
+                     getattr(self.vae, 'use_temporal_tiling', 'N/A'))
 
         # Process single image for I2V (latent dimensions computed but not used directly)
 
@@ -906,15 +917,23 @@ class MatrixGameImageVAEEncodingStage(ImageVAEEncodingStage):
         vae_autocast_enabled = (
             vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
 
-        # Encode Image
-        with torch.autocast(device_type="cuda",
-                            dtype=vae_dtype,
-                            enabled=vae_autocast_enabled):
+        _gpu_mem("before encode")
+        logger.info("video_condition shape=%s dtype=%s device=%s",
+                     video_condition.shape, video_condition.dtype,
+                     video_condition.device)
+
+        # Encode Image (no_grad avoids autograd graph that would retain all
+        # intermediate activations across the feature-cache loop iterations)
+        with torch.no_grad(), torch.autocast(
+                device_type="cuda", dtype=vae_dtype,
+                enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
             if not vae_autocast_enabled:
                 video_condition = video_condition.to(vae_dtype)
+            _gpu_mem("right before vae.encode()")
             encoder_output = self.vae.encode(video_condition)
+            _gpu_mem("after vae.encode()")
 
         # MatrixGame uses deterministic VAE encode for the first-frame conditioning.
         # Sampling would inject random noise into the cond_concat tensor and destroy the action guidance.

fastvideo/pipelines/stages/input_validation.py

@@ -210,9 +210,9 @@ class InputValidationStage(PipelineStage):
                     f"keyboard_cond must have 3 dimensions (B, T, K), but got {batch.keyboard_cond.dim()}"
                 )
             keyboard_dim = batch.keyboard_cond.shape[-1]
-            if keyboard_dim not in {2, 4, 6, 7}:
+            if keyboard_dim not in {2, 3, 4, 6, 7}:
                 raise ValueError(
-                    f"keyboard_cond last dimension must be 2, 4, 6, or 7, but got {keyboard_dim}"
+                    f"keyboard_cond last dimension must be 2, 3, 4, 6, or 7, but got {keyboard_dim}"
                 )
             logger.info(
                 "Action control: keyboard_cond validated - shape %s (dim=%d)",

fastvideo/pipelines/stages/matrixgame_denoising.py

@@ -313,32 +313,43 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
         attention_head_dim = getattr(
             self.transformer, 'attention_head_dim',
             self.transformer.hidden_size // num_attention_heads)
+        # WanGame PRoPE stores cat([rope, prope], dim=-1) in cache
+        head_dim_mult = getattr(self.transformer,
+                                '_kv_cache_head_dim_multiplier', 1)
+        attention_head_dim = attention_head_dim * head_dim_mult
         if self.local_attn_size != -1:
             kv_cache_size = self.local_attn_size * self.frame_seq_length
         else:
             kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
 
+        # WanGame uses None-initialized cache (concat-based); MatrixGame uses
+        # zero-tensor cache with index tracking.
+        use_none_init = head_dim_mult > 1
+
         for _ in range(self.num_transformer_blocks):
-            kv_cache.append({
-                "k":
-                torch.zeros([
-                    batch_size, kv_cache_size, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "v":
-                torch.zeros([
-                    batch_size, kv_cache_size, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "global_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-                "local_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-            })
+            if use_none_init:
+                kv_cache.append({"k": None, "v": None})
+            else:
+                kv_cache.append({
+                    "k":
+                    torch.zeros([
+                        batch_size, kv_cache_size, num_attention_heads,
+                        attention_head_dim
+                    ],
+                                dtype=dtype,
+                                device=device),
+                    "v":
+                    torch.zeros([
+                        batch_size, kv_cache_size, num_attention_heads,
+                        attention_head_dim
+                    ],
+                                dtype=dtype,
+                                device=device),
+                    "global_end_index":
+                    torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index":
+                    torch.tensor([0], dtype=torch.long, device=device),
+                })
 
         return kv_cache
 
@@ -429,6 +440,227 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
             })
         return crossattn_cache
 
+    def _denoise_one_step(
+        self,
+        current_latents: torch.Tensor,
+        noise_latents_btchw: torch.Tensor,
+        batch: ForwardBatch,
+        start_index: int,
+        current_num_frames: int,
+        timestep: torch.Tensor,
+        step_idx: int,
+        next_timestep: torch.Tensor | None,
+        ctx: BlockProcessingContext,
+        action_kwargs: dict[str, Any],
+        noise_generator: Callable[[tuple, torch.dtype, int], torch.Tensor]
+        | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run one denoising step. Returns (updated_latents, updated_noise_latents_btchw)."""
+        prompt_embeds = batch.prompt_embeds
+        t_cur = timestep
+
+        if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep:
+            current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
+        else:
+            current_model = self.transformer
+
+        noise_latents = noise_latents_btchw.clone()
+        latent_model_input = current_latents.to(ctx.target_dtype)
+
+        # For models that don't handle image_latent internally (e.g., WanGame),
+        # concatenate image conditioning along the channel dimension.
+        # MatrixGame handles this inside its forward() via forward_context.
+        if (batch.image_latent is not None
+                and not getattr(current_model, '_concatenates_image_latent',
+                                False)):
+            image_cond = batch.image_latent.to(ctx.target_dtype)
+            end = start_index + current_num_frames
+            if image_cond.shape[2] >= end:
+                image_cond = image_cond[:, :, start_index:end]
+            elif image_cond.shape[2] > start_index:
+                image_cond = image_cond[:, :, start_index:]
+                pad_frames = current_num_frames - image_cond.shape[2]
+                if pad_frames > 0:
+                    pad = torch.zeros(image_cond.shape[0],
+                                      image_cond.shape[1],
+                                      pad_frames,
+                                      image_cond.shape[3],
+                                      image_cond.shape[4],
+                                      device=image_cond.device,
+                                      dtype=image_cond.dtype)
+                    image_cond = torch.cat([image_cond, pad], dim=2)
+            else:
+                image_cond = torch.zeros(image_cond.shape[0],
+                                         image_cond.shape[1],
+                                         current_num_frames,
+                                         image_cond.shape[3],
+                                         image_cond.shape[4],
+                                         device=image_cond.device,
+                                         dtype=image_cond.dtype)
+            latent_model_input = torch.cat([latent_model_input, image_cond],
+                                           dim=1)
+
+        independent_first_frame = getattr(self.transformer,
+                                          'independent_first_frame', False)
+        if batch.image_latent is not None and independent_first_frame and start_index == 0:
+            latent_model_input = torch.cat([
+                latent_model_input,
+                batch.image_latent.to(ctx.target_dtype)
+            ],
+                                           dim=2)
+
+        # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents
+        t_expand = t_cur.repeat(latent_model_input.shape[0] *
+                                current_num_frames)
+
+        # Build attention metadata if VSA is available
+        if vsa_available and self.attn_backend == VideoSparseAttentionBackend:
+            self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
+            )
+            if self.attn_metadata_builder_cls is not None:
+                self.attn_metadata_builder = self.attn_metadata_builder_cls(
+                )
+                h, w = current_latents.shape[-2:]
+                attn_metadata = self.attn_metadata_builder.build(
+                    current_timestep=step_idx,
+                    raw_latent_shape=(current_num_frames, h, w),
+                    patch_size=ctx.fastvideo_args.pipeline_config.
+                    dit_config.patch_size,
+                    STA_param=batch.STA_param,
+                    VSA_sparsity=ctx.fastvideo_args.VSA_sparsity,
+                    device=get_local_torch_device(),
+                )
+                assert attn_metadata is not None, "attn_metadata cannot be None"
+            else:
+                attn_metadata = None
+        else:
+            attn_metadata = None
+
+        with torch.autocast(device_type="cuda",
+                            dtype=ctx.target_dtype,
+                            enabled=ctx.autocast_enabled), \
+            set_forward_context(current_timestep=step_idx,
+                                attn_metadata=attn_metadata,
+                                forward_batch=batch):
+            # Expand timestep to per-frame format [batch, num_frames] for causal model
+            t_expanded_noise = t_cur * torch.ones(
+                (latent_model_input.shape[0], current_num_frames),
+                device=latent_model_input.device,
+                dtype=torch.long)
+
+            model_kwargs = {
+                "kv_cache": ctx.get_kv_cache(t_cur),
+                "crossattn_cache": ctx.crossattn_cache,
+                "current_start": start_index * self.frame_seq_length,
+                "start_frame": start_index,
+            }
+
+            if self.use_action_module and current_model == self.transformer:
+                model_kwargs.update({
+                    "kv_cache_mouse":
+                    ctx.kv_cache_mouse,
+                    "kv_cache_keyboard":
+                    ctx.kv_cache_keyboard,
+                })
+                model_kwargs.update(action_kwargs)
+
+            # For models with discrete action embedding (WanGame-style):
+            # convert keyboard one-hot to scalar action = trans_label * 9
+            # and provide identity camera matrices for PRoPE
+            if getattr(current_model, '_uses_discrete_action', False):
+                dev = latent_model_input.device
+                b = latent_model_input.shape[0]
+                vae_ratio = 4
+                start_pixel = (0 if start_index == 0 else 1 + vae_ratio *
+                               (start_index - 1))
+                if (batch.keyboard_cond is not None
+                        and start_pixel < batch.keyboard_cond.shape[1]):
+                    kbd = batch.keyboard_cond[:, start_pixel, :].to(dev)
+                    has_key = kbd.sum(dim=-1) > 0  # [B]
+                    trans_label = (kbd.argmax(dim=-1) + 1) * has_key.long()
+                else:
+                    trans_label = torch.zeros(b, device=dev, dtype=torch.long)
+                action_value = trans_label * 9  # composite: trans * 9 + rot
+                model_kwargs["action"] = action_value.unsqueeze(1).expand(
+                    -1, current_num_frames).to(dev)
+                # Identity viewmats and default normalized intrinsics for PRoPE
+                model_kwargs["viewmats"] = torch.eye(
+                    4, device=dev, dtype=ctx.target_dtype
+                ).unsqueeze(0).unsqueeze(0).expand(
+                    b, current_num_frames, -1, -1).contiguous()
+                default_K = torch.tensor(
+                    [[0.5052, 0.0, 0.5],
+                     [0.0, 0.8979, 0.5],
+                     [0.0, 0.0, 1.0]],
+                    device=dev, dtype=ctx.target_dtype)
+                model_kwargs["Ks"] = default_K.unsqueeze(0).unsqueeze(
+                    0).expand(b, current_num_frames, -1, -1).contiguous()
+
+            pred_noise_btchw = current_model(
+                latent_model_input,
+                prompt_embeds,
+                t_expanded_noise,
+                **ctx.image_kwargs,
+                **ctx.pos_cond_kwargs,
+                **model_kwargs,
+            ).permute(0, 2, 1, 3, 4)
+
+        if ctx.boundary_timestep is not None and t_cur >= ctx.boundary_timestep:
+            pred_video_btchw = pred_noise_to_x_bound(
+                pred_noise=pred_noise_btchw.flatten(0, 1),
+                noise_input_latent=noise_latents.flatten(0, 1),
+                timestep=t_expand,
+                boundary_timestep=torch.ones_like(t_expand) *
+                ctx.boundary_timestep,
+                scheduler=self.scheduler).unflatten(
+                    0, pred_noise_btchw.shape[:2])
+        else:
+            pred_video_btchw = pred_noise_to_pred_video(
+                pred_noise=pred_noise_btchw.flatten(0, 1),
+                noise_input_latent=noise_latents.flatten(0, 1),
+                timestep=t_expand,
+                scheduler=self.scheduler).unflatten(
+                    0, pred_noise_btchw.shape[:2])
+
+        if next_timestep is not None:
+            next_t = next_timestep * torch.ones(
+                [1], dtype=torch.long, device=pred_video_btchw.device)
+
+            # Use custom noise generator if provided (for streaming), else generate
+            if noise_generator is not None:
+                noise = noise_generator(pred_video_btchw.shape,
+                                        pred_video_btchw.dtype, step_idx)
+            else:
+                noise = torch.randn(
+                    pred_video_btchw.shape,
+                    dtype=pred_video_btchw.dtype,
+                    generator=(batch.generator[0] if isinstance(
+                        batch.generator, list) else batch.generator)).to(
+                            pred_video_btchw.device)
+
+            noise_btchw = noise
+            if ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and step_idx < len(
+                    ctx.high_noise_timesteps) - 1:
+                noise_latents_btchw = self.scheduler.add_noise_high(
+                    pred_video_btchw.flatten(0, 1),
+                    noise_btchw.flatten(0, 1), next_t,
+                    torch.ones_like(next_t) *
+                    ctx.boundary_timestep).unflatten(
+                        0, pred_video_btchw.shape[:2])
+            elif ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and step_idx == len(
+                    ctx.high_noise_timesteps) - 1:
+                noise_latents_btchw = pred_video_btchw
+            else:
+                noise_latents_btchw = self.scheduler.add_noise(
+                    pred_video_btchw.flatten(0,
+                                             1), noise_btchw.flatten(0, 1),
+                    next_t).unflatten(0, pred_video_btchw.shape[:2])
+            current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4)
+        else:
+            current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4)
+
+        return current_latents, noise_latents_btchw
+
     def _process_single_block(
         self,
         current_latents: torch.Tensor,
@@ -442,144 +674,23 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
         | None = None,
         progress_bar: Any | None = None,
     ) -> torch.Tensor:
-        prompt_embeds = batch.prompt_embeds
         noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
 
         for i, t_cur in enumerate(timesteps):
-            if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep:
-                current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer
-            else:
-                current_model = self.transformer
-
-            noise_latents = noise_latents_btchw.clone()
-            latent_model_input = current_latents.to(ctx.target_dtype)
-
-            independent_first_frame = getattr(self.transformer,
-                                              'independent_first_frame', False)
-            if batch.image_latent is not None and independent_first_frame and start_index == 0:
-                latent_model_input = torch.cat([
-                    latent_model_input,
-                    batch.image_latent.to(ctx.target_dtype)
-                ],
-                                               dim=2)
-
-            # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents
-            t_expand = t_cur.repeat(latent_model_input.shape[0] *
-                                    current_num_frames)
-
-            # Build attention metadata if VSA is available
-            if vsa_available and self.attn_backend == VideoSparseAttentionBackend:
-                self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
-                )
-                if self.attn_metadata_builder_cls is not None:
-                    self.attn_metadata_builder = self.attn_metadata_builder_cls(
-                    )
-                    h, w = current_latents.shape[-2:]
-                    attn_metadata = self.attn_metadata_builder.build(
-                        current_timestep=i,
-                        raw_latent_shape=(current_num_frames, h, w),
-                        patch_size=ctx.fastvideo_args.pipeline_config.
-                        dit_config.patch_size,
-                        STA_param=batch.STA_param,
-                        VSA_sparsity=ctx.fastvideo_args.VSA_sparsity,
-                        device=get_local_torch_device(),
-                    )
-                    assert attn_metadata is not None, "attn_metadata cannot be None"
-                else:
-                    attn_metadata = None
-            else:
-                attn_metadata = None
-
-            with torch.autocast(device_type="cuda",
-                                dtype=ctx.target_dtype,
-                                enabled=ctx.autocast_enabled), \
-                set_forward_context(current_timestep=i,
-                                    attn_metadata=attn_metadata,
-                                    forward_batch=batch):
-                # Expand timestep to per-frame format [batch, num_frames] for causal model
-                t_expanded_noise = t_cur * torch.ones(
-                    (latent_model_input.shape[0], current_num_frames),
-                    device=latent_model_input.device,
-                    dtype=torch.long)
-
-                model_kwargs = {
-                    "kv_cache": ctx.get_kv_cache(t_cur),
-                    "crossattn_cache": ctx.crossattn_cache,
-                    "current_start": start_index * self.frame_seq_length,
-                    "start_frame": start_index,
-                }
-
-                if self.use_action_module and current_model == self.transformer:
-                    model_kwargs.update({
-                        "kv_cache_mouse":
-                        ctx.kv_cache_mouse,
-                        "kv_cache_keyboard":
-                        ctx.kv_cache_keyboard,
-                    })
-                    model_kwargs.update(action_kwargs)
-
-                pred_noise_btchw = current_model(
-                    latent_model_input,
-                    prompt_embeds,
-                    t_expanded_noise,
-                    **ctx.image_kwargs,
-                    **ctx.pos_cond_kwargs,
-                    **model_kwargs,
-                ).permute(0, 2, 1, 3, 4)
-
-            if ctx.boundary_timestep is not None and t_cur >= ctx.boundary_timestep:
-                pred_video_btchw = pred_noise_to_x_bound(
-                    pred_noise=pred_noise_btchw.flatten(0, 1),
-                    noise_input_latent=noise_latents.flatten(0, 1),
-                    timestep=t_expand,
-                    boundary_timestep=torch.ones_like(t_expand) *
-                    ctx.boundary_timestep,
-                    scheduler=self.scheduler).unflatten(
-                        0, pred_noise_btchw.shape[:2])
-            else:
-                pred_video_btchw = pred_noise_to_pred_video(
-                    pred_noise=pred_noise_btchw.flatten(0, 1),
-                    noise_input_latent=noise_latents.flatten(0, 1),
-                    timestep=t_expand,
-                    scheduler=self.scheduler).unflatten(
-                        0, pred_noise_btchw.shape[:2])
-
-            if i < len(timesteps) - 1:
-                next_timestep = timesteps[i + 1] * torch.ones(
-                    [1], dtype=torch.long, device=pred_video_btchw.device)
-
-                # Use custom noise generator if provided (for streaming), else generate
-                if noise_generator is not None:
-                    noise = noise_generator(pred_video_btchw.shape,
-                                            pred_video_btchw.dtype, i)
-                else:
-                    noise = torch.randn(
-                        pred_video_btchw.shape,
-                        dtype=pred_video_btchw.dtype,
-                        generator=(batch.generator[0] if isinstance(
-                            batch.generator, list) else batch.generator)).to(
-                                pred_video_btchw.device)
-
-                noise_btchw = noise
-                if ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i < len(
-                        ctx.high_noise_timesteps) - 1:
-                    noise_latents_btchw = self.scheduler.add_noise_high(
-                        pred_video_btchw.flatten(0, 1),
-                        noise_btchw.flatten(0, 1), next_timestep,
-                        torch.ones_like(next_timestep) *
-                        ctx.boundary_timestep).unflatten(
-                            0, pred_video_btchw.shape[:2])
-                elif ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i == len(
-                        ctx.high_noise_timesteps) - 1:
-                    noise_latents_btchw = pred_video_btchw
-                else:
-                    noise_latents_btchw = self.scheduler.add_noise(
-                        pred_video_btchw.flatten(0,
-                                                 1), noise_btchw.flatten(0, 1),
-                        next_timestep).unflatten(0, pred_video_btchw.shape[:2])
-                current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4)
-            else:
-                current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4)
+            next_timestep = timesteps[i + 1] if i < len(timesteps) - 1 else None
+            current_latents, noise_latents_btchw = self._denoise_one_step(
+                current_latents=current_latents,
+                noise_latents_btchw=noise_latents_btchw,
+                batch=batch,
+                start_index=start_index,
+                current_num_frames=current_num_frames,
+                timestep=t_cur,
+                step_idx=i,
+                next_timestep=next_timestep,
+                ctx=ctx,
+                action_kwargs=action_kwargs,
+                noise_generator=noise_generator,
+            )
 
             if progress_bar is not None:
                 progress_bar.update()
@@ -605,6 +716,36 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                                dtype=torch.long) * int(context_noise)
         context_bcthw = current_latents.to(ctx.target_dtype)
 
+        # Channel concatenation for models that don't handle it internally
+        if (batch.image_latent is not None
+                and not getattr(self.transformer, '_concatenates_image_latent',
+                                False)):
+            image_cond = batch.image_latent.to(ctx.target_dtype)
+            end = start_index + current_num_frames
+            if image_cond.shape[2] >= end:
+                image_cond = image_cond[:, :, start_index:end]
+            elif image_cond.shape[2] > start_index:
+                image_cond = image_cond[:, :, start_index:]
+                pad_frames = current_num_frames - image_cond.shape[2]
+                if pad_frames > 0:
+                    pad = torch.zeros(image_cond.shape[0],
+                                      image_cond.shape[1],
+                                      pad_frames,
+                                      image_cond.shape[3],
+                                      image_cond.shape[4],
+                                      device=image_cond.device,
+                                      dtype=image_cond.dtype)
+                    image_cond = torch.cat([image_cond, pad], dim=2)
+            else:
+                image_cond = torch.zeros(image_cond.shape[0],
+                                         image_cond.shape[1],
+                                         current_num_frames,
+                                         image_cond.shape[3],
+                                         image_cond.shape[4],
+                                         device=image_cond.device,
+                                         dtype=image_cond.dtype)
+            context_bcthw = torch.cat([context_bcthw, image_cond], dim=1)
+
         with torch.autocast(device_type="cuda",
                             dtype=ctx.target_dtype,
                             enabled=ctx.autocast_enabled), \
@@ -628,6 +769,37 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                 })
                 context_model_kwargs.update(action_kwargs)
 
+            # Discrete action and camera matrices for WanGame-style models
+            if getattr(self.transformer, '_uses_discrete_action', False):
+                dev = latents_device
+                b = context_bcthw.shape[0]
+                vae_ratio = 4
+                start_pixel = (0 if start_index == 0 else 1 + vae_ratio *
+                               (start_index - 1))
+                if (batch.keyboard_cond is not None
+                        and start_pixel < batch.keyboard_cond.shape[1]):
+                    kbd = batch.keyboard_cond[:, start_pixel, :].to(dev)
+                    has_key = kbd.sum(dim=-1) > 0
+                    trans_label = (kbd.argmax(dim=-1) + 1) * has_key.long()
+                else:
+                    trans_label = torch.zeros(b, device=dev, dtype=torch.long)
+                action_value = trans_label * 9
+                context_model_kwargs["action"] = (
+                    action_value.unsqueeze(1).expand(
+                        -1, current_num_frames).to(dev))
+                context_model_kwargs["viewmats"] = torch.eye(
+                    4, device=dev, dtype=ctx.target_dtype
+                ).unsqueeze(0).unsqueeze(0).expand(
+                    b, current_num_frames, -1, -1).contiguous()
+                default_K = torch.tensor(
+                    [[0.5052, 0.0, 0.5],
+                     [0.0, 0.8979, 0.5],
+                     [0.0, 0.0, 1.0]],
+                    device=dev, dtype=ctx.target_dtype)
+                context_model_kwargs["Ks"] = default_K.unsqueeze(
+                    0).unsqueeze(0).expand(
+                        b, current_num_frames, -1, -1).contiguous()
+
             if ctx.boundary_timestep is not None and self.transformer_2 is not None:
                 self.transformer_2(
                     context_bcthw,
@@ -641,6 +813,10 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                     **ctx.pos_cond_kwargs,
                 )
 
+            # WanGame uses is_cache=True to store KV without prepending old cache
+            if getattr(self.transformer, '_uses_discrete_action', False):
+                context_model_kwargs["is_cache"] = True
+
             self.transformer(
                 context_bcthw,
                 prompt_embeds,

fastvideo/pipelines/stages/multi_user_engine.py

@@ -0,0 +1,695 @@
+"""Multi-user engine implementing ORCA-style batching for streaming generation.
+
+Groups users at the same (block_idx, denoising_step) and batches their forward
+passes together, sharing model weights while maintaining per-user KV caches.
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any, Iterator
+
+import torch
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.matrixgame_denoising import (
+    BlockProcessingContext,
+    MatrixGameCausalDenoisingStage,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class CompletedResult:
+    """Result for a user whose block generation finished."""
+    user_id: str
+    output_batch: ForwardBatch
+
+
+@dataclass
+class UserSession:
+    """Per-user state for multi-user streaming."""
+    user_id: str
+    ctx: BlockProcessingContext
+    vae_cache: list[torch.Tensor | None] | None
+    batch: ForwardBatch
+    fastvideo_args: FastVideoArgs
+
+    # Denoising state (set when a step is in progress)
+    denoising_step: int = -1  # -1 means no pending work
+    current_latents: torch.Tensor | None = None
+    noise_latents_btchw: torch.Tensor | None = None
+    action_kwargs: dict[str, Any] = field(default_factory=dict)
+    block_start_index: int = 0
+    block_num_frames: int = 0
+    dit_elapsed_ms: float = 0.0  # accumulated DiT time across denoising steps
+
+
+class MultiUserEngine:
+    """ORCA-style scheduler that batches multiple users on a single GPU."""
+
+    def __init__(self, pipeline, disable_batching: bool = False):
+        self.pipeline = pipeline
+        self.disable_batching = disable_batching
+        if not pipeline.post_init_called:
+            pipeline.post_init()
+        self.denoiser: MatrixGameCausalDenoisingStage = (
+            pipeline._stage_name_mapping["denoising_stage"]
+        )
+        self.decoder = pipeline._stage_name_mapping["decoding_stage"]
+        self.users: dict[str, UserSession] = {}
+
+    def add_user(
+        self,
+        user_id: str,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> None:
+        """Initialize a new user session (runs preprocessing + streaming_reset)."""
+        def _gpu_mem():
+            a = torch.cuda.memory_allocated() / 1024**3
+            r = torch.cuda.memory_reserved() / 1024**3
+            return f"alloc={a:.2f}GiB, reserved={r:.2f}GiB"
+
+        logger.info("add_user(%s) start: %s", user_id, _gpu_mem())
+        logger.info("add_user(%s) memory summary:\n%s",
+                     user_id, torch.cuda.memory_summary())
+
+        if user_id in self.users:
+            logger.warning("User %s already exists, removing first", user_id)
+            self.remove_user(user_id)
+
+        # Free any existing streaming state FIRST to make room on GPU.
+        # After warmup, the denoiser holds KV caches + noise pool.
+        import gc
+        logger.info("add_user(%s) denoiser._streaming_initialized=%s, ctx=%s",
+                     user_id, self.denoiser._streaming_initialized,
+                     self.denoiser._streaming_ctx is not None)
+        if self.denoiser._streaming_initialized:
+            ctx = self.denoiser._streaming_ctx
+            if ctx is not None:
+                # Explicitly delete large tensors to break any reference cycles
+                ctx.kv_cache1 = None
+                ctx.kv_cache2 = None
+                ctx.kv_cache_mouse = None
+                ctx.kv_cache_keyboard = None
+                ctx.crossattn_cache = None
+                ctx.noise_pool = None
+                if ctx.batch is not None:
+                    ctx.batch.latents = None
+                    ctx.batch.image_latent = None
+                    ctx.batch.prompt_embeds = None
+                ctx.batch = None
+            self.denoiser._streaming_ctx = None
+            self.denoiser._streaming_initialized = False
+            logger.info("add_user(%s) after ctx cleanup: %s", user_id, _gpu_mem())
+            gc.collect()
+            torch.cuda.empty_cache()
+            logger.info("add_user(%s) after gc+empty_cache: %s", user_id, _gpu_mem())
+            logger.info("add_user(%s) post-cleanup summary:\n%s",
+                         user_id, torch.cuda.memory_summary())
+
+        # Run preprocessing stages
+        if not self.pipeline.post_init_called:
+            self.pipeline.post_init()
+
+        # Ensure VAE temporal tiling is enabled to avoid OOM when encoding
+        # the full video_condition [1,3,num_frames,H,W] through 3D conv layers.
+        # Enable on both the pipeline's VAE and the stage's VAE reference.
+        vae = self.pipeline.get_module("vae", None)
+        if vae is not None:
+            vae.enable_tiling(use_temporal_tiling=True)
+            logger.info("Pipeline VAE tiling: use_tiling=%s, temporal=%s, min_frames=%s, id=%s",
+                        vae.use_tiling, vae.use_temporal_tiling,
+                        getattr(vae, 'tile_sample_min_num_frames', 'N/A'), id(vae))
+        img_stage = self.pipeline._stage_name_mapping.get("image_latent_preparation_stage")
+        if img_stage is not None and hasattr(img_stage, 'vae'):
+            img_stage.vae.enable_tiling(use_temporal_tiling=True)
+            logger.info("Stage VAE tiling: use_tiling=%s, temporal=%s, min_frames=%s, id=%s",
+                        img_stage.vae.use_tiling, img_stage.vae.use_temporal_tiling,
+                        getattr(img_stage.vae, 'tile_sample_min_num_frames', 'N/A'),
+                        id(img_stage.vae))
+
+        stages_to_run = [
+            "input_validation_stage", "prompt_encoding_stage",
+            "image_encoding_stage", "conditioning_stage",
+            "latent_preparation_stage", "image_latent_preparation_stage",
+        ]
+        for stage_name in stages_to_run:
+            if stage_name in self.pipeline._stage_name_mapping:
+                batch = self.pipeline._stage_name_mapping[stage_name].forward(
+                    batch, fastvideo_args)
+                logger.info("add_user(%s) after %s: %s",
+                            user_id, stage_name, _gpu_mem())
+
+        # Initialize denoiser state for this user (creates KV caches etc.)
+        # We call streaming_reset to set up the context, then steal it
+        logger.info("add_user(%s) before streaming_reset: %s", user_id, _gpu_mem())
+        self.denoiser.streaming_reset(batch, fastvideo_args)
+        ctx = self.denoiser._streaming_ctx
+        assert ctx is not None
+        # Detach from denoiser so it doesn't interfere with other users
+        self.denoiser._streaming_ctx = None
+        self.denoiser._streaming_initialized = False
+
+        self.users[user_id] = UserSession(
+            user_id=user_id,
+            ctx=ctx,
+            vae_cache=None,
+            batch=batch,
+            fastvideo_args=fastvideo_args,
+        )
+        logger.info("Added user %s (total: %d)", user_id, len(self.users))
+
+    def remove_user(self, user_id: str) -> None:
+        """Remove a user session and free GPU memory."""
+        session = self.users.pop(user_id, None)
+        if session is not None:
+            # Let GC handle tensor cleanup
+            session.ctx = None  # type: ignore
+            session.vae_cache = None
+            session.current_latents = None
+            session.noise_latents_btchw = None
+            logger.info("Removed user %s (remaining: %d)",
+                        user_id, len(self.users))
+
+    def submit_step(
+        self,
+        user_id: str,
+        keyboard_action: torch.Tensor | None,
+        mouse_action: torch.Tensor | None,
+    ) -> None:
+        """Queue a block generation request for a user."""
+        session = self.users.get(user_id)
+        if session is None:
+            raise KeyError(f"User {user_id} not found")
+
+        ctx = session.ctx
+        if ctx.block_idx >= len(ctx.block_sizes):
+            logger.warning("User %s has no more blocks to generate", user_id)
+            return
+
+        batch = session.batch
+        latents = batch.latents
+        assert latents is not None
+
+        current_num_frames = ctx.block_sizes[ctx.block_idx]
+        start_index = ctx.start_index
+
+        current_latents = latents[:, :, start_index:start_index +
+                                  current_num_frames, :, :]
+
+        # Update batch with new actions
+        if keyboard_action is not None or mouse_action is not None:
+            vae_ratio = 4
+            start_frame = 0 if start_index == 0 else 1 + vae_ratio * (
+                start_index - 1)
+
+            if keyboard_action is not None:
+                n = keyboard_action.shape[1]
+                batch.keyboard_cond[:, start_frame:start_frame +
+                                    n] = keyboard_action.to(
+                                        batch.keyboard_cond.device)
+            if mouse_action is not None:
+                n = mouse_action.shape[1]
+                batch.mouse_cond[:, start_frame:start_frame +
+                                 n] = mouse_action.to(batch.mouse_cond.device)
+
+        action_kwargs = self.denoiser._prepare_action_kwargs(
+            batch, start_index, current_num_frames)
+
+        # Set up denoising state
+        session.denoising_step = 0
+        session.dit_elapsed_ms = 0.0
+        session.current_latents = current_latents
+        session.noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
+        session.action_kwargs = action_kwargs
+        session.block_start_index = start_index
+        session.block_num_frames = current_num_frames
+
+    def has_pending_work(self) -> bool:
+        """Check if any user has pending denoising work."""
+        return any(s.denoising_step >= 0 for s in self.users.values())
+
+    def run_iteration(self) -> Iterator[CompletedResult]:
+        """Run one denoising step for all ready users (ORCA iteration-level scheduling).
+
+        Groups users by (block_idx, denoising_step) and batches compatible groups.
+        Yields completed results immediately after each user's VAE decode finishes,
+        so responses can be sent without waiting for all users in the batch.
+        """
+        # Group users with pending work by (block_idx, denoising_step)
+        groups: dict[tuple[int, int], list[UserSession]] = {}
+        for session in self.users.values():
+            if session.denoising_step < 0:
+                continue
+            key = (session.ctx.block_idx, session.denoising_step)
+            groups.setdefault(key, []).append(session)
+
+        for (block_idx, step_idx), sessions in groups.items():
+            timesteps = sessions[0].ctx.timesteps
+            t_cur = timesteps[step_idx]
+            next_timestep = timesteps[step_idx + 1] if step_idx < len(timesteps) - 1 else None
+
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+
+            user_ids = [s.user_id[:8] for s in sessions]
+            if len(sessions) == 1 or self.disable_batching:
+                # Process each user individually
+                for s in sessions:
+                    logger.info("ORCA step: block=%d step=%d single user=%s",
+                                block_idx, step_idx, s.user_id[:8])
+                    self._run_single_user_step(s, t_cur, step_idx, next_timestep)
+            else:
+                # Multiple users at same state - batch them
+                logger.info("ORCA step: block=%d step=%d BATCHED %d users=%s",
+                            block_idx, step_idx, len(sessions), user_ids)
+                self._run_batched_step(sessions, t_cur, step_idx, next_timestep)
+
+            torch.cuda.synchronize()
+            step_ms = (time.perf_counter() - t0) * 1000
+
+            # Advance step counters and finalize completed users
+            for s in sessions:
+                s.dit_elapsed_ms += step_ms
+                s.denoising_step += 1
+                if s.denoising_step >= len(timesteps):
+                    # Denoising complete — VAE decode and yield immediately
+                    result = self._finalize_block(s)
+                    if result is not None:
+                        yield result
+
+    def _run_single_user_step(
+        self,
+        session: UserSession,
+        timestep: torch.Tensor,
+        step_idx: int,
+        next_timestep: torch.Tensor | None,
+    ) -> None:
+        """Run one denoising step for a single user (no batching)."""
+        noise_gen = self._make_noise_generator(session)
+
+        current_latents, noise_latents_btchw = self.denoiser._denoise_one_step(
+            current_latents=session.current_latents,
+            noise_latents_btchw=session.noise_latents_btchw,
+            batch=session.batch,
+            start_index=session.block_start_index,
+            current_num_frames=session.block_num_frames,
+            timestep=timestep,
+            step_idx=step_idx,
+            next_timestep=next_timestep,
+            ctx=session.ctx,
+            action_kwargs=session.action_kwargs,
+            noise_generator=noise_gen,
+        )
+        session.current_latents = current_latents
+        session.noise_latents_btchw = noise_latents_btchw
+
+    def _run_batched_step(
+        self,
+        sessions: list[UserSession],
+        timestep: torch.Tensor,
+        step_idx: int,
+        next_timestep: torch.Tensor | None,
+    ) -> None:
+        """Run one denoising step for multiple users batched together."""
+        n = len(sessions)
+
+        # Concatenate along batch dimension
+        batched_latents = torch.cat(
+            [s.current_latents for s in sessions], dim=0)
+        batched_noise = torch.cat(
+            [s.noise_latents_btchw for s in sessions], dim=0)
+
+        # Build a merged ForwardBatch with concatenated per-user tensors
+        merged_batch = self._merge_batches(sessions)
+
+        # Build merged context with concatenated KV caches
+        merged_ctx = self._merge_contexts(sessions, merged_batch)
+
+        # Merge action kwargs
+        merged_action_kwargs = self._merge_action_kwargs(sessions)
+
+        # Create noise generator that concatenates per-user noise
+        device = batched_latents.device
+
+        def batched_noise_gen(shape, dtype, si):
+            noises = []
+            for s in sessions:
+                gen = self._make_noise_generator(s)
+                if gen is not None:
+                    per_user_shape = (1,) + shape[1:]
+                    noises.append(gen(per_user_shape, dtype, si))
+                else:
+                    noises.append(torch.randn(
+                        (1,) + shape[1:], dtype=dtype).to(device))
+            return torch.cat(noises, dim=0)
+
+        current_latents, noise_latents_btchw = self.denoiser._denoise_one_step(
+            current_latents=batched_latents,
+            noise_latents_btchw=batched_noise,
+            batch=merged_batch,
+            start_index=sessions[0].block_start_index,
+            current_num_frames=sessions[0].block_num_frames,
+            timestep=timestep,
+            step_idx=step_idx,
+            next_timestep=next_timestep,
+            ctx=merged_ctx,
+            action_kwargs=merged_action_kwargs,
+            noise_generator=batched_noise_gen,
+        )
+
+        # Unbatch results back to per-user tensors
+        latents_list = current_latents.split(1, dim=0)
+        noise_list = noise_latents_btchw.split(1, dim=0)
+
+        # Restore per-user KV caches from the batched caches
+        self._split_contexts(merged_ctx, sessions)
+
+        for i, s in enumerate(sessions):
+            s.current_latents = latents_list[i]
+            s.noise_latents_btchw = noise_list[i]
+
+    def _merge_batches(self, sessions: list[UserSession]) -> ForwardBatch:
+        """Create a merged ForwardBatch with concatenated user tensors."""
+        ref = sessions[0].batch
+        # prompt_embeds and image_embeds are shared (same game)
+        # latents, keyboard_cond, mouse_cond differ per user
+        merged = ForwardBatch.__new__(ForwardBatch)
+        merged.__dict__.update(ref.__dict__)
+
+        # image_latent: cat along batch dim if present
+        if ref.image_latent is not None:
+            merged.image_latent = torch.cat(
+                [s.batch.image_latent for s in sessions], dim=0)
+
+        return merged
+
+    def _merge_contexts(
+        self, sessions: list[UserSession],
+        merged_batch: ForwardBatch | None = None,
+    ) -> BlockProcessingContext:
+        """Create a merged BlockProcessingContext with concatenated KV caches."""
+        ref = sessions[0].ctx
+        n = len(sessions)
+
+        # Concatenate KV caches along batch dimension
+        merged_kv1 = self._cat_kv_caches([s.ctx.kv_cache1 for s in sessions])
+        merged_kv2 = None
+        if ref.kv_cache2 is not None:
+            merged_kv2 = self._cat_kv_caches(
+                [s.ctx.kv_cache2 for s in sessions])
+
+        merged_crossattn = self._cat_crossattn_caches(
+            [s.ctx.crossattn_cache for s in sessions])
+
+        merged_kv_mouse = None
+        merged_kv_keyboard = None
+        if ref.kv_cache_mouse is not None:
+            merged_kv_mouse = self._cat_action_kv_caches(
+                [s.ctx.kv_cache_mouse for s in sessions],
+                is_mouse=True)
+        if ref.kv_cache_keyboard is not None:
+            merged_kv_keyboard = self._cat_kv_caches(
+                [s.ctx.kv_cache_keyboard for s in sessions])
+
+        # image_kwargs: cat image embeds
+        merged_image_kwargs = dict(ref.image_kwargs)
+        if "encoder_hidden_states_image" in ref.image_kwargs:
+            embeds = ref.image_kwargs["encoder_hidden_states_image"]
+            if isinstance(embeds, list) and len(embeds) > 0 and torch.is_tensor(embeds[0]):
+                # Cat each tensor in the list across users
+                merged_embeds = []
+                for idx in range(len(embeds)):
+                    merged_embeds.append(torch.cat(
+                        [s.ctx.image_kwargs["encoder_hidden_states_image"][idx]
+                         for s in sessions], dim=0))
+                merged_image_kwargs["encoder_hidden_states_image"] = merged_embeds
+
+        if merged_batch is None:
+            merged_batch = self._merge_batches(sessions)
+
+        return BlockProcessingContext(
+            batch=merged_batch,
+            block_idx=ref.block_idx,
+            start_index=ref.start_index,
+            kv_cache1=merged_kv1,
+            kv_cache2=merged_kv2,
+            kv_cache_mouse=merged_kv_mouse,
+            kv_cache_keyboard=merged_kv_keyboard,
+            crossattn_cache=merged_crossattn,
+            timesteps=ref.timesteps,
+            block_sizes=ref.block_sizes,
+            noise_pool=None,  # noise handled per-user
+            fastvideo_args=ref.fastvideo_args,
+            target_dtype=ref.target_dtype,
+            autocast_enabled=ref.autocast_enabled,
+            boundary_timestep=ref.boundary_timestep,
+            high_noise_timesteps=ref.high_noise_timesteps,
+            context_noise=ref.context_noise,
+            image_kwargs=merged_image_kwargs,
+            pos_cond_kwargs=ref.pos_cond_kwargs,
+        )
+
+    def _cat_kv_caches(
+        self, caches_list: list[list[dict[str, Any]]]
+    ) -> list[dict[str, Any]]:
+        """Concatenate KV caches along batch dimension."""
+        num_layers = len(caches_list[0])
+        merged = []
+        for layer_idx in range(num_layers):
+            merged.append({
+                "k": torch.cat(
+                    [c[layer_idx]["k"] for c in caches_list], dim=0),
+                "v": torch.cat(
+                    [c[layer_idx]["v"] for c in caches_list], dim=0),
+                "global_end_index": caches_list[0][layer_idx]["global_end_index"],
+                "local_end_index": caches_list[0][layer_idx]["local_end_index"],
+            })
+        return merged
+
+    def _cat_action_kv_caches(
+        self,
+        caches_list: list[list[dict[str, Any]]],
+        is_mouse: bool = False,
+    ) -> list[dict[str, Any]]:
+        """Concatenate action KV caches.
+
+        Mouse caches have shape [B*frame_seq_length, ...] so we need
+        to concatenate correctly.
+        """
+        num_layers = len(caches_list[0])
+        merged = []
+        for layer_idx in range(num_layers):
+            if is_mouse:
+                # Mouse: [B*F, cache_size, heads, dim] -> [N*F, cache_size, heads, dim]
+                # Each user has [1*F, ...], so simple cat along dim 0 works
+                merged.append({
+                    "k": torch.cat(
+                        [c[layer_idx]["k"] for c in caches_list], dim=0),
+                    "v": torch.cat(
+                        [c[layer_idx]["v"] for c in caches_list], dim=0),
+                    "global_end_index": caches_list[0][layer_idx]["global_end_index"],
+                    "local_end_index": caches_list[0][layer_idx]["local_end_index"],
+                })
+            else:
+                merged.append({
+                    "k": torch.cat(
+                        [c[layer_idx]["k"] for c in caches_list], dim=0),
+                    "v": torch.cat(
+                        [c[layer_idx]["v"] for c in caches_list], dim=0),
+                    "global_end_index": caches_list[0][layer_idx]["global_end_index"],
+                    "local_end_index": caches_list[0][layer_idx]["local_end_index"],
+                })
+        return merged
+
+    def _cat_crossattn_caches(
+        self, caches_list: list[list[dict[str, Any]]]
+    ) -> list[dict[str, Any]]:
+        """Concatenate cross-attention caches along batch dimension."""
+        num_layers = len(caches_list[0])
+        merged = []
+        for layer_idx in range(num_layers):
+            merged.append({
+                "k": torch.cat(
+                    [c[layer_idx]["k"] for c in caches_list], dim=0),
+                "v": torch.cat(
+                    [c[layer_idx]["v"] for c in caches_list], dim=0),
+                "is_init": caches_list[0][layer_idx]["is_init"],
+            })
+        return merged
+
+    def _split_contexts(
+        self,
+        merged_ctx: BlockProcessingContext,
+        sessions: list[UserSession],
+    ) -> None:
+        """Split merged KV caches back into per-user caches after forward pass."""
+        n = len(sessions)
+
+        self._split_kv_caches(merged_ctx.kv_cache1,
+                              [s.ctx.kv_cache1 for s in sessions])
+        if merged_ctx.kv_cache2 is not None:
+            self._split_kv_caches(merged_ctx.kv_cache2,
+                                  [s.ctx.kv_cache2 for s in sessions])
+
+        self._split_crossattn_caches(merged_ctx.crossattn_cache,
+                                     [s.ctx.crossattn_cache for s in sessions])
+
+        if merged_ctx.kv_cache_mouse is not None:
+            frame_seq_len = self.denoiser.frame_seq_length
+            self._split_action_kv_caches(
+                merged_ctx.kv_cache_mouse,
+                [s.ctx.kv_cache_mouse for s in sessions],
+                is_mouse=True, frame_seq_length=frame_seq_len)
+        if merged_ctx.kv_cache_keyboard is not None:
+            self._split_kv_caches(merged_ctx.kv_cache_keyboard,
+                                  [s.ctx.kv_cache_keyboard for s in sessions])
+
+    def _split_kv_caches(
+        self,
+        merged: list[dict[str, Any]],
+        per_user: list[list[dict[str, Any]]],
+    ) -> None:
+        """Copy data from merged KV cache back to per-user caches."""
+        n = len(per_user)
+        for layer_idx in range(len(merged)):
+            k_split = merged[layer_idx]["k"].split(1, dim=0)
+            v_split = merged[layer_idx]["v"].split(1, dim=0)
+            for i in range(n):
+                per_user[i][layer_idx]["k"].copy_(k_split[i])
+                per_user[i][layer_idx]["v"].copy_(v_split[i])
+                per_user[i][layer_idx]["global_end_index"].copy_(
+                    merged[layer_idx]["global_end_index"])
+                per_user[i][layer_idx]["local_end_index"].copy_(
+                    merged[layer_idx]["local_end_index"])
+
+    def _split_action_kv_caches(
+        self,
+        merged: list[dict[str, Any]],
+        per_user: list[list[dict[str, Any]]],
+        is_mouse: bool = False,
+        frame_seq_length: int = 1,
+    ) -> None:
+        """Split action KV caches back to per-user."""
+        n = len(per_user)
+        for layer_idx in range(len(merged)):
+            if is_mouse:
+                # [N*F, ...] -> N chunks of [F, ...]
+                k_split = merged[layer_idx]["k"].split(frame_seq_length, dim=0)
+                v_split = merged[layer_idx]["v"].split(frame_seq_length, dim=0)
+            else:
+                k_split = merged[layer_idx]["k"].split(1, dim=0)
+                v_split = merged[layer_idx]["v"].split(1, dim=0)
+            for i in range(n):
+                per_user[i][layer_idx]["k"].copy_(k_split[i])
+                per_user[i][layer_idx]["v"].copy_(v_split[i])
+                per_user[i][layer_idx]["global_end_index"].copy_(
+                    merged[layer_idx]["global_end_index"])
+                per_user[i][layer_idx]["local_end_index"].copy_(
+                    merged[layer_idx]["local_end_index"])
+
+    def _split_crossattn_caches(
+        self,
+        merged: list[dict[str, Any]],
+        per_user: list[list[dict[str, Any]]],
+    ) -> None:
+        """Split cross-attention caches back to per-user."""
+        n = len(per_user)
+        for layer_idx in range(len(merged)):
+            k_split = merged[layer_idx]["k"].split(1, dim=0)
+            v_split = merged[layer_idx]["v"].split(1, dim=0)
+            for i in range(n):
+                per_user[i][layer_idx]["k"].copy_(k_split[i])
+                per_user[i][layer_idx]["v"].copy_(v_split[i])
+                per_user[i][layer_idx]["is_init"] = merged[layer_idx]["is_init"]
+
+    def _merge_action_kwargs(
+        self, sessions: list[UserSession]
+    ) -> dict[str, Any]:
+        """Merge action kwargs by concatenating tensors along batch dim."""
+        ref = sessions[0].action_kwargs
+        if not ref:
+            return {}
+
+        merged: dict[str, Any] = {}
+        for key in ref:
+            vals = [s.action_kwargs[key] for s in sessions]
+            if torch.is_tensor(vals[0]):
+                merged[key] = torch.cat(vals, dim=0)
+            else:
+                # Scalars (like num_frame_per_block) - use ref value
+                merged[key] = vals[0]
+        return merged
+
+    def _make_noise_generator(self, session: UserSession):
+        """Create a noise generator using the user's pre-allocated noise pool."""
+        ctx = session.ctx
+        latents = session.batch.latents
+
+        def noise_gen(shape, dtype, step_idx):
+            if ctx.noise_pool is not None and step_idx < len(ctx.noise_pool):
+                return ctx.noise_pool[step_idx][:, :shape[1], :, :, :].to(
+                    latents.device)
+            else:
+                return torch.randn(shape, dtype=dtype).to(latents.device)
+
+        return noise_gen
+
+    def _finalize_block(self, session: UserSession) -> CompletedResult | None:
+        """Finalize a completed block: update context cache + VAE decode."""
+        ctx = session.ctx
+        batch = session.batch
+        latents = batch.latents
+        start_index = session.block_start_index
+        current_num_frames = session.block_num_frames
+
+        # Write denoised latents back
+        latents[:, :, start_index:start_index +
+                current_num_frames, :, :] = session.current_latents
+
+        # Update KV caches with clean context
+        self.denoiser._update_context_cache(
+            current_latents=session.current_latents,
+            batch=batch,
+            start_index=start_index,
+            current_num_frames=current_num_frames,
+            ctx=ctx,
+            action_kwargs=session.action_kwargs,
+            context_noise=ctx.context_noise,
+        )
+
+        # Advance streaming state
+        old_start = ctx.start_index
+        ctx.start_index += current_num_frames
+        ctx.block_idx += 1
+
+        # VAE decode
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+
+        current_latents_for_vae = latents[:, :, start_index:start_index +
+                                          current_num_frames, :, :]
+        decoded_frames, session.vae_cache = self.decoder.streaming_decode(
+            current_latents_for_vae,
+            session.fastvideo_args,
+            cache=session.vae_cache,
+            is_first_chunk=(start_index == 0),
+        )
+        torch.cuda.synchronize()
+        vae_ms = (time.perf_counter() - t0) * 1000
+
+        batch.output = decoded_frames
+        batch.stage_timings = {"dit_ms": session.dit_elapsed_ms, "vae_ms": vae_ms}
+
+        # Reset denoising state
+        session.denoising_step = -1
+        session.current_latents = None
+        session.noise_latents_btchw = None
+
+        return CompletedResult(user_id=session.user_id, output_batch=batch)
+

fastvideo/platforms/cuda.py

@@ -155,7 +155,7 @@ class CudaPlatformBase(Platform):
                 )
         elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE:
             try:
-                from fastvideo.attention.backends.sageattn.api import sageattn_blackwell  # noqa: F401
+                from sageattn3 import sageattn3_blackwell  # noqa: F401
 
                 from fastvideo.attention.backends.sage_attn3 import (  # noqa: F401
                     SageAttention3Backend)

fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp2_hsdp_shard4.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Mismatch compatibility test (direction B):
+  - WORLD_SIZE=4
+  - SP: sp_size=2
+  - HSDP/FSDP2 mesh: replicate_dim=1, shard_dim=4 (mesh_shape=(1,4))
+
+We verify gradients from:
+  (A) full MSE: mean((pred-target)^2)
+  (B) SP-sharded loss: sp_size * local_sse / total_numel
+match under this mismatch.
+
+Run:
+  source /home/hao_lab/miniconda3/etc/profile.d/conda.sh
+  conda activate alexfv
+  torchrun --standalone --nproc_per_node=4 -m pytest -q fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp2_hsdp_shard4.py
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed import cleanup_dist_env_and_memory
+from fastvideo.distributed.parallel_state import (
+    get_sp_group,
+    maybe_init_distributed_environment_and_model_parallel,
+)
+from fastvideo.training.training_utils import shard_latents_across_sp
+
+
+def _world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def _rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+@pytest.fixture(scope="module")
+def dist_setup():
+    if _world_size() != 4:
+        pytest.skip("Designed for torchrun WORLD_SIZE=4")
+    if not torch.cuda.is_available():
+        pytest.skip("Requires CUDA")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank)
+
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=2)
+    assert get_sp_group().world_size == 2
+    yield
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+    cleanup_dist_env_and_memory()
+
+
+class ScaleModule(torch.nn.Module):
+    def __init__(self, c: int, init: torch.Tensor):
+        super().__init__()
+        assert init.shape == (c,)
+        self.scale = torch.nn.Parameter(init.clone())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.scale.view(1, -1, 1, 1, 1)
+
+
+def _broadcast_(t: torch.Tensor, src: int = 0) -> None:
+    dist.broadcast(t, src=src)
+
+
+def _gather_full_vec_allranks(local_shard: torch.Tensor) -> torch.Tensor:
+    gathered = [torch.empty_like(local_shard) for _ in range(dist.get_world_size())]
+    dist.all_gather(gathered, local_shard)
+    return torch.cat(gathered, dim=0)
+
+
+def test_sp_sharded_loss_matches_full_mse_under_mismatch(dist_setup):
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+    torch.manual_seed(2027)
+
+    b, c, t, h, w = 2, 8, 1, 1, 8  # thw=8 divisible by sp=2
+
+    x = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        x.normal_()
+    _broadcast_(x)
+    target = (0.25 * x).contiguous()
+
+    init = torch.empty((c,), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        init.normal_()
+    _broadcast_(init)
+
+    mesh = init_device_mesh(
+        "cuda",
+        mesh_shape=(1, 4),
+        mesh_dim_names=("replicate", "shard"),
+    )
+
+    mp = MixedPrecisionPolicy(
+        param_dtype=torch.float32,
+        reduce_dtype=torch.float32,
+        output_dtype=torch.float32,
+        cast_forward_inputs=False,
+    )
+
+    # (A) Full loss.
+    m_full = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_full, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred_full = m_full(x)
+    loss_full = ((pred_full - target) ** 2).mean()
+    loss_full.backward()
+    g_full = m_full.scale.grad
+    g_full_local = g_full.to_local() if hasattr(g_full, "to_local") else g_full
+    g_full_vec = _gather_full_vec_allranks(g_full_local.flatten())
+
+    # (B) SP-sharded loss.
+    m_sp = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_sp, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred = m_sp(x)
+    sp = get_sp_group().world_size  # 2
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(target)
+    local_sse = ((sharded_pred - sharded_target) ** 2).sum()
+    loss_sp = (sp * local_sse) / pred.numel()
+    loss_sp.backward()
+    g_sp = m_sp.scale.grad
+    g_sp_local = g_sp.to_local() if hasattr(g_sp, "to_local") else g_sp
+    g_sp_vec = _gather_full_vec_allranks(g_sp_local.flatten())
+
+    torch.testing.assert_close(g_sp_vec, g_full_vec, rtol=1e-4, atol=1e-5)
+
+

fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp4_hsdp_shard2.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Mismatch compatibility test (direction A):
+  - WORLD_SIZE=4
+  - SP: sp_size=4
+  - HSDP/FSDP2 mesh: replicate_dim=2, shard_dim=2 (mesh_shape=(2,2))
+
+We verify gradients from:
+  (A) full MSE: mean((pred-target)^2)
+  (B) SP-sharded loss: sp_size * local_sse / total_numel
+match under this mismatch.
+
+Run:
+  source /home/hao_lab/miniconda3/etc/profile.d/conda.sh
+  conda activate alexfv
+  torchrun --standalone --nproc_per_node=4 -m pytest -q fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp4_hsdp_shard2.py
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed import cleanup_dist_env_and_memory
+from fastvideo.distributed.parallel_state import (
+    get_sp_group,
+    maybe_init_distributed_environment_and_model_parallel,
+)
+from fastvideo.training.training_utils import shard_latents_across_sp
+
+
+def _world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def _rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+@pytest.fixture(scope="module")
+def dist_setup():
+    if _world_size() != 4:
+        pytest.skip("Designed for torchrun WORLD_SIZE=4")
+    if not torch.cuda.is_available():
+        pytest.skip("Requires CUDA")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank)
+
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=4)
+    assert get_sp_group().world_size == 4
+    yield
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+    cleanup_dist_env_and_memory()
+
+
+class ScaleModule(torch.nn.Module):
+    def __init__(self, c: int, init: torch.Tensor):
+        super().__init__()
+        assert init.shape == (c,)
+        self.scale = torch.nn.Parameter(init.clone())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.scale.view(1, -1, 1, 1, 1)
+
+
+def _broadcast_(t: torch.Tensor, src: int = 0) -> None:
+    dist.broadcast(t, src=src)
+
+
+def _reconstruct_full_vec_from_replicate0(local_shard: torch.Tensor) -> torch.Tensor:
+    """
+    Mesh layout: [[0,2],[1,3]] with dim_names=("replicate","shard").
+    replicate0 group is ranks [0,2], and shard order inside replicate0 is [0,2].
+    Reconstruct a single full vector from replicate0 only, then broadcast to all ranks.
+    """
+    group_ranks = [0, 2]
+    pg = dist.new_group(ranks=group_ranks)
+    if dist.get_rank() in group_ranks:
+        gathered = [torch.empty_like(local_shard) for _ in range(len(group_ranks))]
+        dist.all_gather(gathered, local_shard, group=pg)
+        full = torch.cat(gathered, dim=0)
+    else:
+        full = torch.empty((local_shard.numel() * len(group_ranks),),
+                           device=local_shard.device,
+                           dtype=local_shard.dtype)
+    dist.broadcast(full, src=0)
+    return full
+
+
+def test_sp_sharded_loss_matches_full_mse_under_mismatch(dist_setup):
+    from torch.distributed.device_mesh import DeviceMesh
+    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+    torch.manual_seed(2026)
+
+    b, c, t, h, w = 2, 8, 1, 1, 8  # thw=8 divisible by sp=4
+
+    x = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        x.normal_()
+    _broadcast_(x)
+    target = (0.5 * x).contiguous()
+
+    init = torch.empty((c,), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        init.normal_()
+    _broadcast_(init)
+
+    mesh = DeviceMesh(
+        "cuda",
+        [[0, 2], [1, 3]],
+        mesh_dim_names=("replicate", "shard"),
+    )
+
+    mp = MixedPrecisionPolicy(
+        param_dtype=torch.float32,
+        reduce_dtype=torch.float32,
+        output_dtype=torch.float32,
+        cast_forward_inputs=False,
+    )
+
+    # (A) Full loss.
+    m_full = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_full, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred_full = m_full(x)
+    loss_full = ((pred_full - target) ** 2).mean()
+    loss_full.backward()
+    g_full = m_full.scale.grad
+    g_full_local = g_full.to_local() if hasattr(g_full, "to_local") else g_full
+    g_full_vec = _reconstruct_full_vec_from_replicate0(g_full_local.flatten())
+
+    # (B) SP-sharded loss.
+    m_sp = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_sp, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred = m_sp(x)
+    sp = get_sp_group().world_size  # 4
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(target)
+    local_sse = ((sharded_pred - sharded_target) ** 2).sum()
+    loss_sp = (sp * local_sse) / pred.numel()
+    loss_sp.backward()
+    g_sp = m_sp.scale.grad
+    g_sp_local = g_sp.to_local() if hasattr(g_sp, "to_local") else g_sp
+    g_sp_vec = _reconstruct_full_vec_from_replicate0(g_sp_local.flatten())
+
+    torch.testing.assert_close(g_sp_vec, g_full_vec, rtol=1e-4, atol=1e-5)
+
+

fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed import (cleanup_dist_env_and_memory,
+                                  get_sp_group,
+                                  maybe_init_distributed_environment_and_model_parallel)
+from fastvideo.training.training_utils import shard_latents_across_sp
+
+
+def _world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def _rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+@pytest.fixture(scope="module")
+def sp_dist():
+    """
+    These tests are intended to be run under torchrun with multiple GPUs, e.g.
+      torchrun --nproc_per_node=2 -m pytest -q fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py
+      torchrun --nproc_per_node=4 -m pytest -q fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py
+      torchrun --nproc_per_node=8 -m pytest -q fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py
+    """
+    ws = _world_size()
+    if ws <= 1:
+        pytest.skip("Requires torchrun with WORLD_SIZE>1")
+    if not torch.cuda.is_available():
+        pytest.skip("Requires CUDA")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank)
+
+    torch.manual_seed(1234)
+    # Initialize FastVideo dist + model-parallel groups.
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=ws)
+    assert get_sp_group().world_size == ws
+
+    yield
+
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+    cleanup_dist_env_and_memory()
+
+
+def _broadcast_tensor_(t: torch.Tensor, src: int = 0) -> torch.Tensor:
+    if dist.is_available() and dist.is_initialized():
+        dist.broadcast(t, src=src)
+    return t
+
+
+def _all_gather_cat(t: torch.Tensor, dim: int) -> torch.Tensor:
+    ws = _world_size()
+    gathered = [torch.empty_like(t) for _ in range(ws)]
+    dist.all_gather(gathered, t)
+    return torch.cat(gathered, dim=dim)
+
+
+def test_shard_latents_t_not_divisible_but_thw_divisible_no_error(sp_dist):
+    """
+    Regression test for the commit "fix splitting on t":
+    - t is NOT divisible by sp_size
+    - (t*h*w) IS divisible by sp_size
+    - sharding should NOT raise, and shards should round-trip to the original.
+    """
+    ws = _world_size()
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+
+    b, c, t, h, w = 2, 3, 3, 2, 4  # t=3 (not divisible by 8), thw=24 (divisible by 8)
+    assert (t * h * w) % ws == 0
+    assert t % ws != 0
+
+    latents = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        latents.normal_()
+    _broadcast_tensor_(latents)
+
+    shard = shard_latents_across_sp(latents)
+    assert shard.shape == (b, c, (t * h * w) // ws)
+
+    gathered = _all_gather_cat(shard, dim=2)
+    torch.testing.assert_close(gathered, latents.reshape(b, c, t * h * w))
+
+    dist.barrier()
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        # No padding needed: thw divisible by 8, but t is not.
+        (2, 1, 3, 2, 4),  # thw=24
+        # Padding needed: thw not divisible by 8.
+        (2, 1, 3, 2, 5),  # thw=30 -> pad to 32
+    ],
+)
+def test_sharded_loss_gradient_matches_full_mse_after_rank_avg(sp_dist, shape):
+    """
+    Validates the math behind the new loss computation:
+
+    If each rank computes:
+      loss_rank = sp_world_size * local_sse / total_numel
+    and the training stack averages gradients across ranks, then the resulting
+    gradient should match the full MSE gradient.
+    """
+    ws = _world_size()
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+    b, c, t, h, w = shape
+
+    # Make inputs identical across ranks for deterministic comparison.
+    init = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    x = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    y = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        init.normal_()
+        x.normal_()
+        y.normal_()
+    _broadcast_tensor_(init)
+    _broadcast_tensor_(x)
+    _broadcast_tensor_(y)
+
+    # Full loss baseline (redundant compute; correct reference).
+    w_full = torch.nn.Parameter(init.clone())
+    pred_full = w_full * x
+    full_loss = ((pred_full - y)**2).mean()
+    full_loss.backward()
+    grad_full = w_full.grad.detach().clone()
+
+    # Sharded loss (compute on a shard only).
+    w_shard = torch.nn.Parameter(init.clone())
+    pred = w_shard * x
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(y)
+    local_sse = ((sharded_pred - sharded_target)**2).sum()
+    shard_loss = ws * local_sse / pred.numel()
+    shard_loss.backward()
+    grad_shard = w_shard.grad.detach().clone()
+
+    # Simulate "gradient averaging across ranks" (DDP/FSDP replicate-dim behavior).
+    dist.all_reduce(grad_shard, op=dist.ReduceOp.SUM)
+    grad_shard /= ws
+
+    # (Optional) also average grad_full, for symmetry.
+    dist.all_reduce(grad_full, op=dist.ReduceOp.SUM)
+    grad_full /= ws
+
+    torch.testing.assert_close(grad_shard, grad_full, rtol=1e-5, atol=1e-6)
+
+    dist.barrier()
+
+
+def test_padding_does_not_change_global_sse(sp_dist):
+    """
+    Padding-specific correctness test.
+
+    When (t*h*w) is NOT divisible by sp_size, shard_latents_across_sp pads with
+    zeros on the flattened axis. This test validates that:
+    1) Summed SSE across SP ranks equals the full (unpadded) SSE.
+    2) Any padded tokens that land on a rank are exactly zero for both pred/target.
+    """
+    from fastvideo.distributed.utils import compute_padding_for_sp
+
+    ws = _world_size()
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+
+    b, c, t, h = 2, 2, 1, 1
+    # Force padding for any ws>1: seq_len = 2*ws + 1  -> remainder 1
+    w = 2 * ws + 1
+    seq_len = t * h * w
+    assert seq_len % ws != 0
+
+    pred = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    target = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        pred.normal_()
+        target.normal_()
+    _broadcast_tensor_(pred)
+    _broadcast_tensor_(target)
+
+    # Full (unpadded) SSE reference.
+    full_sse = ((pred - target)**2).sum()
+
+    # Sharded local SSE (includes padded region, which should contribute 0).
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(target)
+    local_sse = ((sharded_pred - sharded_target)**2).sum()
+
+    global_sse = local_sse.clone()
+    dist.all_reduce(global_sse, op=dist.ReduceOp.SUM)
+    torch.testing.assert_close(global_sse, full_sse, rtol=1e-5, atol=1e-6)
+
+    # Explicitly validate padded tokens are zeros on ranks that cover them.
+    padded_seq_len, padding_amount = compute_padding_for_sp(seq_len, ws)
+    assert padding_amount > 0
+    elements_per_rank = padded_seq_len // ws
+    start = _rank() * elements_per_rank
+    end = (_rank() + 1) * elements_per_rank
+
+    if end > seq_len:
+        # There is padding on this rank: positions [max(seq_len, start), end)
+        pad_start_global = max(seq_len, start)
+        pad_end_global = end
+        pad_start_local = pad_start_global - start
+        pad_end_local = pad_end_global - start
+        assert pad_start_local < pad_end_local
+
+        pad_slice_pred = sharded_pred[:, :, pad_start_local:pad_end_local]
+        pad_slice_target = sharded_target[:, :, pad_start_local:pad_end_local]
+        torch.testing.assert_close(pad_slice_pred, torch.zeros_like(pad_slice_pred))
+        torch.testing.assert_close(pad_slice_target, torch.zeros_like(pad_slice_target))
+
+    dist.barrier()
+
+

fastvideo/tests/ssim/test_matrixgame_similarity.py

@@ -5,7 +5,7 @@ import torch
 import pytest
 
 from fastvideo import VideoGenerator
-from fastvideo.models.dits.matrix_game.utils import create_action_presets
+from fastvideo.models.dits.matrixgame.utils import create_action_presets
 from fastvideo.logger import init_logger
 from fastvideo.tests.utils import (
     compute_video_ssim_torchvision,

fastvideo/training/matrixgame_training_pipeline.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+from typing import Any
+
+import torch
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
+    FlowUniPCMultistepScheduler)
+from fastvideo.pipelines.basic.matrixgame.matrixgame_i2v_pipeline import (
+    MatrixGamePipeline)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch
+from fastvideo.training.training_pipeline import TrainingPipeline
+from fastvideo.utils import is_vsa_available, shallow_asdict
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+class MatrixGameTrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for Matrix-Game-2.0.
+    """
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowUniPCMultistepScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        May be used in future refactors.
+        """
+        pass
+
+    def set_schemas(self):
+        self.train_dataset_schema = pyarrow_schema_matrixgame
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        args_copy.dit_cpu_offload = True
+        # args_copy.pipeline_config.vae_config.load_encoder = False
+        # validation_pipeline = WanImageToVideoValidationPipeline.from_pretrained(
+        self.validation_pipeline = MatrixGamePipeline.from_pretrained(
+            training_args.model_path,
+            args=None,
+            inference_mode=True,
+            loaded_modules={
+                "transformer": self.get_module("transformer"),
+            },
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            dit_cpu_offload=True)
+
+    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
+        batch = next(self.train_loader_iter, None)  # type: ignore
+        if batch is None:
+            self.current_epoch += 1
+            logger.info("Starting epoch %s", self.current_epoch)
+            # Reset iterator for next epoch
+            self.train_loader_iter = iter(self.train_dataloader)
+            # Get first batch of new epoch
+            batch = next(self.train_loader_iter)
+
+        latents = batch['vae_latent']
+        latents = latents[:, :, :self.training_args.num_latent_t]
+        # encoder_hidden_states = batch['text_embedding']
+        # encoder_attention_mask = batch['text_attention_mask']
+        clip_features = batch['clip_feature']
+        image_latents = batch['first_frame_latent']
+        image_latents = image_latents[:, :, :self.training_args.num_latent_t]
+        pil_image = batch['pil_image']
+        infos = batch['info_list']
+
+        training_batch.latents = latents.to(get_local_torch_device(),
+                                            dtype=torch.bfloat16)
+        training_batch.encoder_hidden_states = None
+        training_batch.encoder_attention_mask = None
+        # MatrixGame doesn't use text encoder
+        training_batch.preprocessed_image = pil_image.to(
+            get_local_torch_device())
+        training_batch.image_embeds = clip_features.to(get_local_torch_device())
+        training_batch.image_latents = image_latents.to(
+            get_local_torch_device())
+        training_batch.infos = infos
+
+        # Action conditioning
+        if 'mouse_cond' in batch and batch['mouse_cond'].numel() > 0:
+            training_batch.mouse_cond = batch['mouse_cond'].to(
+                get_local_torch_device(), dtype=torch.bfloat16)
+        else:
+            training_batch.mouse_cond = None
+
+        if 'keyboard_cond' in batch and batch['keyboard_cond'].numel() > 0:
+            training_batch.keyboard_cond = batch['keyboard_cond'].to(
+                get_local_torch_device(), dtype=torch.bfloat16)
+        else:
+            training_batch.keyboard_cond = None
+
+        return training_batch
+
+    def _prepare_dit_inputs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        """Override to properly handle I2V concatenation - call parent first, then concatenate image conditioning."""
+
+        # First, call parent method to prepare noise, timesteps, etc. for video latents
+        training_batch = super()._prepare_dit_inputs(training_batch)
+
+        assert isinstance(training_batch.image_latents, torch.Tensor)
+        image_latents = training_batch.image_latents.to(
+            get_local_torch_device(), dtype=torch.bfloat16)
+
+        temporal_compression_ratio = self.training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        num_frames = (self.training_args.num_latent_t -
+                      1) * temporal_compression_ratio + 1
+        batch_size, num_channels, _, latent_height, latent_width = image_latents.shape
+        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height,
+                                   latent_width)
+        mask_lat_size[:, :, 1:] = 0
+
+        first_frame_mask = mask_lat_size[:, :, :1]
+        first_frame_mask = torch.repeat_interleave(
+            first_frame_mask, dim=2, repeats=temporal_compression_ratio)
+        mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]],
+                                  dim=2)
+        mask_lat_size = mask_lat_size.view(batch_size, -1,
+                                           temporal_compression_ratio,
+                                           latent_height, latent_width)
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(
+            image_latents.device).to(dtype=torch.bfloat16)
+
+        training_batch.noisy_model_input = torch.cat(
+            [training_batch.noisy_model_input, mask_lat_size, image_latents],
+            dim=1)
+
+        return training_batch
+
+    def _build_input_kwargs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+
+        # Image Embeds for conditioning
+        image_embeds = training_batch.image_embeds
+        assert torch.isnan(image_embeds).sum() == 0
+        image_embeds = image_embeds.to(get_local_torch_device(),
+                                       dtype=torch.bfloat16)
+        encoder_hidden_states_image = image_embeds
+
+        # NOTE: noisy_model_input already contains concatenated image_latents from _prepare_dit_inputs
+        training_batch.input_kwargs = {
+            "hidden_states":
+            training_batch.noisy_model_input,
+            "encoder_hidden_states":
+            training_batch.encoder_hidden_states,  # None for MatrixGame
+            "timestep":
+            training_batch.timesteps.to(get_local_torch_device(),
+                                        dtype=torch.bfloat16),
+            # "encoder_attention_mask":
+            # training_batch.encoder_attention_mask,
+            "encoder_hidden_states_image":
+            encoder_hidden_states_image,
+            # Action conditioning
+            "mouse_cond":
+            training_batch.mouse_cond,
+            "keyboard_cond":
+            training_batch.keyboard_cond,
+            "return_dict":
+            False,
+        }
+        return training_batch
+
+    def _prepare_validation_batch(self, sampling_param: SamplingParam,
+                                  training_args: TrainingArgs,
+                                  validation_batch: dict[str, Any],
+                                  num_inference_steps: int) -> ForwardBatch:
+        sampling_param.prompt = validation_batch['prompt']
+        sampling_param.height = training_args.num_height
+        sampling_param.width = training_args.num_width
+        sampling_param.image_path = validation_batch.get(
+            'image_path') or validation_batch.get('video_path')
+        sampling_param.num_inference_steps = num_inference_steps
+        sampling_param.data_type = "video"
+        assert self.seed is not None
+        sampling_param.seed = self.seed
+
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
+                        sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+        temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        num_frames = (training_args.num_latent_t -
+                      1) * temporal_compression_factor + 1
+        sampling_param.num_frames = num_frames
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            latents=None,
+            generator=torch.Generator(device="cpu").manual_seed(self.seed),
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=training_args.VSA_sparsity,
+        )
+        if "image" in validation_batch and validation_batch["image"] is not None:
+            batch.pil_image = validation_batch["image"]
+
+        if "keyboard_cond" in validation_batch and validation_batch[
+                "keyboard_cond"] is not None:
+            keyboard_cond = validation_batch["keyboard_cond"]
+            keyboard_cond = torch.tensor(keyboard_cond, dtype=torch.bfloat16)
+            keyboard_cond = keyboard_cond.unsqueeze(0)
+            batch.keyboard_cond = keyboard_cond
+
+        if "mouse_cond" in validation_batch and validation_batch[
+                "mouse_cond"] is not None:
+            mouse_cond = validation_batch["mouse_cond"]
+            mouse_cond = torch.tensor(mouse_cond, dtype=torch.bfloat16)
+            mouse_cond = mouse_cond.unsqueeze(0)
+            batch.mouse_cond = mouse_cond
+
+        return batch
+
+
+def main(args) -> None:
+    logger.info("Starting training pipeline...")
+
+    pipeline = MatrixGameTrainingPipeline.from_pretrained(
+        args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+    main(args)

fastvideo/training/training_pipeline.py

@@ -443,11 +443,26 @@ class TrainingPipeline(LoRAPipeline, ABC):
             # make sure no implicit broadcasting happens
             assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
 
-            sharded_pred = shard_latents_across_sp(model_pred)
-            sharded_target = shard_latents_across_sp(target)
-            loss = (torch.mean(
-                (sharded_pred.float() - sharded_target.float())**2) /
-                    self.training_args.gradient_accumulation_steps)
+            # Defensive: avoid NaNs/div0 if an upstream bug ever produces empty tensors.
+            # mean(empty) -> NaN; division by 0 -> inf/NaN. Keep a 0 loss with grad history.
+            if model_pred.numel() == 0:
+                loss = (model_pred.sum() *
+                        0.0) / self.training_args.gradient_accumulation_steps
+            else:
+                # Compute MSE on one SP shard. We shard along the flattened token axis
+                # (t*h*w) with optional padding so (t*h*w) need not be divisible by sp_size.
+                sp_world_size = get_sp_group().world_size
+                if sp_world_size > 1:
+                    sharded_pred = shard_latents_across_sp(model_pred)
+                    sharded_target = shard_latents_across_sp(target)
+                    local_sse = ((sharded_pred.float() -
+                                  sharded_target.float())**2).sum()
+                    loss = (sp_world_size * local_sse / model_pred.numel()
+                            ) / self.training_args.gradient_accumulation_steps
+                else:
+                    loss = (torch.mean(
+                        (model_pred.float() - target.float())**2) /
+                            self.training_args.gradient_accumulation_steps)
 
             loss.backward()
 
@@ -654,7 +669,11 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 seq_len = (training_batch.raw_latent_shape[2] // patch_t) * (
                     training_batch.raw_latent_shape[3] //
                     patch_h) * (training_batch.raw_latent_shape[4] // patch_w)
-                context_len = int(training_batch.encoder_hidden_states.shape[1])
+                if training_batch.encoder_hidden_states is not None:
+                    context_len = int(
+                        training_batch.encoder_hidden_states.shape[1])
+                else:
+                    context_len = 0
 
                 metrics["dit_seq_len"] = int(seq_len)
                 metrics["context_len"] = context_len

fastvideo/training/training_utils.py

@@ -20,9 +20,10 @@ from fastvideo.training.checkpointing_utils import (ModelWrapper,
                                                     RandomStateWrapper,
                                                     SchedulerWrapper)
 
-from einops import rearrange
 from fastvideo.distributed.parallel_state import (get_sp_parallel_rank,
                                                   get_sp_world_size)
+from fastvideo.distributed.utils import (compute_padding_for_sp,
+                                         pad_sequence_tensor)
 
 logger = init_logger(__name__)
 
@@ -867,10 +868,25 @@ def shard_latents_across_sp(latents: torch.Tensor) -> torch.Tensor:
     sp_world_size = get_sp_world_size()
     rank_in_sp_group = get_sp_parallel_rank()
     if sp_world_size > 1:
-        latents = rearrange(latents,
-                            "b c (n s) h w -> b c n s h w",
-                            n=sp_world_size).contiguous()
-        latents = latents[:, :, rank_in_sp_group, :, :, :]
+        # Shard on the flattened token axis (t*h*w) rather than raw `t`, so we
+        # don't require t % sp_world_size == 0. Pad on the flattened axis if needed.
+        assert latents.ndim == 5, f"Expected latents [b,c,t,h,w], got {latents.shape}"
+        b, c, t, h, w = latents.shape
+        latents = latents.reshape(b, c, t * h * w)
+
+        original_seq_len = latents.shape[2]
+        padded_seq_len, padding_amount = compute_padding_for_sp(
+            original_seq_len, sp_world_size)
+        if padding_amount > 0:
+            latents = pad_sequence_tensor(latents,
+                                          padded_seq_len,
+                                          seq_dim=2,
+                                          pad_value=0.0)
+
+        elements_per_rank = padded_seq_len // sp_world_size
+        start = rank_in_sp_group * elements_per_rank
+        end = (rank_in_sp_group + 1) * elements_per_rank
+        latents = latents[:, :, start:end].contiguous()
     return latents
 
 

fastvideo/worker/gpu_worker.py

@@ -75,6 +75,10 @@ class Worker:
 
         self.pipeline = build_pipeline(self.fastvideo_args)
 
+    def get_pipeline(self):
+        """Return the loaded pipeline instance."""
+        return self.pipeline
+
     def execute_forward(self, forward_batch: ForwardBatch,
                         fastvideo_args: FastVideoArgs) -> ForwardBatch:
         output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args)

fastvideo/worker/multiproc_executor.py

@@ -39,13 +39,17 @@ logger = init_logger(__name__)
 class StreamingTaskType(str, Enum):
     """
     Enumeration for different streaming task types.
-    
+
     Inherits from str to allow string comparison for backward compatibility.
     """
     RESET = "reset"
     STEP = "step"
     CLEAR = "clear"
     EXIT = "exit"
+    # Multi-user task types
+    USER_JOIN = "user_join"
+    USER_STEP = "user_step"
+    USER_LEAVE = "user_leave"
 
 
 @dataclass
@@ -58,6 +62,8 @@ class StreamingTask:
     # For RESET tasks:
     batch: ForwardBatch | None = None
     fastvideo_args: FastVideoArgs | None = None
+    # For multi-user tasks:
+    user_id: str | None = None
 
 
 @dataclass
@@ -66,6 +72,8 @@ class StreamingResult:
     task_type: StreamingTaskType
     output_batch: ForwardBatch | None = None
     error: Exception | None = None
+    # For multi-user results:
+    user_id: str | None = None
 
 
 class MultiprocExecutor(Executor):
@@ -183,6 +191,13 @@ class MultiprocExecutor(Executor):
         self.collective_rpc("start_streaming_queue_loop")
         self._streaming_enabled = True
 
+    def enable_multi_user_streaming(self) -> None:
+        if self._streaming_enabled:
+            return
+
+        self.collective_rpc("start_multi_user_streaming_loop")
+        self._streaming_enabled = True
+
     def disable_streaming(self) -> None:
         if not self._streaming_enabled:
             return
@@ -225,6 +240,42 @@ class MultiprocExecutor(Executor):
             self._streaming_input_queue.put(
                 StreamingTask(task_type=StreamingTaskType.CLEAR))
 
+    def submit_user_join(self, user_id: str, forward_batch: ForwardBatch,
+                         fastvideo_args: FastVideoArgs) -> None:
+        if not self._streaming_enabled:
+            self.enable_streaming()
+
+        self._streaming_input_queue.put(
+            StreamingTask(
+                task_type=StreamingTaskType.USER_JOIN,
+                user_id=user_id,
+                batch=forward_batch,
+                fastvideo_args=fastvideo_args,
+            ))
+
+    def submit_user_step(self, user_id: str,
+                         keyboard_action: torch.Tensor | None,
+                         mouse_action: torch.Tensor | None) -> None:
+        if not self._streaming_enabled:
+            raise RuntimeError(
+                "Streaming mode not enabled. Call enable_streaming() first.")
+
+        self._streaming_input_queue.put(
+            StreamingTask(
+                task_type=StreamingTaskType.USER_STEP,
+                user_id=user_id,
+                keyboard_action=keyboard_action,
+                mouse_action=mouse_action,
+            ))
+
+    def submit_user_leave(self, user_id: str) -> None:
+        if self._streaming_enabled and self._streaming_input_queue is not None:
+            self._streaming_input_queue.put(
+                StreamingTask(
+                    task_type=StreamingTaskType.USER_LEAVE,
+                    user_id=user_id,
+                ))
+
     def get_result(self,
                    timeout: float | None = None) -> StreamingResult | None:
         if not self._streaming_enabled or self._streaming_output_queue is None:
@@ -642,6 +693,11 @@ class WorkerMultiprocProc:
                             {"status": "streaming_queue_loop_started"})
                         self.streaming_queue_loop()
                         continue
+                    if method == "start_multi_user_streaming_loop":
+                        self.pipe.send(
+                            {"status": "multi_user_streaming_loop_started"})
+                        self.multi_user_streaming_loop()
+                        continue
                     if method == 'execute_forward':
                         forward_batch = kwargs['forward_batch']
                         fastvideo_args = kwargs['fastvideo_args']
@@ -715,11 +771,126 @@ class WorkerMultiprocProc:
                     self.worker.execute_streaming_clear()
                     self.streaming_output_queue.put(
                         StreamingResult(task_type=StreamingTaskType.CLEAR))
+                elif task.task_type in (StreamingTaskType.USER_JOIN,
+                                        StreamingTaskType.USER_STEP,
+                                        StreamingTaskType.USER_LEAVE):
+                    # Switch to multi-user mode
+                    self.multi_user_streaming_loop(first_task=task)
+                    break
             except Exception as e:
                 logger.error("Worker %d queue loop error: %s", self.rank, e)
                 self.streaming_output_queue.put(
                     StreamingResult(task_type=StreamingTaskType.STEP, error=e))
 
+    def multi_user_streaming_loop(
+        self, first_task: StreamingTask | None = None
+    ) -> None:
+        """ORCA-style multi-user streaming loop.
+
+        Processes USER_JOIN/USER_STEP/USER_LEAVE tasks and batches
+        compatible users together for each denoising step.
+        """
+        from fastvideo.pipelines.stages.multi_user_engine import MultiUserEngine
+
+        if self.streaming_input_queue is None or self.streaming_output_queue is None:
+            logger.error("Worker %d: streaming queues not initialized",
+                         self.rank)
+            return
+
+        pipeline = self.worker.get_pipeline()
+        try:
+            from ui.world_model.server.config import DISABLE_BATCHING
+        except ImportError:
+            DISABLE_BATCHING = False
+        engine = MultiUserEngine(pipeline, disable_batching=DISABLE_BATCHING)
+        logger.info("Worker %d: multi-user streaming loop started", self.rank)
+
+        def handle_task(task: StreamingTask) -> bool:
+            """Handle a single task. Returns False if should exit."""
+            if task.task_type == StreamingTaskType.EXIT:
+                return False
+            elif task.task_type == StreamingTaskType.USER_JOIN:
+                try:
+                    engine.add_user(
+                        task.user_id, task.batch, task.fastvideo_args)
+                    self.streaming_output_queue.put(
+                        StreamingResult(
+                            task_type=StreamingTaskType.USER_JOIN,
+                            user_id=task.user_id))
+                except Exception as e:
+                    logger.error("Worker %d user_join error for %s: %s",
+                                 self.rank, task.user_id, e)
+                    self.streaming_output_queue.put(
+                        StreamingResult(
+                            task_type=StreamingTaskType.USER_JOIN,
+                            user_id=task.user_id, error=e))
+            elif task.task_type == StreamingTaskType.USER_STEP:
+                try:
+                    engine.submit_step(
+                        task.user_id, task.keyboard_action, task.mouse_action)
+                except Exception as e:
+                    logger.error("Worker %d user_step submit error for %s: %s",
+                                 self.rank, task.user_id, e)
+                    self.streaming_output_queue.put(
+                        StreamingResult(
+                            task_type=StreamingTaskType.USER_STEP,
+                            user_id=task.user_id, error=e))
+            elif task.task_type == StreamingTaskType.USER_LEAVE:
+                engine.remove_user(task.user_id)
+                self.streaming_output_queue.put(
+                    StreamingResult(
+                        task_type=StreamingTaskType.USER_LEAVE,
+                        user_id=task.user_id))
+            return True
+
+        # Handle the first task that triggered the switch
+        if first_task is not None:
+            if not handle_task(first_task):
+                return
+
+        while True:
+            # 1. Drain input queue (non-blocking) for new requests
+            while True:
+                try:
+                    task = self.streaming_input_queue.get_nowait()
+                    if not handle_task(task):
+                        return
+                except queue.Empty:
+                    break
+
+            # 2. Run one ORCA iteration
+            if engine.has_pending_work():
+                try:
+                    completed = engine.run_iteration()
+                    for result in completed:
+                        self.streaming_output_queue.put(
+                            StreamingResult(
+                                task_type=StreamingTaskType.USER_STEP,
+                                user_id=result.user_id,
+                                output_batch=result.output_batch))
+                except Exception as e:
+                    logger.error("Worker %d ORCA iteration error: %s",
+                                 self.rank, e, exc_info=True)
+                    # Reset all pending users to prevent infinite error loop
+                    for uid, sess in engine.users.items():
+                        if sess.denoising_step >= 0:
+                            sess.denoising_step = -1
+                            sess.current_latents = None
+                            sess.noise_latents_btchw = None
+                            self.streaming_output_queue.put(
+                                StreamingResult(
+                                    task_type=StreamingTaskType.USER_STEP,
+                                    user_id=uid, error=e))
+
+            # 3. If no work, block on queue to avoid busy-waiting
+            if not engine.has_pending_work():
+                try:
+                    task = self.streaming_input_queue.get(timeout=0.1)
+                    if not handle_task(task):
+                        return
+                except queue.Empty:
+                    continue
+
     @staticmethod
     def setup_proc_title_and_log_prefix() -> None:
         dp_size = get_dp_group().world_size

mkdocs.yml

@@ -152,9 +152,10 @@ nav:
     - Overview: attention/index.md
     - Video Sparse Attention: attention/vsa/index.md
     - Sliding Tile Attention: attention/sta/index.md
-    - Adding a New Attention Backend: attention/developer/index.md
+    - Backend Development: contributing/attention_backend.md
   - Utilities:
     - LoRA: utilities/lora.md
+    - Debugging: utilities/debugging.md
   - Design:
     - Overview: design/overview.md
   - Developer Guide:
@@ -165,7 +166,8 @@ nav:
       - RunPod: contributing/developer_env/runpod.md
     - Testing: contributing/testing.md
     - Profiling: contributing/profiling.md
-    - Adding a New Attention Backend: attention/developer/index.md
+    - Coding Agents: contributing/coding_agents.md
+    - Attention Backend Development: contributing/attention_backend.md
   - API Reference:
     - FastVideo: api/fastvideo.md
 

ui/world_model/client/index.html

@@ -0,0 +1,12 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <title>World Model Interface</title>
+</head>
+<body>
+  <div id="app"></div>
+  <script type="module" src="/src/main.js"></script>
+</body>
+</html>

ui/world_model/client/package.json

@@ -0,0 +1,16 @@
+{
+  "name": "wm-interface-frontend",
+  "private": true,
+  "version": "0.0.1",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "vite build",
+    "preview": "vite preview"
+  },
+  "devDependencies": {
+    "@sveltejs/vite-plugin-svelte": "^3.0.1",
+    "svelte": "^4.2.8",
+    "vite": "^5.0.11"
+  }
+}

ui/world_model/client/public/logo.svg

@@ -0,0 +1,18 @@
+<svg width="252" height="105" viewBox="0 0 252 105" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM100.768 13.1217L87.7028 29.4852H103.143L100.768 13.1217Z" fill="#356CFF"/>
+<path d="M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM109.081 90.697L116.802 65.8487C116.802 65.8487 120.959 65.8487 132.242 65.8487C143.525 65.8487 137.586 78.5759 135.211 84.0304C133.307 88.4021 127.491 90.697 122.74 90.697C117.989 90.697 109.081 90.697 109.081 90.697Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944H159.747C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273L125.188 48.273L124 37.97L147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852H131.836C120.142 29.4852 125.897 1.00043 141.337 1.00043L173.188 1.00056Z" fill="#356CFF"/>
+<path d="M179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056Z" fill="#356CFF"/>
+<path d="M161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457Z" fill="#356CFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM237.948 77.9692C239.984 70.6965 240.917 65.242 228.446 65.242C215.975 65.242 211.818 71.9087 210.037 77.9692C208.255 84.0298 208.255 91.3025 219.538 91.3025C230.821 91.3025 235.911 85.2419 237.948 77.9692Z" fill="#356CFF"/>
+<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944M173.188 1.00056C173.188 1.00056 156.777 1.00043 141.337 1.00043M173.188 1.00056L141.337 1.00043M141.337 20.3944C146.088 20.3944 150.839 20.3944 159.747 20.3944M141.337 20.3944H159.747M159.747 20.3944C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273M148.463 48.273C139.556 48.273 125.188 48.273 125.188 48.273M148.463 48.273L125.188 48.273M125.188 48.273L124 37.97M124 37.97C124 37.97 141.931 37.97 147.87 37.97M124 37.97L147.87 37.97M147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852M151.433 29.4852C146.682 29.4852 138.962 29.4852 131.836 29.4852M151.433 29.4852H131.836M131.836 29.4852C120.142 29.4852 125.897 1.00043 141.337 1.00043M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057ZM96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM87.7028 29.4852L100.768 13.1217L103.143 29.4852H87.7028ZM89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457ZM108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM116.802 65.8487L109.081 90.697C109.081 90.697 117.989 90.697 122.74 90.697C127.491 90.697 133.307 88.4021 135.211 84.0304C137.586 78.5759 143.525 65.8487 132.242 65.8487C120.959 65.8487 116.802 65.8487 116.802 65.8487ZM179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056ZM161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457ZM230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM228.446 65.242C240.917 65.242 239.984 70.6965 237.948 77.9692C235.911 85.2419 230.821 91.3025 219.538 91.3025C208.255 91.3025 208.255 84.0298 210.037 77.9692C211.818 71.9087 215.975 65.242 228.446 65.242Z" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M15.2524 55.5451L21.191 100.999L24.7541 100.999L18.8156 55.5451L15.2524 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 55.5451L14.065 100.999L15.2527 100.999L9.31417 55.5451L8.12646 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 55.5451L6.93853 100.999L7.53239 100.999L1.59385 55.5451L1 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M15.2524 48.2724L30.0988 1H33.6619L18.8156 48.2724H15.2524Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M8.12646 48.2724L22.9728 1H24.1605L9.31417 48.2724H8.12646Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
+<path d="M1 48.2724L15.8463 1H16.4402L1.59385 48.2724H1Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
+<path d="M85.3271 55.5457H67.5116L87 12.7363L44.3513 68.2729H58.6038L43.1636 101L85.3271 55.5457Z" fill="#FDC717" stroke="#FDC717" stroke-width="1.18771" stroke-miterlimit="16"/>
+</svg>

ui/world_model/client/src/main.js

@@ -0,0 +1,8 @@
+import App from './App.svelte'
+import './style.css'
+
+const app = new App({
+  target: document.getElementById('app')
+})
+
+export default app

ui/world_model/client/src/style.css

@@ -0,0 +1,20 @@
+* {
+  margin: 0;
+  padding: 0;
+  box-sizing: border-box;
+}
+
+body {
+  font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
+  background: #1a1a1a;
+  color: white;
+  min-height: 100vh;
+  overflow-y: auto;
+}
+
+#app {
+  width: 100%;
+  min-height: 100vh;
+  display: flex;
+  flex-direction: column;
+}

ui/world_model/client/vite.config.js

@@ -0,0 +1,18 @@
+import { defineConfig } from 'vite'
+import { svelte } from '@sveltejs/vite-plugin-svelte'
+
+export default defineConfig({
+  plugins: [svelte()],
+  server: {
+    port: 5173,
+    proxy: {
+      '/ws': {
+        target: 'http://localhost:8000',
+        ws: true,
+      },
+      '/status': {
+        target: 'http://localhost:8000',
+      },
+    },
+  }
+})

ui/world_model/server/config.py

@@ -0,0 +1,61 @@
+from pathlib import Path
+
+# Repo root: ui/world_model/server/config.py -> go up 3 levels
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+
+# Model registry - add new models here
+MODEL_REGISTRY = {
+    "matrix-game-2.0-base": {
+        "name": "Matrix-Game 2.0 Base",
+        "model_path": "FastVideo/Matrix-Game-2.0-Base-Diffusers",
+        "keyboard_dim": 4,
+        "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png",
+    },
+    "wangame-1.3b-mc-w-only-still-7k": {
+        "name": "WANGame 1.3B MC (W-Only, 7k)",
+        "model_path": "weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers",
+        "init_weights_from_safetensors": str(_REPO_ROOT / "checkpoints" / "wangame-1.3b-mc-w-only-still-7k" / "checkpoint-7000" / "transformer"),
+        "override_transformer_cls_name": "WanGameActionTransformer3DModel",
+        "override_pipeline_cls_name": "WanGameCausalDMDPipeline",
+        "keyboard_dim": 4,
+        #"image_url": "/server-assets/mc.png",
+        "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png",
+        # "image_path": str(Path(__file__).resolve().parent / "mc.png"),
+    },
+}
+
+DEFAULT_MODEL_ID = "matrix-game-2.0-base"
+
+# Active model configuration (set by server or user selection)
+MODEL_CONFIG = MODEL_REGISTRY[DEFAULT_MODEL_ID]
+
+# Keyboard mappings (WASD)
+KEYBOARD_MAP = {
+    "w": [1, 0, 0, 0],
+    "s": [0, 1, 0, 0],
+    "a": [0, 0, 1, 0],
+    "d": [0, 0, 0, 1],
+}
+
+# Camera mappings (Arrow keys)
+CAM_VALUE = 0.1
+CAMERA_MAP = {
+    "ArrowUp": [CAM_VALUE, 0],
+    "ArrowDown": [-CAM_VALUE, 0],
+    "ArrowLeft": [0, -CAM_VALUE],
+    "ArrowRight": [0, CAM_VALUE],
+}
+
+# Generation limits
+MAX_BLOCKS = 50
+SESSION_TIMEOUT_SECONDS = 900
+MAX_USERS_PER_GPU = 16
+DISABLE_BATCHING = False  # Process users sequentially (batch_size=1) instead of batching
+
+# Frame settings
+NUM_FRAMES = 597
+FRAME_HEIGHT = 352
+FRAME_WIDTH = 640
+NUM_INFERENCE_STEPS = 4
+JPEG_QUALITY = 85
+BATCH_SIZE = 12

ui/world_model/server/main.py

@@ -0,0 +1,262 @@
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from contextlib import asynccontextmanager
+import asyncio
+import base64
+import cv2
+import os
+import time
+import uuid
+
+from config import (
+    KEYBOARD_MAP, CAMERA_MAP, MAX_BLOCKS, JPEG_QUALITY, BATCH_SIZE,
+    SESSION_TIMEOUT_SECONDS, MODEL_CONFIG, MODEL_REGISTRY, DEFAULT_MODEL_ID
+)
+from gpu_pool import GPUPool, GPUSlot, get_available_gpus
+
+# Global GPU pool
+gpu_pool: GPUPool = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan manager."""
+    global gpu_pool
+
+    print("Starting server...")
+
+    # Get available GPUs
+    gpu_ids = get_available_gpus()
+    print(f"Detected GPUs: {gpu_ids}")
+
+    # Initialize GPU pool (spawns subprocess per GPU)
+    gpu_pool = GPUPool(gpu_ids)
+    await gpu_pool.initialize()
+
+    print("Server ready")
+    yield
+
+    print("Shutting down server...")
+    await gpu_pool.shutdown()
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.get("/status")
+async def get_status():
+    """Get the current status of the GPU pool."""
+    if gpu_pool is None:
+        return {"error": "GPU pool not initialized"}
+    return gpu_pool.get_status()
+
+
+@app.get("/models")
+async def get_models():
+    """Get available models and the currently active model."""
+    models = []
+    for model_id, config in MODEL_REGISTRY.items():
+        models.append({
+            "id": model_id,
+            "name": config["name"],
+        })
+    return {
+        "models": models,
+        "default_model_id": DEFAULT_MODEL_ID,
+    }
+
+
+def encode_frames(frames: list) -> list[str]:
+    """Encode frames to base64 JPEG."""
+    encoded = []
+    for frame in frames:
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        _, buffer = cv2.imencode('.jpg', frame_rgb, [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
+        encoded.append(base64.b64encode(buffer.tobytes()).decode("utf-8"))
+    return encoded
+
+
+async def send_frames(websocket: WebSocket, frames: list):
+    """Encode and send frames in batches."""
+    if not frames:
+        return
+    encoded_frames = encode_frames(frames)
+    for i in range(0, len(encoded_frames), BATCH_SIZE):
+        batch = encoded_frames[i:i + BATCH_SIZE]
+        await websocket.send_json({"type": "frame_batch", "frames": batch})
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+
+    client_id = str(uuid.uuid4())
+    print(f"Client {client_id[:8]} connected")
+
+    # Send initial queue status
+    status = gpu_pool.get_status()
+    await websocket.send_json({
+        "type": "queue_status",
+        "position": status["queue_size"] + 1 if status["available_gpus"] == 0 else 0,
+        "total_gpus": status["total_gpus"],
+        "available_gpus": status["available_gpus"]
+    })
+
+    gpu_id = None
+    slot: GPUSlot = None
+    timeout_task: asyncio.Task = None
+
+    async def session_timeout():
+        """Close the session after timeout."""
+        await asyncio.sleep(SESSION_TIMEOUT_SECONDS)
+        print(f"[GPU {gpu_id}] Session timeout for client {client_id[:8]}")
+        try:
+            await websocket.send_json({
+                "type": "session_timeout",
+                "message": f"Session expired after {SESSION_TIMEOUT_SECONDS} seconds"
+            })
+            await websocket.close(code=1000, reason="Session timeout")
+        except Exception:
+            pass  # WebSocket may already be closed
+
+    try:
+        # Wait for model selection from client
+        model_id = DEFAULT_MODEL_ID
+        model_config = MODEL_CONFIG
+        try:
+            init_data = await asyncio.wait_for(websocket.receive_json(), timeout=5.0)
+            if init_data.get("type") == "select_model":
+                model_id = init_data.get("model_id", DEFAULT_MODEL_ID)
+                if model_id in MODEL_REGISTRY:
+                    model_config = MODEL_REGISTRY[model_id]
+                    print(f"Client {client_id[:8]} selected model: {model_id}")
+        except asyncio.TimeoutError:
+            pass  # Use default model config
+
+        # Acquire a GPU slot (may wait in queue, may share with other users)
+        gpu_id, slot = await gpu_pool.acquire(client_id, websocket)
+
+        # Start session timeout
+        timeout_task = asyncio.create_task(session_timeout())
+
+        # Join the multi-user engine on this GPU (triggers reload if model differs)
+        await slot.join_user(client_id, model_id=model_id)
+
+        # Notify client they're connected to a GPU
+        await websocket.send_json({
+            "type": "gpu_assigned",
+            "gpu_id": gpu_id,
+            "session_timeout": SESSION_TIMEOUT_SECONDS,
+            "image_url": model_config.get("image_url")
+        })
+
+        # Generate and send initial frame
+        block_count = 0
+        await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS})
+
+        try:
+            frames, timings = await slot.user_step(
+                client_id, [0, 0, 0, 0], [0, 0])
+            block_count = 1
+            await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS})
+            await send_frames(websocket, frames)
+            if timings:
+                print(f"[GPU {gpu_id}] Initial ({client_id[:8]}): model={timings.get('model_step_ms', 0):.0f}ms")
+        except Exception as e:
+            print(f"[GPU {gpu_id}] Initial frame error for {client_id[:8]}: {e}")
+
+        # Main event loop
+        while True:
+            data = await websocket.receive_json()
+            message_type = data.get("type", "key")
+
+            if message_type == "reset":
+                print(f"[GPU {gpu_id}] Reset requested by client {client_id[:8]}")
+                await websocket.send_json({"type": "reset_started"})
+
+                try:
+                    # Leave and rejoin to reset user state
+                    await slot.leave_user(client_id)
+                    await slot.join_user(client_id, model_id=model_id)
+
+                    frames, timings = await slot.user_step(
+                        client_id, [0, 0, 0, 0], [0, 0])
+                    block_count = 1
+                    await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS})
+                    await send_frames(websocket, frames)
+                    await websocket.send_json({"type": "reset_complete"})
+                    print(f"[GPU {gpu_id}] Reset complete for {client_id[:8]}")
+                except Exception as e:
+                    print(f"[GPU {gpu_id}] Reset error for {client_id[:8]}: {e}")
+                    await websocket.send_json({"type": "reset_complete"})
+                continue
+
+            # Handle key press
+            key = data.get("key")
+            if key and block_count < MAX_BLOCKS:
+                if key in CAMERA_MAP:
+                    keyboard_vector = [0, 0, 0, 0]
+                    mouse_vector = CAMERA_MAP[key]
+                else:
+                    keyboard_vector = KEYBOARD_MAP.get(key, [0, 0, 0, 0])
+                    mouse_vector = [0, 0]
+
+                print(f"[GPU {gpu_id}] Key '{key}' from {client_id[:8]}, generating block {block_count + 1}...")
+
+                t_start = time.time()
+                frames, timings = await slot.user_step(
+                    client_id, keyboard_vector, mouse_vector)
+                t_generation = time.time() - t_start
+
+                block_count += 1
+                await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS})
+
+                if frames:
+                    t_encode_start = time.time()
+                    await send_frames(websocket, frames)
+                    t_encode = time.time() - t_encode_start
+
+                    dit_ms = timings.get('dit_ms', 0)
+                    vae_ms = timings.get('vae_ms', 0)
+                    print(f"[GPU {gpu_id}] ({client_id[:8]}) DiT: {dit_ms:.0f}ms, VAE: {vae_ms:.0f}ms, Encode+Send: {t_encode*1000:.0f}ms, Total: {t_generation*1000:.0f}ms")
+
+    except WebSocketDisconnect:
+        print(f"Client {client_id[:8]} disconnected")
+    except Exception as e:
+        print(f"Client {client_id[:8]} error: {e}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        # Cancel the timeout task if still running
+        if timeout_task and not timeout_task.done():
+            timeout_task.cancel()
+            try:
+                await timeout_task
+            except asyncio.CancelledError:
+                pass
+        if client_id:
+            await gpu_pool.release(client_id)
+
+
+# Serve server-side assets (images, etc.)
+server_dir = os.path.dirname(os.path.abspath(__file__))
+app.mount("/server-assets", StaticFiles(directory=server_dir), name="server-assets")
+
+# Serve built frontend (must be after API/WebSocket routes)
+static_dir = os.path.join(server_dir, "..", "client", "dist")
+if os.path.isdir(static_dir):
+    app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8001)

ui/world_model/server/mock_server.py

@@ -0,0 +1,332 @@
+"""
+Mock server to test latency perception.
+Generates frames with a moving object based on keyboard input.
+Configurable latency to find the acceptable threshold.
+"""
+
+import asyncio
+import base64
+import io
+import time
+import uuid
+from contextlib import asynccontextmanager
+
+import numpy as np
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from PIL import Image, ImageDraw
+import os
+
+# Configuration
+LATENCY_MS = 200  # Simulated model latency - adjust this to test
+FRAME_WIDTH = 640
+FRAME_HEIGHT = 352
+NUM_FRAMES = 12  # Frames per block
+FPS = 24
+BATCH_SIZE = 4
+JPEG_QUALITY = 85
+SESSION_TIMEOUT_SECONDS = 90
+MAX_BLOCKS = 50
+
+# Mock game state
+class GameState:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.x = FRAME_WIDTH // 2
+        self.y = FRAME_HEIGHT // 2
+        self.size = 40
+        self.color = (100, 200, 255)  # Light blue
+        self.trail = []  # Previous positions for motion blur effect
+
+    def update(self, keyboard_vector, mouse_vector):
+        """Store the movement direction for this block."""
+        self.keyboard_vector = keyboard_vector
+        self.mouse_vector = mouse_vector
+
+        # Arrow keys for camera (just change color as visual feedback)
+        if mouse_vector[0] != 0 or mouse_vector[1] != 0:
+            r = int(100 + mouse_vector[0] * 50) % 256
+            g = int(200 + mouse_vector[1] * 50) % 256
+            self.color = (r, g, 255)
+
+    def render_frame(self, frame_idx: int) -> np.ndarray:
+        """Render a single frame with incremental movement."""
+        speed_per_frame = 8
+
+        # Apply movement for this frame
+        if hasattr(self, 'keyboard_vector'):
+            if self.keyboard_vector[0]:  # W - up
+                self.y -= speed_per_frame
+            if self.keyboard_vector[1]:  # A - left
+                self.x -= speed_per_frame
+            if self.keyboard_vector[2]:  # S - down
+                self.y += speed_per_frame
+            if self.keyboard_vector[3]:  # D - right
+                self.x += speed_per_frame
+
+        # Wrap around screen
+        self.x = self.x % FRAME_WIDTH
+        self.y = self.y % FRAME_HEIGHT
+
+        # Update trail
+        self.trail.append((self.x, self.y))
+        if len(self.trail) > 5:
+            self.trail.pop(0)
+        """Render a single frame."""
+        # Create frame with gradient background
+        img = Image.new('RGB', (FRAME_WIDTH, FRAME_HEIGHT), (20, 20, 30))
+        draw = ImageDraw.Draw(img)
+
+        # Draw grid for visual reference
+        for x in range(0, FRAME_WIDTH, 40):
+            draw.line([(x, 0), (x, FRAME_HEIGHT)], fill=(40, 40, 50), width=1)
+        for y in range(0, FRAME_HEIGHT, 40):
+            draw.line([(0, y), (FRAME_WIDTH, y)], fill=(40, 40, 50), width=1)
+
+        # Interpolate position for smooth animation within block
+        t = frame_idx / NUM_FRAMES
+
+        # Draw trail (motion blur effect)
+        for i, (tx, ty) in enumerate(self.trail):
+            alpha = (i + 1) / len(self.trail) * 0.3
+            trail_size = int(self.size * (0.5 + alpha * 0.5))
+            trail_color = tuple(int(c * alpha) for c in self.color)
+            draw.ellipse(
+                [tx - trail_size//2, ty - trail_size//2,
+                 tx + trail_size//2, ty + trail_size//2],
+                fill=trail_color
+            )
+
+        # Draw main object (circle with glow)
+        # Outer glow
+        for r in range(3, 0, -1):
+            glow_size = self.size + r * 8
+            glow_alpha = 0.2 / r
+            glow_color = tuple(int(c * glow_alpha) for c in self.color)
+            draw.ellipse(
+                [self.x - glow_size//2, self.y - glow_size//2,
+                 self.x + glow_size//2, self.y + glow_size//2],
+                fill=glow_color
+            )
+
+        # Main circle
+        draw.ellipse(
+            [self.x - self.size//2, self.y - self.size//2,
+             self.x + self.size//2, self.y + self.size//2],
+            fill=self.color
+        )
+
+        # Inner highlight
+        highlight_size = self.size // 3
+        highlight_offset = self.size // 6
+        draw.ellipse(
+            [self.x - highlight_offset - highlight_size//2,
+             self.y - highlight_offset - highlight_size//2,
+             self.x - highlight_offset + highlight_size//2,
+             self.y - highlight_offset + highlight_size//2],
+            fill=(255, 255, 255)
+        )
+
+        # Add latency indicator text
+        draw.text((10, 10), f"Latency: {LATENCY_MS}ms", fill=(150, 150, 150))
+        draw.text((10, 30), f"Frame: {frame_idx + 1}/{NUM_FRAMES}", fill=(100, 100, 100))
+
+        return np.array(img)
+
+
+def encode_frame(frame: np.ndarray) -> str:
+    """Encode frame to base64 JPEG."""
+    img = Image.fromarray(frame)
+    buffer = io.BytesIO()
+    img.save(buffer, format='JPEG', quality=JPEG_QUALITY)
+    return base64.b64encode(buffer.getvalue()).decode('utf-8')
+
+
+def generate_frames(game_state: GameState) -> list[str]:
+    """Generate and encode a block of frames."""
+    encoded = []
+    for i in range(NUM_FRAMES):
+        frame = game_state.render_frame(i)
+        encoded.append(encode_frame(frame))
+    return encoded
+
+
+# Keyboard/camera mappings (same as real server)
+KEYBOARD_MAP = {
+    'w': [1, 0, 0, 0],
+    'a': [0, 1, 0, 0],
+    's': [0, 0, 1, 0],
+    'd': [0, 0, 0, 1],
+}
+
+CAMERA_MAP = {
+    'ArrowUp': [0, -1],
+    'ArrowDown': [0, 1],
+    'ArrowLeft': [-1, 0],
+    'ArrowRight': [1, 0],
+}
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    print(f"Mock server starting with {LATENCY_MS}ms simulated latency...")
+    yield
+    print("Mock server shutting down...")
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+async def send_frames(websocket: WebSocket, frames: list[str]):
+    """Send frames in batches."""
+    for i in range(0, len(frames), BATCH_SIZE):
+        batch = frames[i:i + BATCH_SIZE]
+        await websocket.send_json({"type": "frame_batch", "frames": batch})
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+
+    client_id = str(uuid.uuid4())
+    print(f"Client {client_id[:8]} connected")
+
+    game_state = GameState()
+    block_count = 0
+    timeout_task = None
+
+    async def session_timeout():
+        await asyncio.sleep(SESSION_TIMEOUT_SECONDS)
+        try:
+            await websocket.send_json({
+                "type": "session_timeout",
+                "message": f"Session expired after {SESSION_TIMEOUT_SECONDS} seconds"
+            })
+            await websocket.close()
+        except Exception:
+            pass
+
+    try:
+        # Send initial status
+        await websocket.send_json({
+            "type": "queue_status",
+            "position": 0,
+            "total_gpus": 1,
+            "available_gpus": 1
+        })
+
+        # Simulate GPU assignment delay
+        await asyncio.sleep(0.1)
+
+        # Start timeout
+        timeout_task = asyncio.create_task(session_timeout())
+
+        # Send GPU assigned
+        await websocket.send_json({
+            "type": "gpu_assigned",
+            "gpu_id": 0,
+            "session_timeout": SESSION_TIMEOUT_SECONDS,
+            "image_url": None  # No initial image for mock
+        })
+
+        # Generate initial frame
+        await websocket.send_json({"type": "block_count", "count": 0, "max": MAX_BLOCKS})
+
+        # Simulate latency
+        await asyncio.sleep(LATENCY_MS / 1000)
+
+        frames = generate_frames(game_state)
+        block_count = 1
+        await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS})
+        await send_frames(websocket, frames)
+
+        print(f"[Mock] Initial frames sent")
+
+        # Main loop
+        while True:
+            data = await websocket.receive_json()
+            message_type = data.get("type", "key")
+
+            if message_type == "reset":
+                print(f"[Mock] Reset requested")
+                await websocket.send_json({"type": "reset_started"})
+
+                game_state.reset()
+                await asyncio.sleep(LATENCY_MS / 1000)
+
+                frames = generate_frames(game_state)
+                block_count = 1
+                await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS})
+                await send_frames(websocket, frames)
+                await websocket.send_json({"type": "reset_complete"})
+                continue
+
+            key = data.get("key")
+            if key and block_count < MAX_BLOCKS:
+                if key in CAMERA_MAP:
+                    keyboard_vector = [0, 0, 0, 0]
+                    mouse_vector = CAMERA_MAP[key]
+                else:
+                    keyboard_vector = KEYBOARD_MAP.get(key, [0, 0, 0, 0])
+                    mouse_vector = [0, 0]
+
+                print(f"[Mock] Key '{key}' pressed, generating block {block_count + 1}...")
+
+                t_start = time.time()
+
+                # Simulate model latency
+                await asyncio.sleep(LATENCY_MS / 1000)
+
+                # Update game state and generate frames
+                game_state.update(keyboard_vector, mouse_vector)
+                frames = generate_frames(game_state)
+
+                t_total = (time.time() - t_start) * 1000
+
+                block_count += 1
+                await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS})
+                await send_frames(websocket, frames)
+
+                print(f"[Mock] Block generated in {t_total:.0f}ms (target: {LATENCY_MS}ms)")
+
+    except WebSocketDisconnect:
+        print(f"Client {client_id[:8]} disconnected")
+    except Exception as e:
+        print(f"Client {client_id[:8]} error: {e}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        if timeout_task and not timeout_task.done():
+            timeout_task.cancel()
+
+
+# Serve static files
+static_dir = os.path.join(os.path.dirname(__file__), "..", "client", "dist")
+if os.path.isdir(static_dir):
+    app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
+
+
+if __name__ == "__main__":
+    import argparse
+    import uvicorn
+
+    parser = argparse.ArgumentParser(description="Mock server for latency testing")
+    parser.add_argument("--latency", type=int, default=200, help="Simulated latency in ms")
+    parser.add_argument("--port", type=int, default=8001, help="Server port")
+    args = parser.parse_args()
+
+    LATENCY_MS = args.latency
+    print(f"Starting mock server with {LATENCY_MS}ms latency on port {args.port}")
+
+    uvicorn.run(app, host="0.0.0.0", port=args.port)

ui/world_model/server/test_batching.py

@@ -0,0 +1,149 @@
+"""Test ORCA batching: synchronized vs random keypresses."""
+import argparse
+import asyncio
+import json
+import random
+import time
+import websockets
+
+SERVER_URL = "ws://localhost:8000/ws"
+KEYS = ["w", "a", "s", "d"]
+# Think time range for random mode (ms)
+MIN_THINK_MS = 200
+MAX_THINK_MS = 1500
+
+
+async def connect_and_wait(client_id, ws):
+    """Drain messages until initial frame_batch is received."""
+    while True:
+        msg = json.loads(await ws.recv())
+        if msg["type"] == "frame_batch":
+            print(f"  Client {client_id}: got initial frame")
+            return
+
+
+async def synced_client(client_id: int, barrier: asyncio.Barrier, results: list,
+                        num_rounds: int):
+    """Synchronized keypresses — all clients press at the same time."""
+    async with websockets.connect(
+        SERVER_URL, max_size=50_000_000, ping_interval=None,
+    ) as ws:
+        await connect_and_wait(client_id, ws)
+
+        for round_idx in range(num_rounds):
+            await barrier.wait()
+
+            t0 = time.perf_counter()
+            await ws.send(json.dumps({"type": "key", "key": "w"}))
+
+            while True:
+                msg = json.loads(await ws.recv())
+                if msg["type"] == "frame_batch":
+                    break
+
+            elapsed_ms = (time.perf_counter() - t0) * 1000
+            results.append(elapsed_ms)
+            print(f"  Client {client_id} round {round_idx}: {elapsed_ms:.0f}ms")
+
+
+async def random_client(client_id: int, start_event: asyncio.Event, results: list,
+                        num_rounds: int):
+    """Random keypresses with think time between presses."""
+    async with websockets.connect(
+        SERVER_URL, max_size=50_000_000, ping_interval=None,
+    ) as ws:
+        await connect_and_wait(client_id, ws)
+        start_event.set()
+
+        # Small stagger so clients don't all start at once
+        await asyncio.sleep(random.uniform(0, 1.0))
+
+        for _ in range(num_rounds):
+            think_ms = random.uniform(MIN_THINK_MS, MAX_THINK_MS)
+            await asyncio.sleep(think_ms / 1000)
+
+            t0 = time.perf_counter()
+            await ws.send(json.dumps({"type": "key", "key": random.choice(KEYS)}))
+
+            while True:
+                msg = json.loads(await ws.recv())
+                if msg["type"] == "frame_batch":
+                    break
+
+            elapsed_ms = (time.perf_counter() - t0) * 1000
+            results.append(elapsed_ms)
+
+        print(f"  Client {client_id}: done")
+
+
+def print_stats(times: list[float], label: str):
+    """Print summary statistics and histogram."""
+    times.sort()
+    n = len(times)
+    print(f"\n--- {label} ({n} keypresses) ---")
+    print(f"  Min:    {times[0]:.0f}ms")
+    print(f"  p25:    {times[n//4]:.0f}ms")
+    print(f"  Median: {times[n//2]:.0f}ms")
+    print(f"  p75:    {times[3*n//4]:.0f}ms")
+    print(f"  p95:    {times[int(n*0.95)]:.0f}ms")
+    print(f"  Max:    {times[-1]:.0f}ms")
+    print(f"  Avg:    {sum(times)/n:.0f}ms")
+
+    buckets = [0] * 10
+    labels = ["<1s", "<1.5s", "<2s", "<2.5s", "<3s",
+              "<3.5s", "<4s", "<4.5s", "<5s", "5s+"]
+    for ms in times:
+        idx = min(int(ms / 500), 9)
+        buckets[idx] += 1
+    print(f"\n  Distribution:")
+    for lbl, count in zip(labels, buckets):
+        if count > 0:
+            bar = "#" * count
+            pct = count / n * 100
+            print(f"    {lbl:>6s}: {bar} ({count}, {pct:.0f}%)")
+
+
+async def run_synced(num_clients: int, num_rounds: int):
+    print(f"=== SYNCHRONIZED MODE ({num_clients} clients, {num_rounds} rounds) ===")
+    barrier = asyncio.Barrier(num_clients)
+    results = []
+    tasks = [synced_client(i, barrier, results, num_rounds) for i in range(num_clients)]
+    await asyncio.gather(*tasks)
+    print_stats(results, "Synchronized (worst case)")
+
+
+async def run_random(num_clients: int, num_rounds: int):
+    print(f"=== RANDOM MODE ({num_clients} clients, {num_rounds} presses each) ===")
+    print(f"    Think time: {MIN_THINK_MS}-{MAX_THINK_MS}ms")
+    results = []
+    events = [asyncio.Event() for _ in range(num_clients)]
+    tasks = [random_client(i, events[i], results, num_rounds) for i in range(num_clients)]
+    await asyncio.gather(*tasks)
+    print_stats(results, "Random (realistic)")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Test ORCA multi-user batching")
+    parser.add_argument("mode", nargs="?", default="both",
+                        choices=["synced", "random", "both"])
+    parser.add_argument("-c", "--clients", type=int, default=4,
+                        help="Number of concurrent clients (default: 4)")
+    parser.add_argument("-r", "--rounds", type=int, default=10,
+                        help="Number of rounds per client (default: 10)")
+    return parser.parse_args()
+
+
+async def main():
+    args = parse_args()
+    print(f"Connecting {args.clients} clients to {SERVER_URL}...\n")
+
+    if args.mode in ("synced", "both"):
+        await run_synced(args.clients, args.rounds)
+    if args.mode in ("random", "both"):
+        if args.mode == "both":
+            print("\n" + "=" * 60 + "\n")
+        await run_random(args.clients, args.rounds)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())