Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

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()</code></pre></details>


fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend


AttnQatTrainBackend()


Source
class AttnQatTrainBackend(AttentionBackend):

accept_output_buffer: bool = True

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

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

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

@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_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 = &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_train does not support dropout (got dropout_p=&#123;dropout_p&#125;). &quot;
                                  &quot;The QAT training kernel applies no stochastic dropout.&quot;)

def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -&gt; torch.Tensor:
    return attn_qat_train(query, key, value, is_causal=self.causal, sm_scale=self.softmax_scale)</code></pre></details>


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)

Was this page helpful?