torch.nn.functional.scaled_dot_product_attention

torch.nn.functional.scaled_dot_product_attention()

scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,

is_causal=False, scale=None, enable_gqa=False) -> Tensor

在 query、key 和 value 张量上计算 scaled dot product attention，如果传入attention mask，则使用它，并且如果指定了大于 0.0 的 dropout 概率，则应用 dropout。可选的 scale 参数只能作为关键字参数指定。

# Efficient implementation equivalent to the following:
def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
        is_causal=False, scale=None, enable_gqa=False) -> torch.Tensor:
    L, S = query.size(-2), key.size(-2)
    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
    if is_causal:
        assert attn_mask is None
        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
        attn_bias.to(query.dtype)

    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
        else:
            attn_bias = attn_mask + attn_bias

    if enable_gqa:
        key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
        value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)

    attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight += attn_bias
    attn_weight = torch.softmax(attn_weight, dim=-1)
    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
    return attn_weight @ value

警告

此函数处于 Beta 阶段，未来可能会发生变化。

警告

此函数始终根据指定的 dropout_p 参数应用 dropout。为了在评估期间禁用 dropout，请务必在调用此函数的模块不处于训练模式时传入 0.0 的值。

例如

class MyModel(nn.Module):
    def __init__(self, p=0.5):
        super().__init__()
        self.p = p

    def forward(self, ...):
        return F.scaled_dot_product_attention(...,
            dropout_p=(self.p if self.training else 0.0))

注意

目前支持三种 scaled dot product attention 实现

• FlashAttention-2: 更快的注意力，具有更好的并行性和工作划分

• 内存高效注意力 (Memory-Efficient Attention)

• 基于 C++ 实现的 PyTorch 版本，匹配上述公式

使用 CUDA 后端时，函数可能会调用优化的内核以提高性能。对于所有其他后端，将使用 PyTorch 实现。

所有实现默认启用。Scaled dot product attention 会尝试根据输入自动选择最优化的实现。为了更精细地控制使用哪种实现，提供了以下函数用于启用和禁用实现。上下文管理器是首选机制

• torch.nn.attention.sdpa_kernel(): 一个上下文管理器，用于启用或禁用任何实现。

• torch.backends.cuda.enable_flash_sdp(): 全局启用或禁用 FlashAttention。

• torch.backends.cuda.enable_mem_efficient_sdp(): 全局启用或禁用 Memory-Efficient Attention。

• torch.backends.cuda.enable_math_sdp(): 全局启用或禁用 PyTorch C++ 实现。

每个融合内核都有特定的输入限制。如果用户需要使用特定的融合实现，请使用 torch.nn.attention.sdpa_kernel() 禁用 PyTorch C++ 实现。如果融合实现不可用，将发出警告并说明融合实现无法运行的原因。

由于浮点运算融合的特性，此函数的输出可能会因所选择的后端内核而异。C++ 实现支持 torch.float64，当需要更高精度时可以使用。对于 math 后端，如果输入是 torch.half 或 torch.bfloat16，所有中间结果都保持在 torch.float 中。

更多信息请参阅数值精度

Grouped Query Attention (GQA) 是一个实验性特性。目前仅适用于 CUDA 张量上的 Flash_attention 和 math kernel，不支持 Nested tensor。GQA 的限制条件：

• number_of_heads_query % number_of_heads_key_value == 0 且，

• number_of_heads_key == number_of_heads_value

注意

在某些情况下，给定 CUDA 设备上的张量并使用 CuDNN 时，此算子可能会选择非确定性算法以提高性能。如果不需要这种情况，可以通过设置 torch.backends.cudnn.deterministic = True 来尝试使操作确定（可能以性能为代价）。有关更多信息，请参阅可复现性。

参数

• query (Tensor) – Query 张量；形状为$(N, ..., Hq, L, E)$。

• key (Tensor) – Key 张量；形状为$(N, ..., H, S, E)$。

• value (Tensor) – Value 张量；形状为$(N, ..., H, S, Ev)$。

• attn_mask (可选 Tensor) – Attention 掩码；形状必须能够广播到 attention 权重（形状为$(N,..., L, S)$）的形状。支持两种类型的掩码。一种是布尔掩码，其中 True 值表示元素应该参与 attention。另一种是与 query, key, value 类型相同的浮点掩码，它被加到 attention 分数上。

• dropout_p (float) – Dropout 概率；如果大于 0.0，则应用 dropout

• is_causal (bool) – 如果设置为 True，当掩码是方阵时，注意力掩码是一个下三角矩阵。当掩码不是方阵时，由于对齐（参见 torch.nn.attention.bias.CausalBias），注意力掩码的形式为左上角的因果偏置。如果同时设置了 attn_mask 和 is_causal，则会抛出错误。

• scale (可选 python:float, 仅限关键字参数) – 应用于 softmax 之前的缩放因子。如果为 None，默认值设置为$\frac{1}{\sqrt{E}}$

• enable_gqa (bool) – 如果设置为 True，启用 Grouped Query Attention (GQA)，默认为 False。

返回

注意力输出；形状为$(N, ..., Hq, L, Ev)$。

返回类型

output (Tensor)

形状图例

• $N: \text{Batch size} ... : \text{Any number of other batch dimensions (optional)}$

• $S: \text{Source sequence length}$

• $L: \text{Target sequence length}$

• $E: \text{Embedding dimension of the query and key}$

• $Ev: \text{Embedding dimension of the value}$

• $Hq: \text{Number of heads of query}$

• $H: \text{Number of heads of key and value}$

示例

>>> # Optionally use the context manager to ensure one of the fused kernels is run
>>> query = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
>>> key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
>>> value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
>>> with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
>>>     F.scaled_dot_product_attention(query,key,value)

>>> # Sample for GQA for llama3
>>> query = torch.rand(32, 32, 128, 64, dtype=torch.float16, device="cuda")
>>> key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
>>> value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device="cuda")
>>> with sdpa_kernel(backends=[SDPBackend.MATH]):
>>>     F.scaled_dot_product_attention(query,key,value,enable_gqa=True)