--- title: attn_qat_train ---

fastvideo.attention.backends.attn_qat_train.AttentionBackend

AttentionBackend()

Abstract class for attention backends.

fastvideo.attention.backends.attn_qat_train.AttentionBackend.get_name

get_name() -> str

fastvideo.attention.backends.attn_qat_train.AttentionBackend.get_impl_cls

get_impl_cls() -> type[AttentionImpl]

fastvideo.attention.backends.attn_qat_train.AttentionBackend.get_metadata_cls

get_metadata_cls() -> type[AttentionMetadata]

fastvideo.attention.backends.attn_qat_train.AttentionBackend.get_builder_cls

get_builder_cls() -> type[AttentionMetadataBuilder]

fastvideo.attention.backends.attn_qat_train.AttentionImpl

AttentionImpl(num_heads: int, head_size: int, softmax_scale: float, causal: bool = False, num_kv_heads: int | None = None, prefix: str = '', **extra_impl_args)

fastvideo.attention.backends.attn_qat_train.AttentionImpl.preprocess_qkv

preprocess_qkv(qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor

Preprocess QKV tensor before performing attention operation.

Default implementation returns the tensor unchanged. Subclasses can override this to implement custom preprocessing like reshaping, tiling, scaling, or other transformations.

Called AFTER all_to_all for distributed attention

Parameters

qkv: torch.Tensor

The query-key-value tensor

attn_metadata: T

Metadata for the attention operation

Returns

torch.Tensor

Processed QKV tensor

fastvideo.attention.backends.attn_qat_train.AttentionImpl.postprocess_output

postprocess_output(output: torch.Tensor, attn_metadata: T) -> torch.Tensor

Postprocess the output tensor after the attention operation.

Default implementation returns the tensor unchanged. Subclasses can override this to implement custom postprocessing like untiling, scaling, or other transformations.

Called BEFORE all_to_all for distributed attention

Parameters

output: torch.Tensor

The output tensor from the attention operation

attn_metadata: T

Metadata for the attention operation

Returns

torch.Tensor

Postprocessed output tensor

fastvideo.attention.backends.attn_qat_train.AttentionImpl.forward

forward(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: T) -> torch.Tensor

fastvideo.attention.backends.attn_qat_train.AttentionMetadata

AttentionMetadata(current_timestep: int, *, VSA_sparsity: float = 0.0)

Attention metadata for prefill and decode batched together.

fastvideo.attention.backends.attn_qat_train.AttentionMetadata.asdict_zerocopy

asdict_zerocopy(skip_fields: set[str] | None = None) -> dict[str, Any]

Similar to dataclasses.asdict, but avoids deepcopying.

fastvideo.attention.backends.attn_qat_train.AttentionMetadataBuilder

AttentionMetadataBuilder()

Abstract class for attention metadata builders.

Create the builder, remember some configuration and parameters.

fastvideo.attention.backends.attn_qat_train.AttentionMetadataBuilder.prepare

prepare() -> None

Prepare for one batch.

fastvideo.attention.backends.attn_qat_train.AttentionMetadataBuilder.build

build(**kwargs: Any) -> AttentionMetadata

Build attention metadata with on-device tensors.

fastvideo.attention.backends.attn_qat_train.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.attention.backends.attn_qat_train.is_attn_qat_train_available

is_attn_qat_train_available() -> bool
Source
def is_attn_qat_train_available() -> bool:
    return _get_attn_qat_train_attention() is not None

fastvideo.attention.backends.attn_qat_train.attn_qat_train

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
Source
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:
        detail = f" Original import error: {_attn_qat_train_import_error}" if _attn_qat_train_import_error else ""
        raise ImportError("ATTN_QAT_TRAIN requires FastVideo's fastvideo-kernel package. Install it or make "
                          f"fastvideo-kernel/python importable.{detail}")

    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
    # Triton 3.7's NVWS pass aborts while compiling this kernel on Blackwell.
    # The kernel has a supported non-warp-specialized path, so use it on both
    # datacenter (sm_100) and consumer (sm_120) Blackwell GPUs.
    capability_major = torch.cuda.get_device_capability()[0]
    warp_specialize = capability_major not in (10, 12)
    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()

fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend

AttnQatTrainBackend()
Source
class AttnQatTrainBackend(AttentionBackend):

    accept_output_buffer: bool = True

    @staticmethod
    def get_supported_head_sizes() -> list[int]:
        return [128]

    @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

fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend.get_supported_head_sizes

get_supported_head_sizes() -> list[int]
Source
@staticmethod
def get_supported_head_sizes() -> list[int]:
    return [128]

fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend.get_name

get_name() -> str
Source
@staticmethod
def get_name() -> str:
    return "ATTN_QAT_TRAIN"

fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend.get_impl_cls

get_impl_cls() -> type[AttnQatTrainImpl]
Source
@staticmethod
def get_impl_cls() -> type["AttnQatTrainImpl"]:
    return AttnQatTrainImpl

fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend.get_metadata_cls

get_metadata_cls() -> type[AttentionMetadata]
Source
@staticmethod
def get_metadata_cls() -> type["AttentionMetadata"]:
    raise NotImplementedError

fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend.get_builder_cls

get_builder_cls() -> type[AttentionMetadataBuilder[AttentionMetadata]]
Source
@staticmethod
def get_builder_cls() -> type["AttentionMetadataBuilder[AttentionMetadata]"]:
    raise NotImplementedError

fastvideo.attention.backends.attn_qat_train.AttnQatTrainImpl

AttnQatTrainImpl(num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = '', **extra_impl_args)
Source
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)

fastvideo.attention.backends.attn_qat_train.AttnQatTrainImpl.forward

forward(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor
Source
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)