[feat] QAD 5090: Wire the Attn-QAT training attention backend (10/12)

Make AttnQatTrainBackend (deadcode from #1358) selectable for
quantization-aware finetuning, mirroring the inference wiring in #1457. The
fake-quant-in-backward attention path becomes reachable via
FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN.

- platforms/interface.py: add ATTN_QAT_TRAIN to AttentionBackendEnum.
- platforms/cuda.py: dispatch ATTN_QAT_TRAIN -> AttnQatTrainBackend, guarded by
  is_attn_qat_train_available().
- attention/backends/attn_qat_train.py: add is_attn_qat_train_available()
  helper (mirrors is_attn_qat_infer_available()).
- configs/models/dits/base.py: add ATTN_QAT_TRAIN to the default supported set.

Config-driven, no monkey-patch module swapping. The training Triton kernel
(fastvideo_kernel.triton_kernels.attn_qat_train) is not upstreamed yet, so for
now selecting ATTN_QAT_TRAIN logs a warning and falls back to Flash Attention
until a follow-up lands the kernel.

Depends on #1457 (overlaps the enum / dispatch / supported-list files).

Author: alexzms

fastvideo/attention/backends/attn_qat_train.py

@@ -49,6 +49,10 @@ def _get_attn_qat_train_attention() -> Callable[..., torch.Tensor] | None:
     return _attn_qat_train_attention
 
 
+def is_attn_qat_train_available() -> bool:
+    return _get_attn_qat_train_attention() is not None
+
+
 def attn_qat_train(q_BLHD: torch.Tensor,
                    k_BLHD: torch.Tensor,
                    v_BLHD: torch.Tensor,

fastvideo/configs/models/dits/base.py

@@ -24,7 +24,8 @@ class DiTArchConfig(ArchConfig):
                                                  AttentionBackendEnum.TORCH_SDPA,
                                                  AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                                                  AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
-                                                 AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN)
+                                                 AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.SLA_ATTN,
+                                                 AttentionBackendEnum.SAGE_SLA_ATTN)
 
     hidden_size: int = 0
     num_attention_heads: int = 0

fastvideo/platforms/cuda.py

@@ -140,6 +140,14 @@ def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, hea
             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.ATTN_QAT_TRAIN:
+            from fastvideo.attention.backends.attn_qat_train import (  # noqa: F401
+                AttnQatTrainBackend, is_attn_qat_train_available)
+            if is_attn_qat_train_available():
+                logger.info("Using Attn-QAT training (fake-quantized attention) backend.")
+                return "fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend"
+            logger.warning("Attn-QAT training kernel is not built; falling back to Flash Attention "
+                           "(NO fake-quant in the attention path).")
         elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
             try:
                 from fastvideo_kernel import video_sparse_attn  # noqa: F401

fastvideo/platforms/interface.py

@@ -15,6 +15,7 @@ class AttentionBackendEnum(enum.Enum):
     TORCH_SDPA = enum.auto()
     SAGE_ATTN = enum.auto()
     SAGE_ATTN_THREE = enum.auto()
+    ATTN_QAT_TRAIN = enum.auto()
     VIDEO_SPARSE_ATTN = enum.auto()
     BSA_ATTN = enum.auto()
     VMOBA_ATTN = enum.auto()