Skip to content

sage_attn3

Classes

fastvideo.attention.backends.sage_attn3.SageAttention3Impl 

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

Bases: AttentionImpl

Source code in fastvideo/attention/backends/sage_attn3.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
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
    self.dropout = extra_impl_args.get("dropout_p", 0.0)

Methods:

fastvideo.attention.backends.sage_attn3.SageAttention3Impl.forward 
forward(query: Tensor, key: Tensor, value: Tensor, attn_metadata: AttentionMetadata) -> Tensor

Call sageattn3_blackwell directly. Input is already [B, H, L, D] and contiguous from preprocess_qkv.

Source code in fastvideo/attention/backends/sage_attn3.py 
72
73
74
75
76
77
78
79
80
81
82
def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    """Call sageattn3_blackwell directly. Input is already [B, H, L, D]
    and contiguous from preprocess_qkv."""
    output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
    return output
fastvideo.attention.backends.sage_attn3.SageAttention3Impl.postprocess_output 
postprocess_output(output: Tensor, attn_metadata: AttentionMetadata) -> Tensor

Transpose output from [B, H, L, D] back to [B, L, H, D].

Source code in fastvideo/attention/backends/sage_attn3.py 
84
85
86
87
88
89
90
def postprocess_output(
    self,
    output: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    """Transpose output from [B, H, L, D] back to [B, L, H, D]."""
    return output.permute(0, 2, 1, 3).contiguous()
fastvideo.attention.backends.sage_attn3.SageAttention3Impl.preprocess_qkv 
preprocess_qkv(qkv: Tensor, attn_metadata: AttentionMetadata) -> Tensor

Transpose stacked QKV from [3B, L, H, D] to [3B, H, L, D].

Single bulk permute+contiguous on the entire stacked tensor rather than three separate transposed views for Q, K, V. The .contiguous() is required: sageattn_blackwell's fake kernel returns empty_like(q), so the op's output strides must match contiguous q under torch.compile.

Source code in fastvideo/attention/backends/sage_attn3.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
def preprocess_qkv(
    self,
    qkv: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    """Transpose stacked QKV from [3B, L, H, D] to [3B, H, L, D].

    Single bulk permute+contiguous on the entire stacked tensor rather than
    three separate transposed views for Q, K, V. The .contiguous() is
    required: sageattn_blackwell's fake kernel returns empty_like(q), so the
    op's output strides must match contiguous q under torch.compile.
    """
    return qkv.permute(0, 2, 1, 3).contiguous()

Functions: