你的 FlashAttention 真的在跑吗？几个简单方法确认

之前有个朋友在昇腾 NPU 上部署模型按文档开了--enable-flash-attn跑起来也没报错。但他总觉得延迟不对——跟之前没开的时候差不多。他问我怎么确认 FlashAttention 真的生效了不会是静默降级了吧这个问题问得很好。很多推理框架如 vLLM、TGI在 FlashAttention 初始化失败时会静默降级到标准 Attention不会报错但性能直接掉一大截。你要是不知道怎么查就只能蒙在鼓里。今天我把几个简单的确认方法讲清楚——不用读源码不用会调试器几条命令就能确认。方法一看启动日志最简单vLLM 和 TGI 启动的时候会打一行日志说明 FlashAttention 的状态。你要是没看到这行就是没开。vLLM 的日志# 开了 FlashAttention正常INFO: flash_attn: FlashAttention V2 enabled(block_size128)# 没开被降级了INFO: flash_attn: Using standard attention(FlashAttention not available)踩坑预警vLLM 的这行日志在INFO级别。你要是把日志级别设成了WARNING或ERROR就看不到。得把环境变量设对exportVLLM_LOG_LEVELINFOTGI 的日志# 开了 FlashAttention正常INFO text_generation_inference: FlashAttention enabled(via npu_flash_attention)# 没开被降级了WARN text_generation_inference: FlashAttention not available, falling back to standard attentionTGI 的日志在WARN级别也会打比 vLLM 友好一点。方法二看显存占用最直观FlashAttention 最明显的特征是显存占用跟seq_len几乎无关标准 Attention 是O(N2)O(N^2)O(N2)的。你跑两个不同seq_len的请求看显存占用的差值# 第一次请求seq_len512npu-smi inspect-c0|grepMemory Usage# 输出Memory Usage: 3840 MB# 第二次请求seq_len4096同一个进程并发请求npu-smi inspect-c0|grepMemory Usage# FlashAttentionMemory Usage: 3968 MB涨了 128 MB很少# 标准 AttentionMemory Usage: 11264 MB涨了 7424 MB爆炸判断标准seq_len从 512 涨到 40968 倍显存涨 10%→FlashAttention 在跑seq_len从 512 涨到 4096显存涨 50%→标准 Attention 在跑没开 FlashAttention原理标准 Attention 要存O(N2)O(N^2)O(N2)的注意力矩阵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)O(N)O(N)增长标准 Attention 是O(N2)O(N^2)O(N2)增长。你跑几组不同seq_len的延迟画个曲线就知道。importtimeimporttorchimporttorch_npu# 测试不同 seq_len 的延迟seq_lens[512,1024,2048,4096,8192]latencies[]forseq_leninseq_lens:qtorch.randn(1,32,seq_len,128,dtypetorch.float16,devicenpu)ktorch.randn(1,32,seq_len,128,dtypetorch.float16,devicenpu)vtorch.randn(1,32,seq_len,128,dtypetorch.float16,devicenpu)# 预热for_inrange(5):_torch_npu.contrib.functional.npu_flash_attention(q,k,v,head_num32)torch.npu.synchronize()# 计时starttime.time()_torch_npu.contrib.functional.npu_flash_attention(q,k,v,head_num32)torch.npu.synchronize()endtime.time()latencies.append((end-start)*1000)# msprint(fseq_len{seq_len}, latency{latencies[-1]:.2f}ms)# 打印结果forseq_len,latencyinzip(seq_lens,latencies):print(f{seq_len}\t{latency:.2f})判断标准画成双对数坐标Attention 类型双对数坐标里的斜率标准 Attention~2.0O(N2)O(N^2)O(N2)FlashAttention~1.0O(N)O(N)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-n0.5npu-smi inspect -c 0 | grep AI CoreFlashAttention 的 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 环境变量exportASCEND_PROFILING_MODE1exportASCEND_PROFILING_DIR./profiling_output# 2. 跑一次推理python my_inference_script.py# 3. 关闭 ProfilingunsetASCEND_PROFILING_MODEunsetASCEND_PROFILING_DIR# 4. 用 Ascend Profiler 可视化asc-prof ./profiling_output--output./profiling_viz打开./profiling_viz/index.html看 Kernel 列表里有没有flash_attention_v2这个 KernelKernel 列表部分 ✅ flash_attention_v2_0 # FlashAttention 在跑 ✅ flash_attention_v2_1 matmul_qk # 标准 Attention 的矩阵乘法不应该看到 softmax_fwd # 标准 Attention 的 Softmax不应该看到判断标准看到flash_attention_v2→ FlashAttention 在跑看到matmul_qksoftmax_fwdmatmul_pv三个独立 Kernel → 标准 Attention 在跑FlashAttention 没生效踩坑预警Profiling 会拖慢推理速度 10-20 倍只能用来调试不能在生产环境开。常见导致 FlashAttention 没生效的原因你按上面五个方法查了一遍发现 FlashAttention 确实没在跑。常见原因原因1ops-transformer 没编译对# 检查 FlashAttention 算子是否存在ls/usr/local/Ascend/ascend-toolkit/latest/op_api/flash_attention_v2/# 应该能看到 libflash_attention_v2.so# 要是没有重新编译cdops-transformer/src/flash_attention_v2bashbuild.sh--socAscend910--typreleasesudo./output/flash_attention_v2_Ascend910.run原因2推理框架版本太老# vLLM 需要 ≥ v0.4.0python-cimport vllm; print(vllm.__version__)# TGI 需要 ≥ v1.2.0# 看 TGI 的 Cargo.toml 里的版本号原因3CANN 版本太老不支持 FlashAttention V2# CANN 需要 ≥ 8.0npu-smi info|grepVersion# 输出Version: 8.0.RC1 → OK# 输出Version: 7.0.RC1 → 太老不支持 FlashAttention V2原因4head_dim 不是 128 的倍数FlashAttention 的 SRAM 分块要求head_dim是 128 的倍数。你要是用的模型head_dim64比如 GPT-2FlashAttention 会静默降级。# 检查模型的 head_dimfromtransformersimportAutoConfig configAutoConfig.from_pretrained(./models/your-model)print(config.hidden_size/config.num_attention_heads)# head_dim# 输出 128 → OK# 输出 64 → 不支持会降级原因5seq_len 1024阈值以下用标准 Attention 更快之前文章讲过FlashAttention 在seq_len 1024的时候可能更慢所以有些框架的默认行为是seq_len 1024的时候用标准 Attention。# vLLM 里强制开 FlashAttention不管 seq_lenexportVLLM_FORCE_FLASH_ATTN1总结一下确认 FlashAttention 有没有在跑按这个优先级查看启动日志最简单30 秒看显存随 seq_len 的变化最直观5 分钟测延迟随 seq_len 的变化曲线最准确30 分钟用 npu-smi 看 AI Core 利用率最硬核需要权限用 Profiling 工具抓 Timeline最专业适合调试最常见的坑ops-transformer 没编译对占 80% 的案例CANN 版本太老占 15%head_dim 不支持占 5%。