attn_qat_infer

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

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})"

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</code></pre></details>
fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend
AttnQatInferBackend()
Source
class AttnQatInferBackend(AttentionBackend):
accept_output_buffer: bool = True

@staticmethod
def get_supported_head_sizes() -&gt; list&#91;int&#93;:
    return &#91;64, 128&#93;

@staticmethod
def get_name() -&gt; str:
    return &quot;ATTN_QAT_INFER&quot;

@staticmethod
def get_impl_cls() -&gt; type&#91;&quot;AttnQatInferImpl&quot;&#93;:
    return AttnQatInferImpl

@staticmethod
def get_metadata_cls() -&gt; type&#91;&quot;AttentionMetadata&quot;&#93;:
    raise NotImplementedError

@staticmethod
def get_builder_cls() -&gt; type&#91;&quot;AttentionMetadataBuilder&#91;AttentionMetadata&#93;&quot;&#93;:
    raise NotImplementedError</code></pre></details>
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 = &quot;&quot;,
    **extra_impl_args,
) -&gt; None:
    self.causal = causal
    self.softmax_scale = softmax_scale
    dropout_p = extra_impl_args.get(&quot;dropout_p&quot;, 0.0)
    if dropout_p &gt; 0:
        raise NotImplementedError(f&quot;attn_qat_infer does not support dropout (got dropout_p=&#123;dropout_p&#125;). &quot;
                                  &quot;The QAT inference kernel applies no stochastic dropout.&quot;)
    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,
) -&gt; 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 == &quot;fa4_fp4&quot;:
        return self._forward_fa4_fp4(query, key, value)
    if kernel is None:
        raise ImportError(f&quot;attn_qat_infer is not available (&#123;attn_qat_infer_receipt()&#125;). &quot;
                          &quot;Please ensure an ATTN_QAT_INFER kernel is installed for this device.&quot;)

    attn_qat_infer = _get_attn_qat_infer()
    assert attn_qat_infer is not None  # kernel == &quot;cutlass_sm12x&quot; 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,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;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.&quot;&quot;&quot;
    _nvfp4_quantize_for_fa4, flash_attn_fp4_func = _resolve_fa4_route_ops()

    orig_seqlen_q = query.shape&#91;1&#93;
    orig_seqlen_k = key.shape&#91;1&#93;

    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 == &quot;fp8&quot;:
        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&#91;:, :orig_seqlen_q&#93;
    k_fp4 = k_fp4&#91;:, :orig_seqlen_k&#93;

    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&#91;0&#93;
    return output</code></pre></details>
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 == &quot;cutlass_sm12x&quot; 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()</code></pre></details>