[feat]: Attn-QAT inference + training backends (deadcode) (Attn-QAT 4/12)

Add fastvideo/attention/backends/attn_qat_{infer,train}.py extracted from PR hao-ai-lab#1225.
Backends not yet registered in the selector — full deadcode until activation (slice 12).

Attn-QAT-Stack: 4/12

Co-Authored-By: jzhang38 <42993249+jzhang38@users.noreply.github.com>
Co-Authored-By: RandNMR73 <99706358+RandNMR73@users.noreply.github.com>

Author: 3 people

fastvideo/attention/backends/attn_qat_infer.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib
+import sys
+from collections.abc import Callable
+from pathlib import Path
+
+import torch
+
+from fastvideo.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from fastvideo.logger import init_logger
+
+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"
+_attn_qat_infer: Callable[..., torch.Tensor] | None = None
+_attn_qat_infer_import_attempted = False
+
+
+def _ensure_kernel_paths() -> None:
+    for path in (_project_root, _kernel_root, _kernel_python_root):
+        path_str = str(path)
+        if path_str not in sys.path:
+            sys.path.insert(0, path_str)
+
+
+def _get_attn_qat_infer() -> Callable[..., torch.Tensor] | None:
+    global _attn_qat_infer
+    global _attn_qat_infer_import_attempted
+
+    if _attn_qat_infer_import_attempted:
+        return _attn_qat_infer
+
+    _attn_qat_infer_import_attempted = True
+    _ensure_kernel_paths()
+
+    try:
+        # Prefer the in-repo kernel implementation during local development.
+        _attn_qat_infer = importlib.import_module("attn_qat_infer").sageattn_blackwell
+    except ImportError:
+        _attn_qat_infer = None
+
+    return _attn_qat_infer
+
+
+def is_attn_qat_infer_available() -> bool:
+    return _get_attn_qat_infer() is not None
+
+
+class AttnQatInferBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_name() -> str:
+        return "ATTN_QAT_INFER"
+
+    @staticmethod
+    def get_impl_cls() -> type["AttnQatInferImpl"]:
+        return AttnQatInferImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder[AttentionMetadata]"]:
+        raise NotImplementedError
+
+
+class AttnQatInferImpl(AttentionImpl[AttentionMetadata]):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        dropout_p = extra_impl_args.get("dropout_p", 0.0)
+        if dropout_p > 0:
+            raise NotImplementedError(f"attn_qat_infer does not support dropout (got dropout_p={dropout_p}). "
+                                      "The QAT inference kernel applies no stochastic dropout.")
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        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 = attn_qat_infer(
+            query,
+            key,
+            value,
+            attn_mask=None,
+            is_causal=self.causal,
+            sm_scale=self.softmax_scale,
+        )
+        return output.transpose(1, 2).contiguous()

fastvideo/attention/backends/attn_qat_train.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib
+import sys
+from collections.abc import Callable
+from pathlib import Path
+
+import torch
+
+from fastvideo.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from fastvideo.logger import init_logger
+
+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"
+_attn_qat_train_attention: Callable[..., torch.Tensor] | None = None
+_attn_qat_train_import_attempted = False
+
+
+def _ensure_kernel_paths() -> None:
+    for path in (_project_root, _kernel_root, _kernel_python_root):
+        path_str = str(path)
+        if path_str not in sys.path:
+            sys.path.insert(0, path_str)
+
+
+def _get_attn_qat_train_attention() -> Callable[..., torch.Tensor] | None:
+    global _attn_qat_train_attention
+    global _attn_qat_train_import_attempted
+
+    if _attn_qat_train_import_attempted:
+        return _attn_qat_train_attention
+
+    _attn_qat_train_import_attempted = True
+    _ensure_kernel_paths()
+
+    try:
+        _attn_qat_train_attention = importlib.import_module("fastvideo_kernel.triton_kernels.attn_qat_train").attention
+    except ImportError:
+        _attn_qat_train_attention = None
+
+    return _attn_qat_train_attention
+
+
+def attn_qat_train(q_BLHD: torch.Tensor,
+                   k_BLHD: torch.Tensor,
+                   v_BLHD: torch.Tensor,
+                   is_causal: bool = False,
+                   sm_scale: float | None = None) -> torch.Tensor:
+    attention = _get_attn_qat_train_attention()
+    if attention is None:
+        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()
+    v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous()
+
+    use_qat_qkv_backward = True
+    smooth_k = False
+    warp_specialize = True
+    is_qat = True
+    two_level_quant_p_sage3 = False
+    fake_quant_p_bwd = True
+    use_high_prec_o = True
+    smooth_q = False
+    if sm_scale is None:
+        sm_scale = 1.0 / (q_BHLD.shape[-1]**0.5)
+    use_global_sf_qkv = False
+    use_global_sf_p = False
+
+    o_BHLD = attention(
+        q_BHLD,
+        k_BHLD,
+        v_BHLD,
+        is_causal,
+        sm_scale,
+        use_qat_qkv_backward,
+        smooth_k,
+        warp_specialize,
+        is_qat,
+        two_level_quant_p_sage3,
+        fake_quant_p_bwd,
+        use_high_prec_o,
+        smooth_q,
+        use_global_sf_p,
+        use_global_sf_qkv,
+    )
+    return o_BHLD.permute(0, 2, 1, 3).contiguous()
+
+
+class AttnQatTrainBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "ATTN_QAT_TRAIN"
+
+    @staticmethod
+    def get_impl_cls() -> type["AttnQatTrainImpl"]:
+        return AttnQatTrainImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder[AttentionMetadata]"]:
+        raise NotImplementedError
+
+
+class AttnQatTrainImpl(AttentionImpl[AttentionMetadata]):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        dropout_p = extra_impl_args.get("dropout_p", 0.0)
+        if dropout_p > 0:
+            raise NotImplementedError(f"attn_qat_train does not support dropout (got dropout_p={dropout_p}). "
+                                      "The QAT training kernel applies no stochastic dropout.")
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        return attn_qat_train(query, key, value, is_causal=self.causal, sm_scale=self.softmax_scale)