CANN/cann-bench MLA算子API文档

MLA 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介多头潜在注意力 (Multi-Head Latent Attention) 算子仅包含注意力计算部分不含 KV 解压缩是 DeepSeek-V2/V3 等模型的核心注意力机制。主要应用场景大语言模型推理中的 MLA 注意力计算如 DeepSeek-V2、DeepSeek-V3超长序列推理场景需要在推理效率和模型质量之间取得平衡的大规模 Transformer 架构算子特征难度等级L4FusedComposite多输入q_nope, q_rope, k_nope, k_rope, v单输出Q 和 K 均分为 nope 和 rope 两部分传入内部拼接后计算注意力V 与 K_nope 共享 d_nope 维度输出 head dim 为 d_nope支持 GQA 模式N_q 个 query head 共享 N_kv 个 KV headMLA 典型配置 N_kv1支持 BSND 和 BNSD 两种输入输出 layout2. 算子定义数学公式$$ Q \text{concat}(Q_{nope}, Q_{rope}) \quad \text{dim: } d_{nope} d_{rope} $$$$ K \text{concat}(K_{nope}, K_{rope}) \quad \text{dim: } d_{nope} d_{rope} $$$$ y \text{softmax}\left(Q \times K^T \times \text{scaleValue}\right) \times V $$其中$Q_{nope}$、$K_{nope}$ 为 nope 部分维度 $d_{nope}$$Q_{rope}$、$K_{rope}$ 为 rope 部分经过 RoPE 编码维度 $d_{rope}$$V$ 的 head dim 为 $d_{nope}$因此输出的 head dim 也为 $d_{nope}$$\text{scaleValue}$ 为缩放因子0 时自动使用 $1/\sqrt{d_{nope} d_{rope}}$当 $N_q N_{kv}$ 时进行 GQA 扩展每个 KV head 被 $N_q / N_{kv}$ 个 query head 共享具体子步骤Q 拼接$Q \text{concat}(Q_{nope}, Q_{rope})$K 拼接$K \text{concat}(K_{nope}, K_{rope})$GQA 扩展将 KV head 复制 $N_q / N_{kv}$ 次匹配 query head 数缩放点积$\text{scores} Q \times K^T \times \text{scaleValue}$Softmax 归一化$\text{attn_weights} \text{softmax}(\text{scores}, \text{dim}-1)$加权求和$y \text{attn_weights} \times V$3. 接口规范算子原型cann_bench.mla(Tensor q_nope, Tensor q_rope, Tensor k_nope, Tensor k_rope, Tensor v, int numKVHeads1, float scaleValue-1.0, str inputLayoutBSND) - Tensor y输入参数说明参数类型默认值描述q_nopeTensor必选query 的 nope 部分BSND: [B, S, N_q, d_nope]BNSD: [B, N_q, S, d_nope]q_ropeTensor必选query 的 rope 部分BSND: [B, S, N_q, d_rope]BNSD: [B, N_q, S, d_rope]k_nopeTensor必选key 的 nope 部分BSND: [B, S_kv, N_kv, d_nope]BNSD: [B, N_kv, S_kv, d_nope]k_ropeTensor必选key 的 rope 部分BSND: [B, S_kv, N_kv, d_rope]BNSD: [B, N_kv, S_kv, d_rope]vTensor必选值张量BSND: [B, S_kv, N_kv, d_nope]BNSD: [B, N_kv, S_kv, d_nope]numKVHeadsint1KV 头数scaleValuefloat-1.0缩放因子0 时自动使用 1/sqrt(d_nope d_rope)inputLayoutstrBSND输入输出 layoutBSND 或 BNSD输出参数Shapedtype描述yBSND: [B, S, N_q, d_nope]BNSD: [B, N_q, S, d_nope]与输入相同注意力输出张量head dim 为 d_nope数据类型输入 dtype输出 dtypefloat16float16float32float32bfloat16bfloat16规则与约束所有输入 Tensor 的 dtype 必须一致N_q 必须能被 N_kv 整除q_nope 和 k_nope 的 head dim 一致d_nopeq_rope 和 k_rope 的 head dim 一致d_ropev 的 head dim d_nope输出的 head dim 也为 d_nope所有输入和输出遵循相同的 layoutBSND 或 BNSD4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差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 MLA算子Torch Golden参考实现 多头潜在注意力 (Multi-Head Latent Attention)仅包含注意力计算部分 Q 和 K 均分为 nope 和 rope 两部分传入内部拼接后计算注意力 V 与 K_nope 共享 d_nope 维度 支持 BSND 和 BNSD 两种输入 layout 公式: Q concat(Q_nope, Q_rope) dim: d_nope d_rope K concat(K_nope, K_rope) dim: d_nope d_rope y softmax(Q K^T * scaleValue) V def mla( q_nope: torch.Tensor, q_rope: torch.Tensor, k_nope: torch.Tensor, k_rope: torch.Tensor, v: torch.Tensor, numKVHeads: int 1, scaleValue: float -1.0, inputLayout: str BSND, ) - torch.Tensor: 多头潜在注意力 (Multi-Head Latent Attention) Args: q_nope: query 的 nope 部分BSND: [B, S, N_q, d_nope]BNSD: [B, N_q, S, d_nope] q_rope: query 的 rope 部分BSND: [B, S, N_q, d_rope]BNSD: [B, N_q, S, d_rope] k_nope: key 的 nope 部分BSND: [B, S_kv, N_kv, d_nope]BNSD: [B, N_kv, S_kv, d_nope] k_rope: key 的 rope 部分BSND: [B, S_kv, N_kv, d_rope]BNSD: [B, N_kv, S_kv, d_rope] v: 值张量BSND: [B, S_kv, N_kv, d_nope]BNSD: [B, N_kv, S_kv, d_nope] numKVHeads: KV 头数 scaleValue: 缩放因子0 时自动使用 1/sqrt(d_nope d_rope) inputLayout: 输入 layoutBSND 或 BNSD Returns: 输出张量与输入 layout 一致head dim 为 d_nope # 统一转为 BSND 内部计算 if inputLayout BNSD: q_nope q_nope.permute(0, 2, 1, 3) q_rope q_rope.permute(0, 2, 1, 3) k_nope k_nope.permute(0, 2, 1, 3) k_rope k_rope.permute(0, 2, 1, 3) v v.permute(0, 2, 1, 3) B, S, N_q, d_nope q_nope.shape d_rope q_rope.shape[-1] D_qk d_nope d_rope S_kv k_nope.shape[1] N_kv numKVHeads if scaleValue 0: scaleValue 1.0 / (D_qk ** 0.5) # 拼接 Q [Q_nope, Q_rope]: [B, S, N_q, d_nope d_rope] q torch.cat([q_nope, q_rope], dim-1) # 拼接 K [K_nope, K_rope]: [B, S_kv, N_kv, d_nope d_rope] k torch.cat([k_nope, k_rope], dim-1) # GQA 扩展: 每个 KV head 复制 N_q // N_kv 次 G N_q // N_kv if G 1: k k.unsqueeze(3).expand(B, S_kv, N_kv, G, D_qk).reshape(B, S_kv, N_q, D_qk) v v.unsqueeze(3).expand(B, S_kv, N_kv, G, d_nope).reshape(B, S_kv, N_q, d_nope) # 转置为 [B, N, S, D] q q.transpose(1, 2) k k.transpose(1, 2) v v.transpose(1, 2) # 缩放点积注意力 scores torch.matmul(q, k.transpose(-2, -1)) * scaleValue attn_weights torch.nn.functional.softmax(scores, dim-1) out torch.matmul(attn_weights, v) # [B, N_q, S, d_nope] # 转回 BSND: [B, S, N_q, d_nope] out out.transpose(1, 2) # 按输入 layout 输出 if inputLayout BNSD: out out.permute(0, 2, 1, 3) return out6. 额外信息算子调用示例import torch import cann_bench B, S, S_kv 2, 128, 256 N_q, N_kv 128, 1 d_nope, d_rope 512, 64 # BSND layout q_nope torch.randn(B, S, N_q, d_nope, dtypetorch.float16, devicenpu) q_rope torch.randn(B, S, N_q, d_rope, dtypetorch.float16, devicenpu) k_nope torch.randn(B, S_kv, N_kv, d_nope, dtypetorch.float16, devicenpu) k_rope torch.randn(B, S_kv, N_kv, d_rope, dtypetorch.float16, devicenpu) v torch.randn(B, S_kv, N_kv, d_nope, dtypetorch.float16, devicenpu) y cann_bench.mla(q_nope, q_rope, k_nope, k_rope, v, numKVHeadsN_kv, scaleValue-1.0, inputLayoutBSND) # y shape: [B, S, N_q, d_nope] [2, 128, 128, 512]【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考