快捷方式

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

在查询、键和值张量上计算缩放点积注意力,使用可选的注意力掩码(如果传递),并在指定大于 0.0 的概率时应用丢弃。可选的 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)
    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

    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 参数应用丢弃。若要在评估期间禁用丢弃,请确保在模块处于非训练模式时传递 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))

注意

目前支持三种缩放点积注意力的实现方法

当使用 CUDA 后端时,此函数可能会调用经过优化的内核来提高性能。对于所有其他后端,将使用 PyTorch 实现。

默认情况下启用所有实现。缩放点积注意力尝试根据输入自动选择最佳实现。为了提供对使用哪些实现的更细粒度的控制,提供了以下函数来启用和禁用实现。上下文管理器是首选机制

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

由于浮点运算融合的性质,此函数的输出可能因所选后端内核而异。c++ 实现支持 torch.float64,可以在需要更高精度时使用。对于数学后端,如果输入为 torch.half 或 torch.bfloat16,则所有中间结果将保留在 torch.float 中。

有关更多信息,请参见 数值精度

分组查询注意力 (GQA) 是一项实验性功能。目前它仅适用于 CUDA 张量上的 Flash_attention 和数学内核,不支持嵌套张量。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) – 查询张量;形状 (N,...,Hq,L,E)(N, ..., Hq, L, E).

  • key (张量) – 键张量;形状为 (N,...,H,S,E)(N, ..., H, S, E).

  • value (张量) – 值张量;形状为 (N,...,H,S,Ev)(N, ..., H, S, Ev).

  • attn_mask (可选张量) – 注意力掩码;形状必须可广播到注意力权重的形状,即 (N,...,L,S)(N,..., L, S). 支持两种类型的掩码。布尔掩码,其中 True 表示元素 *应该* 参与注意力。与查询、键、值相同类型、添加至注意力得分的浮点掩码。

  • dropout_p (浮点数) – 丢弃概率;如果大于 0.0,则应用丢弃。

  • is_causal (布尔值) – 如果设置为 true,当掩码为方阵时,注意力掩码为下三角矩阵。当掩码为非方阵时,注意力掩码具有由于对齐导致的左上角因果偏差的形式(参见 torch.nn.attention.bias.CausalBias)。如果同时设置了 attn_mask 和 is_causal,则会抛出错误。

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

  • enable_gqa (布尔值) – 如果设置为 True,则启用分组查询注意力 (GQA),默认情况下设置为 False。

返回值

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

返回类型

output (张量)

形状说明
  • N:批次大小...:任何数量的其他批次维度 (可选)N: \text{批次大小} ... : \text{任何数量的其他批次维度 (可选)}

  • S:源序列长度S: \text{源序列长度}

  • L:目标序列长度L: \text{目标序列长度}

  • E:查询和键的嵌入维度E: \text{查询和键的嵌入维度}

  • Ev:值的嵌入维度Ev: \text{值的嵌入维度}

  • Hq:查询的头数Hq: \text{查询的头数}

  • H:键和值的头的数量H: \text{键和值的头的数量}

示例

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

文档

访问 PyTorch 的全面的开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源