FlexAttention：PyTorch 的灵活性与 FlashAttention 的性能

作者：PyTorch 团队：Horace He、Driss Guessous、Yanbo Liang、Joy Dong

理论上，Attention Is All You Need（注意力机制足矣）。但实际上，我们也需要像 FlashAttention 这样经过优化的注意力实现。

尽管这些融合的注意力实现显著提高了性能并支持了长上下文，但这种效率是以牺牲灵活性为代价的。您无法再通过编写一些 PyTorch 运算符来尝试新的注意力变体，您通常需要编写新的自定义内核！这对于 ML 研究人员来说就像一场“软件彩票”——如果您的注意力变体不适合现有的优化内核之一，您注定会面临缓慢的运行时和 CUDA OOM。

关于注意力变体的一些示例，我们有因果 (Causal)、相对位置编码、Alibi、滑动窗口注意力、PrefixLM、文档掩码/样本打包/锯齿张量、Tanh 软上限、PagedAttention 等。更糟糕的是，人们通常想要这些的组合！滑动窗口注意力 + 文档掩码 + 因果 + 上下文并行？或者 PagedAttention + 滑动窗口 + Tanh 软上限又如何？

下面的左图代表了当今世界的状况——掩码 + 偏差 + 设置的某些组合已经实现了内核。但是，各种选项导致了指数级的设置数量，因此总体而言，我们最终得到了相当零星的支持。更糟糕的是，研究人员提出的新的注意力变体将获得零支持。

为了彻底解决这个超立方体问题，我们引入了 FlexAttention，这是一个新的 PyTorch API。

1. 我们提供了一个灵活的 API，允许用几行地道的 PyTorch 代码实现许多注意力变体（包括到目前为止博客文章中提到的所有变体）。
2. 我们通过 torch.compile 将其降低到融合的 FlashAttention 内核中，生成一个 FlashAttention 内核，该内核不会物化任何额外的内存，并且性能与手写内核相当。
3. 我们还利用 PyTorch 的 autograd 机制自动生成反向传播。
4. 最后，我们还可以利用注意力掩码中的稀疏性，从而显著改进标准注意力实现。

借助 FlexAttention，我们希望尝试新的注意力变体将仅受您的想象力限制。

您可以在 Attention Gym 中找到许多 FlexAttention 示例：https://github.com/pytorch-labs/attention-gym。如果您有任何很酷的应用，请随时提交示例！

PS：我们还发现这个 API 非常令人兴奋，因为它以有趣的方式利用了许多现有的 PyTorch 基础设施——最后会有更多介绍。

FlexAttention

这是经典的注意力公式

代码形式

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 之前修改注意力分数。令人惊讶的是，这最终足以满足绝大多数注意力变体（以下示例）！

具体而言，score_mod 的预期签名有些独特。

def score_mod(score: f32[], b: i32[], h: i32[], q_idx: i32[], kv_idx: i32[])
    return score # noop - standard attention

换句话说，score 是一个标量 pytorch 张量，表示查询令牌和键令牌的点积。其余参数告诉您当前正在计算哪个点积——b（批次中的当前元素）、h（当前头）、q_idx（查询中的位置）、kv_idx（键/值张量中的位置）。

要应用此函数，我们可以将其实现为

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，我们会自动将您的函数降低到单个融合的 FlexAttention 内核中——保证有效，否则退款！

这个 API 最终出乎意料地具有表现力。让我们看一些例子。

分数修改示例

完全注意力

我们首先进行“完全注意力”，或标准双向注意力。在这种情况下，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()

相对位置编码

一种常见的注意力变体是 “相对位置编码”。相对位置编码不是对查询和键的绝对距离进行编码，而是根据查询和键之间的“距离”调整分数。

def relative_positional(score, b, h, q_idx, kv_idx):
    return score + (q_idx - kv_idx)

请注意，与典型的实现不同，这不需要物化 SxS 张量。相反，FlexAttention 在内核中“动态”计算偏差值，从而显著提高内存和性能。

ALiBi 偏差

来源：Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

ALiBi 在 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation 中引入，并声称在推理时具有有益的长度外推属性。值得注意的是，MosaicML 指出 “缺乏内核支持” 是他们最终从 ALiBi 切换到旋转嵌入的主要原因。

Alibi 与相对位置编码类似，但有一个例外——它有一个通常预先计算的每个头因子。

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

这演示了 torch.compile 提供的一个有趣的灵活性——我们可以从 alibi_bias 加载，即使它没有显式作为输入传入！生成的 Triton 内核将计算来自 alibi_bias 张量的正确加载并融合它。请注意，您可以重新生成 alibi_bias，我们仍然不需要重新编译。

软上限

软上限是 Gemma2 和 Grok-1 中使用的一种技术，可防止 logits 过度增长。在 FlexAttention 中，它看起来像

softcap = 20
def soft_cap(score, b, h, q_idx, kv_idx):
    score = score / softcap
    score = torch.tanh(score)
    score = score * softcap
    return score

请注意，我们还从前向传播自动生成了反向传播。此外，尽管此实现语义上是正确的，但在这种情况下，出于性能原因，我们可能希望使用 tanh 近似值。有关更多详细信息，请参阅 attention-gym。

因果掩码

虽然双向注意力是最简单的，但最初的 Attention is All You Need 论文和绝大多数 LLM 在仅解码器设置中使用注意力，其中每个令牌只能关注其之前的令牌。人们通常将此视为下三角掩码，但使用 score_mod API，它可以表示为

def causal_mask(score, b, h, q_idx, kv_idx):
    return torch.where(q_idx >= kv_idx, score, -float("inf"))

基本上，如果查询令牌“在”键令牌“之后”，我们保留分数。否则，我们通过将其设置为 -inf 来屏蔽它，从而确保它不会参与 softmax 计算。

但是，与其他修改相比，掩码是特殊的——如果某事物被屏蔽，我们可以完全跳过其计算！在这种情况下，因果掩码具有大约 50% 的稀疏性，因此不利用稀疏性将导致速度减慢 2 倍。尽管此 score_mod 足以正确实现因果掩码，但要获得稀疏性的性能优势，还需要另一个概念——mask_mod。

掩码修改

为了利用掩码的稀疏性，我们需要做更多的工作。具体而言，通过将 mask_mod 传递给 create_block_mask，我们可以创建一个 BlockMask。然后，FlexAttention 可以使用 BlockMask 来利用稀疏性！

mask_mod 的签名与 score_mod 非常相似——只是没有 score。特别是

# returns True if this position should participate in the computation
mask_mod(b, h, q_idx, kv_idx) => bool

请注意，score_mod 在严格意义上比 mask_mod 更具表现力。但是，对于掩码，建议使用 mask_mod 和 create_block_mask，因为它性能更高。请参阅关于为什么 score_mod 和 mask_mod 是分开的常见问题解答。

现在，让我们看一下我们如何使用 mask_mod 实现因果掩码。

因果掩码

from torch.nn.attention.flex_attention import create_block_mask

def causal(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

# Because the sparsity pattern is independent of batch and heads, we'll set them to None (which broadcasts them) 
block_mask = create_block_mask(causal, B=None, H=None, Q_LEN=1024, KV_LEN=1024)
# In this case, we don't need a score_mod, so we won't pass any in.
# However, score_mod can still be combined with block_mask if you need the additional flexibility.
flex_attention(query, key, value, block_mask=block_mask)

请注意，create_block_mask 是一个相对昂贵的操作！ 尽管 FlexAttention 在其更改时不需要重新编译，但如果您不小心缓存它，它可能会导致明显的减速（查看常见问题解答，了解有关最佳实践的建议）。

虽然 TFlops 大致相同，但 mask_mod 版本的执行时间快 2 倍！这表明我们可以利用 BlockMask 为我们提供的稀疏性，而不会损失硬件效率。

滑动窗口 + 因果

来源：Mistral 7B

由 Mistral 推广的滑动窗口注意力（也称为局部注意力）利用了最近的令牌最有用的直觉。特别是，它允许查询令牌仅关注最近的 1024 个令牌（例如）。这通常与因果注意力一起使用。

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

# If you want to be cute...
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

来源：PaliGemma: A versatile 3B VLM for transfer

在 Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer 中提出的 T5 架构描述了一种注意力变体，该变体对“前缀”执行完全双向注意力，对其余部分执行因果注意力。我们再次组合两个掩码函数来实现此目的，一个用于因果掩码，另一个基于前缀长度。

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)
# In this case, our mask is different per sequence so we set B equal to our batch size
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。请参阅重新计算块掩码与重新编译。

但是，作为交换，我们不仅能够为 prefixLM 拥有高效的注意力内核，而且还能够利用输入中存在的任何稀疏性！FlexAttention 将根据 BlockMask 数据动态调整其性能，而无需重新编译内核。

文档掩码/锯齿序列

另一种常见的注意力变体是文档掩码/锯齿序列。想象一下，您有许多长度不同的序列。您想一起训练所有这些序列，但不幸的是，大多数运算符只接受矩形张量。

通过 BlockMask，我们也可以在 FlexAttention 中高效地支持这一点！

1. 首先，我们将所有序列展平为一个总和（序列长度）令牌的单个序列。
2. 然后，我们计算每个令牌所属的 document_id。
3. 最后，在我们的 mask_mod 中，我们只需判断查询和 kv 令牌是否属于同一文档！

# The document that each token belongs to.
# e.g. [0, 0, 0, 1, 1, 2, 2, 2, 2, 2, 2] corresponds to sequence lengths 3, 2, and 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):
    # Get unique document IDs and their counts
    _, counts = torch.unique_consecutive(document_id, return_counts=True)
    # Create cumulative counts (offsets)
    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)

现在，此掩码呈“分块-prefixLM-对角”形状。:)

这就是我们所有的示例！注意力变体的数量远远超过我们可以列出的空间，因此请查看 Attention Gym 以获取更多示例。我们希望社区也能贡献一些他们最喜欢的 FlexAttention 应用。

常见问题解答

问：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) # Compiles the kernel here 

bias_mod2 = create_bias_mod(torch.tensor(2))
flex_attention(..., score_mod=bias_mod2) # Doesn't need to recompile! 

即使更改块稀疏性也不需要重新编译。但是，如果块稀疏性发生变化，我们确实需要重新计算 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):
    # Compute block mask at beginning of forwards
    block_mask = create_block_mask(doc_mask, None, None, S, S)    
    x = self.layer1(x, block_mask)
    x = self.layer2(x, block_mask)
    ...
    # amortize block mask construction cost across all layers
    x = self.layer3(x, block_mask) 
    return x

掩码每层更改（例如数据相关的稀疏性）
这是最难的设置，因为我们无法在多个 FlexAttention 调用中摊销块掩码计算。尽管 FlexAttention 当然仍然可以使这种情况受益，但 BlockMask 的实际好处取决于您的注意力掩码有多稀疏以及我们构建 BlockMask 的速度有多快。这使我们想到…

问：我们如何更快地计算 BlockMask？

create_block_mask 非常昂贵，无论是从内存还是计算角度来看，因为确定块是否完全稀疏需要评估块中每个点的 mask_mod。有几种方法可以解决这个问题

1. 如果您的掩码在批次大小或头之间相同，请确保您在这些掩码上进行广播（即在 create_block_mask 中将它们设置为 None）。
2. 编译 create_block_mask。不幸的是，今天，由于一些不幸的限制，torch.compile 不能直接在 create_block_mask 上工作。但是，您可以设置 _compile=True，这将显著减少峰值内存和运行时（在我们的测试中通常是一个数量级）。

3. 为 BlockMask 编写自定义构造函数。BlockMask 的元数据非常简单（请参阅 文档）。它本质上是两个张量。a. num_blocks：为每个查询块计算的 KV 块的数量。
   b. indices：为每个查询块计算的 KV 块的位置。

   例如，这是 causal_mask 的自定义 BlockMask 构造函数。

def create_causal_mask(S):
    BLOCK_SIZE = 128
    # The first query block computes one block, the second query block computes 2 blocks, etc.
    num_blocks = torch.arange(S // BLOCK_SIZE, device="cuda") + 1
    # Since we're always computing from the left to the right,
    # we can use the indices [0, 1, 2, ...] for every query block.
    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 需要实际的 score 值作为输入，但是当我们预先计算 BlockMask 时，我们没有实际的 score 值。我们可以通过传入全零来伪造值，如果 score_mod 返回 -inf，那么我们认为它被屏蔽了（事实上，我们最初就是这样做的！）。

但是，存在两个问题。首先，这是 hacky 的——如果用户的 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 内核。但是，由于其通用性，我们确实会产生小的性能损失。例如，我们必须产生一些额外的延迟来确定接下来要计算哪个块。在某些情况下，我们提供了一些内核选项，这些选项可能会影响内核的性能，同时改变其行为。它们可以在这里找到：性能旋钮

作为案例研究，让我们探讨旋钮如何影响因果注意力的性能。我们将比较 A100 上 triton 内核与 FlashAttentionv2 的性能。脚本可以在 这里找到。

FlexAttention 在前向传播中实现了 FlashAttention2 性能的 90%，在后向传播中实现了 85%。FlexAttention 当前正在使用一种确定性算法，该算法重新计算的中间值比 FAv2 更多，但我们计划改进 FlexAttention 的后向算法，并希望缩小这一差距！

结论

我们希望您在使用 FlexAttention 时能像我们开发它时一样有趣！在开发过程中，我们最终发现了比我们预期的更多的此 API 应用。我们已经看到它将 torchtune 的 样本打包吞吐量提高了 71%，取代了研究人员花费一周多时间编写自己的自定义 Triton 内核的需求，并提供了与自定义手写注意力变体相当的性能。

实现 FlexAttention 非常有趣的最后一件事是，我们能够以有趣的方式利用许多现有的 PyTorch 基础设施。例如，TorchDynamo（torch.compile 的前端）的独特方面之一是，它不要求编译函数中使用的张量显式作为输入传入。这使我们能够编译诸如文档掩码之类的修改，这些修改需要访问全局变量，而全局变量需要更改！

bias = torch.randn(1024, 1024)
def score_mod(score, b, h, q_idx, kv_idx):
    return score + bias[q_idx][kv_idx] # The bias tensor can change!

此外，torch.compile 是通用图形捕获机制这一事实也使其能够支持更“高级”的转换，例如将任何 mask_mod 转换为适用于锯齿张量的高阶转换。

我们还利用 TorchInductor（torch.compile 的后端）基础设施进行 Triton 模板。这不仅使支持代码生成 FlexAttention 变得容易，而且还自动为我们提供了对动态形状以及后记融合（即在注意力末尾融合运算符）的支持！未来，我们计划扩展此支持，以允许量化版本的注意力或诸如 RadixAttention 之类的东西。

此外，我们还利用了高阶运算、PyTorch 的 autograd 来自动生成反向传播，以及 vmap 来自动应用 score_mod 以创建 BlockMask。

而且，当然，如果没有 Triton 和 TorchInductor 生成 Triton 代码的能力，这个项目是不可能实现的。

我们期待在未来将我们在这里使用的方法应用于更多应用！

限制和未来工作

• FlexAttention 目前在 PyTorch nightly 版本中可用，我们计划在 2.5.0 中将其作为原型功能发布
• 我们在此处未介绍如何将 FlexAttention 用于推理（或如何实现 PagedAttention）——我们将在以后的文章中介绍这些内容。
• 我们正在努力改进 FlexAttention 的性能，以在 H100 GPU 上与 FlashAttention3 相匹配。
• FlexAttention 要求所有序列长度必须是 128 的倍数——这将很快得到解决。
• 我们计划很快添加 GQA 支持——目前，您只需复制 kv 头即可。

致谢

我们要重点介绍一些启发了 FlexAttention 的先前工作（和人员）。

• Tri Dao 在 FlashAttention 上的工作
• Francisco Massa 和 Xformers 团队在 Triton 中的 BlockSparseAttention
• Jax 团队在 SplashAttention 上的工作
• Philippe Tillet 和 Keren Zhou 在 Triton 方面对我们的帮助
• Ali Hassani 关于邻域注意力的讨论
• 所有抱怨注意力内核不支持他们最喜欢的注意力变体的人 :)