refactor sampling pipeline

Debug transformer output misalignment
Fix OCR Rewards
2026-01-20 15:40:53 -08:00 · 2026-01-20 14:38:47 -08:00 · 2026-01-20 14:38:22 -08:00 · 2026-01-20 14:38:22 -08:00 · 2026-01-20 14:38:22 -08:00 · 2026-01-20 14:38:21 -08:00
396 changed files with 586904 additions and 42140 deletions
@@ -45,12 +45,6 @@ jobs:
      - name: Setup Pages
        uses: actions/configure-pages@v4

-      - name: Generate docs examples
-        run: python docs/generate_examples.py
-
-      - name: Check docs links
-        run: python scripts/check_docs_links.py
-
      - name: Build documentation
        run: mkdocs build

@@ -69,4 +63,4 @@ jobs:
    steps:
      - name: Deploy to GitHub Pages
        id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v4
@@ -64,10 +64,7 @@ jobs:
            # - torch-version: '2.7.1'
            #   cuda-version: '12.8.0'
            #   torch-cuda-short: 'cu128'
-            # - torch-version: '2.9.1'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'
-            - torch-version: '2.10.0'
+            - torch-version: '2.9.1'
              cuda-version: '12.8.0'
              torch-cuda-short: 'cu128'

@@ -18,7 +18,6 @@ venv/
 .venv/
 runs/
 samples/
-Miniconda3-latest-Linux-x86_64.sh
 *validation/
 data/
 outputs/
@@ -34,11 +33,6 @@ env
 *.log
 weights/

-# SSIM test outputs
-fastvideo/tests/ssim/generated_videos/
-**/.cache/**
-
-
 # Distribution / packaging
 build/
 dist/
@@ -75,11 +69,6 @@ docs/distillation/examples/
 !docs/assets/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif
-!assets/images/**/*.png
-!assets/images/**/*.jpg
-!assets/images/**/*.jpeg
-!assets/images/**/*.gif
-!assets/videos/**/*.mp4

 dmd_t2v_output/
 preprocess_output_text/
@@ -10,7 +10,7 @@ exclude: |
    demo/.*|
    predict\.py|
    scripts/.*|
-    assets/prompts/.*|
+    prompts/.*|
    fastvideo/data_preprocess/.*|
    fastvideo/dataset/.*|
    fastvideo/models/.*|
@@ -1,42 +0,0 @@
-# Repository Guidelines
-
-## Project Structure & Module Organization
- Core Python package: `fastvideo/` (models, pipelines, training, distributed runtime, CLI entrypoints).
- CUDA/custom kernels: `fastvideo-kernel/` (separate build/test flow).
- Tests:
-  - `fastvideo/tests/` for package-level tests (dataset, encoders, inference, training, SSIM, workflow).
-  - `tests/local_tests/` for additional local/component checks.
- Docs and guides: `docs/` (MkDocs source), with contributor docs in `docs/contributing/`.
- Runnable examples and scripts: `examples/` and `scripts/`.
- Static assets: `assets/` (including `assets/images/`, `assets/videos/`, and `assets/prompts/`) and `comfyui/assets/`.
-
-## Build, Test, and Development Commands
- `uv pip install -e .[dev]`: editable install with lint/test extras.
- `pre-commit install --hook-type pre-commit --hook-type commit-msg`: enable local hooks.
- `pre-commit run --all-files`: run formatter/lint/type/spelling checks.
- `pytest tests/`: run top-level test suite.
- `pytest fastvideo/tests/ -v`: run package tests.
- `pytest fastvideo/tests/ssim/ -vs`: run SSIM regression tests (GPU-heavy).
- `cd fastvideo-kernel && ./build.sh`: build kernel extensions.
-
-## Coding Style & Naming Conventions
- Python 3.10+; 4-space indentation; keep code and imports readable and explicit.
- Style tools are configured in `pyproject.toml` and `.pre-commit-config.yaml`:
-  - `yapf` (format), `ruff` (lint, auto-fix), `mypy` (typing), `codespell`.
- Target line length is 80.
- Naming: `snake_case` for functions/files, `PascalCase` for classes, `UPPER_SNAKE_CASE` for constants.
-
-## Testing Guidelines
- Use `pytest` and place tests near relevant domains (e.g., `fastvideo/tests/encoders/`).
- Prefer descriptive names like `test_<feature>_<expected_behavior>.py`.
- For new pipelines/backends, include at least one regression-oriented test; add SSIM coverage when output quality must be preserved.
- Document GPU assumptions in tests that require specific hardware.
-
-## Commit & Pull Request Guidelines
- Follow existing commit style: short subject with optional tag prefix, e.g. `[bugfix]: ...`, `[feat]: ...`, `[misc]: ...`, and include PR reference like `(#1234)` when applicable.
- Keep commits focused by concern (feature, refactor, fix).
- PRs should include:
-  - clear problem/solution summary,
-  - test evidence (`pytest`/SSIM outputs or rationale if skipped),
-  - linked issue/PR context,
-  - screenshots or sample outputs for UI/demo/docs changes.
@@ -1 +0,0 @@
-@AGENTS.md
@@ -1,101 +1,169 @@
-# Sliding Tile Attention (STA) Branch
+<div align="center">
+<img src=assets/logos/logo.svg width="30%"/>
+</div>

-This branch is a stash/testing branch for people who want to try
-Sliding Tile Attention (STA). The top-level README is intentionally
-STA-only.
+<p align="center">
+     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/sv3MMKyv" target="_blank"> <b> WeChat </b> </a> |
+</p>

-## What is STA
+**FastVideo is a unified post-training and inference framework for accelerated video generation.**

-Sliding Tile Attention is an optimized attention backend for
-window-based video generation.
+## NEWS
+- ```2025/11/19```: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
+- ```2025/08/04```: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

- Blog: https://hao-ai-lab.github.io/blogs/sta/
- Paper: https://arxiv.org/abs/2502.04507
- In-repo STA docs: `docs/attention/sta/index.md`
+<details>
+<summary>More</summary>

-## Setup
+- ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
+- ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
+- ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).

-Install FastVideo from source:
+</details>
+
+## Key Features
+
+FastVideo has the following features:
+- End-to-end post-training support for bidirectional and autoregressive models:
+  - Support full finetuning and LoRA finetuning for state-of-the-art open video DiTs
+  - Data preprocessing pipeline for video, image, and text data
+  - Distribution Matching Distillation (DMD2) stepwise distillation.
+  - Sparse attention with [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
+  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) to achineve >50x denoising speedup
+  - Scalable training with FSDP2, sequence parallelism, and selective activation checkpointing.
+  - Causal distillation through Self-Forcing
+  - See this [page](https://hao-ai-lab.github.io/FastVideo/training/overview/) for full list of supported models and recipes.
+- State-of-the-art performance optimizations for inference
+  - Sequence Parallelism for distributed inference
+  - Multiple state-of-the-art attention backends
+  - User-friendly CLI and Python API
+  - See this [page](https://hao-ai-lab.github.io/FastVideo/inference/optimizations/) for full list of supported optimizations.
+- Diverse hardware and OS support
+  - Support H100, A100, 4090
+  - Support Linux, Windows, MacOS
+  - See this [page](https://hao-ai-lab.github.io/FastVideo/inference/hardware_support/) for full list of supported hardware and OS.
+
+## Getting Started
+We recommend using an environment manager such as `Conda` to create a clean environment:

 ```bash
-uv pip install -e .
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
 ```

-Build the STA kernel package:
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.
+
+## Sparse Distillation
+For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
+
+See below for recipes and datasets:
+
+|                                            Model                                              |                                               Sparse Distillation                                                 |                                                  Dataset                                                  |
+|:-------------------------------------------------------------------------------------------:  |:---------------------------------------------------------------------------------------------------------------:  |:--------------------------------------------------------------------------------------------------------: |
+| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)         |    [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P)      | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)      |
+| [FastWan2.1-T2V-14B-Preview](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-Diffusers)   |                                                   Coming soon!                                                    |   [FastVideo Synthetic Wan2.1 720P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x768x1280_250k)   |
+| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers)           | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free)   | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k)  |
+
+## Inference
+### Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation/). Create a file called `example.py` with the following code:
+
+```python
+import os
+from fastvideo import VideoGenerator
+
+def main():
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:

 ```bash
-cd fastvideo-kernel
-./build.sh
-cd ..
+python example.py
 ```

-## Run STA Inference
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/).

-STA backend:
+### Other docs:

-```bash
-export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
-```
+- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
+- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/)

-Ready-to-run examples:
+## Distillation and Finetuning
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
+<!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->

- HunyuanVideo: `scripts/inference/v1_inference_hunyuan_STA.sh`
- Wan2.1-T2V-14B: `scripts/inference/v1_inference_wan_STA.sh`
+## Awesome work using FastVideo or our research projects

-Run:
+- [SGLang](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen): SGLang's diffusion inference functionality is based  on a fork of FastVideo on Sept. 24, 2025. [![Star](https://img.shields.io/github/stars/sgl-project/sglang.svg?style=social&label=Star)](https://github.com/sgl-project/sglang)

-```bash
-bash scripts/inference/v1_inference_hunyuan_STA.sh
-# or
-bash scripts/inference/v1_inference_wan_STA.sh
-```
+- [DanceGRPO](https://github.com/XueZeyue/DanceGRPO): A  unified framework to adapt Group Relative Policy Optimization (GRPO) to visual generation paradigms. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/XueZeyue/DanceGRPO.svg?style=social&label=Star)](https://github.com/XueZeyue/DanceGRPO)
+- [SRPO](https://github.com/Tencent-Hunyuan/SRPO): A method to directly align the full diffusion trajectory with fine-grained human preference. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/SRPO.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/SRPO)
+- [DCM](https://github.com/Vchitect/DCM): Dual-expert consistency model for efficient and high-quality video generation. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Vchitect/DCM.svg?style=social&label=Star)](https://github.com/Vchitect/DCM)
+- [Hunyuan Video 1.5](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5): A leading lightweight video generation model, where they proposed SSTA based on Sliding Tile Attention. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/HunyuanVideo-1.5.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5)
+- [Kandinsky-5.0](https://github.com/kandinskylab/kandinsky-5): A family of diffusion models for video & image generation, where their NABLA attention includes a Sliding Tile Attention branch. [![Star](https://img.shields.io/github/stars/kandinskylab/kandinsky-5.svg?style=social&label=Star)](https://github.com/kandinskylab/kandinsky-5)
+- [LongCat Video](https://github.com/meituan-longcat/LongCat-Video): A foundational video generation model with 13.6B parameters with block-sparse attention similar to Video Sparse Attention. [![Star](https://img.shields.io/github/stars/meituan-longcat/LongCat-Video.svg?style=social&label=Star)](https://github.com/meituan-longcat/LongCat-Video)

-Both scripts already set STA-related env vars:
+## 🤝 Contributing

- `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
- `FASTVIDEO_ATTENTION_CONFIG` to an STA mask strategy JSON
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview/).
+See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/899).
+## Acknowledgement
+We learned and reused code from the following projects:
+- [Wan-Video](https://github.com/Wan-Video)
+- [ThunderKittens](https://github.com/HazyResearch/ThunderKittens)
+- [Triton](https://github.com/triton-lang/triton)
+- [DMD2](https://github.com/tianweiy/DMD2)
+- [diffusers](https://github.com/huggingface/diffusers)
+- [xDiT](https://github.com/xdit-project/xDiT)
+- [vLLM](https://github.com/vllm-project/vllm)
+- [SGLang](https://github.com/sgl-project/sglang)

-## STA Mask Strategy Files
-
- HunyuanVideo config: `assets/mask_strategy_hunyuan.json`
- Wan config: `assets/mask_strategy_wan.json`
-
-## STA Mask Search (Wan2.1-T2V-14B)
-
-Run mask search + tuning from repo root:
-
-```bash
-bash examples/inference/sta_mask_search/inference_wan_sta.sh
-```
-
-What this script does:
-
- Runs `STA_searching` first.
- Runs `STA_tuning` next (`skip_time_steps=12` by default).
- Uses prompt shards from `assets/prompt_0.txt` to `assets/prompt_7.txt`.
-
-Important notes:
-
- Default script is set to 8 GPUs (`num_gpu=8`). If needed, edit
-  `examples/inference/sta_mask_search/inference_wan_sta.sh`.
- STA searching/tuning currently supports `69x768x1280` (Wan setting).
-
-Generated files:
-
- Search results: `output/mask_search_result_pos_1280x768/`
- Tuned strategy: `output/mask_search_strategy_1280x768/mask_strategy_s12.json`
-
-Use the tuned mask for STA inference:
-
-```bash
-export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
-export FASTVIDEO_ATTENTION_CONFIG=output/mask_search_strategy_1280x768/mask_strategy_s12.json
-python examples/inference/sta_mask_search/wan_example.py --STA_mode STA_inference --num_gpus 1
-```
+We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.

 ## Citation
+If you find FastVideo useful, please considering citing our work:

 ```bibtex
+@software{fastvideo2024,
+  title        = {FastVideo: A Unified Framework for Accelerated Video Generation},
+  author       = {The FastVideo Team},
+  url          = {https://github.com/hao-ai-lab/FastVideo},
+  month        = apr,
+  year         = {2024},
+}
+
+@article{zhang2025vsa,
+  title={Vsa: Faster video diffusion with trainable sparse attention},
+  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
+  journal={arXiv preprint arXiv:2505.13389},
+  year={2025}
+}
+
@article{zhang2025fast,
  title={Fast video generation with sliding tile attention},
  author={Zhang, Peiyuan and Chen, Yongqi and Su, Runlong and Ding, Hangliang and Stoica, Ion and Liu, Zhengzhong and Zhang, Hao},
@@ -1 +0,0 @@
-Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.
@@ -1 +0,0 @@
-A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature.
@@ -1 +0,0 @@
-A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.
@@ -1 +0,0 @@
-A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.
@@ -1 +0,0 @@
-A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1
@@ -1 +0,0 @@
-fox in the forest close-up quickly turned its head to the left
@@ -1 +0,0 @@
-Man walking his dog in the woods on a hot sunny day
@@ -1 +0,0 @@
-A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.
@@ -1,7 +0,0 @@
-# FastVideo/assets/videos
-
-This folder is used to store **video assets for examples**, primarily **input videos** consumed by scripts under `FastVideo/examples/`.
-
- **Typical contents**: short input clips for demos (e.g., video2world / image2video examples).
- **Non-critical**: these assets are for convenience and are not required to use the FastVideo library.
- **Large files**: avoid committing large videos to git; prefer shared storage or download-on-demand.
@@ -0,0 +1,195 @@
+import argparse
+import os
+import tempfile
+
+import gradio as gr
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import export_to_video
+
+from fastvideo.distill.solver import PCMFMScheduler
+from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
+from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
+
+
+def init_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompts", nargs="+", default=[])
+    parser.add_argument("--num_frames", type=int, default=25)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=848)
+    parser.add_argument("--num_inference_steps", type=int, default=8)
+    parser.add_argument("--guidance_scale", type=float, default=4.5)
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--seed", type=int, default=12345)
+    parser.add_argument("--transformer_path", type=str, default=None)
+    parser.add_argument("--scheduler_type", type=str, default="pcm_linear_quadratic")
+    parser.add_argument("--lora_checkpoint_dir", type=str, default=None)
+    parser.add_argument("--shift", type=float, default=8.0)
+    parser.add_argument("--num_euler_timesteps", type=int, default=50)
+    parser.add_argument("--linear_threshold", type=float, default=0.1)
+    parser.add_argument("--linear_range", type=float, default=0.75)
+    parser.add_argument("--cpu_offload", action="store_true")
+    return parser.parse_args()
+
+
+def load_model(args):
+    if args.scheduler_type == "euler":
+        scheduler = FlowMatchEulerDiscreteScheduler()
+    else:
+        linear_quadratic = True if "linear_quadratic" in args.scheduler_type else False
+        scheduler = PCMFMScheduler(
+            1000,
+            args.shift,
+            args.num_euler_timesteps,
+            linear_quadratic,
+            args.linear_threshold,
+            args.linear_range,
+        )
+
+    if args.transformer_path:
+        transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path)
+    else:
+        transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder="transformer/")
+
+    pipe = MochiPipeline.from_pretrained(args.model_path, transformer=transformer, scheduler=scheduler)
+    pipe.enable_vae_tiling()
+    # pipe.to(device)
+    # if args.cpu_offload:
+    pipe.enable_sequential_cpu_offload()
+    return pipe
+
+
+def generate_video(
+    prompt,
+    negative_prompt,
+    use_negative_prompt,
+    seed,
+    guidance_scale,
+    num_frames,
+    height,
+    width,
+    num_inference_steps,
+    randomize_seed=False,
+):
+    if randomize_seed:
+        seed = torch.randint(0, 1000000, (1, )).item()
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    if not use_negative_prompt:
+        negative_prompt = None
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        output = pipe(
+            prompt=[prompt],
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        ).frames[0]
+
+    output_path = os.path.join(tempfile.mkdtemp(), "output.mp4")
+    export_to_video(output, output_path, fps=30)
+    return output_path, seed
+
+
+examples = [
+    "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand’s movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.",
+    "A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.",
+    "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.",
+]
+
+args = init_args()
+pipe = load_model(args)
+print("load model successfully")
+with gr.Blocks() as demo:
+    gr.Markdown("# Fastvideo Mochi Video Generation Demo")
+
+    with gr.Group():
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            run_button = gr.Button("Run", scale=0)
+        result = gr.Video(label="Result", show_label=False)
+
+    with gr.Accordion("Advanced options", open=False):
+        with gr.Group():
+            with gr.Row():
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=1024,
+                    step=32,
+                    value=args.height,
+                )
+                width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=args.width)
+
+            with gr.Row():
+                num_frames = gr.Slider(
+                    label="Number of Frames",
+                    minimum=21,
+                    maximum=163,
+                    value=args.num_frames,
+                )
+                guidance_scale = gr.Slider(
+                    label="Guidance Scale",
+                    minimum=1,
+                    maximum=12,
+                    value=args.guidance_scale,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Inference Steps",
+                    minimum=4,
+                    maximum=100,
+                    value=args.num_inference_steps,
+                )
+
+            with gr.Row():
+                use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=False,
+            )
+
+            seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=args.seed)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            seed_output = gr.Number(label="Used Seed")
+
+    gr.Examples(examples=examples, inputs=prompt)
+
+    use_negative_prompt.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt,
+        outputs=negative_prompt,
+    )
+
+    run_button.click(
+        fn=generate_video,
+        inputs=[
+            prompt,
+            negative_prompt,
+            use_negative_prompt,
+            seed,
+            guidance_scale,
+            num_frames,
+            height,
+            width,
+            num_inference_steps,
+            randomize_seed,
+        ],
+        outputs=[result, seed_output],
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)
@@ -0,0 +1,15 @@
+Fast-Hunyuan comparison with original Hunyuan, achieving an 8X diffusion speed boost with the FastVideo framework.
+
+https://github.com/user-attachments/assets/064ac1d2-11ed-4a0c-955b-4d412a96ef30
+
+Fast-Mochi comparison with original Mochi, achieving an 8X diffusion speed boost with the FastVideo framework.
+
+https://github.com/user-attachments/assets/5fbc4596-56d6-43aa-98e0-da472cf8e26c
+
+Comparison between OpenAI Sora, original Hunyuan and FastHunyuan
+
+https://github.com/user-attachments/assets/d323b712-3f68-42b2-952b-94f6a49c4836
+
+Comparison between original FastHunyuan, LLM-INT8 quantized FastHunyuan and NF4 quantized FastHunyuan
+
+https://github.com/user-attachments/assets/cf89efb5-5f68-4949-a085-f41c1ef26c94
@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp310-cp310-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp310-cp310-linux_x86_64.whl

 COPY . .

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp311-cp311-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp311-cp311-linux_x86_64.whl

 COPY . .

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp312-cp312-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp312-cp312-linux_x86_64.whl

 COPY . .

@@ -121,15 +121,6 @@ This page contains the complete API reference for the FastVideo library.
      show_root_toc_entry: true
      heading_level: 4

-## fastvideo.registry
-
-::: fastvideo.registry
-    options:
-      show_source: true
-      show_root_heading: true
-      show_root_toc_entry: true
-      heading_level: 3
-
 ## fastvideo.pipelines

 ::: fastvideo.pipelines
@@ -6,8 +6,6 @@ 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

@@ -1,176 +0,0 @@
-# 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.
@@ -1,480 +0,0 @@
-# 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/registry.py` using explicit
-  `register_configs(...)` blocks (this file is the single source of truth now).
-
-### 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`
@@ -1,55 +1,16 @@

 # 📦 Developing FastVideo on RunPod

-You can easily use the FastVideo Pod Template on [RunPod](https://www.runpod.io) for development or experimentation.
+You can easily use the FastVideo Docker image as a custom container on [RunPod](https://www.runpod.io) for development or experimentation.

 ## Creating a new pod

- Make sure you are using the correct RunPod account.
-![RunPod Account Selection](../../assets/images/runpod_account.png)
+Choose a GPU that supports CUDA 12.8
+
+Pick 1 or 2 L40S GPU(s)

- 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:

 ```
@@ -63,3 +24,30 @@ 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/
+```
@@ -1,84 +1,76 @@

 # 🛠️ Contributing to FastVideo

-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.
+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!

-## Quick prerequisites
+Our community is open to everyone and welcomes any contributions no matter how large or small.

- **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).
+# 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.

-For a full install checklist, see `docs/getting_started/installation/gpu.md`.
-
-## Local development (Conda + editable install)
+We recommend using a fresh Python 3.10 Conda environment to develop FastVideo:

 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:
+Create and activate a Conda environment for FastVideo:

-```bash
+```
 conda create -n fastvideo python=3.12 -y
 conda activate fastvideo
 ```

 Install `uv` (optional, but recommended):

-```bash
+From instructions on [uv](https://astral.sh/uv/):
+
+```
 curl -LsSf https://astral.sh/uv/install.sh | sh
-# or
+# or 
 wget -qO- https://astral.sh/uv/install.sh | sh
 ```

-Clone the repo:
+Clone the FastVideo repository and go to the FastVideo directory:

-```bash
+```
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+
 ```

-Install FastVideo in editable mode and set up hooks:
+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.

 ```bash
 uv pip install -e .[dev]

-# Optional: FlashAttention (builds native kernels)
-uv pip install flash-attn --no-build-isolation
+# Can also install flash-attn (optional)
+uv pip install flash-attn --no-build-isolation 

-# Linting, formatting, static typing
+# Linting, formatting and static type checking
 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, installing FlashAttention 3 can improve
-performance (see `docs/inference/optimizations.md`).
+If you are on a Hopper GPU, you should also install [FA3](https://github.com/Dao-AILab/flash-attention) for much better performance:

-## Docker development (optional)
+```
+git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention/hopper

-If you prefer a containerized environment, use the dev image documented in
-`docs/contributing/developer_env/docker.md`.
+# make sure you have ninja installed
+uv pip install ninja
+
+python setup.py install
+```

 ## Testing

-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`.
+Please refer to the [Testing Guide](testing.md) for more information on how to add and run tests in FastVideo.
@@ -1,179 +1,424 @@
-# FastVideo Architecture Overview
+# 🔍 FastVideo Overview

-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.
+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.

-## FastVideo structure at a glance
+## Table of Contents - Directory Structure and Files

-FastVideo maps a Diffusers-style repo into a pipeline like this:
+- [`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/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.
+## Core Architecture

-Flow:
-`model_index.json` -> component loaders -> model modules -> pipeline stages ->
-sampling params.
+FastVideo separates model components from execution logic with these principles:

-Minimal usage (from `examples/inference/basic/basic.py`):
+- **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:

 ```python
-from fastvideo import VideoGenerator
-from fastvideo.configs.sample import SamplingParam
+# Load arguments from command line
+fastvideo_args = prepare_fastvideo_args(sys.argv[1:])

-model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # or official_weights/<model_name>/
-generator = VideoGenerator.from_pretrained(model_id, num_gpus=1)
+# Access parameters
+model = load_model(fastvideo_args.model_path)

-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,
+# 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
 )
 ```

-## Configuration system
+### Sequence Parallelism

-FastVideo uses typed configs to keep model definitions, pipeline wiring, and
-runtime parameters consistent:
+Sequence parallelism splits sequences across devices:

- `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/registry.py`: unified registry for pipeline config + sampling
-  defaults and model metadata resolution, defined via explicit
-  `register_configs(...)` blocks (no separate dict registries).
+- **Implementation**: Through `DistributedAttention` and sequence splitting
+- **Use cases**: Long video sequences or high-resolution processing. Used by DiT models.

-`FastVideoArgs` (in `fastvideo/fastvideo_args.py`) provides runtime settings and
-is passed into pipeline construction and stages.
+```python
+# Distributed attention for long sequences
+from fastvideo.attention import DistributedAttention

-## 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
+self.attn = DistributedAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    causal=False,
+    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN)
+)
 ```

-Key points:
+### Communication Primitives

- `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 distributed operations via AllGather, AllReduce, and synchronization mechanisms.

-Note on tensor names:
+Efficient communication primitives minimize distributed overhead:

-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.
+- **Sequence-Parallel AllGather**: Collects sequence chunks
+- **Tensor-Parallel AllReduce**: Combines partial results
+- **Distributed Synchronization**: Coordinates execution

-Example HF repo (Wan 2.1 T2V 1.3B Diffusers):
+## Forward Context Management

-```
-https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/tree/main
+### 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)
 ```

-Example `model_index.json` from that repo:
+## Executor and Worker System

-```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 `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
 ```

-How this maps to FastVideo:
+The platform system is designed to be extensible for future hardware targets.

- `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`
+## Logger

-## Pipeline system
+See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)

- `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.
+*TODO*: (help wanted) Add an environment variable that disables process-aware logging.

-## Model components
+## Contributing to FastVideo

- 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/`
+If you're a new contributor, here are some common areas to explore:

-## Attention and distributed execution
+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 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/`.
+When adding code, follow these practices:

-## Related docs
-
- [Contributing overview](../contributing/overview.md)
- [Coding agents workflow](../contributing/coding_agents.md)
- [Testing guide](../contributing/testing.md)
+- 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
@@ -6,7 +6,7 @@ All documented examples are autogenerated using [generate_examples.py](https://g

 ## Examples

- [Examples Distillation Index](../distillation/examples/examples_distillation_index.md)
- [Examples Training Index](../training/examples/examples_training_index.md)
- [Examples Inference Index](../inference/examples/examples_inference_index.md)
+- [Examples Distillation Index](distillation/examples/examples_distillation_index.md)
+- [Examples Training Index](training/examples/examples_training_index.md)
+- [Examples Inference Index](inference/examples/examples_inference_index.md)

@@ -2,7 +2,6 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py

 import itertools
-import os
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -127,20 +126,8 @@ class Example:
        Raises:
            IndexError: If no Markdown files are found in the directory.
        """ # noqa: E501
-        if self.path.is_file():
-            return self.path
-
-        markdown_files = sorted(self.path.glob("*.md"))
-        if not markdown_files:
-            raise IndexError(f"No Markdown files found in {self.path}")
-
-        readme_files = [
-            f for f in markdown_files if f.name.lower() == "readme.md"
-        ]
-        if readme_files:
-            return readme_files[0]
-
-        return markdown_files[0]
+        return self.path if self.path.is_file() else list(
+            self.path.glob("*.md")).pop()

    def determine_other_files(self) -> list[Path]:
        """
@@ -534,11 +521,11 @@ def generate_examples(generate_main_index: bool = False) -> None:
            # Add to main index if it exists
            if generate_main_index and examples_index:
                main_index_dir = examples_index.path.parent
-                rel_path = os.path.relpath(category_index.path,
-                                           start=main_index_dir)
+                rel_path = category_index.path.relative_to(
+                    main_index_dir.parent)
                examples_index.documents.insert(
                    0,
-                    str(rel_path).replace("\\", "/").replace(".md", ""))
+                    str(rel_path).replace(".md", ""))

            # Write the category index file
            with open(category_index.path, "w+") as f:
@@ -18,25 +18,12 @@ conda activate fastvideo
 pip install fastvideo
 ```

-### Using uv
-
-```bash
-# Create and activate a new uv environment
-uv venv --python 3.12 --seed
-source .venv/bin/activate
-
-uv pip install fastvideo
-```
-
 ### From source

 ```bash
 git clone https://github.com/hao-ai-lab/FastVideo.git
 cd FastVideo
 pip install -e .
-
-# or if you are using uv
-uv pip install -e .
 ```

 Also optionally install flash-attn:
@@ -79,6 +79,5 @@ if __name__ == '__main__':

 - [Installation Guide](installation.md) - Detailed installation instructions
 - [Configuration](../inference/configuration.md) - Learn about configuration options
- [Examples](../inference/examples/examples_inference_index.md) - Explore more
-  examples
+- [Examples](../inference/examples/) - Explore more examples
 - [Optimizations](../inference/optimizations.md) - Performance optimization tips
@@ -101,17 +101,14 @@ Replace standard attention with FastVideo's optimized attention:
 ```python
 # Local attention patterns
 from fastvideo.attention import LocalAttention
-from fastvideo.platforms.interface import AttentionBackendEnum
+from fastvideo.attention.backends.abstract import _Backend
 self.attn = LocalAttention(
    num_heads=num_heads,
    head_size=head_dim,
    dropout_rate=0.0,
    softmax_scale=None,
    causal=False,
-    supported_attention_backends=(
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-    )
+    supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
 )

 # Distributed attention for long sequences
@@ -122,60 +119,29 @@ self.attn = DistributedAttention(
    dropout_rate=0.0,
    softmax_scale=None,
    causal=False,
-    supported_attention_backends=(
-        AttentionBackendEnum.SLIDING_TILE_ATTN,
-        AttentionBackendEnum.FLASH_ATTN,
-        AttentionBackendEnum.TORCH_SDPA,
-    )
+    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
 )
 ```

 #### Define supported backend selection

 ```python
-_supported_attention_backends = (
-    AttentionBackendEnum.FLASH_ATTN,
-    AttentionBackendEnum.TORCH_SDPA,
-)
+   _supported_attention_backends = (_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
 ```

 ### Registering Models

-Register implemented modules for auto‑discovery by adding `EntryClass` in each
-model module (the registry scans for it):
+Register implemented modules in the model registry:

 ```python
-# In fastvideo/models/dits/your_module.py
-class YourTransformerModel(...):
-    ...
+# In fastvideo/models/registry.py
+_TEXT_TO_VIDEO_DIT_MODELS = {
+    "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"),
+}

-# Entry point for model registry
-EntryClass = YourTransformerModel
-```
-
-```python
-# In fastvideo/models/vaes/your_vae.py
-class YourVAEModel(...):
-    ...
-
-# Entry point for model registry
-EntryClass = YourVAEModel
-```
-
-Register pipeline config + sampling defaults in the unified registry:
-
-```python
-# In fastvideo/registry.py
-register_configs(
-    sampling_param_cls=YourSamplingParam,
-    pipeline_config_cls=YourPipelineConfig,
-    hf_model_paths=[
-        "org/your-model-id",
-    ],
-    model_detectors=[
-        lambda path: "your-model" in path.lower(),
-    ],
-)
+_VAE_MODELS = {
+    "YourVAEModel": ("vaes", "yourvae", "YourVAEClass"),
+}
 ```

 ## Step 2: Directory Structure
@@ -1,85 +1,104 @@
 # FastVideo CLI Inference

-The FastVideo CLI exposes the same core inference controls as the Python API.
+The FastVideo CLI provides a quick way to access the FastVideo inference pipeline for video generation. For more advanced usage,
+see the Python interface [here](examples/basic.md).

 ## Basic Usage

-Use either:
-
-1. `--model-path` + `--prompt`
-2. `--model-path` + `--prompt-txt` (batch prompts, one line per prompt)
-3. `--config` (JSON/YAML)
+The basic command to generate a video is:

 ```bash
-fastvideo generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
-  --prompt "A cat playing with a ball of yarn"
+fastvideo generate --model-path {MODEL_PATH} --prompt {PROMPT}
 ```

-```bash
-fastvideo generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
-  --prompt-txt prompts.txt
-```
+### Required Parameters

-You cannot provide both `--prompt` and `--prompt-txt` in the same run.
+- `--model-path {MODEL_PATH}`: Path to the model or model ID
+- `--prompt {PROMPT}`: Text description for the video you want to generate

-## View All Arguments
+## Common Arguments
+
+To see all the options, you can use the `--help` flag:

 ```bash
 fastvideo generate --help
 ```

-Arguments come from:
+### Hardware Configuration

- FastVideo runtime args (`FastVideoArgs`)
- Sampling args (`SamplingParam`)
- Pipeline config args (`PipelineConfig`)
+- `--num-gpus {NUM_GPUS}`: Number of GPUs to use
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder, should not be larger than 1 if text encoder offload is enabled, as layerwise offload + prefetch is faster)
+- `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)

-## Common Arguments
+#### Video Configuration

-### Parallelism
+- `--height {HEIGHT}`: Height of the generated video
+- `--width {WIDTH}`: Width of the generated video
+- `--num-frames {NUM_FRAMES}`: Number of frames to generate
+- `--fps {FPS}`: Frames per second for the saved video

- `--num-gpus`
- `--sp-size`
- `--tp-size`
+#### Generation Parameters

-### Sampling
+- `--num-inference-steps {STEPS}`: Number of denoising steps
+- `--negative-prompt {PROMPT}`: Negative prompt to guide generation away from certain concepts
+- `--seed {SEED}`: Random seed for reproducible generation

- `--num-frames`
- `--height` / `--width`
- `--num-inference-steps`
- `--guidance-scale`
- `--seed`
- `--negative-prompt`
+#### Output Options

-### Output
+- `--output-path {PATH}`: Directory to save the generated video
+- `--save-video`: Whether to save the video to disk
+- `--return-frames`: Whether to return the raw frames

- `--output-path`
- `--save-video` / `--no-save-video`
- `--return-frames`
+## Using Configuration Files

-### Offloading and Performance
-
- `--dit-layerwise-offload`
- `--use-fsdp-inference`
- `--text-encoder-cpu-offload`
- `--image-encoder-cpu-offload`
- `--vae-cpu-offload`
- `--enable-torch-compile`
- `--torch-compile-kwargs`
-
-## Using Config Files
+Instead of specifying all parameters on the command line, you can use a configuration file:

 ```bash
-fastvideo generate --config config.yaml
+fastvideo generate --config {CONFIG_FILE_PATH}
 ```

-Config files can be JSON or YAML. CLI flags override config-file values.
+The config file should be in JSON or YAML format with the same parameter names as the CLI options. Command-line arguments will take precedence over settings in the configuration file, allowing you to override specific values while keeping the rest from the config file.

-Example `config.yaml`:
+Example configuration file (config.json):
+
+```json
+{
+    "model_path": "FastVideo/FastHunyuan-diffusers",
+    "prompt": "A beautiful woman in a red dress walking down a street",
+    "output_path": "outputs/",
+    "num_gpus": 2,
+    "sp_size": 2,
+    "tp_size": 1,
+    "num_frames": 45,
+    "height": 720,
+    "width": 1280,
+    "num_inference_steps": 6,
+    "seed": 1024,
+    "fps": 24,
+    "precision": "bf16",
+    "vae_precision": "fp16",
+    "vae_tiling": true,
+    "vae_sp": true,
+    "vae_config": {
+        "load_encoder": false,
+        "load_decoder": true,
+        "tile_sample_min_height": 256,
+        "tile_sample_min_width": 256
+    },
+    "text_encoder_precisions": [
+        "fp16",
+        "fp16"
+    ],
+    "mask_strategy_file_path": null,
+    "enable_torch_compile": false
+}
+```
+
+Or using YAML format (config.yaml):

 ```yaml
 model_path: "FastVideo/FastHunyuan-diffusers"
-prompt: "A capybara lounging in a hammock"
+prompt: "A beautiful woman in a red dress walking down a street"
 output_path: "outputs/"
 num_gpus: 2
 sp_size: 2
@@ -89,34 +108,44 @@ height: 720
 width: 1280
 num_inference_steps: 6
 seed: 1024
-dit_precision: "bf16"
+fps: 24
+precision: "bf16"
 vae_precision: "fp16"
 vae_tiling: true
 vae_sp: true
+vae_config:
+  load_encoder: false
+  load_decoder: true
+  tile_sample_min_height: 256
+  tile_sample_min_width: 256
+text_encoder_precisions:
+  - "fp16"
+  - "fp16"
+mask_strategy_file_path: null
 enable_torch_compile: false
 ```

-Notes:
-
- Use `dit_precision` / `vae_precision` (not `precision`).
- Nested config objects are supported, for example `vae_config` and
-  `dit_config`.
-
 ## Examples

-Simple generation:
+Generating a simple video:

 ```bash
-fastvideo generate \
-  --model-path FastVideo/FastHunyuan-diffusers \
-  --prompt "A cat playing with a ball of yarn" \
-  --num-frames 45 --height 720 --width 1280 \
-  --num-inference-steps 6 --seed 1024 \
-  --output-path outputs/
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A cat playing with a ball of yarn" --num-frames 45 --height 720 --width 1280 --num-inference-steps 6 --seed 1024 --output-path outputs/
 ```

-Config + CLI override:
+Using a negative prompt to avoid certain elements:

 ```bash
-fastvideo generate --config config.yaml --prompt "A panda skiing at sunset"
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A beautiful forest landscape" --negative-prompt "people, buildings, roads"
 ```
+
+Combining command line arguments and a configuration file:
+
+```bash
+fastvideo generate --config config.json --prompt "A capybara lounging in a hammock"
+```
+
+## Troubleshooting
+
+- If you encounter CUDA out-of-memory errors, try reducing the video dimensions or number of frames, or the number of inference steps.
+- For reproducible results, set the same seed value between runs.
@@ -1,3 +1,4 @@
+
 # Configuration

 ## Multi-GPU Setup
@@ -17,8 +18,7 @@ generator = VideoGenerator.from_pretrained(
 - `PipelineConfig`: Initialization time parameters
 - `SamplingParam`: Generation time parameters

-You can customize generation behavior using `PipelineConfig` and
-`SamplingParam`:
+You can customize various parameters when generating videos using the `PipelineConfig` and `SamplingParam` class:

 ```python
 from fastvideo import VideoGenerator, SamplingParam, PipelineConfig
@@ -27,12 +27,12 @@ def main():
    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
    config = PipelineConfig.from_pretrained(model_name)
    config.vae_precision = "fp16"
+    config.dit_cpu_offload = True

    # Create the generator
    generator = VideoGenerator.from_pretrained(
        model_name,
        num_gpus=1,
-        dit_layerwise_offload=True,  # FastVideoArgs option
        pipeline_config=config
    )

@@ -72,34 +72,6 @@ if __name__ == '__main__':
    main()
 ```

-## JSON/YAML Config Files (CLI)
-
-The CLI supports `--config` with JSON or YAML. Command-line arguments override
-config file values.
-
-```bash
-fastvideo generate --config config.yaml
-```
-
-Use CLI argument names as keys (underscore or hyphen is accepted). Example:
-
-```yaml
-model_path: "FastVideo/FastHunyuan-diffusers"
-prompt: "A capybara relaxing in a hammock"
-num_gpus: 2
-sp_size: 2
-num_frames: 45
-height: 720
-width: 1280
-num_inference_steps: 6
-seed: 1024
-dit_precision: "bf16"
-vae_precision: "fp16"
-vae_tiling: true
-vae_sp: true
-enable_torch_compile: false
-```
-
 ## Performance Optimization

 For configuring optimizations, please see our [optimizations guide](optimizations.md)
@@ -61,8 +61,7 @@ python example.py

 The generated video will be saved in the current directory under `my_videos/`  

-More inference scripts and recipes can be found in `examples/inference/` and
-`scripts/inference/`.
+More inference example scripts can be found in `scripts/inference/`

 ## Available Models

@@ -104,10 +103,11 @@ Common issues and their solutions:
 If you encounter CUDA out of memory errors:

 - Reduce `num_frames` or video resolution
- Enable FastVideo offloading options such as `dit_layerwise_offload=True`
-  (single GPU) or `use_fsdp_inference=True` (multi-GPU)
+- Enable memory optimization with `enable_model_cpu_offload`
 - Try a smaller model or use distilled versions
 - Use `num_gpus` > 1 if multiple GPUs are available
+- Try enabling FSDP inference with `use_fsdp_inference=True` (may slow down generation)
+- Try enabling DiT layerwise offload with `dit_layerwise_offload=True` (now only a few models support this, but may introduce less overhead than FSDP)

 ### Slow Generation

@@ -1,140 +0,0 @@
-# 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)
-```
@@ -25,9 +25,6 @@ This page describes the various options for speeding up generation times in Fast
 - Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
 - Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE`
- Video MoBA Attention: `FASTVIDEO_ATTENTION_BACKEND=VMOBA_ATTN`
- Sparse Linear Attention: `FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN`
- SageSLA Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_SLA_ATTN`

 ### Configuring Backends

@@ -73,15 +70,22 @@ python setup.py install

 **`SLIDING_TILE_ATTN`**

-Sliding Tile Attention is provided by `fastvideo-kernel`.
-See [STA docs](../attention/sta/index.md) for installation details.
+```bash
+pip install st_attn==0.0.4
+```
+
+Please see [this page](../attention/sta/index.md) for more installation instructions.

 ### Video Sparse Attention

 **`VIDEO_SPARSE_ATTN`**

-Video Sparse Attention is provided by `fastvideo-kernel`.
-See [VSA docs](../attention/vsa/index.md) for installation details.
+```bash
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+Please see [this page](../attention/vsa/index.md) for more installation instructions.

 ### Sage Attention

@@ -99,7 +103,7 @@ python setup.py install  # or pip install -e .

 **`SAGE_ATTN_THREE`**

-[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.
+[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.

 #### Hardware Requirements

@@ -109,13 +113,11 @@ 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, follow the `README.md` in the linked repository to install the package from source.
+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:

-### V-MoBA / SLA / SageSLA
-
-These backends are model-specific and require the corresponding kernels and
-dependencies. Use the support matrix and model examples to confirm compatibility
-before enabling them.
+```bash
+python setup.py install
+```

 ## Teacache

@@ -1,10 +1,6 @@
 # Compatibility Matrix

-This page summarizes common model + optimization combinations.
-
-For the canonical, code-level list of model IDs recognized by
-`VideoGenerator.from_pretrained(...)`, see the registrations in
-`fastvideo/registry.py` (`register_configs(...)` entries).
+The table below shows every supported model and optimizations supported for them.

 The symbols used have the following meanings:

@@ -14,9 +10,7 @@ The symbols used have the following meanings:

 ## Models x Optimization

-The `HuggingFace Model ID` can be passed directly to
-`from_pretrained()`. FastVideo then uses model-specific default settings for
-pipeline initialization and sampling.
+The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.

 <style>
  /* Target tables in this section */
@@ -59,9 +53,9 @@ pipeline initialization and sampling.
 | Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
 | Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ | ⭕ |
 | Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720P | ✅ | ✅ | ✅ | ⭕ | ⭕ |
+| StepVideo T2V | `FastVideo/stepvideo-t2v-diffusers` | 768px768px204f<br>544px992px204f<br>544px992px136f | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | TurboWan2.1 T2V 1.3B | `loayrashid/TurboWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | TurboWan2.1 T2V 14B | `loayrashid/TurboWan2.1-T2V-14B-Diffusers` | 480P, 720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
-| TurboWan2.2 I2V A14B | `loayrashid/TurboWan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | LongCat T2V 13.6B | See note** | 480P<br>720P | ❌ | ❌ | ❌ | ⭕ | ✅ |
 | Matrix Game 2.0 Base | `FastVideo/Matrix-Game-2.0-Base-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Matrix Game 2.0 GTA | `FastVideo/Matrix-Game-2.0-GTA-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
@@ -69,14 +63,11 @@ pipeline initialization and sampling.

 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.

-## Canonical Supported IDs
-
-The authoritative source for model-ID recognition is
-`fastvideo/registry.py`. If a model ID is registered there, FastVideo can
-resolve default pipeline and sampling configuration for it.
-
 ## Special requirements

+### StepVideo T2V
+- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
+
 ### Sliding Tile Attention
 - Currently only Hopper GPUs (H100s) are supported.

@@ -1,75 +0,0 @@
-# Debugging
-
-This page collects practical debugging steps for FastVideo inference issues.
-
-## Collect Environment Info
-
-From the repository root, run:
-
-```bash
-python collect_env.py
-```
-
-Attach the output when filing a GitHub issue.
-
-## Increase Logging
-
-FastVideo logging level is controlled by environment variables:
-
-```bash
-FASTVIDEO_LOGGING_LEVEL=DEBUG \
-FASTVIDEO_STAGE_LOGGING=1 \
-python your_script.py
-```
-
-Useful variables:
-
- `FASTVIDEO_LOGGING_LEVEL`: `DEBUG`, `INFO`, `WARNING`, `ERROR`
- `FASTVIDEO_STAGE_LOGGING`: print per-stage timings during pipeline execution
- `FASTVIDEO_ATTENTION_BACKEND`: force an attention backend (for example
-  `TORCH_SDPA` or `FLASH_ATTN`)
-
-## Common Failure Modes
-
-### Out-of-memory
-
-Try, in order:
-
-1. Reduce `height`, `width`, `num_frames`, or `num_inference_steps`.
-2. Enable offloading flags such as `dit_layerwise_offload` (single GPU) or
-   `use_fsdp_inference` (multi-GPU).
-3. Enable `vae_cpu_offload`, `image_encoder_cpu_offload`, and
-   `text_encoder_cpu_offload`.
-
-See [Inference Offloading](../inference/offloading.md) for recommended
-combinations.
-
-### Attention backend import errors
-
-If forcing a backend fails, verify optional dependencies are installed:
-
- `FLASH_ATTN`: `flash-attn`
- `SLIDING_TILE_ATTN` and `VIDEO_SPARSE_ATTN`: `fastvideo-kernel`
- `SAGE_ATTN` / `SAGE_ATTN_THREE`: SageAttention packages
-
-As a fallback, use:
-
-```bash
-export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-```
-
-### Configuration parsing errors
-
-When using `--config`, keep keys aligned with CLI argument names (underscores or
-hyphens are both accepted). For nested config values, use nested objects
-(`vae_config`, `dit_config`) rather than dotted keys.
-
-## Issue Template
-
-When opening an issue, include:
-
- exact command or Python snippet,
- model ID/path,
- full traceback,
- `collect_env.py` output,
- whether the problem reproduces with `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`.
@@ -1,50 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo import VideoGenerator
-from fastvideo.configs.sample import SamplingParam
-
-
-def main():
-    # Point this to your local diffusers model dir (or replace with a HF model ID).
-    model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
-
-    generator = VideoGenerator.from_pretrained(
-        model_path,
-        num_gpus=1,
-        use_fsdp_inference=False,  # set True if GPU is out of memory
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True,
-    )
-
-    sampling_param = SamplingParam.from_pretrained(model_path)
-
-    # image2world example from official repo
-    image_path = "assets/images/bus_terminal.jpg"
-
-    prompt = (
-        "A nighttime city bus terminal gradually shifts from stillness to subtle movement. "
-        "At first, multiple double-decker buses are parked under the glow of overhead lights, "
-        "with a central bus labeled '87D' facing forward and stationary. "
-        "As the video progresses, the bus in the middle moves ahead slowly, its headlights brightening the surrounding area "
-        "and casting reflections onto adjacent vehicles. "
-        "The motion creates space in the lineup, signaling activity within the otherwise quiet station. "
-        "It then comes to a smooth stop, resuming its position in line. "
-        "Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene."
-    )
-
-    generator.generate_video(
-        prompt,
-        sampling_param=sampling_param,
-        image_path=str(image_path),
-        num_cond_frames=1,
-        output_path="outputs_video/cosmos2_5_i2w.mp4",
-        save_video=True,
-    )
-
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
-
@@ -1,5 +1,4 @@
 from fastvideo import VideoGenerator
-from fastvideo.configs.sample import SamplingParam


 def main():
@@ -16,27 +15,19 @@ def main():
        pin_cpu_memory=True,
    )

-    # Load default sampling parameters (negative_prompt, resolution, steps, etc.)
-    sampling_param = SamplingParam.from_pretrained(model_path)
-
    prompt = (
-        "A high-definition video captures the precision of robotic welding in an industrial setting. "
-        "The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. "
-        "The welding process is in full swing, with bright sparks and intense light illuminating the scene, "
-        "creating a vivid display of blue and white hues. "
-        "A significant amount of smoke billows around the welding area, partially obscuring the view but emphasizing the heat and activity. "
-        "The background reveals parts of the workshop environment, including a ventilation system and various pieces of machinery, "
-        "indicating a busy and functional industrial workspace. "
-        "As the video progresses, the robotic arm maintains its steady position, continuing the welding process and moving to its left. "
-        "The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. "
-        "The metal surface beneath the torch shows ongoing signs of heating and melting. "
-        "The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, "
-        "underscoring the ongoing nature of the welding operation."
+       "A high-definition video captures the precision of robotic welding in an industrial setting. The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. The welding process is in full swing, with bright sparks and intense light illuminating the scene, creating a vivid display of blue and white hues. A significant amount of smoke billows around the welding area, partially obscuring the view but emphasizing the heat and activity. The background reveals parts of the workshop environment, including a ventilation system and various pieces of machinery, indicating a busy and functional industrial workspace. As the video progresses, the robotic arm maintains its steady position, continuing the welding process and moving to its left. The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. The metal surface beneath the torch shows ongoing signs of heating and melting. The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, underscoring the ongoing nature of the welding operation."
    )

-    generator.generate_video(
+    video = generator.generate_video(
        prompt,
-        sampling_param=sampling_param,
+        negative_prompt="",
+        height=704,
+        width=1280,
+        num_frames=77,
+        num_inference_steps=35,
+        guidance_scale=7.0,
+        fps=24,
        output_path="outputs_video/cosmos2_5_t2w.mp4",
        save_video=True,
    )
@@ -1,53 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from fastvideo import VideoGenerator
-from fastvideo.configs.sample import SamplingParam
-
-
-def main():
-    # Point this to your local diffusers model dir (or replace with a HF model ID).
-    model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
-
-    generator = VideoGenerator.from_pretrained(
-        model_path,
-        num_gpus=1,
-        use_fsdp_inference=False,  # set True if GPU is out of memory
-        dit_cpu_offload=False,
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True,
-    )
-
-    sampling_param = SamplingParam.from_pretrained(model_path)
-
-    # video2world example from official repo
-    video_path = "assets/videos/robot_pouring.mp4"
-
-    prompt = (
-        "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. "
-        "The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. "
-        "The robotic arm is in the process of pouring a transparent liquid into the glass. "
-        "To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. "
-        "In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. "
-        "The lighting is bright, casting soft shadows on the table. "
-        "The robotic arm's movements are smooth and controlled, demonstrating precision in its task. "
-        "As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. "
-        "The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. "
-        "The other robotic arm on the right side also stays stationary throughout the video. "
-        "The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper."
-    )
-
-    generator.generate_video(
-        prompt,
-        sampling_param=sampling_param,
-        video_path=str(video_path),
-        num_cond_frames=1,
-        output_path="outputs_video/cosmos2_5_v2w.mp4",
-        save_video=True,
-    )
-
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
-
@@ -1,119 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Basic inference script for HunyuanGameCraft video generation.
-
-HunyuanGameCraft generates game-like videos with camera/action control.
-It takes an optional image input and generates video with camera motion
-based on simple action commands (forward, left, right, backward, rotations).
-
-Available actions:
-    - forward (w): Move camera forward
-    - backward (s): Move camera backward  
-    - left (a): Move camera left (strafe)
-    - right (d): Move camera right (strafe)
-    - left_rot: Rotate camera left (pan)
-    - right_rot: Rotate camera right (pan)
-    - up_rot: Rotate camera up (tilt)
-    - down_rot: Rotate camera down (tilt)
-
-T2V vs I2V:
-    - Default: I2V (uses a default reference image). Set GAMECRAFT_I2V_IMAGE to a
-      URL or path to use a different image.
-    - T2V only (no reference image): run with GAMECRAFT_I2V_IMAGE= (empty).
-"""
-import os
-
-import torch
-
-from fastvideo import VideoGenerator
-from fastvideo.models.camera import create_camera_trajectory
-
-# Model configuration (use GAMECRAFT_MODEL_PATH for local weights)
-MODEL_PATH = os.environ.get("GAMECRAFT_MODEL_PATH", "FastVideo/HunyuanGameCraft-Diffusers")
-
-# Default prompts for demo
-DEFAULT_PROMPTS = {
-    "village": "A charming medieval village with cobblestone streets, thatched-roof houses, and vibrant flower gardens under a bright blue sky.",
-    "temple": "A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.",
-    "forest": "A lush green forest with tall trees, dappled sunlight filtering through the leaves, and a winding dirt path.",
-    "beach": "A tropical beach with crystal clear turquoise water, white sand, and palm trees swaying in the breeze.",
-}
-
-# I2V: default reference image (URL). Can override with a local path.
-DEFAULT_I2V_IMAGE_URL = (
-    "https://huggingface.co/datasets/huggingface/documentation-images/"
-    "resolve/main/diffusers/astronaut.jpg"
-)
-DEFAULT_I2V_PROMPT = (
-    "An astronaut hatching from an egg, on the surface of the moon, "
-    "the darkness and depth of space realised in the background."
-)
-
-OUTPUT_PATH = "video_samples_gamecraft"
-
-
-def main():
-    # Initialize generator
-    # FastVideo will automatically download weights from HuggingFace
-    generator = VideoGenerator.from_pretrained(
-        MODEL_PATH,
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True,
-        vae_cpu_offload=True,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True,
-    )
-
-    # Video parameters
-    height = 704
-    width = 1280
-    num_frames = 33
-    action = "forward"
-    action_speed = 0.2
-
-    # Create camera trajectory (Plücker coordinates)
-    camera_states = create_camera_trajectory(
-        action=action,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        action_speed=action_speed,
-        dtype=torch.bfloat16,
-    )
-    print(f"Camera states shape: {camera_states.shape}")
-
-    # I2V vs T2V: unset GAMECRAFT_I2V_IMAGE -> I2V (default image). Set to "" -> T2V.
-    env_image = os.environ.get("GAMECRAFT_I2V_IMAGE")
-    if env_image is None:
-        image_path = DEFAULT_I2V_IMAGE_URL  # default: I2V
-    elif env_image.strip() == "":
-        image_path = None  # T2V
-    else:
-        image_path = env_image.strip()  # I2V with given URL/path
-
-    is_i2v = image_path is not None
-    prompt = DEFAULT_I2V_PROMPT if is_i2v else DEFAULT_PROMPTS["temple"]
-    print(f"Mode: {'I2V' if is_i2v else 'T2V'}, prompt: {prompt[:60]}...")
-
-    gen_kw = dict(
-        prompt=prompt,
-        negative_prompt="",
-        camera_states=camera_states,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        num_inference_steps=50,
-        guidance_scale=6.0,
-        seed=42,
-        fps=24,
-        output_path=OUTPUT_PATH,
-        save_video=True,
-    )
-    if is_i2v:
-        gen_kw["image_path"] = image_path
-    generator.generate_video(**gen_kw)
-
-
-if __name__ == "__main__":
-    main()
@@ -39,4 +39,4 @@ def main():


 if __name__ == "__main__":
-    main()
+    main()
@@ -1,43 +0,0 @@
-from fastvideo import VideoGenerator
-import json
-# from fastvideo.configs.sample import SamplingParam
-
-OUTPUT_PATH = "video_samples_hy15_1080p"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p
-        # or "weizhou03/HunyuanVideo-1.5-Diffusers-1080p" # 720p -> 1080p
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=False, # set to True if GPU is out of memory
-        dit_cpu_offload=True,
-        vae_cpu_offload=True,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
-        # image_encoder_cpu_offload=False,
-    )
-
-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
-
-    prompt2 = (
-        "A majestic lion strides across the golden savanna, its powerful frame "
-        "glistening under the warm afternoon sun. The tall grass ripples gently in "
-        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
-        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
-        "cinematic.")
-
-    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
-
-
-if __name__ == "__main__":
-    main()
@@ -1,66 +0,0 @@
-from fastvideo import VideoGenerator
-from fastvideo.models.dits.hyworld.resolution_utils import get_resolution_from_image
-
-# Default prompt from HY-WorldPlay run.sh
-DEFAULT_PROMPT = 'A paved pathway leads towards a stone arch bridge spanning a calm body of water.  Lush green trees and foliage line the path and the far bank of the water. A traditional-style pavilion with a tiered, reddish-brown roof sits on the far shore. The water reflects the surrounding greenery and the sky.  The scene is bathed in soft, natural light, creating a tranquil and serene atmosphere. The pathway is composed of large, rectangular stones, and the bridge is constructed of light gray stone.  The overall composition emphasizes the peaceful and harmonious nature of the landscape.'
-DEFAULT_IMAGE = 'https://raw.githubusercontent.com/Tencent-Hunyuan/HY-WorldPlay/main/assets/img/test.png'
-
-OUTPUT_PATH = "video_samples_hyworld"
-def main():
-    import argparse
-
-    # pose: (a, w, s, d) - (15, 31)
-    # num_frames: (61, 125)
-    parser = argparse.ArgumentParser(description="HYWorld video generation with FastVideo")
-    parser.add_argument("--prompt", type=str, default=DEFAULT_PROMPT, help="Text prompt for video generation")
-    parser.add_argument("--image", type=str, default=DEFAULT_IMAGE, help="Path or URL to input image")
-    parser.add_argument("--pose", type=str, default='w-31', help="Pose string (e.g., 'a-31', 'w-31', 's-31', 'd-31')")
-    parser.add_argument("--output_path", type=str, default=OUTPUT_PATH, help="Output video path")
-    parser.add_argument("--num-frames", type=int, default=125, help="Number of frames")
-    parser.add_argument("--seed", type=int, default=1, help="Random seed")
-    parser.add_argument("--resolution", type=str, default="480p", help="Only support 480p for now")
-    args = parser.parse_args()
-
-    # Automatically determine resolution from input image
-    HEIGHT, WIDTH = get_resolution_from_image(args.image, args.resolution)
-    print(f"Image: {args.image}")
-    print(f"Pose: {args.pose}")
-    print(f"Resolution: {HEIGHT}x{WIDTH} (from {args.resolution} buckets)")
-    print(f"Num frames: {args.num_frames}")
-    print(f"Output path: {args.output_path}")
-
-    # Initialize generator
-    print("\nInitializing VideoGenerator for HYWorld...")
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/HY-WorldPlay-Bidirectional-Diffusers",
-        num_gpus=1,
-        use_fsdp_inference=True,
-        dit_cpu_offload=True,
-        vae_cpu_offload=True,
-        text_encoder_cpu_offload=True,
-        pin_cpu_memory=True,
-        image_encoder_cpu_offload=True,
-    )
-
-    # Generate video
-    # The pose string is automatically converted to camera matrices by the pipeline
-    print("\nGenerating video...")
-    generator.generate_video(
-        prompt=args.prompt,
-        image_path=args.image,
-        pose=args.pose,  # Camera trajectory control
-        output_path=args.output_path,
-        save_video=True,
-        negative_prompt="",
-        num_frames=args.num_frames,
-        fps=24,
-        height=HEIGHT,
-        width=WIDTH,
-        seed=args.seed,
-    )
-
-    print(f"\nVideo saved to: {args.output_path}")
-
-
-if __name__ == "__main__":
-    main()
@@ -1,49 +0,0 @@
-from fastvideo import VideoGenerator
-from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding
-
-# from fastvideo.configs.sample import SamplingParam
-OUTPUT_PATH = "video_samples_lingbotworld"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-         "FastVideo/LingBot-World-Base-Cam-Diffusers",
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=False, # set to True if GPU is out of memory
-        dit_cpu_offload=True, # DiT need to be offloaded for MoE
-        vae_cpu_offload=False,
-        text_encoder_cpu_offload=True,
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=True,
-        # image_encoder_cpu_offload=False,
-    )
-
-    num_frames = 81
-    prompt = "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls."
-    image_path = "https://raw.githubusercontent.com/Robbyant/lingbot-world/main/examples/00/image.jpg"
-    action_path = "examples/inference/basic/lingbotworld_examples/00"
-    c2ws_plucker_emb, num_frames = prepare_camera_embedding(
-        action_path=action_path,
-        num_frames=num_frames,
-        height=480,
-        width=832,
-        spatial_scale=8,
-    )
-
-    generator.generate_video(
-        prompt,
-        image_path=image_path,
-        output_path=OUTPUT_PATH,
-        save_video=True,
-        num_frames=num_frames,
-        height=480,
-        width=832,
-        c2ws_plucker_emb=c2ws_plucker_emb,
-    )
-
-
-if __name__ == "__main__":
-    main()
@@ -1,39 +0,0 @@
-from fastvideo import VideoGenerator
-
-
-PROMPT = (
-    "A warm sunny backyard. The camera starts in a tight cinematic close-up "
-    "of a woman and a man in their 30s, facing each other with serious "
-    "expressions. The woman, emotional and dramatic, says softly, \"That's "
-    "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly "
-    "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, "
-    "then mutters defensively, \"He's just having fun.\" The camera slowly "
-    "pans right, revealing the grandfather in the garden wearing enormous "
-    "butterfly wings, waving his arms in the air like he's trying to take "
-    "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. "
-    "The woman covers her face, on the verge of tears. The tone is deadpan, "
-    "absurd, and quietly tragic."
-)
-
-
-def main() -> None:
-    # Uses FastVideo default sampling settings for LTX2 base.
-    generator = VideoGenerator.from_pretrained(
-        "Davids048/LTX2-Base-Diffusers",
-        num_gpus=1,
-    )
-
-    output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.1.mp4"
-    generator.generate_video(
-        prompt=PROMPT,
-        output_path=output_path,
-        save_video=True,
-        num_frames=121,
-        height=1088,
-        width=1920,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
@@ -1,35 +0,0 @@
-from fastvideo import VideoGenerator
-
-PROMPT = (
-    "A warm sunny backyard. The camera starts in a tight cinematic close-up "
-    "of a woman and a man in their 30s, facing each other with serious "
-    "expressions. The woman, emotional and dramatic, says softly, \"That's "
-    "it... Dad's lost it. And we've lost Dad.\" The man exhales, slightly "
-    "annoyed: \"Stop being so dramatic, Jess.\" A beat. He glances aside, "
-    "then mutters defensively, \"He's just having fun.\" The camera slowly "
-    "pans right, revealing the grandfather in the garden wearing enormous "
-    "butterfly wings, waving his arms in the air like he's trying to take "
-    "off. He shouts, \"Wheeeew!\" as he flaps his wings with full commitment. "
-    "The woman covers her face, on the verge of tears. The tone is deadpan, "
-    "absurd, and quietly tragic."
-)
-import os
-os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
-
-def main() -> None:
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/LTX2-Distilled-Diffusers",
-        num_gpus=4,
-    )
-
-    output_path = "outputs_video/ltx2_basic/output_ltx2_distilled_t2v.mp4"
-    generator.generate_video(
-        prompt=PROMPT,
-        output_path=output_path,
-        save_video=True,
-    )
-    generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
@@ -1,5 +1,6 @@
 from fastvideo import VideoGenerator
-from fastvideo.models.dits.matrixgame.utils import create_action_presets
+from fastvideo.configs.pipelines.wan import MatrixGameI2V480PConfig
+from fastvideo.models.dits.matrix_game.utils import create_action_presets

 import torch

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

 import torch
 import asyncio
@@ -1,137 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import os
-import re
-from typing import List
-
-
-DEFAULT_PROMPTS = [
-    "a photo of a cat",
-    "a cinematic photo of a red panda wearing a tiny backpack, standing on a rainy neon-lit street at night, shallow depth of field, sharp focus, 35mm, bokeh",
-]
-
-
-def _safe_filename(text: str, max_len: int = 100) -> str:
-    """
-    Make a stable, filesystem-friendly filename base.
-    VideoGenerator uses prompt[:100].strip() internally, so we mirror that,
-    but also remove path separators and other problematic characters.
-    """
-    s = text[:max_len].strip()
-    s = s.replace(os.sep, "_")
-    if os.altsep:
-        s = s.replace(os.altsep, "_")
-    s = re.sub(r"\s+", " ", s)
-    s = re.sub(r"[^A-Za-z0-9 .,_-]", "_", s)
-    s = s.strip(" .")
-    return s or "prompt"
-
-
-def _remove_existing_outputs(out_dir: str, filename_base: str) -> None:
-    """
-    Ensure deterministic naming by deleting any existing outputs that would
-    cause VideoGenerator to append suffixes like _1, _2, etc.
-    """
-    if not os.path.isdir(out_dir):
-        return
-
-    pattern = re.compile(rf"^{re.escape(filename_base)}(_\d+)?\.(mp4|png)$")
-    for fn in os.listdir(out_dir):
-        if pattern.match(fn):
-            try:
-                os.remove(os.path.join(out_dir, fn))
-            except FileNotFoundError:
-                pass
-
-
-def parse_args() -> argparse.Namespace:
-    p = argparse.ArgumentParser(description="Run SD3.5 Medium text-to-image with FastVideo VideoGenerator.")
-    p.add_argument("--model-path", default="stabilityai/stable-diffusion-3.5-medium", help="Path to local diffusers-format SD3.5 weights directory.")
-    p.add_argument(
-        "--out-dir",
-        "--outdir",
-        default="outputs/sd35/samples",
-        help="Output directory for generated mp4 files.",
-    )
-    p.add_argument(
-        "--prompt",
-        action="append",
-        default=None,
-        help="Prompt text. Repeat --prompt multiple times to generate multiple samples.",
-    )
-    p.add_argument("--negative", default="lowres, blurry, jpeg artifacts, watermark, text", help="Negative prompt.")
-    p.add_argument(
-        "--backend",
-        default=None,
-        help="Set FASTVIDEO_ATTENTION_BACKEND (e.g. TORCH_SDPA). If omitted, respects the existing env var.",
-    )
-    p.add_argument("--seed", type=int, default=42, help="Base seed. Each prompt uses seed + prompt_idx.")
-    p.add_argument("--height", type=int, default=768, help="Output height.")
-    p.add_argument("--width", type=int, default=768, help="Output width.")
-    p.add_argument("--steps", type=int, default=28, help="Number of inference steps.")
-    p.add_argument("--guidance", type=float, default=6.0, help="Guidance scale.")
-    p.add_argument("--num-gpus", type=int, default=1, help="Number of GPUs to use.")
-    return p.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    prompts: List[str] = args.prompt if args.prompt else DEFAULT_PROMPTS
-
-    if args.backend:
-        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
-
-    from fastvideo import VideoGenerator
-
-    os.makedirs(args.out_dir, exist_ok=True)
-
-    init_kwargs = {
-        "num_gpus": args.num_gpus,
-        "workload_type": "t2i",
-        "sp_size": 1,
-        "tp_size": 1,
-        "dit_cpu_offload": False,
-        "dit_layerwise_offload": False,
-        "text_encoder_cpu_offload": False,
-        "vae_cpu_offload": False,
-        "image_encoder_cpu_offload": False,
-        "pin_cpu_memory": False,
-        "use_fsdp_inference": False,
-    }
-
-    generator = VideoGenerator.from_pretrained(model_path=args.model_path, **init_kwargs)
-    try:
-        for i, prompt in enumerate(prompts):
-            seed = args.seed + i
-
-            filename_base = f"sd35_{i:02d}_seed{seed}_{_safe_filename(prompt, max_len=80)}"
-            _remove_existing_outputs(args.out_dir, filename_base)
-
-            output_path = os.path.join(args.out_dir, f"{filename_base}.png")
-            print(f"[sd35] prompt_idx={i} seed={seed} output_path={output_path}")
-
-            generation_kwargs = {
-                "output_path": output_path,
-                "height": args.height,
-                "width": args.width,
-                "num_frames": 1,
-                "fps": 1,
-                "num_inference_steps": args.steps,
-                "guidance_scale": args.guidance,
-                "seed": seed,
-                "negative_prompt": args.negative,
-                "save_video": True,
-            }
-
-            generator.generate_video(prompt, **generation_kwargs)
-
-        print(f"[sd35] done. outputs written to: {args.out_dir}")
-    finally:
-        generator.shutdown()
-
-
-if __name__ == "__main__":
-    main()
@@ -31,7 +31,7 @@ def main():
    sampling_param.height = 480
    sampling_param.seed = 1000

-    with open("assets/prompts/mixkit_i2v.jsonl", "r") as f:
+    with open("prompts/mixkit_i2v.jsonl", "r") as f:
        prompt_image_pairs = json.load(f)

    for prompt_image_pair in prompt_image_pairs:
@@ -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.matrixgame.utils import expand_action_to_frames
+from fastvideo.models.dits.matrix_game.utils import expand_action_to_frames


 VARIANT_CONFIG = {
@@ -188,9 +188,8 @@ def load_example_prompts():
    prompt_to_image = {}
    # Try to find the JSON file relative to project root
    possible_json_paths = [
-        Path("assets/prompts/mixkit_i2v.jsonl"),
-        Path(__file__).resolve().parents[4] / "assets" / "prompts" /
-        "mixkit_i2v.jsonl",
+        Path("prompts/mixkit_i2v.jsonl"),
+        Path(__file__).parent.parent.parent.parent / "prompts" / "mixkit_i2v.jsonl",
    ]
    json_path = None
    for path in possible_json_paths:
@@ -202,8 +201,8 @@ def load_example_prompts():
        try:
            with open(json_path, "r", encoding='utf-8') as f:
                data = json.load(f)
-                # Resolve paths relative to repository root.
-                project_root = Path(__file__).resolve().parents[4]
+                # Get the project root directory (parent of prompts directory)
+                project_root = json_path.parent.parent
                for item in data:
                    prompt_text = item.get("prompt", "").strip()
                    image_path = item.get("image_path", "")
@@ -737,8 +736,8 @@ def main():
        allowed_paths=[
            os.path.abspath("outputs"), 
            os.path.abspath("fastvideo-logos"),
-            os.path.abspath("assets/prompts"),
-            os.path.abspath("assets/images"),
+            os.path.abspath("prompts"),
+            os.path.abspath("images"),
            os.path.abspath(tempfile.gettempdir()),
            os.path.abspath(os.path.join(tempfile.gettempdir(), "gradio")),
        ]
@@ -748,4 +747,4 @@ def main():


 if __name__ == "__main__":
-    main() 
+    main() 
@@ -5,7 +5,7 @@ export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers

 base_port=29503
-num_gpu=8
+num_gpu=1
 gpu_ids=$(seq 0 $((num_gpu-1)))
 skip_time_steps=12

@@ -4,7 +4,7 @@ export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export TOKENIZERS_PARALLELISM=false

-MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
 VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
 NUM_GPUS=1
@@ -14,6 +14,7 @@ NUM_GPUS=1
 training_args=(
  --tracker_project_name "wan_ode_init"
  --output_dir "wan_ode_init_crush_smol"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
  --wandb_run_name "wan_ode_init_crush_smol"
  --max_train_steps 6000
  --train_batch_size 1
@@ -33,7 +34,7 @@ parallel_args=(
  --sp_size 1
  --tp_size 1
  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
+  --hsdp_shard_dim 1
 )

 # Model arguments
@@ -50,17 +51,20 @@ dataset_args=(

 # Validation arguments
 validation_args=(
-  --log-visualization
-  --visualization-steps 100
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
 )

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 6e-6
  --mixed_precision "bf16"
  --weight_only_checkpointing_steps 1000
  --training_state_checkpointing_steps 1000
-  --weight_decay 0.01
+  --weight_decay 1e-4
  --max_grad_norm 1.0
 )

@@ -1,93 +0,0 @@
-#!/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[@]}"
@@ -1,27 +0,0 @@
-#!/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
@@ -44,8 +44,8 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

-# Core training arguments
-core_training_args=(
+# Training arguments
+training_args=(
  --tracker_project_name wan_i2v_VSA
  --output_dir "checkpoints/wan_i2v_finetune_VSA"
  --max_train_steps 4000
@@ -53,14 +53,10 @@ core_training_args=(
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 21
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Validation generation shape arguments (used during training validation)
-validation_generation_shape_args=(
  --num_height 720
  --num_width 1280
  --num_frames 81
+  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
@@ -131,9 +127,8 @@ srun torchrun \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
-    "${core_training_args[@]}" \
-    "${validation_generation_shape_args[@]}" \
+    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
+    "${vsa_args[@]}"  
@@ -6,7 +6,9 @@ These are e2e example scripts for finetuning Wan2.1 T2V with VSA to accelerate i

 ## Make sure you have installed VSA

-Go to fastvideo-kernel/README.md for instructions.
+```bash
+pip install vsa
+```

 ### Download the synthetic dataset:

@@ -44,8 +44,8 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

-# Core training arguments
-core_training_args=(
+# Training arguments
+training_args=(
  --tracker_project_name wan_t2v_VSA
  --output_dir "checkpoints/wan_t2v_finetune_VSA"
  --max_train_steps 4000
@@ -53,14 +53,10 @@ core_training_args=(
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 21
-  # --enable_gradient_checkpointing_type "full" # if OOM enable this
-)
-
-# Validation generation shape arguments (used during training validation)
-validation_generation_shape_args=(
  --num_height 480
  --num_width 832
  --num_frames 81
+  # --enable_gradient_checkpointing_type "full" # if OOM enable this
 )

 # Parallel arguments
@@ -131,9 +127,8 @@ srun torchrun \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
-    "${core_training_args[@]}" \
-    "${validation_generation_shape_args[@]}" \
+    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
+    "${vsa_args[@]}"  
@@ -1,114 +0,0 @@
-
-
-# Basic info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-export TRITON_CACHE_DIR="/tmp/triton_cache_${USER}_$$"
-export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
-export MASTER_PORT="${MASTER_PORT:-29500}"
-export NODE_RANK=0
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
-export WANDB_API_KEY="your_wandb_api_key_here" # TODO: Replace with your actual key or load from a secure location
-
-# Configs
-NUM_GPUS=4
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR=data/Wan-Syn_77x448x832_600k
-VALIDATION_DATASET_FILE=examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
-
-# Core training arguments
-core_training_args=(
-  --tracker_project_name wan_t2v_VSA
-  --output_dir "checkpoints/wan_t2v_finetune_VSA"
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 20
-  --enable_gradient_checkpointing_type "full" # if OOM enable this
-)
-
-# Validation generation shape arguments (used during training validation)
-validation_generation_shape_args=(
-  --num_height 448
-  --num_width 832
-  --num_frames 77
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim $NUM_GPUS
-  --hsdp_shard_dim 1
-)
-
-# 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 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 200
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "5.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-6
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.1
-  --dit_precision "fp32"
-  --ema_start_step 0
-  --flow_shift 1
-  --seed 1000
-)
-
-# VSA arguments
-vsa_args=(
-  --VSA_decay_rate 0.03
-  --VSA_decay_interval_steps 50
-  --VSA_sparsity 0.9
-)
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node "$NUM_GPUS" \
-  --node_rank "$NODE_RANK" \
-  --rdzv_backend c10d \
-  --rdzv_endpoint "$MASTER_ADDR:$MASTER_PORT" \
-  fastvideo/training/wan_training_pipeline.py \
-  "${parallel_args[@]}" \
-  "${model_args[@]}" \
-  "${dataset_args[@]}" \
-  "${core_training_args[@]}" \
-  "${validation_generation_shape_args[@]}" \
-  "${optimizer_args[@]}" \
-  "${validation_args[@]}" \
-  "${miscellaneous_args[@]}" \
-  "${vsa_args[@]}"
@@ -44,8 +44,8 @@ VALIDATION_DATASET_FILE=your_validation_dataset_file
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

-# Core training arguments
-core_training_args=(
+# Training arguments
+training_args=(
  --tracker_project_name wan_t2v_VSA
  --output_dir "checkpoints/wan_t2v_finetune_VSA"
  --max_train_steps 4000
@@ -53,14 +53,10 @@ core_training_args=(
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 21
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Validation generation shape arguments (used during training validation)
-validation_generation_shape_args=(
  --num_height 720
  --num_width 1280
  --num_frames 81
+  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
@@ -131,9 +127,8 @@ srun torchrun \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
-    "${core_training_args[@]}" \
-    "${validation_generation_shape_args[@]}" \
+    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${vsa_args[@]}"  
+    "${vsa_args[@]}"  
@@ -1,8 +1,7 @@
 #!/bin/bash
-# Note: for debugging, you do not need to download them all and just control c.
-

 # 480P dataset
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "data/Wan-Syn_77x448x832_600k" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
+
 # 720P dataset
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x768x1280_250k" --local_dir "data/Wan-Syn_77x768x1280_250k" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x768x1280_250k" --local_dir "FastVideo/Wan-Syn_77x768x1280_250k" --repo_type "dataset"
@@ -1,36 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "In the video, a woman is elegantly showcasing her earrings, bringing attention to their intricate design with a gentle touch of her fingers. She is bathed in ambient purple and pink lighting, which casts a soft glow on her delicate features and enhances the vivid tones of her lipstick and eye makeup. Her hair is styled to frame her face smoothly, emphasizing the contours of her jawline and cheekbones. The background features a blurred neon light, adding an artistic and modern touch to the overall aesthetic.",
-      "video_path": "Fashion/mixkit-face-of-an-elegant-and-captivating-woman-41914_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.",
-      "video_path": "Man/mixkit-a-rancher-riding-a-horse-at-sunset-1143_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "In a dimly lit, eerie setting, a mysterious pink bottle labeled \"Authentic 100% organic POISON\" sits prominently in the foreground, casting a menacing aura. The bottle is accentuated by green fog, which swirls lightly around it, enhancing its sinister allure. Behind it, a shadowy golden bottle adorned with a spider emblem subtly emerges, adding an extra layer of mystery to the scene. Dim candles provide faint, flickering light, which complements the dark atmosphere, making the setting ideal for an illusion of hidden dangers.",
-      "video_path": "smoke/mixkit-poison-in-halloween-ritual-33879_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    },
-    {
-      "caption": "The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.",
-      "video_path": "forest/mixkit-view-of-a-river-between-two-old-trees-560_clip_1.mp4",
-      "num_inference_steps":  3,
-      "height": 448,
-      "width": 832,
-      "num_frames": 61
-    }
-  ]
-}
@@ -1,3 +0,0 @@
-*.err
-*.out
-*slurm_logs/*
@@ -1,25 +0,0 @@
-# LTX-2 Crush-Smol Example
-# TODO: Update this doc.
-
-These are e2e example scripts for finetuning LTX-2 on the crush-smol dataset.
-
-## Execute the following commands from `FastVideo/` to run training:
-
-### Download crush-smol dataset:
-
-`bash examples/training/finetune/ltx2/download_dataset.sh`
-
-#### Or use the scripts at scripts/dataset_preparation to download and prepare the dataset
-
-### Preprocess the videos and captions into latents:
-
-`bash examples/training/finetune/ltx2/preprocess_ltx2_data_t2v_new.sh`
-
-### Edit the following file and run finetuning:
-
-`bash examples/training/finetune/ltx2/finetune_t2v.sh`
-
-Notes:
- Update `DATASET_PATH` in the preprocess script to point to your merged dataset root (`videos/` + `videos2caption.json`).
- `MODEL_PATH` should point to a local LTX-2 diffusers-style directory that contains `model_index.json` and `text_encoder/gemma`.
-
@@ -1,3 +0,0 @@
-# #!/bin/bash
-#
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -1,93 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-
-MODEL_PATH="FastVideo/LTX2-Distilled-Diffusers"
-DATA_DIR="data/crush-smol"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-echo  VALIDATION_DATASET_FILE: $VALIDATION_DATASET_FILE
-NUM_GPUS=4
-HEIGHT=1088
-WIDTH=1920
-FRAMES=121
-
-training_args=(
-  --tracker_project_name "ltx2_t2v_finetune"
-  --output_dir "checkpoints/ltx2_t2v_finetune"
-  --max_train_steps 5000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 16
-  --num_height $HEIGHT
-  --num_width $WIDTH
-  --num_frames $FRAMES
-  --enable_gradient_checkpointing_type "full"
-  --mode "finetuning"
-)
-
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 2
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-dataset_args=(
-  --data_path $DATA_DIR
-  --dataloader_num_workers 1
-)
-
-validation_args=(
-  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 20
-  --validation_sampling_steps "8"
-  --validation_guidance_scale "1.0"
-)
-
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-  --lr_scheduler "linear"
-)
-
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --dit_precision "fp32"
-  --dit_cpu_offload False
-  --dit_layerwise_offload False
-  --text_encoder_cpu_offload False
-  --image_encoder_cpu_offload False
-  --vae_cpu_offload False
-  --ltx2-first-frame-conditioning-p 0.0
-)
-
-# NOTE: Setting this environment variable to TORCH_SDPA to avoid the issue of stacking that failed in flash attn. 
-
-
-torchrun \
-  --nnodes 1 \
-  --master_port 29501 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/ltx2_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,80 +0,0 @@
-#!/bin/bash
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export TOKENIZERS_PARALLELISM=false
-
-MODEL_PATH="FastVideo/LTX2-Distilled-Diffusers"
-DATA_DIR="data/crush-smol"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
-NUM_GPUS=1
-
-training_args=(
-  --tracker_project_name "ltx2_t2v_lora_finetune"
-  --output_dir "checkpoints/ltx2_t2v_lora_finetune"
-  --max_train_steps 2000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 8
-  --num_latent_t 10
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-  --ltx2-first-frame-conditioning-p 0.1
-  --enable_gradient_checkpointing_type "full"
-)
-
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size $NUM_GPUS
-  --tp_size 1
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim $NUM_GPUS
-)
-
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-)
-
-dataset_args=(
-  --data_path $DATA_DIR
-  --dataloader_num_workers 1
-)
-
-validation_args=(
-  --log_validation
-  --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "3.0"
-)
-
-optimizer_args=(
-  --learning_rate 2e-4
-  --mixed_precision "bf16"
-  --weight_only_checkpointing_steps 1000
-  --training_state_checkpointing_steps 1000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-  --lora_training True
-  --lora_rank 16
-  --lora_alpha 16
-)
-
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-)
-
-torchrun \
-  --nnodes 1 \
-  --nproc_per_node $NUM_GPUS \
-    fastvideo/training/ltx2_training_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -1,27 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=4
-MODEL_PATH="FastVideo/LTX2-Distilled-Diffusers"
-DATASET_PATH="data/crush-smol"
-OUTPUT_DIR="$DATASET_PATH"
-WITH_AUDIO=true
-
-
-torchrun  --nproc_per_node=$GPU_NUM \
-    --master_port=29513 \
-    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
-    --model_path $MODEL_PATH \
-    --mode preprocess \
-    --workload_type t2v \
-    --preprocess.video_loader_type torchvision \
-    --preprocess.dataset_type merged \
-    --preprocess.dataset_path $DATASET_PATH \
-    --preprocess.dataset_output_dir $OUTPUT_DIR \
-    --preprocess.with_audio $WITH_AUDIO \
-    --preprocess.preprocess_video_batch_size 1 \
-    --preprocess.dataloader_num_workers 0 \
-    --preprocess.max_height 1088  \
-    --preprocess.max_width 1920 \
-    --preprocess.num_frames 121 \
-    --preprocess.train_fps 24 \
-    --preprocess.video_length_tolerance_range 5
@@ -1,16 +0,0 @@
-{
-  "data": [
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press."
-    },
-    {
-      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press."
-    },
-    {
-      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."
-    },
-    {
-      "caption": "A massive steel piston descends onto a stack of chocolate chip cookies, crushing them into crumbs as though they are being compressed by a hydraulic press."
-    }
-  ]
-}
@@ -0,0 +1,129 @@
+#!/bin/bash
+
+# Change to FastVideo root directory (3 levels up from this script)
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+FASTVIDEO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+cd "$FASTVIDEO_ROOT"
+
+# Add FastVideo root to PYTHONPATH so Python can find the fastvideo package
+export PYTHONPATH="$FASTVIDEO_ROOT${PYTHONPATH:+:$PYTHONPATH}"
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+RL_DATASET_DIR="data/ocr/"  # Path to RL prompt dataset directory (should contain train.txt and test.txt)
+VALIDATION_DATASET_FILE="$SCRIPT_DIR/validation.json"
+NUM_GPUS=1
+
+# use GPU 3
+export CUDA_VISIBLE_DEVICES=3
+
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_t2v_grpo"
+  --output_dir "checkpoints/wan_t2v_grpo"
+  --max_train_steps 5000
+  --train_batch_size 4
+  # --train_sp_batch_size 4
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 5
+  --num_height 240
+  --num_width 416
+  --num_frames 33
+  --lora_rank 32
+  --lora_training True
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS 
+  --sp_size $NUM_GPUS 
+  --tp_size $NUM_GPUS
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim $NUM_GPUS
+  # --use-fsdp-inference False
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments (for RL prompt dataset)
+dataset_args=(
+  --data_path $RL_DATASET_DIR  # Used as fallback if rl_dataset_path not set
+  --rl_dataset_path $RL_DATASET_DIR  # RL prompt dataset directory
+  --rl_dataset_type "text"  # "text" or "geneval"
+  --rl_num_image_per_prompt 4  # k parameter (number of samples per prompt)
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation True
+  --validation_dataset_file $VALIDATION_DATASET_FILE
+  --validation_steps 5
+  --validation_sampling_steps "50" 
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-5
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 10
+  --training_state_checkpointing_steps 10
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# RL-specific arguments
+rl_args=(
+  --inference_mode False
+  --rl_mode True
+  --rl_algorithm "grpo"
+  --rl_kl_beta 0.004  # KL regularization coefficient
+  --rl_policy_clip_range 0.2  # Policy clipping range for GRPO
+  --rl_kl_reward 0.0  # KL reward coefficient (typically 0)
+  --rl_global_std False  # Use per-prompt std (recommended for GRPO)
+  --rl_per_prompt_stat_tracking True  # Enable per-prompt stat tracking
+  --rl_warmup_steps 0  # Number of warmup steps (SFT before RL)
+  --reward-models "{\"paddle_ocr\": 1.0}" # use video_ocr reward function
+)
+
+# CFG arguments
+cfg_args=(
+  --guidance_scale 1.0 # use guidance_scale > 1.0 to enable CFG
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0  # No CFG during training (CFG used in sampling)
+  --dit_precision "fp32"
+  # --dit_precision "bf16"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+  # --resume_from_checkpoint "checkpoints/wan_t2v_grpo/checkpoint-XXX"
+  --enable-gradient-checkpointing-type "full"
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+  --master_port 29501 \
+    "$FASTVIDEO_ROOT/fastvideo/training/wan_rl_training_pipeline.py" \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${rl_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -40,17 +40,6 @@ out = video_sparse_attn(q, k, v, block_sizes, block_sizes, topk=5)
 out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...)
 ```

-## Benchmark
-
-### VSA (block-sparse) TFLOPs
-
-After building/installing `fastvideo-kernel`, run:
-
-```bash
-cd fastvideo-kernel
-python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64
-```
-
 ### TurboDiffusion Kernels

 This package also includes kernels from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion), including INT8 GEMM, Quantization, RMSNorm and LayerNorm.
@@ -1,166 +0,0 @@
-#!/usr/bin/env python3
-"""
-Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs.
-
-This script benchmarks the autograd-enabled wrapper:
-  - fastvideo_kernel.block_sparse_attn.block_sparse_attn
-
-So measured time includes wrapper overhead (map->index conversion, dispatch) plus kernel time.
-"""
-
-from __future__ import annotations
-
-import argparse
-import os
-import random
-from typing import Tuple, Callable
-
-import numpy as np
-import torch
-
-try:
-    from triton.testing import do_bench
-except Exception as e:  # pragma: no cover
-    raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e
-
-
-BLOCK_M = 64
-BLOCK_N = 64
-
-
-def set_seed(seed: int = 42) -> None:
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-
-def parse_arguments() -> argparse.Namespace:
-    p = argparse.ArgumentParser(description="Benchmark FastVideo VSA block-sparse attention")
-    p.add_argument("--batch_size", type=int, default=1)
-    p.add_argument("--num_heads", type=int, default=12)
-    p.add_argument("--head_dim", type=int, default=128, choices=[64, 128])
-    p.add_argument("--topk", type=int, default=None, help="KV blocks per Q block (default: ~90%% sparsity)")
-    p.add_argument("--q_seq_lens", type=int, nargs="+", default=[49152], help="Q sequence lengths (must be /64)")
-    p.add_argument("--kv_seq_lens", type=int, nargs="+", default=None, help="KV sequence lengths (defaults to q_seq_len)")
-    p.add_argument("--warmup", type=int, default=5)
-    p.add_argument("--rep", type=int, default=20)
-    p.add_argument("--seed", type=int, default=42)
-    p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"])
-    p.add_argument("--force_triton", action="store_true", help="Force wrapper to use Triton path (if supported by shapes).")
-    return p.parse_args()
-
-
-def create_qkv(batch: int, heads: int, q_len: int, kv_len: int, d: int, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    q = torch.randn(batch, heads, q_len, d, dtype=dtype, device="cuda")
-    k = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda")
-    v = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda")
-    return q, k, v
-
-
-def make_block_map(bs: int, h: int, num_q_blocks: int, num_kv_blocks: int, topk: int) -> torch.Tensor:
-    # block_map: [bs, h, num_q_blocks, num_kv_blocks] bool
-    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device="cuda")
-    topk = min(max(1, topk), num_kv_blocks)
-    idx = torch.topk(scores, topk, dim=-1).indices
-    block_map = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device="cuda")
-    block_map.scatter_(-1, idx, True)
-    return block_map
-
-
-def flops_sparse_attention(bs: int, h: int, d: int, q_len: int, topk_blocks: int, block_n: int) -> float:
-    # Approx: QK^T + PV, each is ~2*bs*h*q_len*(topk_blocks*block_n)*d
-    return 4.0 * bs * h * d * q_len * (topk_blocks * block_n)
-
-
-def bench_ms(fn: Callable[[], object], warmup: int, rep: int) -> float:
-    return do_bench(fn, warmup=warmup, rep=rep, quantiles=None)
-
-
-def main() -> None:
-    args = parse_arguments()
-    set_seed(args.seed)
-
-    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
-
-    if args.force_triton:
-        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
-
-    from fastvideo_kernel.block_sparse_attn import block_sparse_attn
-
-    bs, h, d = args.batch_size, args.num_heads, args.head_dim
-    kv_seq_lens = args.kv_seq_lens
-    if kv_seq_lens is None:
-        kv_seq_lens = args.q_seq_lens
-    if len(kv_seq_lens) != len(args.q_seq_lens):
-        raise ValueError("kv_seq_lens must have the same number of entries as q_seq_lens (or be omitted).")
-
-    print("VSA Block-Sparse Attention Benchmark (WRAPPER)")
-    print(f"device: {torch.cuda.get_device_name(0)}")
-    print(f"batch={bs}, heads={h}, head_dim={d}, dtype={args.dtype}")
-    print(f"BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}")
-    print("NOTE: timings include wrapper overhead (map->index + dispatch).")
-    if args.force_triton:
-        print("dispatch: forced Triton (FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1)")
-    else:
-        print("dispatch: SM90 if available, else Triton")
-
-    for q_len, kv_len in zip(args.q_seq_lens, kv_seq_lens):
-        if q_len % BLOCK_M != 0 or kv_len % BLOCK_N != 0:
-            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by 64")
-            continue
-
-        num_q_blocks = q_len // BLOCK_M
-        num_kv_blocks = kv_len // BLOCK_N
-        topk = args.topk if args.topk is not None else max(1, num_kv_blocks // 10)
-        topk = min(topk, num_kv_blocks)
-
-        print("\n" + "=" * 80)
-        print(f"q_len={q_len}, kv_len={kv_len}, num_q_blocks={num_q_blocks}, num_kv_blocks={num_kv_blocks}, topk={topk}")
-
-        q, k, v = create_qkv(bs, h, q_len, kv_len, d, dtype)
-        block_map = make_block_map(bs, h, num_q_blocks, num_kv_blocks, topk)
-
-        # Variable block sizes: default full blocks (64 tokens per KV block)
-        variable_block_sizes = torch.full((num_kv_blocks,), BLOCK_N, dtype=torch.int32, device="cuda")
-
-        def _fwd():
-            return block_sparse_attn(q, k, v, block_map, variable_block_sizes)
-
-        fwd_ms = bench_ms(_fwd, warmup=args.warmup, rep=args.rep)
-
-        # Backward benchmark (wrapper autograd). We build the graph once, then repeatedly run backward
-        # on the retained graph so bwd timing excludes the forward compute.
-        q_ = q.detach().requires_grad_(True)
-        k_ = k.detach().requires_grad_(True)
-        v_ = v.detach().requires_grad_(True)
-        o_, _aux_ = block_sparse_attn(q_, k_, v_, block_map, variable_block_sizes)
-        og = torch.randn_like(o_)
-        loss = (o_ * og).sum()
-
-        for _ in range(max(1, args.warmup // 2)):
-            torch.autograd.grad(loss, (q_, k_, v_), retain_graph=True)
-        torch.cuda.synchronize()
-
-        bwd_ms = bench_ms(
-            lambda: torch.autograd.grad(loss, (q_, k_, v_), retain_graph=True),
-            warmup=0,
-            rep=max(5, args.rep // 2),
-        )
-
-        flops = flops_sparse_attention(bs, h, d, q_len, topk, BLOCK_N)
-        fwd_tflops = flops / fwd_ms * 1e-12 * 1e3
-        # Rough backward multiplier (attention backward typically ~2-3x forward)
-        bwd_tflops = (2.5 * flops) / bwd_ms * 1e-12 * 1e3
-
-        print(f"fwd(wrapper): {fwd_ms:.3f} ms  | {fwd_tflops:.2f} TFLOPs (approx)")
-        print(f"bwd(wrapper): {bwd_ms:.3f} ms  | {bwd_tflops:.2f} TFLOPs (approx)")
-
-
-if __name__ == "__main__":
-    if not torch.cuda.is_available():
-        raise RuntimeError("CUDA is required for this benchmark.")
-    main()
-
-
@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"

 [project]
 name = "fastvideo-kernel"
-version = "0.2.6"
+version = "0.2.4"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -30,12 +30,13 @@ def _force_triton() -> bool:
    return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1"


-def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+def _map_to_index_torch(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
-    Preferred map->index conversion used by the wrapper.
-
-    This wrapper **requires** the Triton implementation.
-    If Triton (or the Triton map_to_index module) is not available, it raises.
+    Pure-torch (no triton) conversion:
+      block_map: [B, H, Q, KV] bool  (or [H, Q, KV] which will be treated as B=1)
+      returns:
+        index: [B, H, Q, KV] int32 (packed KV indices, -1 padding)
+        num:   [B, H, Q] int32 (#kv blocks per q block)
    """
    if block_map.dim() == 3:
        block_map = block_map.unsqueeze(0)
@@ -44,17 +45,20 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    if block_map.dtype != torch.bool:
        block_map = block_map.to(torch.bool)

-    if not block_map.is_cuda:
-        raise RuntimeError("block_map must be a CUDA tensor (Triton map_to_index required).")
+    B, H, Q, KV = block_map.shape
+    index = torch.full((B, H, Q, KV), -1, dtype=torch.int32, device=block_map.device)
+    num = torch.zeros((B, H, Q), dtype=torch.int32, device=block_map.device)

-    try:
-        from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index  # local import
-    except Exception as e:
-        raise ImportError(
-            "Triton map_to_index is required but not available. "
-            "Ensure Triton is installed and fastvideo_kernel.triton_kernels.index is importable."
-        ) from e
-    return triton_map_to_index(block_map)
+    # Small sizes in practice (B=1, H<=16, Q/KV<=64), so a Python loop is fine.
+    for b in range(B):
+        for h in range(H):
+            for q in range(Q):
+                kv_idx = torch.nonzero(block_map[b, h, q], as_tuple=False).flatten().to(torch.int32)
+                n = int(kv_idx.numel())
+                if n:
+                    index[b, h, q, :n] = kv_idx
+                num[b, h, q] = n
+    return index, num


@torch.library.custom_op(
@@ -73,7 +77,7 @@ def block_sparse_attn_triton(
    k = k.contiguous()
    v = v.contiguous()
    block_map = block_map.to(torch.bool)
-    q2k_idx, q2k_num = _map_to_index(block_map)
+    q2k_idx, q2k_num = _map_to_index_torch(block_map)

    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import (  # local import
        triton_block_sparse_attn_forward,
@@ -83,7 +87,6 @@ def block_sparse_attn_triton(
    return o, M


-
@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_triton")
 def _block_sparse_attn_triton_fake(
    q: torch.Tensor,
@@ -114,8 +117,8 @@ def block_sparse_attn_backward_triton(
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    grad_output = grad_output.contiguous()
    block_map = block_map.to(torch.bool)
-    q2k_idx, q2k_num = _map_to_index(block_map)
-    k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous())
+    q2k_idx, q2k_num = _map_to_index_torch(block_map)
+    k2q_idx, k2q_num = _map_to_index_torch(block_map.transpose(-1, -2).contiguous())

    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import (  # local import
        triton_block_sparse_attn_backward,
@@ -179,7 +182,7 @@ def block_sparse_attn_sm90(
    k_padded = k_padded.contiguous()
    v_padded = v_padded.contiguous()
    block_map = block_map.to(torch.bool)
-    q2k_idx, q2k_num = _map_to_index(block_map)
+    q2k_idx, q2k_num = _map_to_index_torch(block_map)

    o_padded, lse_padded = block_sparse_fwd(
        q_padded, k_padded, v_padded, q2k_idx, q2k_num, variable_block_sizes.int()
@@ -221,7 +224,7 @@ def block_sparse_attn_backward_sm90(

    grad_output_padded = grad_output_padded.contiguous()
    block_map = block_map.to(torch.bool)
-    k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous())
+    k2q_idx, k2q_num = _map_to_index_torch(block_map.transpose(-1, -2).contiguous())

    dq, dk, dv = block_sparse_bwd(
        q_padded,
@@ -287,8 +290,9 @@ def block_sparse_attn(
    block_sparse_fwd, block_sparse_bwd = _get_sm90_ops()
    if (not _force_triton()) and _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
        return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes)
-    # Triton path: supports q_seq_len != kv_seq_len as long as both are padded
-    # to a multiple of the block size (64 tokens).
+    # Triton path: generally assumes q/k/v share the same padded length
+    if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]:
+        raise RuntimeError("Triton fallback requires q/k/v to have the same padded length.")
    return block_sparse_attn_triton(q, k, v, block_map, variable_block_sizes)


@@ -141,6 +141,12 @@ def video_sparse_attn(
        # Use autograd-enabled wrapper so backward works (and still uses SM90 kernel when available)
        out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
    else:
+        if q_seq_len != kv_seq_len:
+            raise RuntimeError(
+                "q/k have different lengths, but the compiled CUDA kernel (block_sparse_fwd) "
+                "is not available. The Triton fallback currently requires q and k/v to have "
+                "the same padded length."
+            )
        # Triton-only forward (kept for environments without the wrapper deps)
        out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes)

@@ -29,7 +29,7 @@ configs = [


 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
-@triton.autotune(configs, key=["N_CTX_Q", "HEAD_DIM"])
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
@triton.jit
 def _attn_fwd_sparse(
        Q,
@@ -60,8 +60,7 @@ def _attn_fwd_sparse(
        stride_on,
        Z,
        H,
-        N_CTX_Q,  #
-        N_CTX_KV,  #
+        N_CTX,  #
        HEAD_DIM: tl.constexpr,  #
        BLOCK_M: tl.constexpr,
        BLOCK_N: tl.constexpr,
@@ -76,29 +75,24 @@ def _attn_fwd_sparse(
    off_hz = tl.program_id(1)  # fused (batch, head)
    b = off_hz // H
    h = off_hz % H
-    q_tiles = N_CTX_Q // BLOCK_M
+    q_tiles = N_CTX // BLOCK_M
    meta_base = ((b * H + h) * q_tiles + q_blk)

    kv_blocks = tl.load(q2k_num + meta_base)  # int32
    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list

    # ----- base pointers -----
-    # Note: when q and kv have different sequence lengths, their per-(batch,head)
-    # strides differ, so we must compute separate base offsets.
-    q_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
-    k_off = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
-    v_off = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
-    o_off = (b.to(tl.int64) * stride_oz + h.to(tl.int64) * stride_oh)
+    qvk_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)

-    Q_ptr = tl.make_block_ptr(base=Q + q_off,
-                              shape=(N_CTX_Q, HEAD_DIM),
+    Q_ptr = tl.make_block_ptr(base=Q + qvk_off,
+                              shape=(N_CTX, HEAD_DIM),
                              strides=(stride_qm, stride_qk),
                              offsets=(q_blk * BLOCK_M, 0),
                              block_shape=(BLOCK_M, HEAD_DIM),
                              order=(1, 0))

-    K_base = tl.make_block_ptr(base=K + k_off,
-                               shape=(HEAD_DIM, N_CTX_KV),
+    K_base = tl.make_block_ptr(base=K + qvk_off,
+                               shape=(HEAD_DIM, N_CTX),
                               strides=(stride_kk, stride_kn),
                               offsets=(0, 0),
                               block_shape=(HEAD_DIM, BLOCK_N),
@@ -106,15 +100,15 @@ def _attn_fwd_sparse(

    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1,
                                                                              0)
-    V_base = tl.make_block_ptr(base=V + v_off,
-                               shape=(N_CTX_KV, HEAD_DIM),
+    V_base = tl.make_block_ptr(base=V + qvk_off,
+                               shape=(N_CTX, HEAD_DIM),
                               strides=(stride_vk, stride_vn),
                               offsets=(0, 0),
                               block_shape=(BLOCK_N, HEAD_DIM),
                               order=v_order)

-    O_ptr = tl.make_block_ptr(base=Out + o_off,
-                              shape=(N_CTX_Q, HEAD_DIM),
+    O_ptr = tl.make_block_ptr(base=Out + qvk_off,
+                              shape=(N_CTX, HEAD_DIM),
                              strides=(stride_om, stride_on),
                              offsets=(q_blk * BLOCK_M, 0),
                              block_shape=(BLOCK_M, HEAD_DIM),
@@ -156,7 +150,7 @@ def _attn_fwd_sparse(
    # ----- epilogue -----
    m_i += tl.math.log2(l_i)
    acc = acc / l_i[:, None]
-    tl.store(M + off_hz * N_CTX_Q + offs_m, m_i)
+    tl.store(M + off_hz * N_CTX + offs_m, m_i)
    tl.store(O_ptr, acc.to(Out.type.element_ty))


@@ -207,7 +201,7 @@ def _attn_bwd_dkdv(
        stride_tok,
        stride_d,  #
        H,
-        N_CTX_KV,
+        N_CTX,
        BLOCK_M1: tl.constexpr,  #
        BLOCK_N1: tl.constexpr,  #
        HEAD_DIM: tl.constexpr,  #
@@ -227,8 +221,8 @@ def _attn_bwd_dkdv(
    off_hz = tl.program_id(2)  # fused (batch, head)
    b = off_hz // H
    h = off_hz % H
-    kv_tiles = N_CTX_KV // BLOCK_N1
-    meta_base = ((b * H + h) * kv_tiles + kv_blk)
+    q_tiles = N_CTX // BLOCK_N1
+    meta_base = ((b * H + h) * q_tiles + kv_blk)

    q_blocks = tl.load(k2q_num + meta_base)  # int32
    q_ptr = k2q_index + meta_base * max_q_blks  # ptr to list
@@ -308,21 +302,16 @@ def _attn_bwd_dq(

    kv_blocks = tl.load(q2k_num + meta_base)  # int32
    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    block_size = tl.load(variable_block_sizes + q_blk)

    for blk_idx in range(kv_blocks * 2):
-        kv_idx = tl.load(kv_ptr + blk_idx // 2).to(tl.int32)
-        # variable_block_sizes is defined per KV block (tile). Mask must therefore
-        # use kv_idx (not q_blk). Also, because we split each 64-token block into
-        # two 32-token halves, the mask must account for the half-block offset.
-        block_size = tl.load(variable_block_sizes + kv_idx).to(tl.int32)
-        half = (blk_idx % 2).to(tl.int32)
-        block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
+        block_sparse_offset = (tl.load(kv_ptr + blk_idx // 2).to(tl.int32) * 2 +
+                               blk_idx % 2) * step_n * stride_tok
        kT = tl.load(kT_ptrs + block_sparse_offset)
        vT = tl.load(vT_ptrs + block_sparse_offset)
        qk = tl.dot(q, kT)
        p = tl.math.exp2(qk - m)
-        offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
-        mask = offs_in_block < block_size
+        mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32)
        p = tl.where(mask[None, :], p, 0.0)
        # Compute dP and dS.
        dp = tl.dot(do, vT).to(tl.float32)
@@ -478,235 +467,19 @@ def _attn_bwd(
    tl.store(dq_ptrs, dq)


-@triton.jit
-def _attn_bwd_dkdv_kernel(
-        Q,
-        K,
-        V,
-        sm_scale,  #
-        DO,  #
-        DK,
-        DV,  #
-        M,
-        D,
-        k2q_index,
-        k2q_num,
-        max_q_blks,
-        variable_block_sizes,
-        # shared token/dim strides (assumed contiguous along token and dim)
-        stride_tok,
-        stride_d,  #
-        # batch/head strides (may differ between Q and KV)
-        stride_qz,
-        stride_qh,
-        stride_kz,
-        stride_kh,
-        stride_vz,
-        stride_vh,
-        stride_doz,
-        stride_doh,
-        stride_dkz,
-        stride_dkh,
-        stride_dvz,
-        stride_dvh,
-        H,
-        N_CTX_Q,
-        N_CTX_KV,
-        BLOCK_M1: tl.constexpr,  #
-        BLOCK_N1: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr):
-    """
-    Backward kernel that computes dK and dV for each KV block (64 tokens).
-    Grid:
-      pid0: kv_blk in [0, N_CTX_KV/BLOCK_N1)
-      pid2: fused (batch, head) in [0, B*H)
-    """
-    bhid = tl.program_id(2)
-    b = bhid // H
-    h = bhid % H
-    kv_blk = tl.program_id(0)
-
-    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
-    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
-    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
-    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
-    dk_adj = (b.to(tl.int64) * stride_dkz + h.to(tl.int64) * stride_dkh)
-    dv_adj = (b.to(tl.int64) * stride_dvz + h.to(tl.int64) * stride_dvh)
-
-    Q = Q + q_adj
-    K = K + kv_adj_k
-    V = V + kv_adj_v
-    DO = DO + do_adj
-    DK = DK + dk_adj
-    DV = DV + dv_adj
-
-    # M and D (delta) are always sized by Q length.
-    M = M + (bhid * N_CTX_Q).to(tl.int64)
-    D = D + (bhid * N_CTX_Q).to(tl.int64)
-
-    offs_k = tl.arange(0, HEAD_DIM)
-    start_n = kv_blk * BLOCK_N1
-    offs_n = start_n + tl.arange(0, BLOCK_N1)
-
-    # load K and V: they stay in SRAM throughout the inner loop.
-    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
-    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
-
-    dv_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
-    dk_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
-
-    num_steps = N_CTX_Q // BLOCK_M1
-    dk_acc, dv_acc = _attn_bwd_dkdv(
-        dk_acc,
-        dv_acc,
-        Q,
-        k,
-        v,
-        sm_scale,
-        DO,
-        M,
-        D,
-        k2q_index,
-        k2q_num,
-        max_q_blks,
-        variable_block_sizes,
-        stride_tok,
-        stride_d,
-        H,
-        N_CTX_KV,
-        BLOCK_M1=BLOCK_M1,
-        BLOCK_N1=BLOCK_N1,
-        HEAD_DIM=HEAD_DIM,
-        start_n=start_n,
-        start_m=0,
-        num_steps=num_steps,
-    )
-
-    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
-    tl.store(dv_ptrs, dv_acc)
-
-    dk_acc *= sm_scale
-    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
-    tl.store(dk_ptrs, dk_acc)
-
-
-@triton.jit
-def _attn_bwd_dq_kernel(
-        Q,
-        K,
-        V,
-        DO,  #
-        DQ,
-        M,
-        D,
-        q2k_index,
-        q2k_num,
-        max_kv_blks,
-        variable_block_sizes,
-        # shared token/dim strides (assumed contiguous along token and dim)
-        stride_tok,
-        stride_d,  #
-        # batch/head strides (may differ between Q and KV)
-        stride_qz,
-        stride_qh,
-        stride_kz,
-        stride_kh,
-        stride_vz,
-        stride_vh,
-        stride_doz,
-        stride_doh,
-        stride_dqz,
-        stride_dqh,
-        H,
-        N_CTX_Q,
-        BLOCK_M2: tl.constexpr,  #
-        BLOCK_N2: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr):
-    """
-    Backward kernel that computes dQ for each Q block (64 tokens).
-    Grid:
-      pid0: q_blk in [0, N_CTX_Q/BLOCK_M2)
-      pid2: fused (batch, head) in [0, B*H)
-    """
-    LN2 = 0.6931471824645996  # = ln(2)
-    bhid = tl.program_id(2)
-    b = bhid // H
-    h = bhid % H
-    q_blk = tl.program_id(0)
-
-    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
-    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
-    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
-    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
-    dq_adj = (b.to(tl.int64) * stride_dqz + h.to(tl.int64) * stride_dqh)
-
-    Q = Q + q_adj
-    K = K + kv_adj_k
-    V = V + kv_adj_v
-    DO = DO + do_adj
-    DQ = DQ + dq_adj
-
-    M = M + (bhid * N_CTX_Q).to(tl.int64)
-    D = D + (bhid * N_CTX_Q).to(tl.int64)
-
-    offs_k = tl.arange(0, HEAD_DIM)
-    start_m = q_blk * BLOCK_M2
-    offs_m = start_m + tl.arange(0, BLOCK_M2)
-
-    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
-    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
-    m = tl.load(M + offs_m)[:, None]
-
-    dq_acc = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
-    num_steps = 0  # unused in _attn_bwd_dq
-    dq_acc = _attn_bwd_dq(
-        dq_acc,
-        q,
-        K,
-        V,
-        do,
-        m,
-        D,
-        q2k_index,
-        q2k_num,
-        max_kv_blks,
-        variable_block_sizes,
-        stride_tok,
-        stride_d,
-        H,
-        N_CTX_Q,
-        BLOCK_M2=BLOCK_M2,
-        BLOCK_N2=BLOCK_N2,
-        HEAD_DIM=HEAD_DIM,
-        start_m=start_m,
-        start_n=0,
-        num_steps=num_steps,
-    )
-
-    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq_acc *= LN2
-    tl.store(dq_ptrs, dq_acc)
-
-
 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
 def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
                                     variable_block_sizes):
-    B, H, Tq, D = q.shape
-    Tkv = k.shape[2]
+    B, H, T, D = q.shape
    sm_scale = 1.0 / math.sqrt(D)
    max_kv_blks = q2k_index.shape[-1]
-    assert Tq % 64 == 0, f"q length must be a multiple of 64, but got {Tq}"
-    assert Tkv % 64 == 0, f"kv length must be a multiple of 64, but got {Tkv}"
+    assert T % 64 == 0, f"T must be a multiple of 64, but got {T}"
    assert q2k_num.shape[
-        -1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
-    assert variable_block_sizes.numel() == Tkv // 64, (
-        f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, "
-        f"got {variable_block_sizes.numel()}"
-    )
+        -1] == T // 64, f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
    o = torch.empty_like(q)
-    M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device)
+    M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)

-    grid = lambda _: (triton.cdiv(Tq, 64), B * H, 1)
+    grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
    _attn_fwd_sparse[grid](q,
                           k,
                           v,
@@ -735,8 +508,7 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
                           o.stride(3),
                           B,
                           H,
-                           Tq,
-                           Tkv,
+                           T,
                           HEAD_DIM=D,
                           STAGE=3)

@@ -746,21 +518,21 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
 def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
                                      k2q_index, k2q_num, variable_block_sizes):
    assert do.is_contiguous()
+    assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()

-    B, H, Tq, D = q.shape
-    Tkv = k.shape[2]
+    B, H, T, D = q.shape
    sm_scale = 1.0 / math.sqrt(D)
    dq = torch.empty_like(q)
    dk = torch.empty_like(k)
    dv = torch.empty_like(v)
-    BATCH, N_HEAD = q.shape[:2]
+    BATCH, N_HEAD, N_CTX = q.shape[:3]
    BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
    arg_k = k
    arg_k = arg_k * (sm_scale * RCP_LN2)
    PRE_BLOCK = 64
-    assert Tq % PRE_BLOCK == 0
-    pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
+    assert N_CTX % PRE_BLOCK == 0
+    pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
    delta = torch.empty_like(M)
    _attn_bwd_preprocess[pre_grid](
        o,
@@ -768,7 +540,7 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
        delta,  #
        BATCH,
        N_HEAD,
-        Tq,  #
+        N_CTX,  #
        BLOCK_M=PRE_BLOCK,
        HEAD_DIM=D  #
    )
@@ -776,75 +548,36 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
    max_q_blks = k2q_index.shape[-1]
    max_kv_blks = q2k_index.shape[-1]

-    # dK/dV kernel: grid over KV blocks
-    grid_kv = (Tkv // BLOCK_N1, 1, BATCH * N_HEAD)
-    _attn_bwd_dkdv_kernel[grid_kv](
+    grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+    _attn_bwd[grid](
        q,
        arg_k,
        v,
        sm_scale,
        do,
+        dq,
        dk,
-        dv,
+        dv,  #
        M,
-        delta,
+        delta,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
        k2q_index,
        k2q_num,
        max_q_blks,
        variable_block_sizes,
-        q.stride(2),
-        q.stride(3),
        q.stride(0),
        q.stride(1),
-        arg_k.stride(0),
-        arg_k.stride(1),
-        v.stride(0),
-        v.stride(1),
-        do.stride(0),
-        do.stride(1),
-        dk.stride(0),
-        dk.stride(1),
-        dv.stride(0),
-        dv.stride(1),
+        q.stride(2),
+        q.stride(3),  #
        N_HEAD,
-        Tq,
-        Tkv,
+        N_CTX,  #
        BLOCK_M1=BLOCK_M1,
-        BLOCK_N1=BLOCK_N1,
-        HEAD_DIM=D,
-    )
-
-    # dQ kernel: grid over Q blocks
-    grid_q = (Tq // BLOCK_M2, 1, BATCH * N_HEAD)
-    _attn_bwd_dq_kernel[grid_q](
-        q,
-        arg_k,
-        v,
-        do,
-        dq,
-        M,
-        delta,
-        q2k_index,
-        q2k_num,
-        max_kv_blks,
-        variable_block_sizes,
-        q.stride(2),
-        q.stride(3),
-        q.stride(0),
-        q.stride(1),
-        arg_k.stride(0),
-        arg_k.stride(1),
-        v.stride(0),
-        v.stride(1),
-        do.stride(0),
-        do.stride(1),
-        dq.stride(0),
-        dq.stride(1),
-        N_HEAD,
-        Tq,
+        BLOCK_N1=BLOCK_N1,  #
        BLOCK_M2=BLOCK_M2,
-        BLOCK_N2=BLOCK_N2,
-        HEAD_DIM=D,
+        BLOCK_N2=BLOCK_N2,  #
+        HEAD_DIM=D  #
    )

    return dq, dk, dv
@@ -1 +1 @@
-__version__ = "0.2.6"
+__version__ = "0.2.4"
@@ -350,7 +350,7 @@ def select_best_mask_strategy(


 def save_mask_search_results(
-        mask_search_final_result: list[Any],
+        mask_search_final_result: list[dict[str, list[float]]],
        prompt: str,
        mask_strategies: list[str],
        output_dir: str = 'output/mask_search_result/') -> str | None:
@@ -358,9 +358,8 @@ def save_mask_search_results(
        print("No mask search results to save")
        return None

-    # Create result dictionary with nested lists:
-    # [timesteps][layers][heads].
-    mask_search_dict: dict[str, dict[str, list[list[list[float]]]]] = {
+    # Create result dictionary with defaultdict for nested lists
+    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
        "L2_loss": defaultdict(list),
        "L1_loss": defaultdict(list)
    }
@@ -372,79 +371,23 @@ def save_mask_search_results(
        masks_list = [int(x) for x in mask.split(',')]
        selected_masks.append(masks_list)

-    def _to_float_list(loss_values: Any, loss_name: str) -> list[float]:
-        if isinstance(loss_values, np.ndarray):
-            loss_values = loss_values.tolist()
-        if not isinstance(loss_values, list | tuple):
-            raise ValueError(
-                f"{loss_name} must be a sequence of numeric values")
-
-        float_values = []
-        for loss in loss_values:
-            try:
-                float_values.append(float(loss))
-            except (TypeError, ValueError) as exc:
-                raise ValueError(
-                    f"Invalid {loss_name} value {loss!r}: expected a number"
-                ) from exc
-        return float_values
-
-    def _extract_timestep_layer_losses(step_data: Any, loss_name: str,
-                                       strategy_idx: int) -> list[list[float]]:
-        if isinstance(step_data, dict):
-            layer_data_list = [step_data]
-        elif isinstance(step_data, list):
-            layer_data_list = step_data
-        else:
-            return []
-
-        timestep_layer_losses: list[list[float]] = []
-        for layer_data in layer_data_list:
-            if not isinstance(layer_data, dict) or loss_name not in layer_data:
-                continue
-
-            raw_losses = layer_data[loss_name]
-            if isinstance(raw_losses, np.ndarray):
-                raw_losses = raw_losses.tolist()
-            if not isinstance(raw_losses, list | tuple):
-                raise ValueError(f"{loss_name} must be a list or tuple")
-            raw_losses = list(raw_losses)
-
-            if raw_losses and isinstance(raw_losses[0], list | tuple
-                                         | np.ndarray):
-                if strategy_idx >= len(raw_losses):
-                    raise ValueError(
-                        f"Missing strategy index {strategy_idx} in {loss_name}")
-                strategy_losses = raw_losses[strategy_idx]
-            else:
-                if strategy_idx > 0:
-                    continue
-                strategy_losses = raw_losses
-
-            timestep_layer_losses.append(
-                _to_float_list(strategy_losses, loss_name))
-
-        return timestep_layer_losses
-
    # Process each mask strategy
    for i, mask_strategy in enumerate(selected_masks):
        mask_strategy_str = str(mask_strategy)
-        l2_step_results: list[list[list[float]]] = []
-        l1_step_results: list[list[list[float]]] = []
-
+        # Process L2 loss
+        step_results: list[list[float]] = []
        for step_data in mask_search_final_result:
-            l2_layer_losses = _extract_timestep_layer_losses(
-                step_data, "L2_loss", i)
-            if l2_layer_losses:
-                l2_step_results.append(l2_layer_losses)
+            if isinstance(step_data, dict) and "L2_loss" in step_data:
+                layer_losses = [float(loss) for loss in step_data["L2_loss"]]
+                step_results.append(layer_losses)
+        mask_search_dict["L2_loss"][mask_strategy_str] = step_results

-            l1_layer_losses = _extract_timestep_layer_losses(
-                step_data, "L1_loss", i)
-            if l1_layer_losses:
-                l1_step_results.append(l1_layer_losses)
-
-        mask_search_dict["L2_loss"][mask_strategy_str] = l2_step_results
-        mask_search_dict["L1_loss"][mask_strategy_str] = l1_step_results
+        step_results = []
+        for step_data in mask_search_final_result:
+            if isinstance(step_data, dict) and "L1_loss" in step_data:
+                layer_losses = [float(loss) for loss in step_data["L1_loss"]]
+                step_results.append(layer_losses)
+        mask_search_dict["L1_loss"][mask_strategy_str] = step_results

    # Create the output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
@@ -92,59 +92,11 @@ class FlashAttentionImpl(AttentionImpl):
        value: torch.Tensor,
        attn_metadata: FlashAttnMetadata,
    ):
-
-        def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor,
-                                             key_len: int) -> torch.Tensor:
-            # Normalize attn_mask to [B, key_len] where True means valid token.
-            if attn_mask.dim() == 4:
-                attn_mask = attn_mask[:, 0, 0, :]
-            elif attn_mask.dim() == 3:
-                attn_mask = attn_mask[:, 0, :]
-            elif attn_mask.dim() != 2:
-                raise ValueError(
-                    f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}"
-                )
-
-            if attn_mask.dtype == torch.bool:
-                key_padding_mask = attn_mask
-            else:
-                # SDPA additive mask convention: valid=0, masked=-inf/large negative.
-                key_padding_mask = attn_mask >= 0
-
-            if key_padding_mask.shape[-1] != key_len:
-                raise ValueError(
-                    "Invalid key padding mask length for FLASH_ATTN: "
-                    f"expected {key_len}, got {key_padding_mask.shape[-1]}")
-            return key_padding_mask
-
        if attn_metadata is not None and hasattr(
                attn_metadata,
                "attn_mask") and attn_metadata.attn_mask is not None:
-            from fastvideo.attention.utils.flash_attn_no_pad import (
-                flash_attn_no_pad, flash_attn_varlen_qk_no_pad)
+            from fastvideo.attention.utils.flash_attn_no_pad import flash_attn_no_pad
            attn_mask = attn_metadata.attn_mask
-
-            # flash_attn_no_pad packs q/k/v as one tensor and assumes equal q/k
-            # sequence lengths. Cross-attention can violate this.
-            if query.shape[1] != key.shape[1]:
-                query_padding_mask = torch.ones(
-                    (query.shape[0], query.shape[1]),
-                    dtype=torch.bool,
-                    device=query.device,
-                )
-                key_padding_mask = _key_padding_mask_from_attn_mask(
-                    attn_mask, key.shape[1]).to(device=key.device)
-                return flash_attn_varlen_qk_no_pad(
-                    query,
-                    key,
-                    value,
-                    query_padding_mask=query_padding_mask,
-                    key_padding_mask=key_padding_mask,
-                    causal=self.causal,
-                    dropout_p=0.0,
-                    softmax_scale=self.softmax_scale,
-                )
-
            qkv = torch.stack([query, key, value], dim=2)

            attn_mask = F.pad(attn_mask, (qkv.shape[1] - attn_mask.shape[1], 0),
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0

 import torch
-from sageattn3 import sageattn3_blackwell
+from fastvideo.attention.backends.sageattn.api import sageattn_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 = sageattn3_blackwell(query, key, value, is_causal=self.causal)
+        output = sageattn_blackwell(query, key, value, is_causal=self.causal)
        output = output.transpose(1, 2)
        return output
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Will Lin	141a1140f6	refactor sampling pipeline	2026-01-20 15:40:53 -08:00
Shijie Wang	6294015389	Debug transformer output misalignment	2026-01-20 14:38:47 -08:00
Shijie Wang	e31b6c9e90	Fix OCR Rewards	2026-01-20 14:38:22 -08:00
Tamoghno Kandar	bf0ff21eeb	Fix OCR Rewards	2026-01-20 14:38:22 -08:00
Shijie Wang	6f937102ad	Enable validation videos	2026-01-20 14:38:22 -08:00
Tamoghno Kandar	0164e93019	Add Validation Loop	2026-01-20 14:38:21 -08:00
Shijie Wang	67e457aa92	resolved cuda OOM error	2026-01-20 14:38:21 -08:00
Shijie Wang	d795f0c443	remove additional sampling pipeline	2026-01-20 14:38:21 -08:00
loaydatrain	02452dd6e7	fixed dtype mismatch	2026-01-20 14:38:21 -08:00
Shijie Wang	91ef24bc14	update run script	2026-01-20 14:38:20 -08:00
Shijie Wang	e76e9fda15	minor fix	2026-01-20 14:38:20 -08:00
Shijie Wang	3b17f5a621	fix trajectory collection & reward computation	2026-01-20 14:38:20 -08:00
Shijie Wang	d758878705	minor fix	2026-01-20 14:38:19 -08:00
Shijie Wang	689e629420	Add entry point script	2026-01-20 14:38:19 -08:00
Shijie Wang	873dc9695f	Complete train_one_step and grpo policy loss	2026-01-20 14:38:19 -08:00
Shijie Wang	bfc0f46d61	Implement trajectories collection, reward and advantage computing	2026-01-20 14:38:18 -08:00
Shijie Wang	39907dbe4d	Port per-prompt stat tracker	2026-01-20 14:38:18 -08:00
Shijie Wang	abdd0c9b6a	Implement SDE step & SDE pipeline with log prob	2026-01-20 14:38:18 -08:00
Shijie (Jacob) Wang	f32a12200d	Refactor and trim down unnecessary RL args	2026-01-20 14:38:18 -08:00
Shijie Wang	f1d2c9e6b7	Add RL dataset & dataloader	2026-01-20 14:38:18 -08:00
Jiali Chen	450579cb42	init algorithm backbone and refactor rl_pipeline	2026-01-20 14:38:17 -08:00
Jiali Chen	26d7d6cc08	minor bug fix	2026-01-20 14:38:17 -08:00
Jiali Chen	44f0124eaa	refactor and add ocr reward model	2026-01-20 14:38:17 -08:00
Jiali Chen	d3ace51394	Phase 1 minor fixes	2026-01-20 14:38:17 -08:00
Jiali Chen	58954c660b	implement Phase 1 backbone code	2026-01-20 14:38:16 -08:00
				`@@ -1 +0,0 @@`
				`Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.`
				`@@ -1 +0,0 @@`
				`A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature.`
				`@@ -1 +0,0 @@`
				`A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.`
				`@@ -1 +0,0 @@`
				`A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.`
				`@@ -1 +0,0 @@`
				A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1
				`@@ -1 +0,0 @@`
				`fox in the forest close-up quickly turned its head to the left`
				`@@ -1 +0,0 @@`
				`Man walking his dog in the woods on a hot sunny day`
				`@@ -1 +0,0 @@`
				`A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic.`