深入理解 PyTorch FlexAttention：从 Score Mod 到 Mask Mod 的稀疏加速之道

1. 引言：为什么需要 FlexAttention

在 Transformer 的世界里，注意力机制几乎是所有模型的性能核心。但真实场景中的注意力往往不是"标准"的：有的需要因果掩码（causal mask），有的需要滑动窗口（sliding window），有的需要文档级隔离（document mask），还有的需要 ALiBi、相对位置偏置、甚至自定义的注意力偏置。

传统做法是：为每一种变体手写一个融合 kernel（比如 FlashAttention 的各种变体）。这带来两个问题：

开发成本高：每加一种 mask 就要写一个 CUDA kernel，还要保证数值正确、性能达标。
组合爆炸：当"因果 + 滑动窗口 + 文档隔离"叠加时，几乎无法穷举。

PyTorch 2.5 引入的 torch.nn.attention.flex_attention 正是为了解决这个痛点。它的核心思想是：

把"注意力变体"抽象成用户可编程的 Python 函数，由编译器自动生成并融合出高效的 kernel。

你只需要用普通的 Python 函数描述"哪些位置要参与计算、分数要怎么改"，剩下的交给 torch.compile 去生成 Triton kernel。这就是 FlexAttention 名字中 "Flex" 的含义------灵活可编程。

2. 核心思想：可编程的注意力

标准注意力可以写成：

Attn(Q,K,V)=softmax(QK⊤d)V \text{Attn}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right) V Attn(Q,K,V)=softmax(d QK⊤)V

FlexAttention 把中间那个"分数矩阵"的构造过程开放给用户，允许在 softmax 之前插入两个可编程的钩子：

score_mod：对注意力分数做逐元素的数值变换，例如加偏置、乘缩放、加 ALiBi 斜率等。
mask_mod：决定某个 (query, key) 位置是否参与计算，返回布尔值，用于稀疏化。

这两个钩子都是纯 Python 函数 ，签名固定，由 torch.compile 编译成融合 kernel。关键点在于：编译器会分析这些函数，把它们"内联"进生成的 Triton kernel 中，从而避免物化完整的 N×NN \times NN×N 分数矩阵，保持 FlashAttention 级别的内存效率。

3. 编程接口速览

FlexAttention 的主要 API 位于 torch.nn.attention.flex_attention：

python 复制代码

from torch.nn.attention.flex_attention import (
    flex_attention,
    create_block_mask,
    create_mask,
    and_masks,
    or_masks,
)

核心函数 flex_attention 的签名（简化）：

python 复制代码

flex_attention(
    query, key, value,
    score_mod=None,      # 可选：分数变换函数
    block_mask=None,     # 可选：块级掩码（稀疏加速）
    scale=None,          # 可选：缩放因子，默认 1/sqrt(d)
    enable_gqa=False,    # 是否启用分组查询注意力
    return_lse=False,    # 是否返回 log-sum-exp
)

其中 score_mod 和 block_mask 就是本文要重点剖析的两个钩子。注意 block_mask 是通过 create_block_mask 从 mask_mod 生成的，而不是直接传 mask_mod。

4. 示例：因果掩码注意力

下面用因果掩码（causal mask）作为完整示例，展示 FlexAttention 的典型用法。

4.1 定义 score_mod 与 mask_mod

python 复制代码

import torch
from torch.nn.attention.flex_attention import (
    flex_attention,
    create_block_mask,
)

# score_mod：对分数做数值变换
# 这里演示一个简单的"按 query 位置缩放"的变换
def my_score_mod(score, b, h, q_idx, kv_idx):
    # score: 当前 (b, h, q_idx, kv_idx) 位置的分数标量
    # 返回变换后的分数
    return score * 1.0  # 恒等变换，实际可替换为任意数值操作

# mask_mod：返回布尔值，True 表示该位置参与计算
def causal_mask_mod(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

score_mod 和 mask_mod 的签名都是 (b, h, q_idx, kv_idx)，区别在于 score_mod 额外接收当前分数 score 并返回新分数，而 mask_mod 只返回布尔值。

4.2 构建 block_mask 并调用

python 复制代码

import torch

device = "cuda"
dtype = torch.float16

B, H, S, D = 2, 8, 1024, 64

q = torch.randn(B, H, S, D, device=device, dtype=dtype)
k = torch.randn(B, H, S, D, device=device, dtype=dtype)
v = torch.randn(B, H, S, D, device=device, dtype=dtype)

# 从 mask_mod 生成块级掩码
block_mask = create_block_mask(
    causal_mask_mod,
    B=B, H=H, Q_LEN=S, KV_LEN=S,
    device=device,
)

# 编译并调用
flex_attention_compiled = torch.compile(flex_attention)

out = flex_attention_compiled(
    q, k, v,
    score_mod=my_score_mod,
    block_mask=block_mask,
)

print(out.shape)  # torch.Size([2, 8, 1024, 64])

4.3 与标准实现对比验证

python 复制代码

import torch.nn.functional as F

# 标准因果注意力参考实现
def reference_causal_attention(q, k, v):
    scale = 1.0 / (q.size(-1) ** 0.5)
    scores = torch.matmul(q, k.transpose(-2, -1)) * scale
    causal = torch.tril(torch.ones(q.size(-2), k.size(-2), device=q.device, dtype=torch.bool))
    scores = scores.masked_fill(~causal, float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    return torch.matmul(attn, v)

ref = reference_causal_attention(q, k, v)
print(torch.allclose(out, ref, atol=1e-2, rtol=1e-2))  # True

这个例子展示了 FlexAttention 的完整工作流：定义 mask_mod → 生成 block_mask → 传入 flex_attention。整个过程无需手写任何 CUDA 代码。

5. 内部实现：score_mod 与 mask_mod 如何处理

这是本文的重点。要理解两者的关系，必须先理解 FlexAttention 的编译流程。

5.1 整体编译流程

#mermaid-svg-jcCzK9N60Bqhr6oi{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-jcCzK9N60Bqhr6oi .error-icon{fill:#552222;}#mermaid-svg-jcCzK9N60Bqhr6oi .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-jcCzK9N60Bqhr6oi .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-jcCzK9N60Bqhr6oi .marker{fill:#333333;stroke:#333333;}#mermaid-svg-jcCzK9N60Bqhr6oi .marker.cross{stroke:#333333;}#mermaid-svg-jcCzK9N60Bqhr6oi svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-jcCzK9N60Bqhr6oi p{margin:0;}#mermaid-svg-jcCzK9N60Bqhr6oi .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-jcCzK9N60Bqhr6oi .cluster-label text{fill:#333;}#mermaid-svg-jcCzK9N60Bqhr6oi .cluster-label span{color:#333;}#mermaid-svg-jcCzK9N60Bqhr6oi .cluster-label span p{background-color:transparent;}#mermaid-svg-jcCzK9N60Bqhr6oi .label text,#mermaid-svg-jcCzK9N60Bqhr6oi span{fill:#333;color:#333;}#mermaid-svg-jcCzK9N60Bqhr6oi .node rect,#mermaid-svg-jcCzK9N60Bqhr6oi .node circle,#mermaid-svg-jcCzK9N60Bqhr6oi .node ellipse,#mermaid-svg-jcCzK9N60Bqhr6oi .node polygon,#mermaid-svg-jcCzK9N60Bqhr6oi .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-jcCzK9N60Bqhr6oi .rough-node .label text,#mermaid-svg-jcCzK9N60Bqhr6oi .node .label text,#mermaid-svg-jcCzK9N60Bqhr6oi .image-shape .label,#mermaid-svg-jcCzK9N60Bqhr6oi .icon-shape .label{text-anchor:middle;}#mermaid-svg-jcCzK9N60Bqhr6oi .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-jcCzK9N60Bqhr6oi .rough-node .label,#mermaid-svg-jcCzK9N60Bqhr6oi .node .label,#mermaid-svg-jcCzK9N60Bqhr6oi .image-shape .label,#mermaid-svg-jcCzK9N60Bqhr6oi .icon-shape .label{text-align:center;}#mermaid-svg-jcCzK9N60Bqhr6oi .node.clickable{cursor:pointer;}#mermaid-svg-jcCzK9N60Bqhr6oi .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-jcCzK9N60Bqhr6oi .arrowheadPath{fill:#333333;}#mermaid-svg-jcCzK9N60Bqhr6oi .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-jcCzK9N60Bqhr6oi .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-jcCzK9N60Bqhr6oi .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-jcCzK9N60Bqhr6oi .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-jcCzK9N60Bqhr6oi .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-jcCzK9N60Bqhr6oi .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-jcCzK9N60Bqhr6oi .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-jcCzK9N60Bqhr6oi .cluster text{fill:#333;}#mermaid-svg-jcCzK9N60Bqhr6oi .cluster span{color:#333;}#mermaid-svg-jcCzK9N60Bqhr6oi div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-jcCzK9N60Bqhr6oi .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-jcCzK9N60Bqhr6oi rect.text{fill:none;stroke-width:0;}#mermaid-svg-jcCzK9N60Bqhr6oi .icon-shape,#mermaid-svg-jcCzK9N60Bqhr6oi .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-jcCzK9N60Bqhr6oi .icon-shape p,#mermaid-svg-jcCzK9N60Bqhr6oi .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-jcCzK9N60Bqhr6oi .icon-shape .label rect,#mermaid-svg-jcCzK9N60Bqhr6oi .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-jcCzK9N60Bqhr6oi .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-jcCzK9N60Bqhr6oi .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-jcCzK9N60Bqhr6oi :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 用户定义 score_mod / mask_mod
create_block_mask 生成 BlockMask
flex_attention 调用
torch.compile 捕获计算图
分析 score_mod 的 Python 字节码
读取 BlockMask 的稀疏结构
生成 Triton kernel
融合执行：加载 Q/K/V 块
逐元素应用 score_mod
按 BlockMask 跳过全零块
softmax + 乘 V
输出

关键点：score_mod 和 mask_mod 走的是两条不同的路径 。score_mod 被内联进 kernel 的逐元素计算中；mask_mod 则在 create_block_mask 阶段就被"预计算"成块级稀疏结构，kernel 只读取这个结构。

5.2 score_mod 的处理：内联为逐元素数值变换

score_mod 是一个 Python 函数，torch.compile 会对其做字节码级分析（通过 TorchDynamo 和 FX 图捕获），把它转换成一系列逐元素操作，然后内联到生成的 Triton kernel 中。

具体来说，在 kernel 内部，对于每个 (q_idx, kv_idx) 位置：

计算原始分数 score = Q[q_idx] · K[kv_idx] * scale
调用 score_mod(score, b, h, q_idx, kv_idx)，得到变换后的分数
继续后续的 softmax 流程

由于 score_mod 是逐元素的，它不会改变计算量 ------无论你做什么数值变换，kernel 仍然要遍历所有 (q_idx, kv_idx) 位置。它的作用是改变数值 ，而不是减少计算。

一个典型的 score_mod 例子是 ALiBi：

python 复制代码

def alibi_score_mod(score, b, h, q_idx, kv_idx):
    slope = 2 ** (-8 * (h + 1) / H)
    return score + slope * (q_idx - kv_idx)

这里 slope 依赖 head 索引 h，(q_idx - kv_idx) 是相对距离。编译器会把这些操作全部内联，生成高效的逐元素计算。

5.3 mask_mod 的处理：预计算为块级稀疏结构

mask_mod 的处理路径完全不同。它不会被内联进 kernel 的逐元素循环，而是在 create_block_mask 阶段被"采样"成块级结构。

create_block_mask 的工作方式是：

把 Q×KVQ \times KVQ×KV 的分数矩阵按 BLOCK_SIZE（通常是 128×128 或 64×64）划分成块。
对每个块，采样 mask_mod 在块内若干位置（或全部位置）的返回值。
如果整个块全为 False，标记为"跳过块"；如果全为 True，标记为"全计算块"；否则标记为"部分块"。

最终生成一个 BlockMask 对象，包含：

kv_num_blocks / kv_indices：每个 query 块需要访问的 key 块索引（稀疏结构）。
full_kv_num_blocks / full_kv_indices：完全计算的块。
partial_kv_num_blocks / partial_kv_indices：部分计算的块（需要逐元素判断）。

在 kernel 执行时，对于"跳过块"，直接不加载 K/V，不计算；对于"全计算块"，直接算；对于"部分块"，才逐元素应用 mask_mod 判断。

这就是稀疏加速的来源：因果掩码下，约一半的块被跳过，计算量减半。

5.4 两者的关系与作用对比

维度	score_mod	mask_mod
输入	score, b, h, q_idx, kv_idx	b, h, q_idx, kv_idx
输出	变换后的分数（浮点）	布尔值（是否参与）
处理阶段	kernel 内逐元素内联	create_block_mask 阶段预计算
是否改变计算量	否，纯数值变换	是，可跳过整块
是否影响稀疏性	否	是，决定块级稀疏结构
典型用途	ALiBi、相对位置偏置、温度缩放	因果掩码、滑动窗口、文档隔离
性能影响	轻微（逐元素开销）	显著（减少 K/V 加载和计算）

一句话总结：

score_mod 是"数值变换"，mask_mod 是"稀疏加速"。前者改变分数的值，后者改变计算的范围。

5.5 为什么 mask_mod 能加速而 score_mod 不能

根本原因在于信息结构 。mask_mod 返回布尔值，具有"块级可聚合性"------一个块内如果全是 False，整个块就可以跳过。这种聚合可以在 kernel 执行前完成，形成稀疏索引。

而 score_mod 返回的是任意浮点数，没有"块级可跳过"的语义。即使某个块的分数都很小，softmax 仍然需要它们参与归一化（除非你显式做 top-k 稀疏，那是另一回事）。所以 score_mod 只能逐元素计算，无法跳过。

这也解释了为什么 API 设计上 mask_mod 要先经过 create_block_mask 转换，而 score_mod 直接传入------前者需要预计算稀疏结构，后者只需要内联。

5.6 组合使用：两者可以叠加

实际场景中，score_mod 和 mask_mod 经常一起使用。例如"因果掩码 + ALiBi 偏置"：

python 复制代码

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

def alibi_score_mod(score, b, h, q_idx, kv_idx):
    slope = 2 ** (-8 * (h + 1) / H)
    return score + slope * (q_idx - kv_idx)

block_mask = create_block_mask(causal_mask_mod, B=B, H=H, Q_LEN=S, KV_LEN=S, device=device)

out = flex_attention_compiled(
    q, k, v,
    score_mod=alibi_score_mod,
    block_mask=block_mask,
)

此时 kernel 会：先用 block_mask 跳过因果掩码下的上三角块（稀疏加速），再对剩余块逐元素应用 ALiBi 偏置（数值变换）。两者各司其职，互不干扰。

6. 进阶：mask_mod 的组合与动态性

FlexAttention 还提供了 and_masks / or_masks 来组合多个 mask_mod：

python 复制代码

from torch.nn.attention.flex_attention import and_masks

def sliding_window_mod(b, h, q_idx, kv_idx):
    return (q_idx - kv_idx) <= 512

combined = and_masks(causal_mask_mod, sliding_window_mod)
block_mask = create_block_mask(combined, B=B, H=H, Q_LEN=S, KV_LEN=S, device=device)

这实现了"因果 + 滑动窗口"的复合掩码，而无需手写任何 kernel。

另外，create_block_mask 支持动态形状：当序列长度变化时，可以传入 _compile=True 让块掩码的生成也被编译，避免重复开销。

7. 性能实践建议

优先用 block_mask 而非在 score_mod 里做掩码 ：如果你在 score_mod 里返回 -inf 来模拟掩码，虽然数值正确，但不会获得稀疏加速 ，因为 kernel 仍然遍历所有位置。正确做法是用 mask_mod + create_block_mask。
合理选择 BLOCK_SIZE：块太小稀疏粒度细但索引开销大，块太大稀疏效果差。默认 128 通常是不错的起点。
缓存 block_mask ：create_block_mask 有一定开销，固定形状下应缓存复用。
用 torch.compile 包裹 ：flex_attention 必须经过 torch.compile 才能生成融合 kernel，否则会走 fallback 路径，性能大幅下降。

8. 总结

FlexAttention 的核心价值在于把"注意力变体"从"手写 kernel"降维成"写 Python 函数"。理解它的关键在于区分两个钩子：

score_mod：逐元素数值变换，内联进 kernel，不改变计算量，用于 ALiBi、偏置、缩放等。
mask_mod ：布尔掩码，经 create_block_mask 预计算成块级稀疏结构，可跳过整块，用于因果、滑动窗口、文档隔离等，是稀疏加速的来源。

记住这句话：score_mod 改"值"，mask_mod 改"范围"。前者让你灵活调整注意力分布，后者让你在保持正确性的前提下大幅减少计算。两者组合，几乎可以覆盖所有常见的注意力变体，而无需写一行 CUDA。