
理论上,“Attention is All You Need”(注意力即你所需)。然而,在实践中,我们还需要优化的注意力实现,例如 FlashAttention。
尽管这些融合的注意力实现显著提高了性能并支持长上下文,但这种效率是以牺牲灵活性为代价的。您不能再通过编写几个 PyTorch 运算符来尝试新的注意力变体——您通常需要编写一个新的自定义内核!这对于机器学习研究人员来说就像是一种“软件彩票”——如果您的注意力变体不适合现有的优化内核,您就会注定遇到运行时缓慢和 CUDA 内存不足的问题。
关于注意力变体的例子,我们有因果注意力、相对位置编码、Alibi、滑动窗口注意力、PrefixLM、文档掩码/样本打包/锯齿张量、Tanh 软限幅、分页注意力等。更糟糕的是,人们经常想要这些组合!滑动窗口注意力 + 文档掩码 + 因果注意力 + 上下文并行?或者分页注意力 + 滑动窗口 + Tanh 软限幅呢?
下面左图代表了当今世界的状况——一些掩码 + 偏差 + 设置的组合有现有的内核实现。但各种选项导致设置的数量呈指数级增长,因此总的来说,我们最终的支持非常零散。更糟糕的是,研究人员提出的新注意力变体将零支持。

为了彻底解决这个超立方体问题,我们引入了FlexAttention,一个新的 PyTorch API。
- 我们提供了一个灵活的 API,允许用几行地道的 PyTorch 代码实现许多注意力变体(包括博客文章中提到的所有变体)。
- 我们通过
torch.compile
将其转换为融合的 FlashAttention 内核,生成的 FlashAttention 内核不占用任何额外内存,并且性能可与手写内核媲美。 - 我们还利用 PyTorch 的自动求导机制自动生成反向传播。
- 最后,我们还可以利用注意力掩码中的稀疏性,从而显著优于标准注意力实现。
有了 FlexAttention,我们希望尝试新的注意力变体将只受限于您的想象力。
您可以在注意力健身房(Attention Gym)找到许多 FlexAttention 示例:https://github.com/pytorch-labs/attention-gym。如果您有任何很酷的应用,欢迎提交示例!
附:我们还发现这个 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 偏差

ALiBi 在短训长测:带有线性偏差的注意力支持输入长度外推中引入,并声称在推理时具有对长度外推有益的特性。值得注意的是,MosaicML 指出“缺乏内核支持”是他们最终从 ALiBi 切换到旋转嵌入的主要原因。
Alibi 类似于相对位置编码,但有一个例外——它有一个通常预先计算的每个头因子。
alibi_bias = generate_alibi_bias() # [num_heads]
def alibi(score, b, h, q_idx, kv_idx):
bias = alibi_bias[h] * (kv_idx - q_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
,因为它的性能更好。请参阅 FAQ,了解为什么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 不需要在其更改时重新编译,但如果您不小心缓存它,它可能会导致显著的减速(查看 FAQ 以获取最佳实践建议)。

虽然 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 and_masks
def sliding_window(b, h, q_idx, kv_idx)
return q_idx - kv_idx <= SLIDING_WINDOW
sliding_window_causal = and_masks(causal_mask, sliding_window)
我们将其与带有滑动窗口掩码的F.scaled_dot_product_attention
以及带有因果掩码的 FA2(作为性能参考点)进行基准测试。我们不仅比F.scaled_dot_product_attention
显著更快,而且也比带有因果掩码的 FA2 显著更快,因为这种掩码具有显著更高的稀疏性。

PrefixLM

探索用统一的文本到文本转换器进行迁移学习的极限中提出的 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
并将其重新用于模型中的所有注意力调用。请参阅重新计算块掩码与重新编译。
然而,作为交换,我们不仅能够为 prefixLM 提供高效的注意力内核,我们还能够利用输入中存在的任何稀疏性!FlexAttention 将根据 BlockMask 数据动态调整其性能,无需重新编译内核。
文档掩码/锯齿序列
另一种常见的注意力变体是文档掩码/锯齿序列。想象一下,您有许多长度可变的序列。您想将它们一起训练,但不幸的是,大多数操作符只接受矩形张量。
通过BlockMask
,我们也可以在 FlexAttention 中高效地支持这一点!
- 首先,我们将所有序列展平为单个序列,其中包含 sum(序列长度) 个令牌。
- 然后,我们计算每个令牌所属的 document_id。
- 最后,在我们的
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)

现在,这个掩码是“块-前缀LM-对角”形状的。🙂
这就是我们所有的例子!注意力变体比我们有空间列出的要多得多,所以请查看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
。有几种方法可以解决这个问题
- 如果您的掩码在批处理大小或头部之间相同,请确保您正在广播这些(即在
create_block_mask
中将它们设置为None
)。 - 编译
create_block_mask
。不幸的是,目前,由于一些不幸的限制,torch.compile
不能直接在create_block_mask
上工作。但是,您可以设置_compile=True
,这将显著降低峰值内存和运行时(在我们的测试中通常是一个数量级)。 - 为 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
,那么我们认为它被屏蔽了(事实上,我们最初就是这样做的!)。
然而,存在两个问题。首先,这是一种权宜之计——如果用户的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 将导致此元数据降至不到一兆字节。
问:数值比较如何?
尽管结果并非位级相同,但我们相信 FlexAttention 在数值上与 FlashAttention 一样准确。我们对因果和非因果注意力变体在大量输入上比较 FlashAttention 和 FlexAttention 的差异,生成以下分布。误差几乎相同。

性能
总的来说,FlexAttention 的性能几乎与手写的 Triton 内核相当,这并不令人意外,因为我们大量利用了手写的 Triton 内核。然而,由于其通用性,我们确实会产生一些小的性能损失。例如,我们必须承担一些额外的延迟来确定接下来要计算哪个块。在某些情况下,我们提供了一些内核选项,它们可以在不改变其行为的情况下影响内核的性能。它们可以在这里找到:性能旋钮
作为一个案例研究,让我们探讨这些旋钮如何影响因果注意力的性能。我们将比较 Triton 内核与 A100 上的 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 的自动求导来自动生成反向传播,以及 vmap 来自动应用score_mod
以创建 BlockMask。
当然,如果没有 Triton 和 TorchInductor 生成 Triton 代码的能力,这个项目就不可能实现。
我们期待未来将这里使用的方法应用于更多应用!
局限性与未来工作
- FlexAttention 目前已在 PyTorch 每夜版中提供,我们计划在 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 关于邻域注意力的讨论
- 所有抱怨注意力内核不支持他们最喜欢的注意力变体的人 🙂