Shmoosh 更新:
这是迄今为止最提气的速度结果。
交代一下背景:Shmoosh 是我当前的研究项目,目标是让扩散模型跑得更快——通过压缩注意力的某些部分,同时保住模型的行为。要找的是那些在 Underpaint、Comfy 这类本地创作工具里真正管用的技术:人们在消费级 GPU 上跑图像生成和编辑工作流。
而今天,内核终于动起来了。
我推了下一片:
1f0ee4e 加入 TF32 打包注意力快速模式。
简短说,我为融合打包注意力加上了可自选的 `dot_precision` 支持。
默认仍是 `ieee`,因为保真度要紧,我不想悄悄改掉高质量策略。但现在多了一个显式的 `tf32` 快速模式,给你要速度、又容忍极小质量代价的场合。
实测数据
在 RTX 4070 上跑合成数据的 1024px K7/no-QJL packed_t 测试,打包注意力从:
0.6232ms 降到 0.2831ms
也就是打包注意力内核阶段提速两倍多。
encode+attention 全路径从:
0.8177ms 降到 0.4761ms
相对 RMSE 几乎没动:
0.023998 到 0.024101
就是这种结果——让你不再眯眼盯着 profiler,而是开始认真起来。
1024px 图像测试集也扛住了:
平均提速:1.068x
平均 PSNR:54.17 dB
最低 PSNR:51.70 dB
图像测试集中的打包注意力阶段从:
1.2806ms 降到 0.7293ms
合成结果有用,但会说谎。看到打包注意力阶段在真实图像生成里也实打实地动了,才更有意思。
显式的快速模式策略
我没有把 TF32 提拔进高保真默认策略,因为有一个测试用例 reading-nook 交了一小笔质量税:
52.07 dB 降到 51.70 dB
依然非常高,但差别是真的,它理应被当成一个策略选择来对待。
所以我加了一个独立的快速模式配置:
configs/underpaint-juggernaut-sdxl-up0-self-attn1-firstblocks-gated70pct-k7-noqjl-tf32-policy.json
验证很干净:
pytest:87 项通过
compileall:通过
jq:通过
git diff --check:通过
仓库干净且已同步
系统的下一个形状
我在慢慢收拢到对的形状。
只压 K 的压缩,一直比动 V 更有希望。
打包表示赢了直接旋转 K 的路径,因为密度才是关键。晚门控自注意力继续看着像那条清醒的道。
K7/no-QJL 扛得够好,值得继续。
而现在,TF32 给打包注意力内核来了一记实打实的助推。
在此之前,Shmoosh 待在「科学上有指望,快一点点」的地带。可本地创作工具要的,是经得起真实图像工作流、真 GPU 和质量预期接触后仍然成立的胜利。
仍在做。
这是打包注意力路径第一次实打实地快了起来。
我很兴奋。