用于 Blackwell 的 FlashAttention-4 后端
作者: tolleybot创建于 2026年9月2日更新于 2026年9月18日
标签enhancement ✨
`fa3_backend.py` 说明 FA3 仅支持 sm_90,Blackwell 需要使用 FA4。FA4 现已可用(`pip install flash-attn-4`,CuTeDSL,涵盖 Hopper 和 Blackwell),我们希望为它贡献一个后端。这意味着需要一个新模块,即 `fa3_backend.py`,实现 #1165 中的 `AttentionBackend` 协议,对架构本身没有任何更改。形状一致:v3 的 ICL head_dim 为 64,其 fp16/bf16 数据类型位于 FA4 在 SM100 上支持的范围内(根据 `flash_attn/cute/interface.py` 中的 `_validate_head_dims`,版本为 `fa4-v4.0.0.beta28`),`flash_attn.cute.flash_attn_func` 使用 q/k/v 采用与 FA3 相同的 `(B, S, H, D)` 布局,因此 `run()` 非常接近相同。一个注意事项: 64 来自 v3 配置默认值,而不是已发布的检查点,因此如果发布的模型不同,值得了解。
内容来源: PriorLabs/TabPFN