Closed
[Requirement|需求建议]: flash_attn算子在arch92架构上支持sink功能 #6530
comic97Add on 24 days agoClosed at 23 days ago
周云飞
24 days ago Comments:
24 days ago Comments:
感谢提出 flash_attn 在 arch92 架构上支持 sink 功能的需求。
目前提供的信息尚未明确 sink 的具体计算语义。请补充所参考的模型或接口说明,并说明 sink 参数的形状、数据类型及预期计算方式;若有参考实现或最小输入输出示例,也可直接提供。这些信息有助于明确需要支持的行为,避免将不同的 sink 机制混为一谈。
arch92 当前存在 ND 和 DN_UNMERGED 两类计算模板,选择会受到头维度、查询序列长度和掩码模式等条件影响。若需要新增 sink 计算,一个可行的方向是先明确接口参数和校验规则,再打通参数到计算内核的传递,并评估两类模板中相应的计算实现。调整时需要保持不启用 sink 的现有调用行为不变,同时评估新增计算对数值精度和性能的影响。
现有代码片段不足以确认 sink 的完整支持情况。明确上述语义后,才能进一步判断需要补充的实现范围及适用限制。


24 days ago Assign comic97 as responsible person
23 days ago Close issue
23 days ago Added label: resolved
Backgroud(背景信息)
1. attention sink 与 flash_attn 的 sinks 参数
StreamingLLM 等工作发现,softmax 注意力对序列初始 token 存在显著的"汇聚"效应(attention sink)。长上下文推理若直接滑窗丢弃早期 KV,模型精度会明显劣化;引入可学习 sink 项可同时改善数值稳定性与长序列外推能力。
为此,flash_attn 接口已定义 sinks 参数(可选,float32,shape 为 (Q_N,),与 q/k/v 的 dtype 解耦),在每个 query 头的 softmax 归一化中引入 sink 项(见 docs/torchapi_flash_attn.md 的 Sinks 参数组):
2. 现状:arch92 平台缺失 sink 支持
3. common 实现无法直接复用的架构原因
common attn_sink_gs1.h 基于"单窗口 64 lane reg Gather + A/B 双窗口宽模式 + chunk 循环"设计,与 arch92 架构特点不匹配:
Benefit / Necessity (价值/作用)
1. 功能补齐(必要性)
2. 性能收益(针对 arch92 架构特点的专用实现)
3. 可复用性
arch92 sink 逻辑收敛为特性级公共头文件 attn_sink_arch92.h,CopyIn 与直写按 GS1/S1G 布局自适应,arch92 平台其他具有相同 softmax 统计结构的注意力类算子可直接复用。
Design(设计方案)
整体数学方案(Scheme B,prologue 注入)
在首 FdBlock(learnableSinkFlag && isFirstFdBlock)的 softmax buffer 复位阶段完成注入,等效于在每个 query 头序列头部预置一个 score = sink、value = 0 的虚拟 token:
代码改动(关联 PR #14083,5 个文件)
新增 attention/flash_attn/op_kernel/utils/attn_sink_arch92.h(arch92 sink 特性级公共实现):
attention/flash_attn/op_kernel/arch92/flash_attn_block_vec_nd_arch92.h(ND 布局接入):ResetSoftmaxBuffer 在首 FdBlock 且 learnableSinkFlag 时:sum 置 1;按 IS_GS1_LAYOUT(BNSD)与 S1G(BSND/TND)自适应选择 SinkCopyInGS1/SinkCopyInS1G(MTE2 锁内完成 GM→UB,值区借用 vec1Res 槽);V 锁内调用带编译期 MBASE 模板参数的 SinkExpandMaxVf121 展开 max;非 sink 场景维持原 Duplicate(0/−inf) 逻辑,不改变未启用 sink 时的现有行为。
attention/flash_attn/op_kernel/arch92/flash_attn_block_vec_dn_unmerged_arch92.h(DN 不合轴接入):S1-only 场景(每任务仅一个 query 头)sink 退化为单值标量,sinkGm_.GetValue 直接读取后经 FaVectorApi::ResetSoftmaxBuffer 以(sum=1, max=sink)复位,省去 GS1 展开通用路径的 MTE2/V 双锁与 UB 临时搬运。
attention/common/op_kernel/arch92/vf/vf_mul_sel_softmaxflashv2_cast_nz_dn_arch92.h:新增 ResetSoftmaxBuffer/ResetSoftmaxBufferVF——VF 向量化按 64 元素/寄存器分组 StoreAlign,一次性复位 sum/max 两个 buffer(替代两次全量 Duplicate)。
attention/flash_attn/op_kernel/utils/attenmask_gs1.h(配套修改):补充 kernel_operator.h/const_def.h 头文件包含;MergeBand/MaskUbCopyS1GVF 等 VF 函数补充 AscendC::Reg:: 命名空间显式限定,保证 arch92 编译通过。
验收方式
flash_attn sinks 功能用例(如 tests/pytests/test_cases/functional_sinks.py,覆盖 BNSD/BSND/TND × PA 各布局、变长/截断/回绕/空 batch 等场景)在 arch92(ascend960dt)平台精度对齐 golden,CI 流水线通过。