--- title: attn_qat_infer ---

fastvideo.attention.backends.attn_qat_infer.AttentionBackend

AttentionBackend()

Abstract class for attention backends.

fastvideo.attention.backends.attn_qat_infer.AttentionBackend.get_name

get_name() -> str

fastvideo.attention.backends.attn_qat_infer.AttentionBackend.get_impl_cls

get_impl_cls() -> type[AttentionImpl]

fastvideo.attention.backends.attn_qat_infer.AttentionBackend.get_metadata_cls

get_metadata_cls() -> type[AttentionMetadata]

fastvideo.attention.backends.attn_qat_infer.AttentionBackend.get_builder_cls

get_builder_cls() -> type[AttentionMetadataBuilder]

fastvideo.attention.backends.attn_qat_infer.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_infer.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_infer.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_infer.AttentionImpl.forward

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

fastvideo.attention.backends.attn_qat_infer.AttentionMetadata

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

Attention metadata for prefill and decode batched together.

fastvideo.attention.backends.attn_qat_infer.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_infer.AttentionMetadataBuilder

AttentionMetadataBuilder()

Abstract class for attention metadata builders.

Create the builder, remember some configuration and parameters.

fastvideo.attention.backends.attn_qat_infer.AttentionMetadataBuilder.prepare

prepare() -> None

Prepare for one batch.

fastvideo.attention.backends.attn_qat_infer.AttentionMetadataBuilder.build

build(**kwargs: Any) -> AttentionMetadata

Build attention metadata with on-device tensors.

fastvideo.attention.backends.attn_qat_infer.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_infer.validate_fa4_pv_mode

validate_fa4_pv_mode(mode: str) -> str
Source
def validate_fa4_pv_mode(mode: str) -> str:
    # Fail fast at impl construction so a typo never survives to the first
    # forward on a Blackwell box.
    if mode not in _FA4_PV_MODES:
        raise ValueError(f"fa4_pv_mode must be one of {_FA4_PV_MODES}, got {mode!r}")
    return mode

fastvideo.attention.backends.attn_qat_infer.attn_qat_infer_receipt

attn_qat_infer_receipt() -> str

One-line receipt of the resolution decision (arch + kernel + quant knobs), for the selection log and for tooling. qk_mode=nvfp4 (per-16 E4M3 SFs) is the repo's tuned default passed through verbatim; pv_mode is the last configured fa4_pv_mode (process-wide, recorded at impl construction -- not a per-layer value), and the dtype actually fed to the kernel is logged once on the first FA4 forward -- see flash_attn/cute/README.md in the kernel repo.

Source
def attn_qat_infer_receipt() -> str:
    """One-line receipt of the resolution decision (arch + kernel + quant
    knobs), for the selection log and for tooling. qk_mode=nvfp4 (per-16 E4M3
    SFs) is the repo's tuned default passed through verbatim; pv_mode is the
    *last configured* fa4_pv_mode (process-wide, recorded at impl
    construction -- not a per-layer value), and the dtype actually fed to the
    kernel is logged once on the first FA4 forward -- see
    flash_attn/cute/README.md in the kernel repo."""
    cap = _active_capability()
    arch = f"sm_{cap[0]}{cap[1]}" if cap is not None else "no-cuda"
    kernel = _resolved_kernel()
    if kernel == "cutlass_sm12x":
        return f"arch={arch} kernel=fastvideo-kernel-cutlass scheme=sage3-fp4-sm120"
    if kernel == "fa4_fp4":
        return (f"arch={arch} kernel=flash-attention-fp4 qk_mode=nvfp4(per-16-e4m3-sf) "
                f"pv_mode={_current_fa4_pv_mode()} train_sim_mismatch=measured")
    supported = "sm_120a/sm_121a via fastvideo-kernel build.sh; sm_100a/sm_103a via flash-attention-fp4"
    if cap is not None and cap in _FA4_FP4_CAPABILITIES:
        return f"arch={arch} kernel=none (flash_attn.cute not importable -- {_FA4_INSTALL_HINT})"
    return f"arch={arch} kernel=none (supported: {supported})"

fastvideo.attention.backends.attn_qat_infer.is_attn_qat_infer_available

is_attn_qat_infer_available() -> bool

True only when the active device has a built ATTN_QAT_INFER kernel.

The import check alone is not sufficient: CUDA 13 wheel builds can carry the sm_12x extension on any host (e.g. H100, GB200), where the import succeeds, backend selection picks this backend, and the first kernel call then fails with an unsupported-capability error instead of ever reaching the documented FlashAttention fallback in fastvideo.platforms.cuda. Gating on the active device's capability keeps that fallback working on every unsupported GPU, while sm_100a/sm_103a now resolve to the FP4 FA4 kernel (#1221).

Source
def is_attn_qat_infer_available() -> bool:
    """True only when the active device has a built ATTN_QAT_INFER kernel.

    The import check alone is not sufficient: CUDA 13 wheel builds can
    carry the sm_12x extension on any host (e.g. H100, GB200), where the
    import succeeds, backend selection picks this backend, and the first
    kernel call then fails with an unsupported-capability error instead of
    ever reaching the documented FlashAttention fallback in
    fastvideo.platforms.cuda. Gating on the active device's capability
    keeps that fallback working on every unsupported GPU, while
    sm_100a/sm_103a now resolve to the FP4 FA4 kernel (#1221).
    """
    return _resolved_kernel() is not None

fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend

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

fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend.get_supported_head_sizes

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

fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend.get_name

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

fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend.get_impl_cls

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

fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend.get_metadata_cls

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

fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend.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_infer.AttnQatInferImpl

AttnQatInferImpl(num_heads: int, head_size: int, causal: bool, softmax_scale: float, num_kv_heads: int | None = None, prefix: str = '', **extra_impl_args)
Source
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.")
        self.fa4_pv_mode = _resolve_fa4_pv_mode(extra_impl_args)
        # Kernel resolution is per-forward, not per-construction: callers
        # (the validation swap, backend selection) gate on
        # is_attn_qat_infer_available() first, and constructing an impl on a
        # host without the kernel must stay legal (pre-existing contract the
        # validation-swap test pins).
        _log_receipt_once()

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
        # Dispatch on the single per-arch resolution: importability of the
        # bundled sm_12x extension is NOT sufficient (CUDA 13 wheels carry it
        # on unsupported hosts, where calling it is the wrong binary).
        kernel = _resolved_kernel()
        if kernel == "fa4_fp4":
            return self._forward_fa4_fp4(query, key, value)
        if kernel is None:
            raise ImportError(f"attn_qat_infer is not available ({attn_qat_infer_receipt()}). "
                              "Please ensure an ATTN_QAT_INFER kernel is installed for this device.")

        attn_qat_infer = _get_attn_qat_infer()
        assert attn_qat_infer is not None  # kernel == "cutlass_sm12x" implies the import succeeded

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

    def _forward_fa4_fp4(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
    ) -> torch.Tensor:
        """sm_100a/sm_103a path: FP4 FA4 (NVFP4 per-16 block-scaled Q/K; V in
        BF16 by default or fp8 e4m3 per the fa4_pv_mode knob) -- mirrors
        FlashAttentionImpl._forward_nvfp4 (#1221). Inputs/outputs are
        (batch, seqlen, nheads, headdim); no transpose."""
        _nvfp4_quantize_for_fa4, flash_attn_fp4_func = _resolve_fa4_route_ops()

        orig_seqlen_q = query.shape[1]
        orig_seqlen_k = key.shape[1]

        q_fp4, q_sf = _nvfp4_quantize_for_fa4(query)
        k_fp4, k_sf = _nvfp4_quantize_for_fa4(key)

        # fp8 PV: unscaled e4m3 cast (no mSFV/v_descale); output stays BF16.
        if self.fa4_pv_mode == "fp8":
            value = value.to(torch.float8_e4m3fn)
        # The once-log runs on the compiled path too: dynamo cannot trace
        # `logging.Logger` methods, so the helper is registered as a
        # reorderable logging function (see _log_pv_dtype_once).
        _log_pv_dtype_once(value.dtype)

        # FP4/SF buffers are padded to a 128 multiple; FA4 masks to the
        # original lengths so padding never biases the softmax.
        q_fp4 = q_fp4[:, :orig_seqlen_q]
        k_fp4 = k_fp4[:, :orig_seqlen_k]

        output = flash_attn_fp4_func(
            q_fp4,
            k_fp4,
            value,
            q_sf,
            k_sf,
            softmax_scale=self.softmax_scale,
            causal=self.causal,
        )
        if isinstance(output, tuple):
            output = output[0]
        return output

fastvideo.attention.backends.attn_qat_infer.AttnQatInferImpl.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:
    # Dispatch on the single per-arch resolution: importability of the
    # bundled sm_12x extension is NOT sufficient (CUDA 13 wheels carry it
    # on unsupported hosts, where calling it is the wrong binary).
    kernel = _resolved_kernel()
    if kernel == "fa4_fp4":
        return self._forward_fa4_fp4(query, key, value)
    if kernel is None:
        raise ImportError(f"attn_qat_infer is not available ({attn_qat_infer_receipt()}). "
                          "Please ensure an ATTN_QAT_INFER kernel is installed for this device.")

    attn_qat_infer = _get_attn_qat_infer()
    assert attn_qat_infer is not None  # kernel == "cutlass_sm12x" implies the import succeeded

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