FlashAttention的水印攻击：怎么知道你的模型被偷用或篡改了？

之前有个公司发现他们的Llama-2-7B模型被人克隆了一份部署在了另一个云服务上。巧的是那个克隆模型的输出跟他们的一模一样——连生成风格都一样。他们去查代码发现对方的代码里也用了npu_flash_attention。他们想知道能不能从FlashAttention的执行行为里找到证据证明对方用了他们的模型这个问题很有意思。答案是能——FlashAttention的执行行为包含了模型的指纹可以用来检测模型是否被篡改或偷用。今天把这个技术讲清楚。先打个比方每把锁的钥匙痕都不一样你找锁匠配了一把钥匙锁匠的机器在钥匙上磨出了独特的痕迹。另一把钥匙如果是从同一台机器磨出来的痕迹会一模一样——就算你换了钥匙的材料痕迹也一模一样。FlashAttention也是这样——你的模型在昇腾NPU上跑FlashAttention每个分块的大小、对齐方式、执行顺序都会留下独特的痕迹。别人克隆了你的模型就算改了代码风格、加了水印FlashAttention的执行行为也会出卖他们。FlashAttention的执行指纹是什么FlashAttention在昇腾NPU上执行时有几个关键参数会影响执行行为指纹1block_size的余数FlashAttention的分块大小默认是128。如果你的seq_len4096block_size128余数0。但如果你的模型是Llama-2seq_len通常pad到4096的倍数。问题如果你把seq_len改成了4100不是128的倍数FlashAttention会怎么处理方案Apad到4160128×32.5不行方案Bpad到4224128×334224方案C报错不同的实现方案对HBM带宽的影响不一样。如果对方克隆了你的模型也会用同样的padding方案——这是第一个指纹。指纹2head_dim的对齐方式FlashAttention要求head_dim是32的倍数。但不同模型可能有不同的padding策略head_dim128直接用不需要paddinghead_dim96padding到128head_dim100padding到128或96两种方案不同的padding策略会影响算子执行的效率。如果对方克隆了你的模型也会用同样的padding策略——这是第二个指纹。指纹3SRAM的tile策略FlashAttention的SRAM分配策略分多少给Q、分多少给K、分多少给V会影响执行效率。昇腾NPU的ops-transformer有默认的tile策略但如果你自定义了tile参数执行行为会不一样。怎么检测用npu-smi监控SRAM利用率看每个分块的处理时间是否一致。# 检测SRAM tile策略是否一致importtimedefcheck_tile_consistency(q,k,v,head_num,num_iterations100):检测FlashAttention的tile策略是否一致times[]for_inrange(num_iterations):torch.npu.synchronize()starttime.perf_counter()_npu_flash_attention(q,k,v,head_numhead_num)torch.npu.synchronize()times.append((time.perf_counter()-start)*1000)# 计算时间方差mean_timesum(times)/len(times)variancesum((t-mean_time)**2fortintimes)/len(times)std_devvariance**0.5# 判断一致性cvstd_dev/mean_time# 变异系数print(f平均时间{mean_time:.4f}ms)print(f标准差{std_dev:.4f}ms)print(f变异系数{cv:.4f})ifcv0.05:print(✅ tile策略一致执行指纹稳定)else:print(⚠️ tile策略有变化执行指纹不稳定)# 测试qtorch.randn(1,32,4096,128,devicenpu,dtypetorch.float16)ktorch.randn(1,32,4096,128,devicenpu,dtypetorch.float16)vtorch.randn(1,32,4096,128,devicenpu,dtypetorch.float16)check_tile_consistency(q,k,v,head_num32)怎么用FlashAttention指纹检测模型克隆方法1对比执行时间序列同一个模型在同样的硬件上执行时间序列应该高度相似。如果对方克隆了你的模型FlashAttention的执行时间序列也会相似——就算他们改了模型权重权重缩放的方式也会影响执行时间。importnumpyasnpfromscipy.statsimportpearsonrdefgenerate_execution_fingerprint(q,k,v,head_num,num_iterations50):生成FlashAttention执行时间序列times[]for_inrange(num_iterations):torch.npu.synchronize()starttime.perf_counter()_npu_flash_attention(q,k,v,head_numhead_num)torch.npu.synchronize()times.append((time.perf_counter()-start)*1000)returnnp.array(times)defcompare_fingerprints(fp1,fp2):对比两个执行指纹# Pearson相关系数corr,p_valuepearsonr(fp1,fp2)# 平均时间差异mean_diffabs(fp1.mean()-fp2.mean())/fp1.mean()# 时间方差差异var_diffabs(fp1.var()-fp2.var())/fp1.var()print(f相关系数{corr:.4f}p值{p_value:.6f})print(f平均时间差异{mean_diff:.4f})print(f时间方差差异{var_diff:.4f})# 判断ifcorr0.95andp_value0.001:print(✅ 执行指纹高度相似很可能是同一模型)elifcorr0.8:print(⚠️ 执行指纹相似可能是同一模型架构)else:print(❌ 执行指纹不相似不是同一模型)# 生成两个模型的指纹your_model_fpgenerate_execution_fingerprint(q,k,v,head_num32)cloned_model_fpgenerate_execution_fingerprint(q,k,v,head_num32)compare_fingerprints(your_model_fp,cloned_model_fp)方法2对比HBM访问模式FlashAttention的HBM访问模式读写了多少数据、访问频率也是指纹。如果对方克隆了你的模型HBM访问模式也会相似。importsubprocessdefget_hbm_access_stats():获取HBM访问统计resultsubprocess.run([npu-smi,dump,-m,0,-t,hbm,-c,1],capture_outputTrue,textTrue)# 解析输出linesresult.stdout.strip().split(\n)stats{read_bytes:int(lines[1].split()[3]),write_bytes:int(lines[1].split()[5]),read_bandwidth:float(lines[1].split()[7]),write_bandwidth:float(lines[1].split()[9])}returnstatsdefcompare_hbm_pattern(fp1_stats,fp2_stats):对比HBM访问模式read_ratiofp1_stats[read_bytes]/fp2_stats[read_bytes]write_ratiofp1_stats[write_bytes]/fp2_stats[write_bytes]print(f读字节比{read_ratio:.4f}理想1.0)print(f写字节比{write_ratio:.4f}理想1.0)if0.95read_ratio1.05and0.95write_ratio1.05:print(✅ HBM访问模式高度相似)else:print(❌ HBM访问模式不同)总结一下FlashAttention的执行指纹可以用来检测模型是否被克隆或篡改block_size的padding策略不同的padding方案会影响执行行为head_dim的对齐方式不同的padding策略会影响算子执行效率SRAM的tile策略不同的tile分配会影响执行时间执行时间序列相关系数0.95说明很可能是同一模型HBM访问模式读写字节比接近1.0说明很可能是同一模型⚠️ 踩坑预警执行指纹只能证明很可能是同一模型不能作为法律证据。要真正证明对方侵权还需要更多的证据比如代码相似度、训练数据来源等。代码和文档https://atomgit.com/cann/ops-transformer