1. 从显存瓶颈说起：Attention 到底卡在哪

做 Transformer 推理或训练的朋友，大概率都遇到过这样的场景：模型参数量看着还能塞进显卡，但序列长度一往上加，显存就炸了。很多人第一反应是“参数量太大”，可实际上，真正吃掉显存的往往不是权重，而是 Attention 计算过程中产生的那一堆中间矩阵。

标准 Attention 的计算公式很简洁：Q 乘 K 的转置得到注意力分数矩阵 S，经过 softmax 归一化后得到概率矩阵 P，再与 V 相乘得到输出 O。问题就出在 S 和 P 这两个矩阵上。它们的形状都是序列长度乘以序列长度，也就是 N×N。当序列长度是 4096 时，单头单层的 S 矩阵就有 1600 万个元素；如果是 8192，直接变成 6700 万个。再乘上头数、层数、批次大小，显存占用会以平方级别膨胀。

我最早接触这个问题是在一个长文本处理项目里，序列长度设到 6144，单卡 80G 显存，模型本身只占 20G 左右，但前向传播跑到一半就 OOM。用显存分析工具一看，Attention 中间矩阵占了将近 40G。当时第一反应是减小 batch size，但 batch 已经降到 1 了，再降就没法跑。后来才意识到，问题的根源不是 batch，而是 Attention 本身的实现方式——它把完整的 S 和 P 矩阵都物化到了显存里。

这就是 FlashAttention 要解决的核心问题。它做的事情说起来很简单：不存中间矩阵，把 Attention 计算拆成小块，在 SRAM 里算完就丢掉，只保留最终输出。但真正把它讲清楚、用起来、调好，里面有不少门道。这篇文章就从原理到 Triton 实战，把这条链路完整走一遍。

2. FlashAttention 的核心思路拆解

2.1 为什么标准 Attention 必须存中间矩阵

要理解 FlashAttention 的优化逻辑，得先搞清楚标准实现为什么非要存 S 和 P。标准 Attention 的计算流程是这样的：先算 S = QK^T，然后对 S 做 softmax 得到 P，最后算 O = PV。这三步在数学上是串行的，softmax 需要拿到完整的 S 才能计算归一化因子，而 P 又需要完整的 softmax 结果才能和 V 相乘。

在 GPU 上，这个串行流程意味着数据必须在显存和计算单元之间来回搬运。S 算出来后写到显存，softmax 读回来算 P，P 再写回显存，最后和 V 相乘时再读一次。每次搬运都是带宽消耗，而 GPU 的显存带宽虽然高，但架不住数据量大。更关键的是，S 和 P 的显存占用是 O(N²)，序列长度翻倍，显存占用翻四倍。

这里有个容易被忽略的点：softmax 的数值稳定性要求。标准实现里，softmax 需要减去最大值防止指数爆炸，这个最大值是全局的，必须等所有 S 算完才能确定。这就导致无法在计算过程中提前释放 S 的内存，必须等整个 softmax 完成。

2.2 分块计算如何打破显存墙

FlashAttention 的核心洞察是：softmax 虽然需要全局信息，但可以通过在线更新的方式分块计算。具体来说，把 Q、K、V 沿序列维度切成小块，每次只加载一块 Q、一块 K、一块 V 到 SRAM 里，在 SRAM 内完成局部的 S 计算、局部 softmax、局部输出累加。关键在于，softmax 的归一化因子可以增量更新。

假设当前处理到第 j 个 K 块，已经累积了前 j-1 个块的输出 O_prev 和归一化统计量。新块进来后，先算局部 S_j，然后更新全局最大值 m_new = max(m_prev, rowmax(S_j))，接着用新的最大值重新缩放之前的累积结果，再把当前块的贡献加进去。这样就不需要存储完整的 S 和 P，只需要在 SRAM 里保留当前块的数据和累积的输出。

这个思路的妙处在于，它把显存占用从 O(N²) 降到了 O(N)。因为每个块的大小是固定的，SRAM 里同时只存在常数个块的数据。序列长度增加时，只是循环次数变多，显存占用不变。实测下来，序列长度 8192 时，FlashAttention 的显存占用只有标准实现的十分之一左右。

2.3 重计算换显存的工程取舍

FlashAttention 在前向传播时不存中间矩阵，但反向传播需要用到 P 矩阵来计算梯度。如果前向不存，反向就得重算。这就是典型的“重计算换显存”策略：前向传播时只存最终输出和少量统计量，反向传播时重新计算 S 和 P，然后立即用于梯度计算并释放。

这个取舍在工程上非常划算。因为 Attention 的计算量本身是 O(N²)，重算一遍只是把计算量翻倍，但显存占用降了一个数量级。对于大序列场景，显存是瓶颈，计算反而有富余。我实测过一个 4096 序列的模型，标准实现显存占用 38G，FlashAttention 只有 6G，训练速度反而快了 15%，因为减少了显存搬运的开销。

当然，重计算也有代价。如果序列很短，比如 512 以下，计算量本身就不大，重计算带来的额外开销可能超过显存节省的收益。所以 FlashAttention 更适合长序列场景，短序列用标准实现反而更直接。

3. Triton 实现 FlashAttention 的关键细节

3.1 Triton 为什么适合写 FlashAttention

Triton 是 OpenAI 开源的一个 GPU 编程语言，定位在 CUDA 和 PyTorch 之间。用 CUDA 写 FlashAttention 需要手动管理共享内存、寄存器分配、线程同步，代码量大且容易出错。用 PyTorch 原生算子又无法精细控制 SRAM 的使用。Triton 刚好卡在中间：它提供了块级编程抽象，可以显式指定块大小、加载到 SRAM、在块内做计算，同时编译器会自动处理线程映射和内存合并。

用 Triton 写 FlashAttention 的核心优势是能精确控制数据流。你可以决定每次加载多大的块、什么时候释放、怎么在 SRAM 里做归约。这些在 CUDA 里也能做，但 Triton 的代码量大概只有 CUDA 的三分之一，调试也更容易。我对比过两种实现，Triton 版本在 A100 上的性能能达到 CUDA 版本的 90% 以上，但开发时间少了一半。

3.2 块大小选择与 SRAM 容量计算

块大小是 Triton 实现 FlashAttention 最关键的参数。块太大，SRAM 放不下，会溢出到显存，失去优化意义；块太小，循环次数多，计算效率低。A100 的 SRAM 每 SM 是 192KB，H100 是 256KB。假设用 fp16 精度，每个元素 2 字节，Q 块大小是 BLOCK_M，K 块大小是 BLOCK_N，头维度是 D。

SRAM 里需要同时存放 Q 块、K 块、V 块、S 块、O 累积块，以及一些统计量。粗略估算：Q 块占 BLOCK_M × D × 2 字节，K 和 V 各占 BLOCK_N × D × 2 字节，S 块占 BLOCK_M × BLOCK_N × 2 字节，O 累积块占 BLOCK_M × D × 2 字节。以 D=64、BLOCK_M=64、BLOCK_N=64 为例，总占用约 64×64×2×5 = 40KB，远小于 192KB，可以放心用。

但实际选择时还要考虑寄存器压力。Triton 会把块内数据映射到寄存器，块太大寄存器不够会 spill 到本地内存，性能急剧下降。我的经验是 BLOCK_M 和 BLOCK_N 取 64 或 128 比较稳，D 小于等于 128 时用 128 的块，D 大于 128 时用 64 的块。如果序列长度不是块大小的整数倍，需要做边界处理，Triton 的 mask 机制可以处理，但会稍微增加开销。

3.3 在线 softmax 的数值稳定性处理

在线 softmax 是 FlashAttention 的数学核心，也是最容易写错的地方。标准 softmax 是 exp(x - max) / sum(exp(x - max))，在线版本需要维护两个累积量：当前最大值 m 和当前指数和 l。每来一个新块，先算局部最大值 m_j，然后更新全局最大值 m_new = max(m, m_j)，接着用 exp(m - m_new) 缩放之前的 l 和 O，再把当前块的 exp(S_j - m_new) 累加进去。

这里有个细节：缩放因子 exp(m - m_new) 在 m 和 m_new 相等时是 1，不需要缩放；只有 m_new 更大时才需要。Triton 里可以用 where 操作避免不必要的乘法。另外，l 的初始值应该是 0，m 的初始值应该是负无穷，这样第一个块进来时能正确初始化。

数值稳定性方面，fp16 的指数范围有限，exp 的参数不能太大。减去最大值后，参数范围在负无穷到 0 之间，exp 结果在 0 到 1 之间，不会溢出。但如果最大值本身很大，比如 1000，减去后还是可能有问题。实际中 Q 和 K 的点积通常不会那么大，但为了保险，可以在算 S 之前对 Q 和 K 做缩放，除以 sqrt(D)。

4. 完整 Triton 实现与逐段解析

4.1 环境准备与依赖安装

先确认环境。Triton 需要 Python 3.8 以上，PyTorch 1.12 以上，CUDA 11.6 以上。安装命令很简单：

pip install triton torch

如果要用到最新的特性，建议从源码安装 Triton：

git clone https://github.com/openai/triton.git cd triton/python pip install -e .

验证安装是否成功：

import triton import triton.language as tl print(triton.__version__)

我建议在 Docker 里跑，因为 Triton 对 CUDA 版本比较敏感，宿主机环境复杂时容易出问题。用 NVIDIA 官方的 PyTorch 镜像，里面已经配好了 CUDA 和 cuDNN，再装 Triton 就行。

4.2 前向传播 kernel 编写

前向 kernel 是整个实现的核心。先定义函数签名：

@triton.jit def flash_attn_fwd( Q, K, V, O, L, M, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh, stride_om, stride_on, Z, H, N_CTX, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, ):

参数说明：Q、K、V、O 是输入输出指针，L 是 softmax 的指数和，M 是最大值。stride 参数描述张量的内存布局，Z 是 batch 大小，H 是头数，N_CTX 是序列长度。BLOCK_M、BLOCK_N、BLOCK_DMODEL 是编译期常量，Triton 会根据这些常量做优化。

接下来是程序 ID 和块索引：

start_m = tl.program_id(0) off_hz = tl.program_id(1) off_z = off_hz // H off_h = off_hz % H

每个 program 处理一个 Q 块和一个 (batch, head) 组合。然后计算 Q 块的偏移：

q_offset = off_z * stride_qz + off_h * stride_qh Q_block_ptr = tl.make_block_ptr( base=Q + q_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_qm, stride_qk), offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0), ) q = tl.load(Q_block_ptr)

这里用make_block_ptr而不是手动算偏移，好处是 Triton 能自动处理边界和内存合并。q 加载到 SRAM 后，初始化累积量：

m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) l_i = tl.zeros([BLOCK_M], dtype=tl.float32) acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)

m_i 初始化为负无穷，l_i 和 acc 初始化为 0。然后进入 K、V 的循环：

for start_n in range(0, N_CTX, BLOCK_N): k = tl.load(K_block_ptr) qk = tl.dot(q, tl.trans(k)) m_ij = tl.max(qk, 1) m_new = tl.maximum(m_i, m_ij) alpha = tl.exp(m_i - m_new) p = tl.exp(qk - m_new[:, None]) l_i = l_i * alpha + tl.sum(p, 1) acc = acc * alpha[:, None] v = tl.load(V_block_ptr) acc = tl.dot(p.to(tl.float16), v, acc) m_i = m_new

这段是在线 softmax 的核心。tl.dot(q, tl.trans(k))算局部注意力分数，tl.max(qk, 1)算每行的最大值，tl.maximum更新全局最大值，alpha是缩放因子，p是局部 softmax 概率，l_i和acc分别更新指数和与输出累积。注意p要转成 fp16 才能和 v 做点积，因为 Triton 的 dot 要求两个操作数类型一致。

循环结束后，把累积结果写回：

acc = acc / l_i[:, None] O_block_ptr = tl.make_block_ptr( base=O + q_offset, shape=(N_CTX, BLOCK_DMODEL), strides=(stride_om, stride_on), offsets=(start_m * BLOCK_M, 0), block_shape=(BLOCK_M, BLOCK_DMODEL), order=(1, 0), ) tl.store(O_block_ptr, acc.to(tl.float16)) tl.store(L + off_hz * N_CTX + start_m * BLOCK_M + tl.arange(0, BLOCK_M), l_i) tl.store(M + off_hz * N_CTX + start_m * BLOCK_M + tl.arange(0, BLOCK_M), m_i)

最后除以 l_i 做归一化，存回 O，同时保存 l_i 和 m_i 供反向传播使用。

4.3 反向传播 kernel 要点

反向传播比前向复杂，因为需要计算 Q、K、V 的梯度。核心思路是：重新计算前向的 S 和 P，然后用链式法则求导。Triton 官方教程里有完整的反向 kernel，这里说几个关键点。

第一，反向需要用到前向保存的 l 和 m，以及输出 O 的梯度 dO。第二，dQ、dK、dV 的计算涉及多个矩阵乘法，需要仔细安排块的加载顺序。第三，softmax 的反向需要用到 P 矩阵，而 P 是重算出来的，不能存。第四，数值稳定性同样重要，dS 的计算要减去行和。

我踩过的一个坑是：反向 kernel 里 dQ 的累积需要原子加操作，因为多个 K 块都会贡献到同一个 Q 块。Triton 的tl.atomic_add可以处理，但性能会受影响。优化方法是把 dQ 的累积放在寄存器里，最后一次性写回，但这要求 Q 块足够小，能完整放在 SRAM 里。

4.4 启动配置与性能调优

kernel 写完后，需要配置 grid 和 block 大小：

def flash_attn(q, k, v): batch, heads, seq_len, dim = q.shape o = torch.empty_like(q) l = torch.empty((batch, heads, seq_len), device=q.device, dtype=torch.float32) m = torch.empty((batch, heads, seq_len), device=q.device, dtype=torch.float32) BLOCK_M = 64 BLOCK_N = 64 grid = (triton.cdiv(seq_len, BLOCK_M), batch * heads) flash_attn_fwd[grid]( q, k, v, o, l, m, q.stride(0), q.stride(1), q.stride(2), q.stride(3), k.stride(0), k.stride(1), k.stride(2), k.stride(3), v.stride(0), v.stride(1), v.stride(2), v.stride(3), o.stride(0), o.stride(1), o.stride(2), o.stride(3), batch, heads, seq_len, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=dim, ) return o

grid 的第一个维度是序列长度除以 BLOCK_M 向上取整，第二个维度是 batch 乘 heads。BLOCK_M 和 BLOCK_N 的选择前面说过，64 或 128 比较稳。如果 dim 大于 128，BLOCK_DMODEL 要设成 128，然后在 kernel 里用循环处理剩余的维度。

性能调优方面，有几个参数可以调：num_warps控制每个 program 的线程束数量，默认是 4，可以试 8；num_stages控制流水线深度，默认是 3，可以试 4 或 5。这两个参数对性能影响很大，建议用 autotune 自动搜索：

@triton.autotune( configs=[ triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64}, num_warps=4, num_stages=3), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64}, num_warps=8, num_stages=4), triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128}, num_warps=4, num_stages=4), ], key=['N_CTX'], )

autotune 会在第一次运行时测试所有配置，选最快的。注意 key 要包含影响性能的参数，比如序列长度。

5. 常见问题与排查技巧实录

5.1 数值不稳定导致输出 NaN

这是最常见的问题，表现是输出里出现 NaN 或 Inf。原因通常是 softmax 的数值处理有问题。排查步骤：先检查 m_i 的初始化是不是负无穷，l_i 是不是 0；再检查 alpha 的计算，exp(m_i - m_new) 在 m_i 是负无穷时应该是 0，但有些实现会算出 NaN；最后检查 p 的计算，exp(qk - m_new) 在 qk 远小于 m_new 时应该接近 0，不会溢出。

我遇到过一次 NaN，查了半天发现是 BLOCK_DMODEL 设成了 64，但实际 dim 是 80，导致加载 Q 块时越界读了脏数据。Triton 的 block_ptr 在 shape 不匹配时不会报错，只会静默读错。所以一定要确保 BLOCK_DMODEL 大于等于实际 dim，如果不等，要在 kernel 里做 mask。

另一个常见原因是 fp16 的精度不够。qk 的点积在 fp16 下累加会丢精度，建议在 dot 之前把 q 和 k 转成 fp32，算完再转回 fp16。Triton 的tl.dot支持输入 fp32 输出 fp32，但性能会降一些。折中方案是用tl.dot(q, k, allow_tf32=True)，在 A100 上 TF32 的精度比 fp16 好，速度比 fp32 快。

5.2 块大小不匹配引发越界

序列长度不是 BLOCK_M 或 BLOCK_N 的整数倍时，最后一个块会越界。Triton 的 block_ptr 有边界检查，但需要显式设置boundary_check。或者手动算 mask：

offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) mask_m = offs_m < N_CTX q = tl.load(Q_block_ptr, boundary_check=(0,), padding_option="zero")

boundary_check=(0,)表示检查第 0 维的边界，越界部分填 0。填 0 对 Q 和 K 是安全的，因为 0 乘任何数都是 0，不会影响 softmax 结果。但对 V 要注意，如果 V 越界填 0，输出会偏小，需要在最后除以有效的 l_i。

我建议在数据准备阶段就把序列长度 padding 到块大小的整数倍，这样 kernel 里不用做边界检查，性能更好。padding 的开销很小，但能省掉不少麻烦。

5.3 性能不如预期时的排查路径

如果 FlashAttention 跑起来比标准实现还慢，按以下顺序排查：

第一，检查块大小。BLOCK_M 和 BLOCK_N 太小会导致循环次数多，太大导致 SRAM 溢出。用triton.testing.do_bench测不同块大小的性能，选最快的。

第二，检查 num_warps 和 num_stages。这两个参数对性能影响很大，默认值不一定最优。用 autotune 自动搜索，或者手动试几组。

第三，检查内存访问模式。Q、K、V 的 stride 如果是非连续的，加载效率会很低。确保输入张量是连续的，或者用.contiguous()转一下。

第四，检查是否用了 TF32。A100 以上默认开启 TF32，如果没开，fp32 的 dot 会慢很多。在 kernel 里显式设置allow_tf32=True。

第五，检查是否有寄存器溢出。用triton.compile编译后看生成的 PTX，如果看到大量 local memory 操作，说明寄存器不够，需要减小块大小。

5.4 与 PyTorch 原生实现的精度对比

FlashAttention 的输出和标准实现有微小差异，这是正常的。差异来源有三个：一是分块计算导致的浮点累加顺序不同；二是在线 softmax 的缩放因子引入的舍入误差；三是 fp16 和 fp32 混合精度的影响。

实测下来，相对误差通常在 1e-3 到 1e-4 之间，对模型最终效果没有可见影响。但如果你的任务对精度极其敏感，比如科学计算，建议用 fp32 跑 FlashAttention，误差能降到 1e-6 以下。

对比方法很简单：用相同的 Q、K、V 分别跑两个实现，算输出的最大绝对误差和相对误差。如果误差超过 1e-2，说明实现有问题，需要检查 softmax 的数值处理。

6. 实际项目中的落地经验

6.1 什么场景该用 FlashAttention

不是所有场景都适合上 FlashAttention。我的判断标准是：序列长度超过 1024，且显存是瓶颈，就值得用。序列长度小于 512 时，标准实现更快，因为重计算的开销超过显存节省的收益。序列长度在 512 到 1024 之间，看具体硬件，A100 以上用 FlashAttention 有优势，V100 以下可能差不多。

另外，如果模型本身很小，比如参数量不到 1B，显存本来就不紧张，用 FlashAttention 的收益有限。但如果模型很大，比如 7B 以上，即使序列长度只有 1024，FlashAttention 也能省下可观的显存，让你能开更大的 batch。

还有一个场景是推理。推理时没有反向传播，不需要重计算，FlashAttention 的显存优势更明显。而且推理的序列长度往往更长，比如对话模型动辄 4096 以上，这时候 FlashAttention 几乎是必选项。

6.2 与现有训练框架的集成方式

集成到现有框架里，最直接的方式是替换torch.nn.functional.scaled_dot_product_attention。PyTorch 2.0 以上已经内置了 FlashAttention 后端，只要设置torch.backends.cuda.enable_flash_sdp(True)就能启用。但内置版本的可定制性差，比如你想改块大小、调 num_warps，就得自己写 kernel。

自己写 kernel 集成时，建议封装成一个nn.Module，forward 里调 Triton kernel，backward 里调反向 kernel。注意要正确处理 autograd，可以用torch.autograd.Function自定义前向和反向：

class FlashAttn(torch.autograd.Function): @staticmethod def forward(ctx, q, k, v): o = flash_attn_fwd_launch(q, k, v) ctx.save_for_backward(q, k, v, o) return o @staticmethod def backward(ctx, do): q, k, v, o = ctx.saved_tensors dq, dk, dv = flash_attn_bwd_launch(q, k, v, o, do) return dq, dk, dv

这样就能无缝接入现有的训练流程。注意保存的张量要控制数量，q、k、v、o 四个就够了，l 和 m 可以在反向时重算。

6.3 长序列场景的显存实测数据

我在一个长文本分类项目里做过对比测试，硬件是单卡 A100 80G，模型是 12 层 Transformer，隐藏维度 768，12 个头。序列长度从 1024 到 8192，分别测标准实现和 FlashAttention 的显存占用和吞吐量。

数据很直观：序列越长，FlashAttention 的优势越大。4096 时显存省了 70%，吞吐量反而高了 43%。8192 时标准实现直接 OOM，FlashAttention 还能跑。这个项目最后把序列长度从 2048 提到了 4096，模型效果提升了 3 个百分点，靠的就是 FlashAttention 省出来的显存。

6.4 踩过的坑与避坑清单

第一个坑是块大小和头维度不匹配。有一次 dim 是 96，我设 BLOCK_DMODEL=64，结果只加载了前 64 维，后 32 维被丢了，输出完全不对。Triton 不会报错，只会静默算错。所以 BLOCK_DMODEL 一定要大于等于实际 dim，如果不等，要么 padding 到 128，要么在 kernel 里用循环处理。

第二个坑是 num_stages 设太大导致 SRAM 溢出。num_stages 控制流水线深度，设成 5 时 SRAM 占用翻倍，结果溢出到显存，性能反而下降。A100 上 num_stages 设 3 或 4 比较稳，H100 可以试 5。

第三个坑是忘了设置allow_tf32。A100 默认不开 TF32，fp32 的 dot 慢得离谱。加上allow_tf32=True后，速度直接翻倍。但要注意 TF32 的精度比 fp32 低，如果任务对精度敏感，还是用 fp32。

第四个坑是反向传播的原子加冲突。多个 K 块同时往同一个 Q 块累加 dQ，用 atomic_add 会有竞争，性能下降明显。解决方案是把 dQ 的累积放在寄存器里，最后一次性写回。但这要求 Q 块能完整放在 SRAM 里，块太大就不行。

第五个坑是序列长度不是块大小整数倍时的边界处理。我一开始没加 boundary_check，结果最后一个块读了越界内存，有时候是 0，有时候是脏数据，导致输出随机波动。加上 boundary_check 后问题消失，但性能降了 5% 左右。后来改成数据准备阶段 padding，性能恢复。

提示：Triton kernel 调试时，可以用triton.testing.do_bench测性能，用torch.allclose对比精度。如果输出不对，先检查块大小和边界，再检查 softmax 的数值处理。

7. 进一步优化的方向

FlashAttention 本身已经很快了，但还有优化空间。一个方向是结合 PagedAttention 做 KV Cache 管理，在推理场景下进一步省显存。另一个方向是用 FP8 精度，H100 上 FP8 的吞吐量是 FP16 的两倍，但精度损失需要评估。还有一个方向是稀疏 Attention，只计算重要的注意力块，把 O(N²) 降到 O(N log N)。

Triton 本身也在快速迭代，新版本对 FlashAttention 的支持越来越好。比如 Triton 3.0 引入了新的内存布局优化，能自动选择最优的块大小。建议保持关注，及时升级。

我在实际项目里的体会是，FlashAttention 不是银弹，它解决的是显存瓶颈，不是计算瓶颈。如果你的场景是计算受限而不是显存受限，比如小模型短序列，用 FlashAttention 可能得不偿失。判断标准很简单：看显存分析工具，如果 Attention 中间矩阵占了显存的大头，那就值得上；如果权重占了大部分，那优化 Attention 收益有限。

最后分享一个小技巧：调试 Triton kernel 时，先用小规模数据跑通，比如序列长度 128、batch 1、单头，确认输出和 PyTorch 一致后再放大。这样能快速定位问题，避免在大规模数据上浪费时间。