fix

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.

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

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:
 

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

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

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"]

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

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()
 

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 = [

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

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),

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)

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

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, 

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

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

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

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),
 

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:

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(

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:

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

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()

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)