快捷方式

torch.nn.attention.flex_attention

torch.nn.attention.flex_attention.flex_attention(query, key, value, score_mod=None, block_mask=None, scale=None, enable_gqa=False, return_lse=False, kernel_options=None)[source]

此函数使用任意注意分数修改函数实现缩放点积注意。

此函数计算查询、键和值张量之间的缩放点积注意,使用用户定义的注意分数修改函数。注意分数修改函数将在查询和键张量之间计算注意分数后应用。注意分数按如下方式计算

以下签名应具有 score_mod 函数:

def score_mod(
    score: Tensor,
    batch: Tensor,
    head: Tensor,
    q_idx: Tensor,
    k_idx: Tensor
) -> Tensor:
其中
  • score: 表示注意分数的标量张量,其数据类型和设备与查询、键和值张量相同。

  • batch, head, q_idx, k_idx: 标量张量分别指示批次索引、查询头索引、查询索引和键/值索引。这些应具有 torch.int 数据类型,并位于与分数张量相同的设备上。

参数
  • query (Tensor) – 查询张量;形状 (B,Hq,L,E)(B, Hq, L, E).

  • key (Tensor) – 键张量;形状 (B,Hkv,S,E)(B, Hkv, S, E).

  • value (Tensor) – 值张量;形状 (B,Hkv,S,Ev)(B, Hkv, S, Ev).

  • score_mod (Optional[Callable]) – 修改注意分数的函数。默认情况下,不会应用 score_mod。

  • block_mask (Optional[BlockMask]) – 控制注意的块稀疏模式的 BlockMask 对象。

  • scale (可选[float]) – 在 softmax 之前应用的缩放因子。如果未指定,则默认值为 1E\frac{1}{\sqrt{E}}.

  • enable_gqa (bool) – 如果设置为 True,则启用分组查询注意力 (GQA) 并将键/值头广播到查询头。

  • return_lse (bool) – 是否返回注意力分数的 logsumexp。默认值为 False。

  • kernel_options (可选[Dict[str, Any]]) – 传递给 Triton 内核的选项。

返回值

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

返回类型

output (Tensor)

形状图例
  • N:批次大小...:其他任何批次维度(可选)N: \text{批次大小} ... : \text{其他任何批次维度(可选)}

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

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

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

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

警告

torch.nn.attention.flex_attention 是 PyTorch 中的原型功能。敬请期待 PyTorch 未来版本中更稳定的实现。有关功能分类的更多信息,请参阅:https://pytorch.ac.cn/blog/pytorch-feature-classification-changes/#prototype

块掩码实用程序

torch.nn.attention.flex_attention.create_block_mask(mask_mod, B, H, Q_LEN, KV_LEN, device='cuda', BLOCK_SIZE=128, _compile=False)[source]

此函数从 mask_mod 函数创建块掩码元组。

参数
  • mask_mod (Callable) – mask_mod 函数。这是一个可调用函数,定义了注意力机制的掩码模式。它接受四个参数:b(批次大小)、h(头数)、q_idx(查询索引)和 kv_idx(键/值索引)。它应该返回一个布尔张量,指示哪些注意力连接是允许的 (True) 或被掩码 (False)。

  • B (int) – 批次大小。

  • H (int) – 查询头的数量。

  • Q_LEN (int) – 查询的序列长度。

  • KV_LEN (int) – 键/值的序列长度。

  • device (str) – 运行掩码创建的设备。

  • KV_BLOCK_SIZE (int) – 每个查询的块掩码的块大小。

  • Q_BLOCK_SIZE (int) – 每个键/值的块掩码的块大小。

  • _compile (bool) – 是否编译掩码创建。

返回值

包含块掩码信息的 BlockMask 对象。

返回类型

BlockMask

示例用法
def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

block_mask = create_block_mask(causal_mask, 1, 1, 8192, 8192, device="cuda")
query = torch.randn(1, 1, 8192, 64, device="cuda", dtype=torch.float16)
key = torch.randn(1, 1, 8192, 64, device="cuda", dtype=torch.float16)
value = torch.randn(1, 1, 8192, 64, device="cuda", dtype=torch.float16)
output = flex_attention(query, key, value, block_mask=block_mask)
torch.nn.attention.flex_attention.create_mask(mod_fn, B, H, Q_LEN, KV_LEN, device='cuda', _compile=False)[source]

此函数从 mod_fn 函数创建掩码张量。

参数
  • mod_fn (Union[_score_mod_signature, _mask_mod_signature]) – 修改注意力分数的函数。

  • B (int) – 批次大小。

  • H (int) – 查询头的数量。

  • Q_LEN (int) – 查询的序列长度。

  • KV_LEN (int) – 键/值的序列长度。

  • device (str) – 运行掩码创建的设备。

返回值

形状为 (B, H, M, N) 的掩码张量。

返回类型

mask (Tensor)

torch.nn.attention.flex_attention.and_masks(*mask_mods)[source]

返回一个 mask_mod,它是提供的 mask_mods 的交集。

返回类型

Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]

torch.nn.attention.flex_attention.or_masks(*mask_mods)[source]

返回一个 mask_mod,它是所有提供的 mask_mods 的并集。

返回类型

Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]

torch.nn.attention.flex_attention.noop_mask(batch, head, token_q, token_kv)[source]

返回一个 noop mask_mod。

返回类型

Tensor

BlockMask

class torch.nn.attention.flex_attention.BlockMask(kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, q_num_blocks, q_indices, full_q_num_blocks, full_q_indices, BLOCK_SIZE, mask_mod)[source]

BlockMask 是我们用于表示块稀疏注意力掩码的格式。它有点像是 BCSR 和非稀疏格式的结合。

基础

块稀疏掩码意味着,我们不是表示掩码中单个元素的稀疏性,而是考虑 KV_BLOCK_SIZE x Q_BLOCK_SIZE 块仅在该块中所有元素都是稀疏的情况下才稀疏。这与硬件很吻合,硬件通常希望执行连续的加载和计算。

此格式主要针对 1. 简单性和 2. 内核效率进行了优化。值得注意的是,它没有针对大小进行优化,因为此掩码始终会缩小 KV_BLOCK_SIZE * Q_BLOCK_SIZE 的倍数。如果大小是一个问题,可以通过增加块大小来缩小张量的大小。

我们的格式的基本要素是

num_blocks_in_row: Tensor[ROWS]: 描述每行中存在的块数。

col_indices: Tensor[ROWS, MAX_BLOCKS_IN_COL]: col_indices[i] 是第 i 行的块位置序列。该行在 col_indices[i][num_blocks_in_row[i]] 之后的值是未定义的。

例如,要从这种格式重建原始张量

dense_mask = torch.zeros(ROWS, COLS)
for row in range(ROWS):
    for block_idx in range(num_blocks_in_row[row]):
        dense_mask[row, col_indices[row, block_idx]] = 1

值得注意的是,此格式使沿着掩码的实现缩减变得更容易。

详情

我们的格式的基本要素只需要 kv_num_blocks 和 kv_indices。但是,我们在这个对象上最多有 8 个张量。这代表了 4 对

1. (kv_num_blocks, kv_indices): 用于注意力的前向传递,因为我们沿着 KV 维度进行缩减。

2. [可选] (full_kv_num_blocks, full_kv_indices): 这是可选的,纯粹是优化。事实证明,将掩码应用于每个块非常昂贵!如果我们明确知道哪些块是“完整的”,并且根本不需要掩码,那么我们可以跳过将 mask_mod 应用于这些块。这需要用户从 score_mod 中分离出一个单独的 mask_mod。对于因果掩码,这大约可以提高 15% 的速度。

3. [生成] (q_num_blocks, q_indices): 用于反向传递,因为计算 dKV 需要沿着 Q 维度迭代掩码。这些是根据 1 自动生成的。

4. [生成] (full_q_num_blocks, full_q_indices): 与上面相同,但用于反向传递。这些是根据 2 自动生成的。

BLOCK_SIZE: Tuple[int, int]
as_tuple(flatten=True)[source]

返回 BlockMask 属性的元组。

参数

flatten (bool) – 如果为 True,它将展平 (KV_BLOCK_SIZE, Q_BLOCK_SIZE) 的元组

classmethod from_kv_blocks(kv_num_blocks, kv_indices, full_kv_num_blocks=None, full_kv_indices=None, BLOCK_SIZE=128, mask_mod=None)[source]

从键值块信息创建一个 BlockMask 实例。

参数
  • kv_num_blocks (Tensor) – 每个 Q_BLOCK_SIZE 行块中的 kv_块数。

  • kv_indices (Tensor) – 每个 Q_BLOCK_SIZE 行块中键值块的索引。

  • full_kv_num_blocks (Optional[Tensor]) – 每个 Q_BLOCK_SIZE 行块中完整 kv_块数。

  • full_kv_indices (Optional[Tensor]) – 每个 Q_BLOCK_SIZE 行块中完整键值块的索引。

  • BLOCK_SIZE (Union[int, Tuple[int, int]]) – KV_BLOCK_SIZE x Q_BLOCK_SIZE 块的大小。

  • mask_mod (Optional[Callable]) – 用于修改掩码的函数。

返回值

通过 _transposed_ordered 生成的具有完整 Q 信息的实例

返回类型

BlockMask

引发
full_kv_indices: Optional[Tensor]
full_kv_num_blocks: Optional[Tensor]
full_q_indices: Optional[Tensor]
full_q_num_blocks: Optional[Tensor]
kv_indices: Tensor
kv_num_blocks: Tensor
mask_mod: Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]
numel()[source]

返回掩码中元素的数量(不考虑稀疏性)。

q_indices: Optional[Tensor]
q_num_blocks: Optional[Tensor]
property shape

返回掩码的形状。

sparsity()[source]

计算稀疏(即未计算)块的百分比。

返回类型

float

to(device)[source]

将 BlockMask 移动到指定的设备。

参数

device (torch.device or str) – 要将 BlockMask 移动到的目标设备。可以是 torch.device 对象或字符串(例如,‘cpu’,‘cuda:0’)。

返回值

一个新的 BlockMask 实例,其中所有张量组件都移动到指定的设备。

返回类型

BlockMask

注意

此方法不会就地修改原始 BlockMask。相反,它返回一个新的 BlockMask 实例,其中各个张量属性可能会或可能不会移动到指定的设备,具体取决于它们当前的设备放置。

to_dense()[source]

返回一个与块掩码等效的密集块。

返回类型

Tensor

to_string(grid_size=(20, 20), limit=4)[source]

返回块掩码的字符串表示。相当漂亮。

如果 grid_size 为 None,则打印出未压缩的版本。警告,它可能非常大!

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源