背景
-
类似 FlashAttention 这样优化过的 Attention 实现虽然极大提升了性能,但是这种效率的提升却牺牲了灵活性。无法再通过编写一些 PyTorch 操作符来尝试新的 Attention 变体,通常需要编写一个新的自定义内核。
- 一些 attention 变体包含:causal、相对位置嵌入、Alibi、滑动窗口 attention、PrefixLM、Document Masking/Sample Packing/Jagged Tensors, Tanh Soft-Capping, PagedAttention。更麻烦的时候有时候人们还希望组合上面这些变体。
-
下图左侧是现在的状态,一些掩码 + 偏置 + 设置的组合已经有了现成的内核实现。但各种选项导致了设置数量的指数增长,因此整体支持非常不均衡。更糟糕的是,研究人员提出的新 Attention 变体将根本得不到支持。为了彻底解决这个超立方体问题,引入了 FlexAttention。
- FlexAttention 提供了一个灵活的 API,允许使用几行惯用的 PyTorch 代码实现多种 Attention 变体(包括博客文章中提到的所有变体)。
- 通过 torch.compile 将其降低为一个融合的 FlashAttention 内核,生成一个不会产生额外内存且性能与手写内核相当的 FlashAttention 内核。
- 自动生成反向传播,利用 PyTorch 的 autograd 机制。
- 可以利用 Attention 掩码中的稀疏性,相比标准的 Attention 实现显著提高了性能。
flexattention 相关文档
- https://pytorch.org/docs/main/nn.attention.flex_attention.html#module-torch.nn.attention.flex_attention
- https://pytorch.org/blog/flexattention/
- https://github.com/pytorch-labs/attention-gym
FlexAttention 介绍
FlexAttention 的实现形式
这是经典的 Attention 方程:
代码形式:
Q, K, V: Tensor[batch_size, num_heads, sequence_length, head_dim]
score: Tensor[batch_size, num_heads, sequence_length, sequence_length] = (Q @ K) / sqrt(head_dim)
probabilities = softmax(score, dim=-1)
output: Tensor[batch_size, num_heads, sequence_length, head_dim] = probabilities @ V
FlexAttention 允许用户定义一个函数 score_mod
代码形式:
Q, K, V: Tensor[batch_size, num_heads, sequence_length, head_dim]
score: Tensor[batch_size, num_heads, sequence_length, sequence_length] = (Q @ K) / sqrt(head_dim)
modified_scores: Tensor[batch_size, num_heads, sequence_length, sequence_length] = score_mod(score)
probabilities = softmax(modified_scores, dim=-1)
output: Tensor[batch_size, num_heads, sequence_length, head_dim] = probabilities @ V
这个函数允许你在 softmax 之前修改 Attention 分数。这对于绝大多数 Attention 变体来说已经足够了。
具体来说,输入的 score 是一个表示查询 token 和键 token 点积的 PyTorch 标量。其余参数告诉你当前正在计算哪个点积——b(批量中的当前元素),h(当前头),q_idx(查询中的位置),kv_idx(键/值张量中的位置)。
def score_mod(score: f32[], b: i32[], h: i32[], q_idx: i32[], kv_idx: i32[]):
return score # 无操作 - 标准 Attention
为了应用这个函数,我们可以将其实现为:
for b in range(batch_size):
for h in range(num_heads):
for q_idx in range(sequence_length):
for kv_idx in range(sequence_length):
modified_scores[b, h, q_idx, kv_idx] = score_mod(scores[b, h, q_idx, kv_idx], b, h, q_idx, kv_idx)
这并不是 FlexAttention 底层的实现方式。通过利用 torch.compile,torch 会自动将你的函数降低为单个融合的 FlexAttention 内核。
SCORE MOD 式样样例
完全注意力 (Full Attention)
完全注意力是标准的双向注意力。在这种情况下,score_mod 是一个无操作的函数——它接收输入的分数,然后按原样返回它们。
def noop(score, b, h, q_idx, kv_idx):
return score
要将其端到端使用(包括前向和后向):
from torch.nn.attention.flex_attention import flex_attention
flex_attention(query, key, value, score_mod=noop).sum().backward()
相对位置编码 (Relative Position Encodings)
一个常见的 Attention 变体是“相对位置编码”,它不是对查询和键进行绝对距离编码,而是根据查询和键之间的“距离”调整分数。
def relative_positional(score, b, h, q_idx, kv_idx):
return score + (q_idx - kv_idx)
ALiBi 偏置
ALiBi 在《Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation》一文中被提出,并声称在推理时具有有益的长度外推属性。
Alibi 类似于相对位置编码,唯一的例外是它具有一个通常预先计算的 per-head factor。
alibi_bias = generate_alibi_bias() # [num_heads]
def alibi(score, b, h, q_idx, kv_idx):
bias = alibi_bias[h] * (q_idx - kv_idx)
return score + bias
Soft-capping
Soft-capping 是一种在 Gemma2 和 Grok-1 中使用的技术,可以防止 logits 过度增大。
softcap = 20
def soft_cap(score, b, h, q_idx, kv_idx):
score = score / softcap
score = torch.tanh(score)
score = score * softcap
return score
因果掩码
- 《Attention is All You Need》论文以及绝大多数 LLMs 使用的是一种仅解码的设置,其中每个标记只能关注之前的标记。人们通常将其视为下三角掩码。
def causal_mask(score, b, h, q_idx, kv_idx):
return torch.where(q_idx >= kv_idx, score, -float("inf"))
基本上,如果查询标记在键标记“之后”,我们保留分数。否则,我们通过将其设置为 -inf 来屏蔽它,从而确保它不会参与 softmax 计算。
因果掩码优化版
- 如果某些东西被屏蔽了,我们可以完全跳过它的计算!在这种情况下,因果掩码具有大约 50% 的稀疏性,因此不利用稀疏性会导致 2 倍的速度减慢。
- 通过 mask_mod 可以实现带 mask 的稀疏性
# returns True if this position should participate in the computation
mask_mod(b, h, q_idx, kv_idx) => bool
- 通过
create_block_mask
来实现高效版本 causal attention
from torch.nn.attention.flex_attention import create_block_mask
def causal(b, h, q_idx, kv_idx):
return q_idx >= kv_idx
# 由于稀疏模式与批次和头无关,我们将它们设置为 None(这会广播它们)
block_mask = create_block_mask(causal, B=None, H=None, Q_LEN=1024, KV_LEN=1024)
# 在这种情况下,我们不需要 score_mod,因此不会传入它。不过,如果你需要额外的灵活性,score_mod 仍然可以与 block_mask 结合使用。
flex_attention(query, key, value, block_mask=block_mask)
请注意,create_block_mask
是一个相对耗时的操作!尽管 FlexAttention 不会因为掩码变化而重新编译,但如果你不小心缓存它,可能会导致显著的速度减慢
尽管 TFlops 大致相同,但 mask_mod 版本的执行时间快了两倍!这表明我们可以在不损失硬件效率的情况下,利用 BlockMask 提供的稀疏性。
滑动窗口 + 因果掩码
- 由 Mistral 推广的滑动窗口注意力(也称为局部注意力)利用了最近的 token 最有用的直觉。特别地,它允许查询 token 仅关注最近的 1024 个 token。这通常与因果注意力一起使用。
- 代码实现
SLIDING_WINDOW = 1024
def sliding_window_causal(b, h, q_idx, kv_idx):
causal_mask = q_idx >= kv_idx
window_mask = q_idx - kv_idx <= SLIDING_WINDOW
return causal_mask & window_mask
# 以下是更简洁的实现方式...
from torch.nn.attention import or_masks
def sliding_window(b, h, q_idx, kv_idx)
return q_idx - kv_idx <= SLIDING_WINDOW
sliding_window_causal = or_masks(causal_mask, sliding_window)
- 测试结果
- 我们将其与使用滑动窗口掩码的
F.scaled_dot_product_attention
以及使用因果掩码的 FA2 进行了性能基准测试(作为性能的参考点)。我们不仅比F.scaled_dot_product_attention
显著更快,而且也显著快于带因果掩码的 FA2,因为这种掩码具有显著更多的稀疏性。
- 我们将其与使用滑动窗口掩码的
PrefixLM
- T5 架构,提出于《Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer》一文中,描述了一种注意力变体:在“前缀”上执行全双向注意力,在其余部分执行因果注意力。我们同样通过组合两个掩码函数来实现这一点,一个用于因果掩码,另一个基于前缀长度。
- 代码实现
prefix_length: [B]
def prefix_mask(b, h, q_idx, kv_idx):
return kv_idx <= prefix_length[b]
prefix_lm_causal = or_masks(prefix_mask, causal_mask)
# 在这种情况下,我们的掩码会根据每个序列的不同而有所不同,因此我们将 B 设置为批量大小
block_mask = create_block_mask(prefix_lm_causal, B=B, H=None, S, S)
就像 score_mod 一样,mask_mod 允许我们引用函数输入中未显式包含的其他张量!然而,使用 PrefixLM 时,稀疏模式会随着输入的变化而变化。这意味着对于每个新的输入批次,我们需要重新计算 BlockMask。常见的做法是,在模型开始时调用 create_block_mask 并在模型的所有注意力调用中重复使用该 block_mask。
文档掩码/不规则序列
-
另一种常见的注意力变体是文档掩码/不规则序列。假设你有许多长度不同的序列。你想要一起训练它们,但不幸的是,大多数运算符只接受矩形张量。
-
通过 BlockMask,我们也可以在 FlexAttention 中高效地支持这一点!
- 首先,我们将所有序列展平成一个包含 sum(sequence lengths) 个 token 的序列。
- 然后,我们计算每个 token 所属的文档 ID。
- 最后,在我们的 mask_mod 中,我们只需判断查询 token 和键值 token 是否属于同一个文档!
-
代码实现
# 每个标记所属的文档。
# 例如,[0, 0, 0, 1, 1, 2, 2, 2, 2, 2, 2] 对应序列长度 3、2 和 6。
document_id: [SEQ_LEN]
def document_masking(b, h, q_idx, kv_idx):
return document_id[q_idx] == document_id[kv_idx]
块对角掩码
关于文档掩码的一个有趣方面是,容易看出它如何与其他任意组合的掩码结合。例如,我们在上一节中已经定义了 prefixlm_mask。现在我们是否需要定义一个 prefixlm_document_mask 函数呢?
在这些情况下,我们发现一种非常有用的模式是所谓的“更高级别的修改”。在这种情况下,我们可以将现有的 mask_mod 自动转换为适用于不规则序列的版本!
def generate_doc_mask_mod(mask_mod, document_id):
# 获取唯一文档 ID 及其计数
_, counts = torch.unique_consecutive(document_id, return_counts=True)
# 创建累计计数(偏移量)
offsets = torch.cat([torch.tensor([0], device=document_id.device), counts.cumsum(0)[:-1]])
def doc_mask_wrapper(b, h, q_idx, kv_idx):
same_doc = document_id[q_idx] == document_id[kv_idx]
q_logical = q_idx - offsets[document_id[q_idx]]
kv_logical = kv_idx - offsets[document_id[kv_idx]]
inner_mask = mask_mod(b, h, q_logical, kv_logical)
return same_doc & inner_mask
return doc_mask_wrapper
例如,使用上面提到的 prefix_lm_causal 掩码,我们可以将其转换为适用于打包文档的掩码,如下所示:
prefix_length = torch.tensor(2, dtype=torch.int32, device="cuda")
def prefix_mask(b, h, q_idx, kv_idx):
return kv_idx < prefix_length
prefix_lm_causal = or_masks(prefix_mask, causal_mask)
doc_prefix_lm_causal_mask = generate_doc_mask_mod(prefix_lm_causal, document_id)
现在,这个掩码是“块-前缀LM-对角”形状的。
常见问题解答
问:FlexAttention 何时需要重新编译?
由于 FlexAttention 利用 torch.compile
来捕获图形,它实际上在很广泛的情况下可以避免重新编译。特别是,即使捕获的张量的值发生变化,也不需要重新编译!
flex_attention = torch.compile(flex_attention)
def create_bias_mod(bias):
def bias_mod(score, b, h, q_idx, kv_idx):
return score + bias
return bias_mod
bias_mod1 = create_bias_mod(torch.tensor(0))
flex_attention(..., score_mod=bias_mod1) # 在这里编译内核
bias_mod2 = create_bias_mod(torch.tensor(2))
flex_attention(..., score_mod=bias_mod2) # 无需重新编译!
即使改变了块稀疏性(block-sparsity),也不需要重新编译。但是,如果块稀疏性发生变化,我们确实需要重新计算 BlockMask。
问:何时需要重新计算 BlockMask?
每当块稀疏性发生变化时,我们需要重新计算 BlockMask。尽管计算 BlockMask 比重新编译便宜得多(通常是几百微秒,而不是几秒),但您仍然应该注意不要过度重新计算 BlockMask。
以下是一些常见的模式和一些建议。
掩码从不变化(例如因果掩码)
在这种情况下,您可以简单地预计算块掩码并全局缓存,重复使用它进行所有注意力调用。
block_mask = create_block_mask(causal_mask, 1, 1, S, S)
causal_attention = functools.partial(flex_attention, block_mask=block_mask)
掩码每批次变化(例如文档掩码)
在这种情况下,我们建议在模型开始时计算 BlockMask,并在模型中传递它 - 对所有层重复使用 BlockMask。
def forward(self, x, doc_mask):
# 在前向传递的开始计算块掩码
block_mask = create_block_mask(doc_mask, None, None, S, S)
x = self.layer1(x, block_mask)
x = self.layer2(x, block_mask)
...
# 将块掩码构造成本摊到所有层
x = self.layer3(x, block_mask)
return x
掩码每层变化(例如数据依赖的稀疏性)
这是最难的情况,因为我们无法将块掩码计算的成本摊到多个 FlexAttention 调用上。尽管 FlexAttention 在这种情况下仍然有优势,但实际的 BlockMask 效益取决于您的注意力掩码有多稀疏以及我们能多快地构建 BlockMask。这就引出了一个问题…
问:我们如何更快地计算 BlockMask?
create_block_mask
在内存和计算方面都相当昂贵,因为要确定一个块是否完全稀疏,必须在块中的每个点评估 mask_mod
。有几种方法可以解决这个问题:
- 如果您的掩码在批次大小或头之间是相同的,请确保您对这些维度进行广播(即在
create_block_mask
中将它们设置为None
)。 - 编译
create_block_mask
。不幸的是,目前torch.compile
不能直接在create_block_mask
上工作,因为存在一些限制。但是,您可以设置_compile=True
,这将显著减少峰值内存和运行时间(在我们的测试中通常减少一个数量级)。 - 编写自定义的 BlockMask 构造器。BlockMask 的元数据非常简单(请参阅文档)。它本质上是两个张量。a.
num_blocks
:为每个查询块计算的 KV 块数。b.indices
:为每个查询块计算的 KV 块的位置。
例如,这里有一个针对因果掩码的自定义 BlockMask 构造器。
def create_causal_mask(S):
BLOCK_SIZE = 128
# 第一个查询块计算一个块,第二个查询块计算两个块,依此类推。
num_blocks = torch.arange(S // BLOCK_SIZE, device="cuda") + 1
# 由于我们总是从左向右计算,
# 我们可以对每个查询块使用 [0, 1, 2, ...] 的索引。
indices = torch.arange(S // BLOCK_SIZE, device="cuda").expand(
S // BLOCK_SIZE, S // BLOCK_SIZE
)
num_blocks = num_blocks[None, None, :]
indices = indices[None, None, :]
return BlockMask(num_blocks, indices, BLOCK_SIZE=BLOCK_SIZE, mask_mod=causal_mask)
问:为什么 score_mod
和 mask_mod
不同?mask_mod
不就是 score_mod
的一种特殊情况吗?
这是一个非常敏锐的问题!实际上,任何 mask_mod
都可以轻松转换为 score_mod
(我们不建议在实际中使用这个函数!)
def mask_mod_as_score_mod(b, h, q_idx, kv_idx):
return torch.where(mask_mod(b, h, q_idx, kv_idx), score, -float("inf"))
那么,如果 score_mod
能实现 mask_mod
的所有功能,为什么还要保留 mask_mod
呢?
一个直接的挑战是:score_mod
需要实际的分数值作为输入,但在预计算 BlockMask 时,我们并没有实际的分数值。我们可以通过传递全零值来伪造这些值,如果 score_mod
返回 -inf
,那么我们认为它是被掩码了(实际上,我们最初是这么做的!)。
然而,这有两个问题。首先,这种方法是很投机的——如果用户的 score_mod
在输入为 0
时返回 -inf
,怎么办?或者,如果用户的 score_mod
用一个大的负值而不是 -inf
进行掩码,怎么办?看起来我们是在用方形钉子填补圆形孔。但是,分离 mask_mod
和 score_mod
还有一个更重要的原因——它在本质上更有效率!
事实证明,对每一个计算出的元素应用掩码实际上是相当昂贵的——我们的基准测试显示性能下降约 15-20%!所以,尽管通过跳过一半的计算我们可以获得显著的速度提升,但由于需要对每个元素进行掩码,我们会损失一部分速度提升!
幸运的是,如果我们可视化因果掩码,我们会发现绝大多数块根本不需要“因果掩码”——它们是完全计算的!只有在对角线上的块是部分计算和部分掩码的,才需要应用掩码。
块对角掩码
BlockMask 之前告诉我们哪些块需要计算,哪些块可以跳过。现在,我们进一步增强了这个数据结构,告诉我们哪些块是“完全计算”的(即可以跳过掩码)和哪些块是“部分计算”的(即需要应用掩码)。然而,需要注意的是,尽管在“完全计算”的块上可以跳过掩码,但像相对位置嵌入这样的 score_mod
仍然需要应用。
仅仅通过 score_mod
,我们无法可靠地判断它的哪些部分是“掩码”。因此,用户必须自己将这些部分分离到 mask_mod
中。
问:BlockMask 需要多少额外的内存?
BlockMask 的元数据大小为 [BATCH_SIZE, NUM_HEADS, QUERY_LEN//BLOCK_SIZE, KV_LEN//BLOCK_SIZE]
。如果掩码在批次或头部维度上是相同的,可以在该维度上进行广播以节省内存。
在默认的 BLOCK_SIZE
为 128 的情况下,我们预计内存使用在大多数情况下是微不足道的。例如,对于 100 万的序列长度,BlockMask 只会使用 60MB 的额外内存。如果这是一个问题,您可以增加块大小:create_block_mask(..., BLOCK_SIZE=1024)
。例如,将 BLOCK_SIZE
增加到 1024 将使元数据减少到不到 1MB。
问:数值比较如何?
尽管结果不是位相同的,但我们有信心 FlexAttention 的数值精度与 FlashAttention 一样高。我们在广泛的输入范围内比较了 FlashAttention 与 FlexAttention 的差异,涵盖了因果和非因果注意力变体。误差几乎是相同的。
分布图
性能
总的来说,FlexAttention 的性能几乎与手写的 Triton 内核一样高效,这不足为奇,因为我们在很大程度上利用了手写的 Triton 内核。然而,由于其通用性,我们确实付出了一些性能代价。例如,我们必须支付一些额外的延迟来确定下一个要计算的块。在某些情况下,我们提供了一些内核选项,可以在改变内核行为的同时影响其性能。
- FlexAttention 实现了 FlashAttention2 性能的 90%,而在反向传播中达到了 85%。目前,FlexAttention 使用了一种确定性算法,该算法比 FAv2 重新计算了更多的中间变量,但我们计划改进 FlexAttention 的反向算法,并希望缩小这一差距!
局限性和未来工作
- FlexAttention 目前在 PyTorch 的夜间版本中可用,我们计划在 2.5.0 中将其作为原型功能发布。
- 我们在此没有涉及如何将 FlexAttention 用于推理(或如何实现 PagedAttention)——我们将在后续文章中介绍这些内容。
- 我们正在努力提高 FlexAttention 的性能,以在 H100 GPU 上匹配 FlashAttention3 的表现。
- FlexAttention 要求所有序列长度必须是 128 的倍数——我们将很快解决这个问题。
- 我们计划很快添加 GQA 支持——目前,你可以通过复制 kv 头来实现。
总结
- FlexAttention 提供了比 FlashAttention 更加灵活的接口,同时保持了几乎相同的性能。虽然在某些特定情况下 FlexAttention 可能略逊于 FlashAttention,但在大多数情况下,两者之间的差距是微不足道的。