之前有个朋友在昇腾 NPU 上部署模型，按文档开了 --enable-flash-attn，跑起来也没报错。但他总觉得延迟不对——跟之前没开的时候差不多。他问我：怎么确认 FlashAttention 真的生效了？不会是静默降级了吧？

这个问题问得很好。很多推理框架（如 vLLM、TGI）在 FlashAttention 初始化失败时，会静默降级到标准 Attention，不会报错，但性能直接掉一大截。你要是不知道怎么查，就只能蒙在鼓里。

今天我把几个简单的确认方法讲清楚——不用读源码，不用会调试器，几条命令就能确认。

方法一：看启动日志（最简单）

vLLM 和 TGI 启动的时候，会打一行日志说明 FlashAttention 的状态。你要是没看到这行，就是没开。

vLLM 的日志

# 开了 FlashAttention（正常）
INFO: flash_attn: FlashAttention V2 enabled (block_size=128)

# 没开（被降级了）
INFO: flash_attn: Using standard attention (FlashAttention not available)

踩坑预警：vLLM 的这行日志在 INFO 级别。你要是把日志级别设成了 WARNING 或 ERROR，就看不到。得把环境变量设对：

export VLLM_LOG_LEVEL=INFO

TGI 的日志

# 开了 FlashAttention（正常）
INFO  text_generation_inference: FlashAttention enabled (via npu_flash_attention)

# 没开（被降级了）
WARN  text_generation_inference: FlashAttention not available, falling back to standard attention

TGI 的日志在 WARN 级别也会打，比 vLLM 友好一点。

方法二：看显存占用（最直观）

FlashAttention 最明显的特征是显存占用跟 seq_len 几乎无关（标准 Attention 是 $O(N^2)$ 的）。

你跑两个不同 seq_len 的请求，看显存占用的差值：

# 第一次请求：seq_len=512
npu-smi inspect -c 0 | grep "Memory Usage"
# 输出：Memory Usage: 3840 MB

# 第二次请求：seq_len=4096（同一个进程，并发请求）
npu-smi inspect -c 0 | grep "Memory Usage"
# FlashAttention：Memory Usage: 3968 MB（涨了 128 MB，很少）
# 标准 Attention：Memory Usage: 11264 MB（涨了 7424 MB，爆炸）

判断标准：

• seq_len 从 512 涨到 4096（8 倍），显存涨 < 10% → FlashAttention 在跑
• seq_len 从 512 涨到 4096，显存涨 > 50% → 标准 Attention 在跑（没开 FlashAttention）

原理：标准 Attention 要存 $O(N^2)$ 的注意力矩阵，seq_len 涨 8 倍，注意力矩阵涨 64 倍。FlashAttention 不存注意力矩阵，seq_len 涨 8 倍，显存只涨一点点（KV Cache 的部分）。

踩坑预警：这个方法只适用于没有 KV Cache 的场景（比如你每次都新开一个对话）。如果你用的是 vLLM 或 TGI 的 KV Cache 功能，seq_len 涨的时候 KV Cache 也会涨，会干扰判断。得看增量显存占用（第二次请求减去第一次请求的显存），把 KV Cache 的部分去掉。

方法三：看延迟随 seq_len 的变化曲线（最准确）

FlashAttention 的延迟随 seq_len 是 $O(N)$ 增长，标准 Attention 是 $O(N^2)$ 增长。你跑几组不同 seq_len 的延迟，画个曲线就知道。

import time
import torch
import torch_npu

# 测试不同 seq_len 的延迟
seq_lens = [512, 1024, 2048, 4096, 8192]
latencies = []

for seq_len in seq_lens:
    q = torch.randn(1, 32, seq_len, 128, dtype=torch.float16, device='npu')
    k = torch.randn(1, 32, seq_len, 128, dtype=torch.float16, device='npu')
    v = torch.randn(1, 32, seq_len, 128, dtype=torch.float16, device='npu')
    
    # 预热
    for _ in range(5):
        _ = torch_npu.contrib.functional.npu_flash_attention(q, k, v, head_num=32)
    torch.npu.synchronize()
    
    # 计时
    start = time.time()
    _ = torch_npu.contrib.functional.npu_flash_attention(q, k, v, head_num=32)
    torch.npu.synchronize()
    end = time.time()
    
    latencies.append((end - start) * 1000)  # ms
    print(f"seq_len={seq_len}, latency={latencies[-1]:.2f} ms")

# 打印结果
for seq_len, latency in zip(seq_lens, latencies):
    print(f"{seq_len}\t{latency:.2f}")

判断标准（画成双对数坐标）：

Attention 类型	双对数坐标里的斜率
标准 Attention	~2.0（$O(N^2)$）
FlashAttention	~1.0（$O(N)$）

你要是看到斜率接近 2.6，就是标准 Attention 在跑，FlashAttention 没生效。

方法四：用 npu-smi 看 AI Core 利用率（最硬核）

你要是想确认 FlashAttention 的 Kernel 真的在跑，可以用 npu-smi 看 AI Core 的利用率。

# 开两个终端
# 终端1：跑推理
python my_inference_script.py

# 终端2：实时监控 AI Core 利用率
watch -n 0.5 "npu-smi inspect -c 0 | grep 'AI Core'"

FlashAttention 的 AI Core 利用率特征：

• Cube Core 利用率：60-80%（高，说明矩阵乘法在跑）
• Vector Core 利用率：50-70%（高，说明 Softmax 在跑）
• Scalar Core 利用率：10-20%（低，说明控制开销小）

标准 Attention 的 AI Core 利用率特征：

• Cube Core 利用率：30-40%（低，因为经常等 HBM 数据）
• Vector Core 利用率：20-30%（低，同理）
• Scalar Core 利用率：5-10%（低）

判断标准：Cube Core 利用率 > 50% → FlashAttention 在跑。Cube Core 利用率 < 40% → 标准 Attention 在跑（在等 HBM）。

踩坑预警：npu-smi inspect 的输出格式跟驱动版本有关。你要是 grep 不到 AI Core，用 npu-smi inspect -c 0 看一下原始输出，找到 AI Core 利用率对应的行，再 grep。

方法五：用昇腾的 Profiling 工具（最专业）

你要是想 100% 确认，可以用昇腾自带的 Profiling 工具抓一次推理的 Timeline。

# 1. 设置 Profiling 环境变量
export ASCEND_PROFILING_MODE=1
export ASCEND_PROFILING_DIR=./profiling_output

# 2. 跑一次推理
python my_inference_script.py

# 3. 关闭 Profiling
unset ASCEND_PROFILING_MODE
unset ASCEND_PROFILING_DIR

# 4. 用 Ascend Profiler 可视化
asc-prof ./profiling_output --output ./profiling_viz

打开 ./profiling_viz/index.html，看 Kernel 列表里有没有 flash_attention_v2 这个 Kernel：

Kernel 列表（部分）：
✅ flash_attention_v2_0      # FlashAttention 在跑
✅ flash_attention_v2_1
   matmul_qk                 # 标准 Attention 的矩阵乘法（不应该看到）
   softmax_fwd              # 标准 Attention 的 Softmax（不应该看到）

判断标准：

• 看到 flash_attention_v2 → FlashAttention 在跑
• 看到 matmul_qk + softmax_fwd + matmul_pv 三个独立 Kernel → 标准 Attention 在跑（FlashAttention 没生效）

踩坑预警：Profiling 会拖慢推理速度 10-20 倍，只能用来调试，不能在生产环境开。

常见导致 FlashAttention 没生效的原因

你按上面五个方法查了一遍，发现 FlashAttention 确实没在跑。常见原因：

原因1：ops-transformer 没编译对

# 检查 FlashAttention 算子是否存在
ls /usr/local/Ascend/ascend-toolkit/latest/op_api/flash_attention_v2/
# 应该能看到 libflash_attention_v2.so

# 要是没有，重新编译
cd ops-transformer/src/flash_attention_v2
bash build.sh --soc Ascend910 --typ release
sudo ./output/flash_attention_v2_Ascend910.run

原因2：推理框架版本太老

# vLLM 需要 ≥ v0.4.0
python -c "import vllm; print(vllm.__version__)"

# TGI 需要 ≥ v1.2.0
# 看 TGI 的 Cargo.toml 里的版本号

原因3：CANN 版本太老（不支持 FlashAttention V2）

# CANN 需要 ≥ 8.0
npu-smi info | grep "Version"
# 输出：Version: 8.0.RC1 → OK
# 输出：Version: 7.0.RC1 → 太老，不支持 FlashAttention V2

原因4：head_dim 不是 128 的倍数
FlashAttention 的 SRAM 分块要求 head_dim 是 128 的倍数。你要是用的模型 head_dim=64（比如 GPT-2），FlashAttention 会静默降级。

# 检查模型的 head_dim
from transformers import AutoConfig
config = AutoConfig.from_pretrained("./models/your-model")
print(config.hidden_size / config.num_attention_heads)  # head_dim
# 输出 128 → OK
# 输出 64 → 不支持，会降级

原因5：seq_len < 1024（阈值以下用标准 Attention 更快）
之前文章讲过，FlashAttention 在 seq_len < 1024 的时候可能更慢，所以有些框架的默认行为是 seq_len < 1024 的时候用标准 Attention。

# vLLM 里强制开 FlashAttention（不管 seq_len）
export VLLM_FORCE_FLASH_ATTN=1

总结一下

确认 FlashAttention 有没有在跑，按这个优先级查：

1. 看启动日志（最简单，30 秒）
2. 看显存随 seq_len 的变化（最直观，5 分钟）
3. 测延迟随 seq_len 的变化曲线（最准确，30 分钟）
4. 用 npu-smi 看 AI Core 利用率（最硬核，需要权限）
5. 用 Profiling 工具抓 Timeline（最专业，适合调试）

最常见的坑：ops-transformer 没编译对（占 80% 的案例），CANN 版本太老（占 15%），head_dim 不支持（占 5%）。