add readme

README.md

@@ -1,153 +1,101 @@
-<div align="center">
-<img src=assets/logos/logo.svg width="30%"/>
-</div>
+# Sliding Tile Attention (STA) Branch
 
-<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://github.com/hao-ai-lab/FastVideo/discussions/1097" target="_blank"> <b> WeChat </b> </a> |
-</p>
+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.
 
-**FastVideo is a unified post-training and inference framework for accelerated video generation.**
+## What is STA
 
-## NEWS
+Sliding Tile Attention is an optimized attention backend for
+window-based video generation.
 
-- `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`
 
-### More News
+## 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/).
-
-## 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 achieve >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/support_matrix/) for full list of supported models, hardware assumptions, and optimization compatibility.
-
-## Getting Started
-
-We recommend using an environment manager such as `Conda` to create a clean environment:
+Install FastVideo from source:
 
 ```bash
-# Create and activate a new conda environment
-conda create -n fastvideo python=3.12
-conda activate fastvideo
-
-# Install FastVideo
-pip install fastvideo
+uv pip install -e .
 ```
 
-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.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/attention/vsa/#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:
+Build the STA kernel package:
 
 ```bash
-python example.py
+cd fastvideo-kernel
+./build.sh
+cd ..
 ```
 
-For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/).
+## Run STA Inference
 
-## More Guides
+STA backend:
 
-- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
-- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
-- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)
+```bash
+export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
+```
 
-## Awesome work using FastVideo or our research projects
+Ready-to-run examples:
 
-- [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.
-- [DanceGRPO](https://github.com/XueZeyue/DanceGRPO): A unified framework to adapt Group Relative Policy Optimization (GRPO) to visual generation paradigms. Code based on FastVideo.
-- [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.
-- [DCM](https://github.com/Vchitect/DCM): Dual-expert consistency model for efficient and high-quality video generation. Code based on FastVideo.
-- [HY-WorldPlay](https://github.com/Tencent-Hunyuan/HY-WorldPlay): An action-conditioned world model model trained using FastVideo framework.
-- [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.
-- [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.
-- [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.
+- HunyuanVideo: `scripts/inference/v1_inference_hunyuan_STA.sh`
+- Wan2.1-T2V-14B: `scripts/inference/v1_inference_wan_STA.sh`
 
-## 🤝 Contributing
+Run:
 
-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).
+```bash
+bash scripts/inference/v1_inference_hunyuan_STA.sh
+# or
+bash scripts/inference/v1_inference_wan_STA.sh
+```
 
-## Acknowledgement
+Both scripts already set STA-related env vars:
 
-We learned the design and reused code from the following projects: [Wan-Video](https://github.com/Wan-Video), [ThunderKittens](https://github.com/HazyResearch/ThunderKittens), [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). 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.
+- `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
+- `FASTVIDEO_ATTENTION_CONFIG` to an STA mask strategy JSON
+
+## 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
+```
 
 ## Citation
 
-If you find FastVideo useful, please consider citing our research work:
-
 ```bibtex
-@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},

assets/prompt_0.txt

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

assets/prompt_1.txt

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

assets/prompt_2.txt

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

assets/prompt_3.txt

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

assets/prompt_4.txt

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

assets/prompt_5.txt

@@ -0,0 +1 @@
+fox in the forest close-up quickly turned its head to the left

assets/prompt_6.txt

@@ -0,0 +1 @@
+Man walking his dog in the woods on a hot sunny day

assets/prompt_7.txt

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

examples/inference/sta_mask_search/inference_wan_sta.sh

@@ -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=1
+num_gpu=8
 gpu_ids=$(seq 0 $((num_gpu-1)))
 skip_time_steps=12
 

fastvideo/attention/backends/STA_configuration.py

@@ -350,7 +350,7 @@ def select_best_mask_strategy(
 
 
 def save_mask_search_results(
-        mask_search_final_result: list[dict[str, list[float]]],
+        mask_search_final_result: list[Any],
         prompt: str,
         mask_strategies: list[str],
         output_dir: str = 'output/mask_search_result/') -> str | None:
@@ -358,8 +358,9 @@ def save_mask_search_results(
         print("No mask search results to save")
         return None
 
-    # Create result dictionary with defaultdict for nested lists
-    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
+    # Create result dictionary with nested lists:
+    # [timesteps][layers][heads].
+    mask_search_dict: dict[str, dict[str, list[list[list[float]]]]] = {
         "L2_loss": defaultdict(list),
         "L1_loss": defaultdict(list)
     }
@@ -371,23 +372,79 @@ 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)
-        # Process L2 loss
-        step_results: list[list[float]] = []
-        for step_data in mask_search_final_result:
-            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
+        l2_step_results: list[list[list[float]]] = []
+        l1_step_results: list[list[list[float]]] = []
 
-        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
+            l2_layer_losses = _extract_timestep_layer_losses(
+                step_data, "L2_loss", i)
+            if l2_layer_losses:
+                l2_step_results.append(l2_layer_losses)
+
+            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
 
     # Create the output directory if it doesn't exist
     os.makedirs(output_dir, exist_ok=True)

fastvideo/pipelines/stages/denoising.py

@@ -722,20 +722,14 @@ class DenoisingStage(PipelineStage):
         from fastvideo.attention.backends.STA_configuration import save_mask_search_results
         if batch.mask_search_final_result_pos is not None and batch.prompt is not None:
             save_mask_search_results(
-                [
-                    dict(layer_data)
-                    for layer_data in batch.mask_search_final_result_pos
-                ],
+                batch.mask_search_final_result_pos,
                 prompt=str(batch.prompt),
                 mask_strategies=sparse_mask_candidates_searching,
                 output_dir=f'output/mask_search_result_pos_{size[0]}x{size[1]}/'
             )
         if batch.mask_search_final_result_neg is not None and batch.prompt is not None:
             save_mask_search_results(
-                [
-                    dict(layer_data)
-                    for layer_data in batch.mask_search_final_result_neg
-                ],
+                batch.mask_search_final_result_neg,
                 prompt=str(batch.prompt),
                 mask_strategies=sparse_mask_candidates_searching,
                 output_dir=f'output/mask_search_result_neg_{size[0]}x{size[1]}/'