CANN Bench ESA Select TopK算子

AiInfraEsaSelectTopk 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介AiInfraEsaSelectTopk 是 ESAEfficient Sparse Attention索引选择算子用于在长序列 attention 前为每个 query token 选择需要参与计算的 key block。算子将 key 序列划分为 Initial Tokens、Middle Tokens 和 Local Tokens并在压缩后的 key block 上计算相关性输出 top-k block 索引。主要应用场景长上下文 prefill/chunk prefill 中的稀疏 key block 选择decode 场景下对历史 KV cache 的快速 top-k block 检索ESA 稀疏 attention 前置索引生成算子特征难度等级L3SortSelect输入布局为input_layoutTND支持 float16 / bfloat16 输入固定常用 block 语义blk_size64compress_blk_size162. 算子定义数学公式对每个 query token 和 KV head在压缩 key block 上计算相关性分数$$ score(q, B_j) \max_{k \in B_j} q \cdot k $$输出由 Initial block、Local block、Middle top-k block 以及一个边界补充位置组成$$ selected InitialBlocks \cup TopK(MiddleScores) \cup LocalBlocks \cup ExtraBlock $$输出最后一维长度为$$ select_count init_blk_num local_blk_num topk 1 $$步骤说明按actual_seq_q_len_optional将 TND query 切分为 batch/chunk。按actual_seq_k_len_optional获取每个 batch 的原始 key 长度。按actual_cmp_seq_k_len_optional获取压缩后的 key 长度若未传入则由原始 key 长度和compress_blk_size推导。将压缩 key 进一步按blk_size / compress_blk_size个压缩 token 聚合为候选 block。对 Initial block 保留优先级对未来 block 施加 causal 约束。对 Middle block 做 top-k拼接 Initial、Local 和补充 block输出 int32 block 索引无效位置填 -1。3. 接口规范算子原型cann_bench.ai_infra_esa_select_topk( Tensor query, Tensor key, int blk_size64, int init_blk_num2, int local_blk_num4, int topk4, str input_layoutTND, Tensor|list|None actual_seq_q_len_optionalNone, Tensor|list|None actual_seq_k_len_optionalNone, Tensor|list|None actual_cmp_seq_k_len_optionalNone, int compress_blk_size16, ) - Tensor topk_indices输入参数说明参数类型必需或默认值dtypeshape描述queryTensor必选float16 / bfloat16[Tq,Nq,D]Query 输入TND layoutkeyTensor必选float16 / bfloat16[B,Scmp,Nkv,D]压缩后的 Key 输入Scmp为压缩 key token 数blk_sizeint默认64-标量原始 key block 大小init_blk_numint默认2-标量强制保留的开头 block 数local_blk_numint默认4-标量强制保留的局部 block 数topkint默认4-标量Middle block 中选择的 top-k 数input_layoutstr默认TND--输入布局仅支持TNDactual_seq_q_len_optionalTensor/list?默认Noneint32 / int64[B]query 序列长度或 TND 前缀和actual_seq_k_len_optionalTensor/list?默认Noneint32 / int64[B]原始 key 序列长度或前缀和actual_cmp_seq_k_len_optionalTensor/list?默认Noneint32 / int64[B]压缩后 key 长度非累加和通常约为ceil(key_len / compress_blk_size)compress_blk_sizeint默认16-标量每个压缩 key token 对应的原始 token 数输出参数dtypeshape描述topk_indicesint32[Tq,Nkv,init_blk_numlocal_blk_numtopk1]每个 query token、KV head 对应的候选 key block 索引无效位置为 -1规则与约束标准 Golden 只支持input_layoutTND。blk_size应为compress_blk_size的整数倍例如默认配置下64 / 16 4个压缩 token 组成一个 ESA block。actual_cmp_seq_k_len_optional表示压缩后的 K 长度不是原始 token 长度且按每个 batch 的长度传入而非前缀和。decode 场景每个 batch 的 query length 通常为 1prefill/chunk prefill 支持一个 chunk 内多个 query token。当 key 长度小于需要选择的 block 覆盖范围时输出会按 causal 阈值将未来 block 置为 -1。4. 精度要求本算子输出为离散 int32 索引验证时应逐元素比较输出索引是否一致。通过标准输出判定方式topk_indices与 golden 逐元素一致允许无效位置均为 -15. 标准 Golden 代码标准 Golden 参考实现位于同目录golden.py的ai_infra_esa_select_topk函数使用 Torch 模拟 ESA 压缩 key block 选择、causal 约束和 top-k 索引生成。6. 额外信息算子调用示例import torch import cann_bench tq, b, scmp, nq, nkv, d 1024, 1, 512, 32, 2, 128 query torch.randn(tq, nq, d, dtypetorch.float16, devicenpu) key torch.randn(b, scmp, nkv, d, dtypetorch.float16, devicenpu) indices cann_bench.ai_infra_esa_select_topk( query, key, blk_size64, init_blk_num2, local_blk_num4, topk4, input_layoutTND, actual_seq_q_len_optional[tq], actual_seq_k_len_optional[8192], actual_cmp_seq_k_len_optional[512], compress_blk_size16, )【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考