diff --git a/docs/design/overview.md b/docs/design/overview.md index 437348e24..e26e19749 100644 --- a/docs/design/overview.md +++ b/docs/design/overview.md @@ -169,8 +169,8 @@ How this maps to FastVideo: `FASTVIDEO_ATTENTION_BACKEND`. - SageAttention3 is split into two selectable backends: `SAGE_ATTN_THREE` for the regular upstream package and - `MODIFIED_SAGE_ATTN_THREE` for the FastVideoKernel-backed modified variant. -- `QAT_ATTN` is a separate FastVideoKernel Triton backend for the QAT attention + `ATTN_QAT_INFER` for the FastVideoKernel-backed inference variant. +- `ATTN_QAT_TRAIN` is a separate FastVideoKernel Triton backend for the QAT attention path. - `LocalAttention` is used for cross-attention and most attention layers. - `DistributedAttention` is used for full-sequence self-attention in the DiT. diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 13a477acc..fd73759aa 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -21,8 +21,8 @@ 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` -- Modified Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=MODIFIED_SAGE_ATTN_THREE` -- QAT Attention: `FASTVIDEO_ATTENTION_BACKEND=QAT_ATTN` +- Attn QAT Infer: `FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER` +- Attn QAT Train: `FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` - Video MoBA Attention: `FASTVIDEO_ATTENTION_BACKEND=VMOBA_ATTN` - Sparse Linear Attention: `FASTVIDEO_ATTENTION_BACKEND=SLA_ATTN` - SageSLA Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_SLA_ATTN` @@ -110,8 +110,8 @@ environment variable values: - `SAGE_ATTN_THREE`: the regular upstream SageAttention3 backend imported from the `sageattn3` package. -- `MODIFIED_SAGE_ATTN_THREE`: the modified SageAttention3 CUDA-kernel backend - imported from the in-repo `modified_sageattn` package. +- `ATTN_QAT_INFER`: the inference CUDA-kernel backend imported from the + in-repo `attn_qat_infer` package. **`SAGE_ATTN_THREE`** @@ -129,7 +129,7 @@ To use Sage Attention 3 in FastVideo, follow the `README.md` in the linked repos ### Modified Sage Attention 3 -**`MODIFIED_SAGE_ATTN_THREE`** +**`ATTN_QAT_INFER`** This backend uses the modified SageAttention3 implementation that lives in the `fastvideo-kernel` repository alongside the `fastvideo_kernel` Triton kernels. @@ -139,22 +139,22 @@ directly during inference. This backend currently assumes access to the in-repo `fastvideo-kernel` checkout or an equivalent editable/source install that exposes both: -- `modified_sageattn` +- `attn_qat_infer` Example: ```python -os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "MODIFIED_SAGE_ATTN_THREE" +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_INFER" ``` ### QAT Attention -**`QAT_ATTN`** +**`ATTN_QAT_TRAIN`** -This backend uses the FastVideoKernel Triton QAT attention implementation from -`fastvideo_kernel.triton_kernels.qat_attn`. Use it when you specifically want -the QAT/Triton attention path rather than the modified SageAttention3 CUDA -kernel path. +This backend uses the FastVideoKernel Triton attention implementation from +`fastvideo_kernel.triton_kernels.attn_qat_train`. Use it when you specifically +want the training-oriented Triton attention path rather than the +`attn_qat_infer` CUDA kernel path. This backend currently assumes access to an install that exposes: @@ -163,7 +163,7 @@ This backend currently assumes access to an install that exposes: Example: ```python -os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "QAT_ATTN" +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_TRAIN" ``` ### V-MoBA / SLA / SageSLA diff --git a/docs/utilities/debugging.md b/docs/utilities/debugging.md index 82be6ecf0..aeeaddd6d 100644 --- a/docs/utilities/debugging.md +++ b/docs/utilities/debugging.md @@ -28,7 +28,7 @@ Useful variables: - `FASTVIDEO_STAGE_LOGGING`: print per-stage timings during pipeline execution - `FASTVIDEO_ATTENTION_BACKEND`: force an attention backend (for example `TORCH_SDPA`, `FLASH_ATTN`, `SAGE_ATTN_THREE`, or - `MODIFIED_SAGE_ATTN_THREE`, or `QAT_ATTN`) + `ATTN_QAT_INFER`, or `ATTN_QAT_TRAIN`) ## Common Failure Modes @@ -55,9 +55,9 @@ If forcing a backend fails, verify optional dependencies are installed: `sta_do_not_delete` + `fastvideo-kernel` - `SAGE_ATTN`: SageAttention package - `SAGE_ATTN_THREE`: upstream `sageattn3` package -- `MODIFIED_SAGE_ATTN_THREE`: `fastvideo-kernel` checkout/source install that - exposes `modified_sageattn` -- `QAT_ATTN`: `fastvideo-kernel` install exposing `fastvideo_kernel` +- `ATTN_QAT_INFER`: `fastvideo-kernel` checkout/source install that exposes + `attn_qat_infer` +- `ATTN_QAT_TRAIN`: `fastvideo-kernel` install exposing `fastvideo_kernel` As a fallback, use: diff --git a/fastvideo-kernel/MANIFEST.in b/fastvideo-kernel/MANIFEST.in index 630963a5d..aa69fbbd7 100644 --- a/fastvideo-kernel/MANIFEST.in +++ b/fastvideo-kernel/MANIFEST.in @@ -2,6 +2,6 @@ include LICENSE include README.md include pyproject.toml recursive-include python/fastvideo_kernel *.py -recursive-include modified_sageattn *.py *.cu *.cuh *.cpp *.h +recursive-include attn_qat_infer *.py *.cu *.cuh *.cpp *.h recursive-include csrc *.cu *.cuh *.cpp *.h recursive-include include/tk *.cu *.cuh *.cpp *.h *.src diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 327692b6f..662881f0a 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -21,7 +21,7 @@ cd fastvideo-kernel ``` On supported Blackwell environments, the same install also packages -`modified_sageattn` and builds its `fp4attn_cuda` / `fp4quant_cuda` +`attn_qat_infer` and builds its `fp4attn_cuda` / `fp4quant_cuda` extensions automatically. This path requires CUDA Toolkit 12.8+ and targets `sm_120a`. diff --git a/fastvideo-kernel/modified_sageattn/__init__.py b/fastvideo-kernel/attn_qat_infer/__init__.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/__init__.py rename to fastvideo-kernel/attn_qat_infer/__init__.py diff --git a/fastvideo-kernel/modified_sageattn/api.py b/fastvideo-kernel/attn_qat_infer/api.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/api.py rename to fastvideo-kernel/attn_qat_infer/api.py diff --git a/fastvideo-kernel/modified_sageattn/blackwell/__init__.py b/fastvideo-kernel/attn_qat_infer/blackwell/__init__.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/__init__.py rename to fastvideo-kernel/attn_qat_infer/blackwell/__init__.py diff --git a/fastvideo-kernel/modified_sageattn/blackwell/api.cu b/fastvideo-kernel/attn_qat_infer/blackwell/api.cu similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/api.cu rename to fastvideo-kernel/attn_qat_infer/blackwell/api.cu diff --git a/fastvideo-kernel/modified_sageattn/blackwell/block_config.h b/fastvideo-kernel/attn_qat_infer/blackwell/block_config.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/block_config.h rename to fastvideo-kernel/attn_qat_infer/blackwell/block_config.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/block_info.h b/fastvideo-kernel/attn_qat_infer/blackwell/block_info.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/block_info.h rename to fastvideo-kernel/attn_qat_infer/blackwell/block_info.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/blockscaled_layout.h b/fastvideo-kernel/attn_qat_infer/blackwell/blockscaled_layout.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/blockscaled_layout.h rename to fastvideo-kernel/attn_qat_infer/blackwell/blockscaled_layout.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/cute_extension.h b/fastvideo-kernel/attn_qat_infer/blackwell/cute_extension.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/cute_extension.h rename to fastvideo-kernel/attn_qat_infer/blackwell/cute_extension.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/epilogue_tma_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/epilogue_tma_ws.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/epilogue_tma_ws.h rename to fastvideo-kernel/attn_qat_infer/blackwell/epilogue_tma_ws.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/kernel_traits.h b/fastvideo-kernel/attn_qat_infer/blackwell/kernel_traits.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/kernel_traits.h rename to fastvideo-kernel/attn_qat_infer/blackwell/kernel_traits.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/kernel_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/kernel_ws.h rename to fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/launch.h b/fastvideo-kernel/attn_qat_infer/blackwell/launch.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/launch.h rename to fastvideo-kernel/attn_qat_infer/blackwell/launch.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/mainloop_tma_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/mainloop_tma_ws.h rename to fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/named_barrier.h b/fastvideo-kernel/attn_qat_infer/blackwell/named_barrier.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/named_barrier.h rename to fastvideo-kernel/attn_qat_infer/blackwell/named_barrier.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/params.h b/fastvideo-kernel/attn_qat_infer/blackwell/params.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/params.h rename to fastvideo-kernel/attn_qat_infer/blackwell/params.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/softmax_fused.h b/fastvideo-kernel/attn_qat_infer/blackwell/softmax_fused.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/softmax_fused.h rename to fastvideo-kernel/attn_qat_infer/blackwell/softmax_fused.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/static_switch.h b/fastvideo-kernel/attn_qat_infer/blackwell/static_switch.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/static_switch.h rename to fastvideo-kernel/attn_qat_infer/blackwell/static_switch.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/tile_scheduler.h b/fastvideo-kernel/attn_qat_infer/blackwell/tile_scheduler.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/tile_scheduler.h rename to fastvideo-kernel/attn_qat_infer/blackwell/tile_scheduler.h diff --git a/fastvideo-kernel/modified_sageattn/blackwell/utils.h b/fastvideo-kernel/attn_qat_infer/blackwell/utils.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/blackwell/utils.h rename to fastvideo-kernel/attn_qat_infer/blackwell/utils.h diff --git a/fastvideo-kernel/modified_sageattn/quantization/__init__.py b/fastvideo-kernel/attn_qat_infer/quantization/__init__.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/quantization/__init__.py rename to fastvideo-kernel/attn_qat_infer/quantization/__init__.py diff --git a/fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_k.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_k.py rename to fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py diff --git a/fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_q.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_q.py rename to fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py diff --git a/fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_v.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_v.py rename to fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py diff --git a/fastvideo-kernel/modified_sageattn/quantization/bench/bench_utils.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py similarity index 100% rename from fastvideo-kernel/modified_sageattn/quantization/bench/bench_utils.py rename to fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py diff --git a/fastvideo-kernel/modified_sageattn/quantization/cuda_utils.h b/fastvideo-kernel/attn_qat_infer/quantization/cuda_utils.h similarity index 100% rename from fastvideo-kernel/modified_sageattn/quantization/cuda_utils.h rename to fastvideo-kernel/attn_qat_infer/quantization/cuda_utils.h diff --git a/fastvideo-kernel/modified_sageattn/quantization/fp4_quantization_4d.cu b/fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu similarity index 100% rename from fastvideo-kernel/modified_sageattn/quantization/fp4_quantization_4d.cu rename to fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu diff --git a/fastvideo-kernel/pyproject.toml b/fastvideo-kernel/pyproject.toml index 7bbd1af2a..923d1b41f 100644 --- a/fastvideo-kernel/pyproject.toml +++ b/fastvideo-kernel/pyproject.toml @@ -32,4 +32,4 @@ dependencies = [ [tool.scikit-build] cmake.build-type = "Release" minimum-version = "build-system.requires" -wheel.packages = ["python/fastvideo_kernel", "modified_sageattn"] +wheel.packages = ["python/fastvideo_kernel", "attn_qat_infer"] diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/qat_attn.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py similarity index 99% rename from fastvideo-kernel/python/fastvideo_kernel/triton_kernels/qat_attn.py rename to fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py index 159f586c3..32be67e51 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/qat_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py @@ -9,7 +9,7 @@ import triton.language as tl from triton.tools.tensor_descriptor import TensorDescriptor from .quant_utils import fake_quantize_q, fake_quantize_kv, fake_quantize -# from modified_sageattn.api import triton_group_mean +# from attn_qat_infer.api import triton_group_mean def is_cuda(): diff --git a/fastvideo-kernel/tests/test_qat_attn.py b/fastvideo-kernel/tests/test_attn_qat_train.py similarity index 97% rename from fastvideo-kernel/tests/test_qat_attn.py rename to fastvideo-kernel/tests/test_attn_qat_train.py index 07a6e51fc..20783f306 100644 --- a/fastvideo-kernel/tests/test_qat_attn.py +++ b/fastvideo-kernel/tests/test_attn_qat_train.py @@ -1,11 +1,11 @@ #!/usr/bin/env python3 """ -Simple test for QAT attention implementation. +Simple test for attn_qat_train. Tests forward and backward passes with and without QAT enabled. """ import torch -from fastvideo_kernel.triton_kernels.qat_attn import _attention +from fastvideo_kernel.triton_kernels.attn_qat_train import _attention from fastvideo_kernel.triton_kernels.fused_attention import attention as fused_attention from math import sqrt @@ -13,9 +13,9 @@ attention = _attention.apply DEVICE = torch.device("cuda") -def qat_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None): +def attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None): """ - Wrapper function that mimics qat_attn from sage_attn3.py. + Wrapper function that mimics attn_qat_train from the backend wrapper. Converts from BLHD format to BHLD format, calls attention, then converts back. Args: @@ -158,7 +158,7 @@ def test_qat_attention_forward(): v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) # Test with QAT (using wrapper that does permute/contiguous) - out_qat_BLHD = qat_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() @@ -209,7 +209,7 @@ def test_qat_attention_backward(): v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) # Forward pass with QAT (using wrapper) - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) @@ -288,7 +288,7 @@ def test_qat_attention_different_shapes(): v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) sm_scale = 1.0 / sqrt(HEAD_DIM) - out_qat_BLHD = qat_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() @@ -328,7 +328,7 @@ def test_qat_attention_non_causal(): k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - out_qat_BLHD = qat_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() @@ -368,7 +368,7 @@ def test_qat_attention_causal(): k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - out_qat_BLHD = qat_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() @@ -414,7 +414,7 @@ def test_qat_attention_causal_backward(): v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) # Forward pass with QAT - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) @@ -476,7 +476,7 @@ def test_qat_attention_different_seq_lengths(): v_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE) # Test with QAT (using wrapper that does permute/contiguous) - out_qat_BLHD = qat_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() @@ -530,7 +530,7 @@ def test_qat_attention_different_seq_lengths_backward(): v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) # Forward pass with QAT - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) # Forward pass with naive out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) @@ -923,7 +923,7 @@ def test_fused_vs_qat_attention_forward(): q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) # Convert QAT output back to float16 for comparison out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) @@ -971,7 +971,7 @@ def test_fused_vs_qat_attention_backward(): out_fused_BLHD = fused_attn_wrapper(q_fused_BLHD, k_fused_BLHD, v_fused_BLHD, causal, sm_scale) # Forward pass with QAT - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) # Create dummy gradient (same for both, in BLHD format) dout = torch.randn_like(out_fused_BLHD).contiguous() @@ -1053,7 +1053,7 @@ def test_fused_vs_qat_attention_different_shapes(): q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) @@ -1093,7 +1093,7 @@ def test_fused_vs_qat_attention_non_causal(): q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) @@ -1133,7 +1133,7 @@ def test_fused_vs_qat_attention_causal(): q_qat_BLHD = q_BLHD.clone().to(torch.bfloat16) k_qat_BLHD = k_BLHD.clone().to(torch.bfloat16) v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) @@ -1175,7 +1175,7 @@ def test_fused_vs_qat_attention_causal_backward(): out_fused_BLHD = fused_attn_wrapper(q_fused_BLHD, k_fused_BLHD, v_fused_BLHD, causal, sm_scale) # Forward pass with QAT - out_qat_BLHD = qat_attn_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) # Create dummy gradient dout = torch.randn_like(out_fused_BLHD).contiguous() @@ -1234,7 +1234,7 @@ def test_qat_attention_wan_shape_forward(): v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) # Test with QAT (using wrapper that does permute/contiguous) - out_qat_BLHD = qat_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() @@ -1293,7 +1293,7 @@ def test_qat_attention_wan_shape_backward(): # ----------------------------- # FORWARD — QAT # ----------------------------- - out_qat_BLHD = qat_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, causal, sm_scale) # ----------------------------- # FORWARD — NAIVE @@ -1395,7 +1395,7 @@ def test_qat_attention_wan_shape_non_divisible_64(): # ----------------------------- # FORWARD # ----------------------------- - out_qat_BLHD = qat_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, causal, sm_scale) + out_qat_BLHD = attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, causal, sm_scale) out_naive_BHLD = naive_attention(q_BHLD, k_BHLD, v_BHLD, causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0,2,1,3).contiguous() diff --git a/fastvideo-kernel/tests/test_fake_quant.py b/fastvideo-kernel/tests/test_fake_quant.py index 00be67f8b..2396c7cf5 100644 --- a/fastvideo-kernel/tests/test_fake_quant.py +++ b/fastvideo-kernel/tests/test_fake_quant.py @@ -62,7 +62,7 @@ def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr, dst_dtype: tl.constexpr, mx_tensor_dtype: tl.constexpr = tl.uint8): """ - Fake quantize function - matches API from qat_attn.py. + Fake quantize function - matches API from attn_qat_train.py. """ high_prec_src_tensor = src_tensor src_tensor, src_scale, src_s_dec = _compute_quant_and_scale( @@ -449,7 +449,7 @@ def test_fake_quantize_fp8(): def test_fake_quantize_attention_shapes(): - """Test fake_quantize with attention-like shapes (similar to test_qat_attn.py).""" + """Test fake_quantize with attention-like shapes (similar to test_attn_qat_train.py).""" torch.manual_seed(42) test_configs = [ diff --git a/fastvideo/attention/backends/modified_sage_attn3.py b/fastvideo/attention/backends/attn_qat_infer.py similarity index 63% rename from fastvideo/attention/backends/modified_sage_attn3.py rename to fastvideo/attention/backends/attn_qat_infer.py index f9b201269..c2f8472b6 100644 --- a/fastvideo/attention/backends/modified_sage_attn3.py +++ b/fastvideo/attention/backends/attn_qat_infer.py @@ -20,8 +20,8 @@ logger = init_logger(__name__) _project_root = Path(__file__).resolve().parent.parent.parent.parent _kernel_root = _project_root / "fastvideo-kernel" _kernel_python_root = _kernel_root / "python" -_sageattn_blackwell: Callable[..., torch.Tensor] | None = None -_sageattn_import_attempted = False +_attn_qat_infer: Callable[..., torch.Tensor] | None = None +_attn_qat_infer_import_attempted = False def _ensure_kernel_paths() -> None: @@ -31,30 +31,30 @@ def _ensure_kernel_paths() -> None: sys.path.insert(0, path_str) -def _get_sageattn_blackwell() -> Callable[..., torch.Tensor] | None: - global _sageattn_blackwell - global _sageattn_import_attempted +def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None: + global _attn_qat_infer + global _attn_qat_infer_import_attempted - if _sageattn_import_attempted: - return _sageattn_blackwell + if _attn_qat_infer_import_attempted: + return _attn_qat_infer - _sageattn_import_attempted = True + _attn_qat_infer_import_attempted = True _ensure_kernel_paths() try: # Prefer the in-repo kernel implementation during local development. - _sageattn_blackwell = importlib.import_module("modified_sageattn").sageattn_blackwell + _attn_qat_infer = importlib.import_module("attn_qat_infer").sageattn_blackwell except ImportError: - _sageattn_blackwell = None + _attn_qat_infer = None - return _sageattn_blackwell + return _attn_qat_infer -def is_modified_sageattn_available() -> bool: - return _get_sageattn_blackwell() is not None +def is_attn_qat_infer_available() -> bool: + return _get_attn_qat_infer() is not None -class ModifiedSageAttention3Backend(AttentionBackend): +class AttnQatInferBackend(AttentionBackend): accept_output_buffer: bool = True @@ -64,11 +64,11 @@ class ModifiedSageAttention3Backend(AttentionBackend): @staticmethod def get_name() -> str: - return "MODIFIED_SAGE_ATTN_THREE" + return "ATTN_QAT_INFER" @staticmethod - def get_impl_cls() -> type["ModifiedSageAttention3Impl"]: - return ModifiedSageAttention3Impl + def get_impl_cls() -> type["AttnQatInferImpl"]: + return AttnQatInferImpl @staticmethod def get_metadata_cls() -> type["AttentionMetadata"]: @@ -79,7 +79,7 @@ class ModifiedSageAttention3Backend(AttentionBackend): raise NotImplementedError -class ModifiedSageAttention3Impl(AttentionImpl): +class AttnQatInferImpl(AttentionImpl): def __init__( self, @@ -102,16 +102,16 @@ class ModifiedSageAttention3Impl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: - sageattn_blackwell = _get_sageattn_blackwell() - if sageattn_blackwell is None: - raise ImportError("modified_sageattn is not available. Please ensure the " - "modified SageAttention3 kernel is installed.") + attn_qat_infer = _get_attn_qat_infer() + if attn_qat_infer is None: + raise ImportError("attn_qat_infer is not available. Please ensure the " + "attn_qat_infer kernel package is installed.") query = query.transpose(1, 2).contiguous() key = key.transpose(1, 2).contiguous() value = value.transpose(1, 2).contiguous() - output = sageattn_blackwell( + output = attn_qat_infer( query, key, value, @@ -119,7 +119,3 @@ class ModifiedSageAttention3Impl(AttentionImpl): is_causal=self.causal, ) return output.transpose(1, 2).contiguous() - - -SageAttention3Backend = ModifiedSageAttention3Backend -SageAttention3Impl = ModifiedSageAttention3Impl diff --git a/fastvideo/attention/backends/modified_sageattn_3_setup.py b/fastvideo/attention/backends/attn_qat_infer_setup.py similarity index 95% rename from fastvideo/attention/backends/modified_sageattn_3_setup.py rename to fastvideo/attention/backends/attn_qat_infer_setup.py index 6081b3957..d4b7ecb02 100644 --- a/fastvideo/attention/backends/modified_sageattn_3_setup.py +++ b/fastvideo/attention/backends/attn_qat_infer_setup.py @@ -16,7 +16,7 @@ this_dir = os.path.dirname(os.path.abspath(__file__)) repo_dir = Path(this_dir).resolve() package_root = repo_dir.parents[2] / "fastvideo-kernel" -PACKAGE_NAME = "modified_sageattn" +PACKAGE_NAME = "attn_qat_infer" # FORCE_BUILD: Force a fresh build locally, instead of attempting to find prebuilt wheels # SKIP_CUDA_BUILD: Intended to allow CI to use a simple `python setup.py sdist` run to copy over raw files, without any cuda compilation @@ -72,7 +72,7 @@ if not SKIP_CUDA_BUILD: # https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920 if FORCE_CXX11_ABI: torch._C._GLIBCXX_USE_CXX11_ABI = True - modified_sageattn_dir = (package_root / "modified_sageattn") + attn_qat_infer_dir = (package_root / "attn_qat_infer") cutlass_dir = repo_dir / "csrc" / "cutlass" (repo_dir / "csrc").mkdir(parents=True, exist_ok=True) if not cutlass_dir.exists(): @@ -103,7 +103,7 @@ if not SKIP_CUDA_BUILD: "-DDQINRMEM", ] include_dirs = [ - modified_sageattn_dir, + attn_qat_infer_dir, cutlass_dir / "include", cutlass_dir / "tools" / "util" / "include", ] @@ -111,7 +111,7 @@ if not SKIP_CUDA_BUILD: ext_modules.append( CUDAExtension( name="fp4attn_cuda", - sources=[str(modified_sageattn_dir / "blackwell/api.cu")], + sources=[str(attn_qat_infer_dir / "blackwell/api.cu")], extra_compile_args={ "cxx": ["-O3", "-std=c++17"], "nvcc": append_nvcc_threads(nvcc_flags + ["-DEXECMODE=0"] + cc_flag), @@ -122,7 +122,7 @@ if not SKIP_CUDA_BUILD: ext_modules.append( CUDAExtension( name="fp4quant_cuda", - sources=[str(modified_sageattn_dir / "quantization/fp4_quantization_4d.cu")], + sources=[str(attn_qat_infer_dir / "quantization/fp4_quantization_4d.cu")], extra_compile_args={ "cxx": ["-O3", "-std=c++17"], "nvcc": append_nvcc_threads(nvcc_flags + ["-DEXECMODE=0"] + cc_flag), diff --git a/fastvideo/attention/backends/qat_attn.py b/fastvideo/attention/backends/attn_qat_train.py similarity index 68% rename from fastvideo/attention/backends/qat_attn.py rename to fastvideo/attention/backends/attn_qat_train.py index e09b7d2bc..271b78dbe 100644 --- a/fastvideo/attention/backends/qat_attn.py +++ b/fastvideo/attention/backends/attn_qat_train.py @@ -20,8 +20,8 @@ logger = init_logger(__name__) _project_root = Path(__file__).resolve().parent.parent.parent.parent _kernel_root = _project_root / "fastvideo-kernel" _kernel_python_root = _kernel_root / "python" -_qat_attention: Callable[..., torch.Tensor] | None = None -_qat_import_attempted = False +_attn_qat_train_attention: Callable[..., torch.Tensor] | None = None +_attn_qat_train_import_attempted = False def _ensure_kernel_paths() -> None: @@ -31,29 +31,32 @@ def _ensure_kernel_paths() -> None: sys.path.insert(0, path_str) -def _get_qat_attention() -> Callable[..., torch.Tensor] | None: - global _qat_attention - global _qat_import_attempted +def _get_attn_qat_train_attention() -> Callable[..., torch.Tensor] | None: + global _attn_qat_train_attention + global _attn_qat_train_import_attempted - if _qat_import_attempted: - return _qat_attention + if _attn_qat_train_import_attempted: + return _attn_qat_train_attention - _qat_import_attempted = True + _attn_qat_train_import_attempted = True _ensure_kernel_paths() try: - _qat_attention = importlib.import_module("fastvideo_kernel.triton_kernels.qat_attn").attention + _attn_qat_train_attention = importlib.import_module("fastvideo_kernel.triton_kernels.attn_qat_train").attention except ImportError: - _qat_attention = None + _attn_qat_train_attention = None - return _qat_attention + return _attn_qat_train_attention -def qat_attn(q_BLHD: torch.Tensor, k_BLHD: torch.Tensor, v_BLHD: torch.Tensor, is_causal: bool = False) -> torch.Tensor: - attention = _get_qat_attention() +def attn_qat_train(q_BLHD: torch.Tensor, + k_BLHD: torch.Tensor, + v_BLHD: torch.Tensor, + is_causal: bool = False) -> torch.Tensor: + attention = _get_attn_qat_train_attention() if attention is None: - raise ImportError("fastvideo_kernel.triton_kernels.qat_attn is not available. Please ensure the " - "FastVideo kernel package is installed.") + raise ImportError("fastvideo_kernel.triton_kernels.attn_qat_train is not available. " + "Please ensure the FastVideo kernel package is installed.") q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() @@ -91,7 +94,7 @@ def qat_attn(q_BLHD: torch.Tensor, k_BLHD: torch.Tensor, v_BLHD: torch.Tensor, i return o_BHLD.permute(0, 2, 1, 3).contiguous() -class QATAttentionBackend(AttentionBackend): +class AttnQatTrainBackend(AttentionBackend): accept_output_buffer: bool = True @@ -101,11 +104,11 @@ class QATAttentionBackend(AttentionBackend): @staticmethod def get_name() -> str: - return "QAT_ATTN" + return "ATTN_QAT_TRAIN" @staticmethod - def get_impl_cls() -> type["QATAttentionImpl"]: - return QATAttentionImpl + def get_impl_cls() -> type["AttnQatTrainImpl"]: + return AttnQatTrainImpl @staticmethod def get_metadata_cls() -> type["AttentionMetadata"]: @@ -116,7 +119,7 @@ class QATAttentionBackend(AttentionBackend): raise NotImplementedError -class QATAttentionImpl(AttentionImpl): +class AttnQatTrainImpl(AttentionImpl): def __init__( self, @@ -139,4 +142,4 @@ class QATAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: - return qat_attn(query, key, value, is_causal=self.causal) + return attn_qat_train(query, key, value, is_causal=self.causal) diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py index fe487de45..45cd4aae3 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py @@ -3,14 +3,14 @@ import sys import os import traceback -from modified_sageattn.api import ( +from attn_qat_infer.api import ( blockscaled_fp4_attn, preprocess_qkv, scale_and_quant_fp4, scale_and_quant_fp4_permute, scale_and_quant_fp4_transpose, ) -from modified_sageattn.quantization.bench.bench_utils import bench +from attn_qat_infer.quantization.bench.bench_utils import bench def calculate_attention_flops(batch_size, num_heads, seq_len_q, seq_len_k, head_dim, is_causal=False): diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py index e7af8d6d4..40fe11d2f 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py @@ -9,10 +9,10 @@ import numpy as np from typing import Dict, List, Tuple, Optional from flash_attn import flash_attn_func -from modified_sageattn.quantization.bench.bench_utils import bench +from attn_qat_infer.quantization.bench.bench_utils import bench # Import SageAttn components for direct control -from modified_sageattn.api import ( +from attn_qat_infer.api import ( preprocess_qkv, scale_and_quant_fp4, scale_and_quant_fp4_permute, diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py index fd79eabcc..462fe518e 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py @@ -4,7 +4,7 @@ import os import traceback from flash_attn import flash_attn_func -from modified_sageattn.quantization.bench.bench_utils import bench +from attn_qat_infer.quantization.bench.bench_utils import bench def calculate_attention_flops(batch_size, num_heads, seq_len_q, seq_len_k, head_dim, is_causal=False): diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py index 1649da084..34e3ce3a3 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py @@ -3,8 +3,8 @@ import sys import os import traceback -from modified_sageattn.api import sageattn_blackwell -from modified_sageattn.quantization.bench.bench_utils import bench +from attn_qat_infer.api import sageattn_blackwell +from attn_qat_infer.quantization.bench.bench_utils import bench def calculate_attention_flops(batch_size, num_heads, seq_len_q, seq_len_k, head_dim, is_causal=False): diff --git a/fastvideo/configs/models/dits/base.py b/fastvideo/configs/models/dits/base.py index 4fb8f80b9..315d4eb57 100644 --- a/fastvideo/configs/models/dits/base.py +++ b/fastvideo/configs/models/dits/base.py @@ -19,8 +19,8 @@ class DiTArchConfig(ArchConfig): AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE, - AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE, - AttentionBackendEnum.QAT_ATTN, AttentionBackendEnum.SLA_ATTN, + AttentionBackendEnum.ATTN_QAT_INFER, + AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN) hidden_size: int = 0 diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 1ccbd1d47..0f3efa0cb 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -201,8 +201,8 @@ environment_variables: dict[str, Callable[[], Any]] = { # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention # - "SAGE_ATTN": use Sage Attention # - "SAGE_ATTN_THREE": use Sage Attention 3 - # - "MODIFIED_SAGE_ATTN_THREE": use the in-repo modified Sage Attention 3 - # - "QAT_ATTN": use the FastVideoKernel Triton QAT attention backend + # - "ATTN_QAT_INFER": use the in-repo attn_qat_infer inference backend + # - "ATTN_QAT_TRAIN": use the FastVideoKernel Triton attn_qat_train backend "FASTVIDEO_ATTENTION_BACKEND": lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None), diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 4c028ef71..1f6d417de 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -200,7 +200,6 @@ class FastVideoArgs: logger.error("Failed to load V-MoBA config from %s: %s", self.moba_config_path, e) raise self._apply_ltx2_vae_overrides() - self._apply_ltx2_vae_overrides() self.check_fastvideo_args() def _apply_ltx2_vae_overrides(self) -> None: diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 7306860a3..eb09266a6 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -64,7 +64,7 @@ class DenoisingStage(PipelineStage): supported_attention_backends=( AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE, - AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE, AttentionBackendEnum.QAT_ATTN) # hack + AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.ATTN_QAT_TRAIN) # hack ) def forward( diff --git a/fastvideo/pipelines/stages/sr_denoising.py b/fastvideo/pipelines/stages/sr_denoising.py index aa62b2162..02c57a591 100644 --- a/fastvideo/pipelines/stages/sr_denoising.py +++ b/fastvideo/pipelines/stages/sr_denoising.py @@ -61,10 +61,10 @@ class SRDenoisingStage(PipelineStage): self.attn_backend = get_attn_backend( head_size=attn_head_size, dtype=torch.float16, # TODO(will): hack - supported_attention_backends=( - AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, - AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE, - AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE, AttentionBackendEnum.QAT_ATTN) # hack + supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, + AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.SAGE_ATTN_THREE, AttentionBackendEnum.ATTN_QAT_INFER, + AttentionBackendEnum.ATTN_QAT_TRAIN) # hack ) def add_noise_to_lq(self, lq_latents: torch.Tensor, strength: float = 0.7) -> torch.Tensor: diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index caf092201..a687824b5 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -140,31 +140,31 @@ class CudaPlatformBase(Platform): except ImportError as e: logger.info(e) logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.") - elif selected_backend == AttentionBackendEnum.MODIFIED_SAGE_ATTN_THREE: + elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER: try: - from fastvideo.attention.backends.modified_sage_attn3 import ( # noqa: F401 - ModifiedSageAttention3Backend, is_modified_sageattn_available, + from fastvideo.attention.backends.attn_qat_infer import ( # noqa: F401 + AttnQatInferBackend, is_attn_qat_infer_available, ) - if not is_modified_sageattn_available(): - raise ImportError("modified_sageattn could not be imported.") - logger.info("Using modified Sage Attention 3 backend.") + if not is_attn_qat_infer_available(): + raise ImportError("attn_qat_infer could not be imported.") + logger.info("Using attn_qat_infer backend.") - return "fastvideo.attention.backends.modified_sage_attn3.ModifiedSageAttention3Backend" + return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend" except ImportError as e: logger.info(e) - logger.info("Modified Sage Attention 3 backend is not installed. Fall back to Flash Attention.") - elif selected_backend == AttentionBackendEnum.QAT_ATTN: + logger.info("attn_qat_infer backend is not installed. Fall back to Flash Attention.") + elif selected_backend == AttentionBackendEnum.ATTN_QAT_TRAIN: try: - from fastvideo_kernel.triton_kernels.qat_attn import attention # noqa: F401 + from fastvideo_kernel.triton_kernels.attn_qat_train import attention # noqa: F401 - from fastvideo.attention.backends.qat_attn import ( # noqa: F401 - QATAttentionBackend) - logger.info("Using QAT attention backend.") + from fastvideo.attention.backends.attn_qat_train import ( # noqa: F401 + AttnQatTrainBackend) + logger.info("Using attn_qat_train backend.") - return "fastvideo.attention.backends.qat_attn.QATAttentionBackend" + return "fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend" except ImportError as e: logger.info(e) - logger.info("QAT attention backend is not installed. Fall back to Flash Attention.") + logger.info("attn_qat_train backend is not installed. Fall back to Flash Attention.") elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: try: from fastvideo_kernel import video_sparse_attn # noqa: F401 diff --git a/fastvideo/platforms/interface.py b/fastvideo/platforms/interface.py index 53e3169a0..a451de9c8 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/platforms/interface.py @@ -15,8 +15,8 @@ class AttentionBackendEnum(enum.Enum): TORCH_SDPA = enum.auto() SAGE_ATTN = enum.auto() SAGE_ATTN_THREE = enum.auto() - MODIFIED_SAGE_ATTN_THREE = enum.auto() - QAT_ATTN = enum.auto() + ATTN_QAT_INFER = enum.auto() + ATTN_QAT_TRAIN = enum.auto() VIDEO_SPARSE_ATTN = enum.auto() BSA_ATTN = enum.auto() VMOBA_ATTN = enum.auto() diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 94c167fc8..31b852a45 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -56,21 +56,21 @@ except Exception: logger = init_logger(__name__) -def _get_generator_qat_attn_impl_cls() -> type[Any]: - from fastvideo.attention.backends.qat_attn import QATAttentionImpl +def _get_generator_attn_qat_train_impl_cls() -> type[Any]: + from fastvideo.attention.backends.attn_qat_train import AttnQatTrainImpl - return QATAttentionImpl + return AttnQatTrainImpl -def swap_sage_attn3(obj: Any, obj_path: str) -> int: +def swap_attn_qat_train(obj: Any, obj_path: str) -> int: """ - If `obj` has an attribute named `attn_impl`, replace it with the QAT - attention implementation used by generator 4-bit attention, carrying over - `.causal` and `.softmax_scale`. + If `obj` has an attribute named `attn_impl`, replace it with the + attn_qat_train implementation used by generator 4-bit attention, carrying + over `.causal` and `.softmax_scale`. Returns number of swaps performed (0 or 1). """ - impl_cls = _get_generator_qat_attn_impl_cls() + impl_cls = _get_generator_attn_qat_train_impl_cls() # Quick reject: no attribute if not hasattr(obj, "attn_impl"): @@ -139,31 +139,8 @@ def swap_fp4_linear(obj: Any, obj_path: str) -> int: out_features = attr_value.output_size if should_replace and in_features is not None and out_features is not None: - # Create new LinearFWD4BWD16 layer with same dimensions - # new_layer = LinearFWD4BWD16( - # in_features=in_features, - # out_features=out_features, - # bias=bias, - # backend="cutlass", # Default backend - # block_size=16, # Default block size - # use_128x4_sf_layout=True # Default layout - # ) - - # Copy weights if possible - # if hasattr(attr_value, 'weight') and attr_value.weight is not None: - # with torch.no_grad(): - # new_layer.weight.copy_(attr_value.weight.data.float()) - - # if bias and hasattr(attr_value, 'bias') and attr_value.bias is not None: - # with torch.no_grad(): - # new_layer.bias.copy_(attr_value.bias.data.float()) layer = getattr(obj, attr_name) # Replace the attribute layer.forward = types.MethodType(fp4_linear_forward, layer) - - # from fpdb import ForkedPdb; ForkedPdb().set_trace() - # setattr(obj, attr_name, new_layer) - # setattr(new_layer, "weight", old_layer.weight) - # setattr(new_layer, "bias", old_layer.bias) swaps_performed += 1 return swaps_performed @@ -230,8 +207,8 @@ class TrainingPipeline(LoRAPipeline, ABC): self.transformer_2, checkpointing_type=training_args.enable_gradient_checkpointing_type) if training_args.generator_4bit_attn: - num_swaps = traverse_swap_module(self.transformer, swap_fn=swap_sage_attn3) - logger.info("Swapped %s attn_impl to QATAttentionImpl instances in self.transformer", num_swaps) + num_swaps = traverse_swap_module(self.transformer, swap_fn=swap_attn_qat_train) + logger.info("Swapped %s attn_impl to AttnQatTrainImpl instances in self.transformer", num_swaps) if training_args.generator_4bit_linear: num_swaps = traverse_swap_module(self.transformer, swap_fn=swap_fp4_linear) logger.info("Swapped %s linear layers to LinearFWD4BWD16 instances in self.transformer", num_swaps)