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 两种张量布局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]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) - 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 不同sparseIndicesTensor是int32BSND: [B, S1, N2, topK]BNSD: [B, N2, S1, topK]按 KV 头分组的稀疏索引前三维布局与 inputLayout 一致scaleValuefloat是-标量缩放因子通常为 1/sqrt(Dk)inputLayoutstr否--张量布局格式BSND默认或 BNSD输出参数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 和输出的布局保持一致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, ) - torch.Tensor: 稀疏 FlashAttention支持 GQA、不同 head dim 和 BSND/BNSD 布局 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] sparseIndices: 稀疏索引int32BSND: [B, S1, N2, topK]BNSD: [B, N2, S1, topK] scaleValue: 缩放因子 inputLayout: 张量布局BSND 或 BNSD 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 attn_weights torch.softmax(scores, dim-1) 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) # y.shape: [B, S1, N1, Dv]【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考