This commit is contained in:
RandNMR73
2026-04-09 09:32:53 +00:00
parent a3a4e0494a
commit 88f9ac1db4
4 changed files with 232 additions and 101 deletions
+221 -99
View File
@@ -1,59 +1,33 @@
# Attn QAT
# Attention QAT
Attn QAT in FastVideo currently has two related backends:
Attention QAT in FastVideo covers two related, but different, backends:
- `ATTN_QAT_TRAIN`: the training-oriented Triton attention path.
- `ATTN_QAT_INFER`: the inference-oriented CUDA kernel path.
- `ATTN_QAT_INFER`: the inference-oriented CUDA kernel path
- `ATTN_QAT_TRAIN`: the training-oriented Triton attention path
Both are selected with `FASTVIDEO_ATTENTION_BACKEND`, but they come from
different codepaths and serve different purposes.
Both are selected with `FASTVIDEO_ATTENTION_BACKEND`, but they are not
interchangeable. The main practical split is:
## Backend Overview
- use `ATTN_QAT_INFER` for standalone inference with the dedicated inference
kernel
- use `ATTN_QAT_TRAIN` for finetuning, validation during training, or when you
specifically want to reproduce the training-side attention path
| Backend | Best for | Package requirement | Primary kernel location |
|---------|----------|---------------------|-------------------------|
| `ATTN_QAT_TRAIN` | finetuning, validation during training, reproducing the training attention path | `fastvideo_kernel` | `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` |
| `ATTN_QAT_INFER` | standalone inference with the `attn_qat_infer` CUDA kernel | `attn_qat_infer` from the in-repo `fastvideo-kernel` checkout | `fastvideo-kernel/attn_qat_infer/` |
## Quick Start
In FastVideo itself, backend selection is routed through:
If your goal is "run Wan 2.1 14B with Attention QAT inference weights", this is
the shortest path:
- `fastvideo/envs.py`
- `fastvideo/platforms/cuda.py`
- `fastvideo/attention/backends/attn_qat_train.py`
- `fastvideo/attention/backends/attn_qat_infer.py`
1. Build the in-repo kernel package so FastVideo can import `attn_qat_infer`.
2. Download the Wan 2.1 14B QAT checkpoint.
3. Edit the provided inference example to point at the 14B base model and the
downloaded QAT safetensors.
4. Run the example with `ATTN_QAT_INFER`.
The legacy training pipeline also contains explicit Attn QAT integration:
### Step 1. Build the kernel package
- `fastvideo/training/training_pipeline.py`
That pipeline forces generator loading through `ATTN_QAT_TRAIN` when
`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` or
`--generator_4bit_attn` is enabled.
## Where The Code Lives
Use these paths when you want to trace or modify the Attn QAT flow:
| Location | Purpose |
|----------|---------|
| `fastvideo/attention/backends/attn_qat_train.py` | FastVideo wrapper that imports and calls the Triton training kernel |
| `fastvideo/attention/backends/attn_qat_infer.py` | FastVideo wrapper that imports and calls the inference kernel |
| `fastvideo-kernel/CMakeLists.txt` | Kernel build definition that compiles the `attn_qat_infer` inference extensions |
| `fastvideo/platforms/cuda.py` | Chooses the concrete attention backend at runtime |
| `fastvideo/envs.py` | Documents supported `FASTVIDEO_ATTENTION_BACKEND` values |
| `fastvideo/training/training_pipeline.py` | Training-time forcing logic for the generator attention backend |
| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` | Triton implementation for `ATTN_QAT_TRAIN` |
| `fastvideo-kernel/attn_qat_infer/api.py` | Python API entrypoint for the inference kernel |
| `fastvideo-kernel/benchmarks/benchmark_*.py` | Kernel-side attention benchmark scripts for FlashAttn2, SageAttn3, FP4, and comparison plots |
| `fastvideo-kernel/attn_qat_infer/blackwell/api.cu` | CUDA implementation behind `ATTN_QAT_INFER` |
| `fastvideo-kernel/tests/test_attn_qat_train.py` | Kernel-level test coverage for the training path |
| `examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 1.3B Attn QAT finetune launcher |
| `examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 14B Attn QAT finetune launcher |
## Build And Install
Before using either backend, make sure the in-repo kernel package is built from
source.
Before using either Attention QAT backend, build the in-repo
`fastvideo-kernel` package from source:
```bash
git submodule update --init --recursive
@@ -61,80 +35,125 @@ cd fastvideo-kernel
./build.sh
```
For local development, the important result is:
After a successful build:
- `ATTN_QAT_TRAIN` can import `fastvideo_kernel`
- `ATTN_QAT_INFER` can import `attn_qat_infer`
- `ATTN_QAT_TRAIN` should be able to import `fastvideo_kernel`
- `ATTN_QAT_INFER` should be able to import `attn_qat_infer`
Backend-specific notes:
`ATTN_QAT_INFER` currently targets the Blackwell CUDA path under
`fastvideo-kernel/attn_qat_infer/` and requires CUDA 12.8+.
- `ATTN_QAT_TRAIN` is the Triton training path shipped through
`fastvideo-kernel`.
- `ATTN_QAT_INFER` currently targets the Blackwell CUDA path in
`fastvideo-kernel/attn_qat_infer/`, is built by `fastvideo-kernel/CMakeLists.txt`,
and requires CUDA 12.8+.
### Step 2. Download the Wan 2.1 14B QAT checkpoint
For the full kernel build requirements and CUDA notes, see:
FastVideo includes a helper script:
- [Attention Overview](../index.md)
- `fastvideo-kernel/README.md`
- `examples/inference/optimizations/download_14B_qat.sh`
## Running Training
By default it downloads:
Today the checked-in Attn QAT training launchers use the legacy training
pipeline in `fastvideo/training/wan_training_pipeline.py`.
- Hugging Face repo: `FastVideo/14B_qat_400`
- local directory: `checkpoints/14B_qat_400`
### Ready-made launchers
Prerequisites:
Use the provided SLURM scripts directly:
- `huggingface_hub` installed, for example:
`uv pip install huggingface_hub`
- access to the model repo if it is private or gated:
`huggingface-cli login`
Run the downloader:
```bash
sbatch examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh
sbatch examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh
bash examples/inference/optimizations/download_14B_qat.sh
```
Both scripts already set:
To download into a custom directory:
```bash
export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
bash examples/inference/optimizations/download_14B_qat.sh /path/to/14B_qat_400
```
Before launching, update the script-local values that are environment-specific:
The script prints a ready-to-copy `init_weights_from_safetensors=...` value at
the end.
- `WANDB_API_KEY`
- `MODEL_PATH`
- `DATA_DIR`
- `VALIDATION_DATASET_FILE`
- output directory and SLURM resource requests
### Step 3. Edit the provided inference example
### What the launchers run
The example to start from is:
The training scripts eventually invoke:
- `examples/inference/optimizations/attn_qat_inference_example.py`
Open that file and update these two values:
1. Change the base model from `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` to
`Wan-AI/Wan2.1-T2V-14B-Diffusers`
2. Replace
`init_weights_from_safetensors="safetensors_path"` with the directory that
contains the downloaded `.safetensors` files
Example:
```python
import os
from fastvideo import VideoGenerator
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_INFER"
generator = VideoGenerator.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
num_gpus=1,
use_fsdp_inference=True,
dit_cpu_offload=False,
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
init_weights_from_safetensors="checkpoints/14B_qat_400",
)
```
Important:
- the checked-in example currently uses the `1.3B` base model until you edit it
- do not load the 14B QAT weights on top of the `1.3B` base model; the weights
and model config will not match
### Step 4. Run the inference example
```bash
torchrun fastvideo/training/wan_training_pipeline.py ...
python examples/inference/optimizations/attn_qat_inference_example.py
```
If you are adapting the workflow to your own cluster or running outside SLURM,
the key Attn QAT requirement is still the same:
Generated videos are written to `video_samples/` by default.
```bash
export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
```
## Backend Overview
Then launch the normal Wan training pipeline with your preferred `torchrun`
arguments and training flags.
| Backend | Best for | Package requirement | Primary kernel location |
|---------|----------|---------------------|-------------------------|
| `ATTN_QAT_TRAIN` | finetuning, training-time validation, reproducing the training path | `fastvideo_kernel` | `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` |
| `ATTN_QAT_INFER` | standalone inference with the dedicated CUDA kernel | `attn_qat_infer` from the in-repo `fastvideo-kernel` checkout | `fastvideo-kernel/attn_qat_infer/` |
## Running Inference
FastVideo routes backend selection through:
For standalone inference, prefer `ATTN_QAT_INFER` when that kernel is available.
Use `ATTN_QAT_TRAIN` only if you specifically want to exercise the training-side
attention path during validation or debugging.
- `fastvideo/envs.py`
- `fastvideo/platforms/cuda.py`
- `fastvideo/attention/backends/attn_qat_train.py`
- `fastvideo/attention/backends/attn_qat_infer.py`
### Python API
The legacy training pipeline also contains explicit Attention QAT integration:
This follows the same pattern as `examples/inference/basic/basic.py`, but with
the Attn QAT backend forced through the environment:
- `fastvideo/training/training_pipeline.py`
That pipeline forces generator loading through `ATTN_QAT_TRAIN` when
`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` or `--generator_4bit_attn` is
enabled.
## Inference Workflows
For standalone inference, prefer `ATTN_QAT_INFER` when the CUDA kernel is
available. Use `ATTN_QAT_TRAIN` for inference only if you intentionally want to
exercise the training-side attention path for debugging or parity checks.
### Minimal Python example
```python
import os
@@ -155,17 +174,35 @@ generator.generate_video(
)
```
To run inference with a finetuned checkpoint instead of the base model, replace
the model ID with your checkpoint path.
### Loading custom safetensors during inference
### CLI
FastVideo supports loading custom transformer weights through
`init_weights_from_safetensors`.
You can also force the backend when using the FastVideo CLI:
This value can point to either:
- a directory containing one or more `.safetensors` files
- a single `.safetensors` file
For Wan 2.1 14B QAT inference, the common pattern is:
```python
generator = VideoGenerator.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
num_gpus=1,
use_fsdp_inference=True,
init_weights_from_safetensors="checkpoints/14B_qat_400",
)
```
### CLI example
You can also force the backend from the command line:
```bash
FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER \
fastvideo generate \
--model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
--model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
--num-gpus 1 \
--sp-size 1 \
--tp-size 1 \
@@ -178,8 +215,87 @@ fastvideo generate \
--output-path outputs_video/
```
This is the same basic flow used by `scripts/inference/v1_inference_wan.sh`,
just with `FASTVIDEO_ATTENTION_BACKEND` set explicitly.
If you want to use custom QAT transformer weights from the CLI, pass the same
custom weight override that the Python API uses:
```bash
FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER \
fastvideo generate \
--model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \
--init-weights-from-safetensors checkpoints/14B_qat_400 \
--num-gpus 1 \
--output-path outputs_video/ \
--prompt "A cinematic close-up of rain on a neon street at night."
```
## Training Workflows
Today the checked-in Attention QAT training launchers use the legacy training
pipeline in `fastvideo/training/wan_training_pipeline.py`.
### Ready-made launchers
Use the provided SLURM scripts directly:
```bash
sbatch examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh
sbatch examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh
```
Both scripts already set:
```bash
export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
```
Before launching, update the script-local values that depend on your
environment:
- `WANDB_API_KEY`
- `MODEL_PATH`
- `DATA_DIR`
- `VALIDATION_DATASET_FILE`
- output directory and SLURM resource requests
### What the launchers run
The training scripts eventually invoke:
```bash
torchrun fastvideo/training/wan_training_pipeline.py ...
```
If you are adapting the workflow to your own cluster or running outside SLURM,
the main Attention QAT requirement is still:
```bash
export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
```
Then launch the normal Wan training pipeline with your preferred `torchrun`
arguments and training flags.
## Where The Code Lives
Use these paths when you want to trace or modify the Attention QAT flow:
| Location | Purpose |
|----------|---------|
| `fastvideo/attention/backends/attn_qat_train.py` | FastVideo wrapper that imports and calls the Triton training kernel |
| `fastvideo/attention/backends/attn_qat_infer.py` | FastVideo wrapper that imports and calls the inference kernel |
| `fastvideo-kernel/CMakeLists.txt` | Kernel build definition that compiles the `attn_qat_infer` inference extensions |
| `fastvideo/platforms/cuda.py` | Chooses the concrete attention backend at runtime |
| `fastvideo/envs.py` | Documents supported `FASTVIDEO_ATTENTION_BACKEND` values |
| `fastvideo/training/training_pipeline.py` | Training-time forcing logic for the generator attention backend |
| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` | Triton implementation for `ATTN_QAT_TRAIN` |
| `fastvideo-kernel/attn_qat_infer/api.py` | Python API entrypoint for the inference kernel |
| `fastvideo-kernel/benchmarks/benchmark_*.py` | Kernel-side benchmark scripts for FlashAttn2, SageAttention3, FP4, and comparison plots |
| `fastvideo-kernel/attn_qat_infer/blackwell/api.cu` | CUDA implementation behind `ATTN_QAT_INFER` |
| `fastvideo-kernel/tests/test_attn_qat_train.py` | Kernel-level test coverage for the training path |
| `examples/inference/optimizations/attn_qat_inference_example.py` | Ready-to-edit inference example for custom Attention QAT weights |
| `examples/inference/optimizations/download_14B_qat.sh` | Helper script for downloading the Wan 2.1 14B QAT checkpoint |
| `examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 1.3B Attention QAT finetune launcher |
| `examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 14B Attention QAT finetune launcher |
## Troubleshooting
@@ -187,10 +303,16 @@ just with `FASTVIDEO_ATTENTION_BACKEND` set explicitly.
successfully and exposes `fastvideo_kernel`.
- If `ATTN_QAT_INFER` fails to import, verify that the local build exposes the
`attn_qat_infer` package.
- If the Wan 2.1 14B example fails after you changed only the checkpoint path,
make sure you also changed the base model to
`Wan-AI/Wan2.1-T2V-14B-Diffusers`.
- If you hit issues with CPU memory pressure or obscure CUDA argument errors in
the example script, try setting `pin_cpu_memory=False`.
- If you want a known-safe fallback for debugging, use
`FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`.
See also:
## Related Pages
- [Attention Overview](../index.md)
- [Inference Optimizations](../../inference/optimizations.md)
- [Debugging](../../utilities/debugging.md)
+2 -1
View File
@@ -5,7 +5,8 @@ FastVideo provides highly optimized custom attention kernels to accelerate video
## Supported Kernels
* **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks.
* **[Attn QAT](attn_qat/index.md)**: Training and inference backends for the Attn QAT workflow.
* **[Attention QAT](attn_qat/index.md)**: Dedicated guide for Attention QAT
inference, training, checkpoint loading, and troubleshooting.
* **[Sliding Tile Attention (STA)](sta/index.md)**: STA kernel support is kept in
`fastvideo-kernel`; full FastVideo STA pipeline workflow is archived in
`sta_do_not_delete`.
+8
View File
@@ -136,6 +136,10 @@ This backend uses the `attn_qat_infer` implementation that lives in the
Use this backend when you want to run the dedicated FP4 inference CUDA kernel
directly during inference.
For the full Attention QAT guide, including Wan 2.1 14B checkpoint download,
example editing steps, training launchers, and troubleshooting, see
[Attention QAT](../attention/attn_qat/index.md).
This backend currently assumes access to the in-repo `fastvideo-kernel`
checkout or an equivalent editable/source install that exposes:
@@ -156,6 +160,10 @@ This backend uses the FastVideoKernel Triton attention implementation from
want the training-oriented Triton attention path rather than the
`attn_qat_infer` CUDA kernel path.
The dedicated [Attention QAT](../attention/attn_qat/index.md) page covers when
to use `ATTN_QAT_TRAIN` versus `ATTN_QAT_INFER`, the ready-made training
launchers, and the end-to-end Wan 2.1 14B inference workflow.
This backend currently assumes access to an install that exposes:
- `fastvideo_kernel`
+1 -1
View File
@@ -152,7 +152,7 @@ nav:
- DMD: distillation/dmd.md
- Attention:
- Overview: attention/index.md
- Attn QAT: attention/attn_qat/index.md
- Attention QAT: attention/attn_qat/index.md
- Video Sparse Attention: attention/vsa/index.md
- Sliding Tile Attention (Archived): attention/sta/index.md
- Backend Development: contributing/attention_backend.md