SparseFlashAttention 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介大序列长度推理场景的高效稀疏注意力计算算子。与标准 FlashAttention 对全部 KV 位置计算注意力不同SparseFlashAttention 通过sparseIndices指定每个 query 位置需要关注的 KV 子集仅对选中的 KV 位置执行缩放点积注意力从而降低计算和内存开销。支持 Grouped Query AttentionGQAquery 头数 N1 可以大于 KV 头数 N2要求 N1 整除 N2每 N1/N2 个 query 头共享同一组 KV 头及其稀疏索引。同时支持 query/key 和 value 使用不同的 head dimDk 和 Dv。主要应用场景大语言模型长序列推理中的高效注意力计算长文本理解与生成任务中降低注意力计算复杂度从 O(S1×S2) 降至 O(S1×topK)需要稀疏注意力模式的 Transformer 推理加速GQA / MLA 架构下的稀疏注意力算子特征难度等级L4FusedComposite四输入query, key, value, sparseIndices单输出融合稀疏索引 gather、缩放点积注意力与 softmax 计算支持 GQAN1 ≥ N2N1 % N2 0和不同 head dimDk、Dv支持 BSND 和 BNSD 两种张量布局支持is_causal因果掩码在稀疏 gather 之上再叠加 [S1, S2] attention 矩阵的右下角对角线掩码sparseIndices 选中但落在对角线之上即 KV 序列位置满足 idx s1 (S2 - S1)的条目在 softmax 前置 -infvalue 与 key 在数值上同源MLA 推理 latent KV cache 语义当Dk Dv时value keyK 与 V 完全相同当Dk Dv典型 MLA: Dk576512_nope64_rope, Dv512value key[..., :Dv]V 是 K 的前 Dv 维前缀切片接口保留 value 为独立入参以兼容通用 attention API调用方需保证此关系2. 算子定义稀疏机制对于每个 KV 头组(b, n2, s1)sparseIndices中对应位置给出该 query 需要关注的 topK 个 KV 序列位置索引值域[0, S2)。该组内的 N1/N2 个 query 头共享相同的稀疏索引attention 计算仅在这 topK 个选中的 KV 位置上进行。sparseIndices的前三维布局与inputLayout一致BSND 布局sparseIndices shape 为[B, S1, N2, topK]BNSD 布局sparseIndices shape 为[B, N2, S1, topK]数学公式$$ K_{sel} \text{gather}(K, \text{sparseIndices}) \quad V_{sel} \text{gather}(V, \text{sparseIndices}) $$$$ y \text{softmax}\left(Q \times K_{sel}^T \times \text{scaleValue}\right) \times V_{sel} $$具体子步骤稀疏 Gather根据sparseIndices从 key/value 中提取选中的 KV 子集以 BNSD 为例$K_{sel}[b, n2, s1, i, :] K[b, n2, \text{sparseIndices}[b, n2, s1, i], :]$shape[B, N2, S1, topK, Dk]$V_{sel}$ 同理shape[B, N2, S1, topK, Dv]GQA 扩展将K_sel和V_sel沿 head 维度复制 GN1/N2 次扩展到[B, N1, S1, topK, ...]缩放点积$\text{scores} Q \cdot K_{sel}^T \times \text{scaleValue}$shape[B, N1, S1, topK]因果掩码可选当is_causalTrue时对每个 query 位置s1和每个 topK 槽位i若sparseIndices[..., s1, i] s1 (S2 - S1)即所选 KV 位置落在 [S1, S2] attention 矩阵右下角对角线之上则该位置 scores 置为 $-\infty$Softmax 归一化$\text{attn_weights} \text{softmax}(\text{scores}, \text{dim}-1)$在 topK 维上归一化加权求和$y \text{attn_weights} \times V_{sel}$布局说明inputLayoutquerykeyvaluesparseIndicesoutputBSND[B, S1, N1, Dk][B, S2, N2, Dk][B, S2, N2, Dv][B, S1, N2, topK][B, S1, N1, Dv]BNSD[B, N1, S1, Dk][B, N2, S2, Dk][B, N2, S2, Dv][B, N2, S1, topK][B, N1, S1, Dv]与标准 FlashAttention 的区别项目FlashAttentionSparseFlashAttention注意力范围全部 S2 个 KV 位置sparseIndices 指定的 topK 个 KV 位置计算复杂度O(S1 × S2 × D)O(S1 × topK × D)topK S2额外输入无sparseIndicessoftmax 范围S2 维topK 维仅选中的 KV 位置GQA 支持是是N1 % N2 0head dimQ/K 共享 DkV 可用不同 Dv同左3. 接口规范算子原型sparse_flash_attention(Tensor query, Tensor key, Tensor value, Tensor sparseIndices, float scaleValue, str inputLayoutBSND, bool is_causalFalse) - Tensor y输入参数说明参数类型必需dtypeshape描述queryTensor是float16/bfloat16BSND: [B, S1, N1, Dk]BNSD: [B, N1, S1, Dk]查询张量keyTensor是float16/bfloat16BSND: [B, S2, N2, Dk]BNSD: [B, N2, S2, Dk]键张量head dim 与 query 一致valueTensor是float16/bfloat16BSND: [B, S2, N2, Dv]BNSD: [B, N2, S2, Dv]值张量head dim 可与 key 不同。语义上要求value key[..., :Dv]同一份 latent KV cache 的前缀视图DkDv 时退化为value keysparseIndicesTensor是int32BSND: [B, S1, N2, topK]BNSD: [B, N2, S1, topK]按 KV 头分组的稀疏索引前三维布局与 inputLayout 一致scaleValuefloat是-标量缩放因子通常为 1/sqrt(Dk)inputLayoutstr否--张量布局格式BSND默认或 BNSDis_causalbool否--是否启用因果掩码右下角对齐默认 False。True 时在稀疏 gather 之上额外屏蔽sparseIndices[..., s1, i] s1 (S2 - S1)的位置输出参数dtypeshape描述y与 query 一致BSND: [B, S1, N1, Dv]BNSD: [B, N1, S1, Dv]稀疏注意力输出布局与输入一致数据类型query/key/value dtypesparseIndices dtype输出 dtypebfloat16int32bfloat16float16int32float16规则与约束query、key、value 的 dtype 必须一致query 和 key 的 head dim 必须一致Dkvalue 的 head dimDv可以不同N1 必须整除 N2GQA 分组约束N1 N2 时退化为 MHAsparseIndices 的值域为[0, S2)即 KV 序列长度范围内的有效索引topK每个 query 关注的 KV 数量可任意取值1 ≤ topK ≤ S2scaleValue 通常设置为 $1/\sqrt{Dk}$inputLayout 必须为 BSND 或 BNSD所有张量含 sparseIndices 和输出的布局保持一致is_causalTrue时要求 $S1 \le S2$若某个 query 行的全部 topK 槽位都被因果掩码屏蔽极端情况该行 softmax 输出按全 0 处理语义约束value key[..., :Dv]value 必须等于 key 的前 Dv 维前缀同一份 latent KV cache 的不同视图。当Dk Dv时退化为value key当Dk Dv 64MLA 典型 Dk576/Dv512时 value 即 key_nope。算子接口保留 value 为独立入参以兼容通用 attention API调用方有责任保证此关系本规范的 Golden 实现不强制检查支持范围输入 tensor 各维度与参数的支持范围维度 / 参数范围备注Bbatch1 ~ 64cases.csv 实测 1 ~ 32S1query 序列长度1 ~ 2048cases.csv 实测 1 ~ 1024S11 对应 decodeS12/4 对应 MTPS1 较大对应 prefillS2KV 序列长度1 ~ 16384cases.csv 实测 1024 ~ 8192is_causalTrue时要求S1 ≤ S2N1query 头数1 ~ 256cases.csv 实测 32 ~ 128N2KV 头数1 ~ 64cases.csv 实测 1 / 8必须满足N1 % N2 0GQA 分组约束N1 N2时退化为 MHADkquery/key head dim64 对齐cases.csv 实测 128 / 192 / 512 / 576典型 MHA/GQA Dk128 或 192典型 MLA Dk512nope或 Dk576nope512rope64Dvvalue head dim64 对齐Dv ≤ Dkcases.csv 实测 128 / 512语义上value key[..., :Dv]topK稀疏选中数量1 ~ 2048cases.csv 实测 512 / 1024要求1 ≤ topK ≤ S2sparseIndices取值[0, S2)int32必须为 KV 序列长度范围内的有效索引scaleValue任意正浮点cases.csv 实测 0.04167 ~ 0.08838约1/sqrt(Dk)inputLayout{BSND, BNSD}cases.csv 实测两种均覆盖所有张量含 sparseIndices 和输出的布局保持一致is_causal{True, False}cases.csv 实测两种均覆盖True 时屏蔽sparseIndices[..., s1, i] s1 (S2 - S1)的位置dtypefloat16 / bfloat16cases.csv 实测两种均覆盖query/key/value dtype 必须一致sparseIndices 固定 int32约束N1 % N2 0Dv ≤ Dk且语义上要求value key[..., :Dv]is_causalTrue时要求S1 ≤ S2若某 query 行的全部 topK 槽位都被因果掩码屏蔽该行 softmax 输出按全 0 处理。4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch def sparse_flash_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, sparseIndices: torch.Tensor, scaleValue: float, inputLayout: str BSND, is_causal: bool False, ) - torch.Tensor: 稀疏 FlashAttention支持 GQA、不同 head dim 和 BSND/BNSD 布局 语义约束: value 与 key 在数值上同源 (MLA latent KV cache): - Dk Dv 时: value key - Dk Dv 时: value key[..., :Dv] (典型 MLA: Dk576, Dv512) 本 Golden 不做强制检查调用方需保证此关系。 Args: query: 查询张量BSND: [B, S1, N1, Dk]BNSD: [B, N1, S1, Dk] key: 键张量BSND: [B, S2, N2, Dk]BNSD: [B, N2, S2, Dk] value: 值张量BSND: [B, S2, N2, Dv]BNSD: [B, N2, S2, Dv]语义上等于 key[..., :Dv] sparseIndices: 稀疏索引int32BSND: [B, S1, N2, topK]BNSD: [B, N2, S1, topK] scaleValue: 缩放因子 inputLayout: 张量布局BSND 或 BNSD is_causal: 是否启用因果掩码右下角对齐True 时在稀疏 gather 之上额外屏蔽 sparseIndices[..., s1, i] s1 (S2 - S1) 的位置 Returns: 注意力输出布局与输入一致 # 统一转为 BNSD 内部计算 if inputLayout BSND: q query.permute(0, 2, 1, 3) # [B, N1, S1, Dk] k key.permute(0, 2, 1, 3) # [B, N2, S2, Dk] v value.permute(0, 2, 1, 3) # [B, N2, S2, Dv] si sparseIndices.permute(0, 2, 1, 3) # [B, N2, S1, topK] else: # BNSD q, k, v, si query, key, value, sparseIndices B, N1, S1, Dk q.shape N2 k.shape[1] S2 k.shape[2] Dv v.shape[-1] topK si.shape[-1] G N1 // N2 # sparseIndices 为 int32转为 long 用于 gather si si.long() # Gather 选中的 KV: si [B, N2, S1, topK] idx_k si.reshape(B, N2, S1 * topK).unsqueeze(-1).expand(-1, -1, -1, Dk) idx_v si.reshape(B, N2, S1 * topK).unsqueeze(-1).expand(-1, -1, -1, Dv) k_sel k.gather(2, idx_k).reshape(B, N2, S1, topK, Dk) # [B, N2, S1, topK, Dk] v_sel v.gather(2, idx_v).reshape(B, N2, S1, topK, Dv) # [B, N2, S1, topK, Dv] # GQA: 将 KV 头扩展到 N1 个 query 头 k_sel k_sel.unsqueeze(2).expand(-1, -1, G, -1, -1, -1).reshape(B, N1, S1, topK, Dk) v_sel v_sel.unsqueeze(2).expand(-1, -1, G, -1, -1, -1).reshape(B, N1, S1, topK, Dv) # Attention: Q K_sel^T - softmax - V_sel scores torch.einsum(bnsd,bnskd-bnsk, q, k_sel) * scaleValue if is_causal: # 右下角对齐sparseIndices 选中的 KV 位置 idx s1 (S2 - S1) 时需屏蔽 s1_idx torch.arange(S1, devicescores.device).view(1, 1, S1, 1) # [1,1,S1,1] # si 已为 [B, N2, S1, topK]需扩展到 N1 si_n1 si.unsqueeze(2).expand(-1, -1, G, -1, -1).reshape(B, N1, S1, topK) causal_mask si_n1 (s1_idx (S2 - S1)) scores scores.masked_fill(causal_mask, float(-inf)) attn_weights torch.softmax(scores, dim-1) # 处理整行全 -inf 导致的 NaN置 0 attn_weights torch.nan_to_num(attn_weights, nan0.0) out torch.einsum(bnsk,bnskd-bnsd, attn_weights, v_sel) # [B, N1, S1, Dv] # 转回原始布局 if inputLayout BSND: return out.permute(0, 2, 1, 3) # [B, S1, N1, Dv] else: return out # [B, N1, S1, Dv]6. 额外信息算子调用示例import torch B, S1, S2, N1, N2, Dk, Dv, topK 2, 1024, 8192, 32, 8, 128, 128, 512 # BSND 布局 query torch.randn(B, S1, N1, Dk, dtypetorch.float16, devicenpu) key torch.randn(B, S2, N2, Dk, dtypetorch.float16, devicenpu) value torch.randn(B, S2, N2, Dv, dtypetorch.float16, devicenpu) sparseIndices torch.stack([ torch.randperm(S2)[:topK] for _ in range(B * N2 * S1) ]).reshape(B, S1, N2, topK).to(dtypetorch.int32, devicenpu) y sparse_flash_attention(query, key, value, sparseIndices, scaleValue1.0 / (Dk ** 0.5), inputLayoutBSND, is_causalFalse) # y.shape: [B, S1, N1, Dv]【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
