算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载GenericBlockSparseAttention 是 CANN ops-transformer 中基于 CATLASS 模板库实现的高性能块稀疏注意力算子支持沿 S 轴序列维任意粒度的块稀疏模式并通过 Paged KV CachePA_BBND与 Packed GQA 特性服务于长序列大模型推理与训练场景。阅读本文后你将掌握该算子的稀疏分块原理、metadata 前置调度流程、aclnn 两段式 C 接口与 PyTorch 扩展接口的完整调用方法以及全部参数、约束与量化/掩码配置的实操细节。产品支持情况GenericBlockSparseAttention 算子及配套的 metadata 前置算子在以下产品上获得支持产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×产品支持信息同时体现在算子定义中在 generic_block_sparse_attention_def.cpp 中通过AICore().AddConfig(ascend910b)、AddConfig(ascend910_93)、AddConfig(ascend950)注册了对应芯片的 AICore 配置。功能说明稀疏注意力计算机制GenericBlockSparseAttention 沿序列维S 轴按块做稀疏计算。Q 侧按blockShapeX、KV 侧按blockShapeY划分稀疏块稀疏块大小为$$blockShapeX \times blockShapeY$$sparseBlockIdx指定每个 Q 块实际选择的 KV 块索引sparseBlockCount指定每个 Q 块实际保留的 KV 块数量。计算时只对选中的块执行 $qk^{T}$、Softmax 以及与 $v$ 的乘积计算公式为$$ attentionOut Softmax(scaleValue \cdot query \cdot key_{sparse}^{T} atten_mask) \cdot value_{sparse} $$其中 $key_{sparse}$、$value_{sparse}$ 为按sparseBlockIdx/sparseBlockCount从 Paged KV Cache 中选取的 KV 块。这种每个 Q 块只取少量 KV 块的方式相比稠密 Attention 大幅减少了参与计算的数据量是长序列场景下控制计算与访存开销的关键手段。两段式调用metadata 前置算子该算子采用metadata 生成 主算子执行的两段式架构这是其与普通 Flash Attention 算子最大的差异点准备query、key、value、sparseBlockIdx、sparseBlockCount等输入先调用aclnnGenericBlockSparseAttentionMetadataPyTorch 侧为generic_block_sparse_attention_metadata生成metadataOptional再调用aclnnGenericBlockSparseAttentionPyTorch 侧为generic_block_sparse_attention将上一步得到的metadataOptional传入主算子。metadata 记录的是 AICore/AIVCore 的任务切分结果即负载均衡调度信息主算子传入后可以优化任务调度。从 generic_block_sparse_attention.py 的源码可以看到metadata 是一个 shape 固定为(1024,)的 int32 TensorGBSA_METADATA_SIZE 1024。实现架构与 Kernel 侧支撑从源码结构看算子内核在 op_kernel 目录下按芯片架构分目录组织arch22/Atlas A2 系列AICore 220的内核实现arch35/Atlas A3/Ascend 950 系列AICore 310的内核实现包含完整量化路径generic_block_sparse_attention_kernel_arch35_full_quant.hattn_infra/跨架构复用的注意力基础设施包括 gemmQK/PV 分块计算、epilogueonline softmax、rescale、结果写出、layout 与坐标管理等子模块。主内核入口 generic_block_sparse_attention.cpp 通过TILING_KEY在编译期展开不同组合FP16/BF16/FP8、LSE 是否输出、softmaxPrecision 等。例如 arch22 上GBSA_FP16_TND_PA_BBND_TILING对应softmaxPrecision0fp32 Softmax RescaleGBSA_FP16_TND_PA_BBND_HALFSM_TILING对应softmaxPrecision1half Softmax fp32 Rescale这与下文的精度配置一一对应。参数说明说明参数维度含义——B 表示 Batch SizeT 表示 Total tokensN 表示 Head NumD 表示 Head DimtopK 表示sparseBlockIdx最后一维maxSparseBlockCount。TND 中的 N 为 query 的 headNum记为 N1PA_BBND 中的 N 为 key/value 的 headNum记为 N2GQA 下 N1 与 N2 可以不同约束见下文。参数名输入/输出/属性描述数据类型数据格式query输入公式中的 query。layoutQ 为 TND 时shape 为 [T, N, D]N 为 query 的 headNumN1FLOAT16、BFLOAT16、FLOAT8_E4M3FNNDkey输入公式中的 key。layoutKv 为 PA_BBND 时shape 为 [numBlocks, blockSize, N, D]N 为 kv 的 headNumN2FLOAT16、BFLOAT16、FLOAT8_E4M3FNNDvalue输入公式中的 valueshape 与 key 一致FLOAT16、BFLOAT16、FLOAT8_E4M3FNNDsparseBlockIdx输入稀疏块索引。TND isPackedGQA1 时shape 为 [N, totalQBlocks, topK]N 为 kv 的 headNumN2无效位置可用 -1 填充有效值须落在前 sparseBlockCount 个位置INT32NDsparseBlockCount输入每个 Q 块实际选择的 KV 块数量。TND isPackedGQA1 时shape 为 [N, totalQBlocks]INT32NDcuSeqLengthsQOptional输入各 batch 中 query 序列长度前缀和layoutQ 为 TND 时必传shape 为 [B1]第 0 个元素为 0最后一个元素等于 totalQTokens相邻差分得到各 batch 的存储长度INT64NDcuSeqLengthsKvOptional输入各 batch 中 key/value 序列长度前缀和layoutKv 为 TND 时必传非 TND如 PA_BBND时不传shape 为 [B1]INT64NDsequsedQOptional可选输入各 batch 中 query 实际有效长度不传时按 cu 前缀和差分得到的存储长度处理shape 为 [B]INT32NDsequsedKvOptional输入各 batch 中 kv 实际有效长度layoutKv 为 PA_BBND 时必传shape 为 [B]INT32NDblockTableOptional输入PagedAttention 页表shape 为 [B, maxNumBlocksPerBatch]值只能为正整数INT32NDblockShape属性稀疏块形状 [blockShapeX, blockShapeY]当前仅支持 [1, 128]blockShapeX 支持任意值blockShapeY 支持按 16 对齐的任意值均不超过 int64 范围INT64-isPackedGQA属性同 group 内 qHead 是否共享稀疏 pattern当前仅支持 1TrueINT64-layoutQ属性query 数据排布格式目标支持 TND/BNSD/BSND当前仅支持 TNDSTRING-layoutKv属性key/value 数据排布格式目标支持 TND/BNSD/BSND/PA_BBND/PA_BNBD当前仅支持 PA_BBNDSTRING-layoutSparsePattern属性sparseBlockIdx、sparseBlockCount 的数据排布格式当前仅支持取 4INT64-scaleValue属性缩放系数传 0 时算子内按 $1/\sqrt{D}$ 处理一般设置为 D^-0.5DOUBLE-maskType属性掩码类型取值 0~5当前仅支持 1内置 causal maskINT64-quantType属性量化类型当前支持 0Ascend 950 上可选 5取值 1~4 传入将校验失败INT64-dstTypeMax属性MXFP4 CX 量化时传入的自定义量化量程当前版本不支持自定义量程必须传入 0.0DOUBLE-softmaxPrecision属性Softmax 计算精度级别取值 0 或 1详见下文Softmax 精度INT64-winLeft / winRight属性滑窗 attention 场景的前向/后向窗口 token 数当前不支持滑窗只支持传入 -1INT64-residualBlockMode属性KV 序列按 blockShapeY 稀疏后尾部不完整块的状态仅支持 0 或 1见下文约束INT64-isConsistentTopK属性同一 batch 同一 head 内每个 Q 块选择的 KV 块最大数量是否一致仅支持 0 或 1BOOL-returnSoftmaxlse属性是否输出 softmaxLse当前仅支持 0INT64-attentionOut输出公式中的 attentionOut数据类型和 shape 与 query 保持一致FP8 输入时由本 tensor 指定输出 dtypeFLOAT16、BFLOAT16NDsoftmaxLseOptional输出Softmax log-sum-exp 中间结果当前不支持须传入 nullptrreturnSoftmaxlse 须为 0FLOATND在算子定义文件 generic_block_sparse_attention_def.cpp 中可以看到这些属性的默认值block_shape默认[1, 128]、is_packed_gqa默认 1、layout_q默认 TND、scale_value默认 0.0、mask_type默认 0、quant_type默认 0、win_left/win_right默认 -1与上文表格一致。同时该文件还揭示了非连续 Tensor 的处理策略key/value使用IgnoreContiguous()允许非连续输入而query、sparseBlockIdx等使用AutoContiguous()自动做连续化。约束说明通用约束metadata 前置要求调用前须先执行aclnnGenericBlockSparseAttentionMetadata生成metadataOptional再调用本接口metadata 须与当前输入/属性配套每次调用须重新生成。metadata 接口与主算子的 sparseBlockIdx、sparseBlockCount、blockShape、mask_mode、quant_mode、is_packed_gqa 等参数须完全一致否则产生未定义行为精度问题或非法内存访问导致的崩溃。维度约束query/key/value 的 headDimD当前仅支持 128KV 页 blockSize 当前仅支持 128且须等于 blockShapeY。分块与寻址TND isPackedGQA1 时totalQBlocks 按 cuSeqLengthsQ 差分得到的存储长度分块$\sum_i \mathrm{ceilDiv}(qStorageLen_i, blockShapeX)$sparse 分块与 QKV 寻址均按该存储长度不以 seqused 重切分topK 须 ≥ sparseBlockCount 中所有元素的最大值当前上限为 256。双长度语义seqused 与对应 cu 前缀和同时传入时Q 侧及 layoutKv 为 TND 时的 KV 侧分核/任务空间按各 batch 实际有效长度seqused累加各 batch 的 seqused 元素须 ≤ 对应 cu 存储长度且须与 Metadata 侧完全一致。PA_BBND 特有约束layoutKv 为 PA_BBND 时须传sequsedKvOptional不传cuSeqLengthsKvOptional。数据类型一致性输入 query、key、value 的数据类型必须一致。GQA 约束query 的 headNum 为 N1key/value 的 headNum 为 N2则 N1 ≥ N2 且 N1 % N2 0groupSize N1/N2 当前须 ≤ 128。非连续约束PA_BBND 下 key/value 仅 dim0物理页轴可非连续页内 blockSize × N2 × D 须连续且 stride0 ≥ blockSize × N2 × D 并按 N2 × D 对齐。确定性aclnnGenericBlockSparseAttention 默认确定性实现。空 Tensor必选输入和输出 shape 中任意轴为 0 的空 Tensor 用例将全部拦截报错。Tiling 期不校验值cu_seqlens、seqused、sparseBlockIdx、sparseBlockCount 及 blockTable 等 Tensor 在 Tiling 阶段无法获取具体数值tiling 侧不对其值进行校验正确性需要用户自行保证。Softmax 精度softmaxPrecisionsoftmaxPrecision 控制 online softmax 阶段以及 rescale 阶段运算使用的数据类型0online softmax 和 rescale 全部采取 fp32适合追求计算精度的场景1混合精度online softmax 采取 fp16/bf16与 attentionOut 相同rescale 采取 fp32online softmax 阶段可能数值溢出。芯片约束Ascend 950 仅支持 1Atlas A2/A3 上 FLOAT16 可配置 0 或 1BFLOAT16 仅支持 0FP8 路径仅支持 1。内核侧的对应关系可参见 generic_block_sparse_attention.cpp 中TILING_KEY的注释。掩码说明maskTypemaskType含义attentionMaskOptionalwinLeft/winRight0不加 mask不传-1/-11causal mask不传 attenMaskOptional内置 causal-1/-12window mask规划BOOL [2048,2048] 下三角与 winLeft/winRight 配合实际 window 包含的向前/向后看 token 数3~5各类特化 mask后续补充 mask 描述-1/-1当前仅支持 maskType1算子内置 causalattenMaskOptional 须为 nullptrwinLeft/winRight 为 -1maskType 为 0/2/3~5 当前不支持。量化说明quantType量化配置按 Ascend 950A5代际描述完整配置如下quantTypeQKV 数据类型对称/非对称P 量化动态/静态量化粒度量化参数 shape量化参数 dType0非量化QKV 直接作为输入计算---q/k/vDequantScaleOptional、pQuantScaleOptional 均不传-1FLOAT8_E4M3对称静态perGroupQKV 均沿 S 维度分组group 大小和稀疏块尺寸必须相同KV 为 paged cache 时 blockSize 需为 blockShapeY 的整数倍q/k/vDequantScaleOptional 必选pQuantScaleOptional 可选传入为 [1] 静态系数nullptr 时默认 448.0FLOAT322FLOAT8_E4M3对称动态micro scalingQKV 沿矩阵乘累加轴按固定大小 32 分组KV 为 paged cache 时 blockSize 需为 64 的整数倍q/k/vDequantScaleOptional 必选FLOAT8_E4M33FLOAT4_E2M1对称动态 OCP同 quantType2同 quantType2FLOAT8_E4M34FLOAT4_E2M1对称动态 CX同 quantType2同 quantType2FLOAT8_E4M35FLOAT8_E4M3对称静态不传入量化系数算子内直接将 P cast 成 fp8q/k/vDequantScaleOptional、pQuantScaleOptional 均不传FLOAT32当前可用 quantType0Ascend 950 上可选 quantType5。quantType1~4 当前不支持传入将校验失败。quantType0 时 q/k/v_dequant_scale 须为 NonequantType≠0 时 attention_out_dtype 必须传入量化场景输出 dtype 须为 float16 或 bfloat16。在 arch35 内核中量化路径由GbsaInferInterfaceFullQuantfp8_e4m3fn_t, ...模板实例化对应 generic_block_sparse_attention_kernel_arch35_full_quant.h实现了 FP8 输入的完整反量化与计算流程。调用说明GenericBlockSparseAttention 提供两套调用方式调用方式样例代码说明aclnn APItest_aclnn_generic_block_sparse_attention.cpp通过 aclnnGenericBlockSparseAttention 两段式接口调用 GenericBlockSparseAttention 算子PyTorch API见下文 Python 示例通过 generic_block_sparse_attention 接口调用 generic_block_sparse_attention 算子aclnn API两段式接口调用aclnn 接口为两段式设计先调用aclnnGenericBlockSparseAttentionGetWorkspaceSize获取 workspace 大小与执行器再调用aclnnGenericBlockSparseAttention执行计算。第一段接口原型aclnnStatus aclnnGenericBlockSparseAttentionGetWorkspaceSize( const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *sparseBlockIdx, const aclTensor *sparseBlockCount, const aclTensor *metadataOptional, const aclTensor *attenMaskOptional, const aclTensor *qDequantScaleOptional, const aclTensor *kDequantScaleOptional, const aclTensor *vDequantScaleOptional, const aclTensor *pQuantScaleOptional, const aclTensor *cuSeqLengthsQOptional, const aclTensor *cuSeqLengthsKvOptional, const aclTensor *sequsedQOptional, const aclTensor *sequsedKvOptional, const aclTensor *blockTableOptional, const aclIntArray *blockShape, char *layoutQ, char *layoutKv, int64_t layoutSparsePattern, double scaleValue, int64_t maskType, int64_t quantType, double dstTypeMax, int64_t softmaxPrecision, int64_t winLeft, int64_t winRight, int64_t returnSoftmaxlse, int64_t residualBlockMode, bool isConsistentTopK, aclTensor *attentionOut, aclTensor *softmaxLseOptional, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnGenericBlockSparseAttention( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)典型的 C 调用流程完整示例见 examples/test_aclnn_generic_block_sparse_attention.cpp 与 docs/aclnnGenericBlockSparseAttention.md#include acl/acl.h #include aclnnop/aclnn_generic_block_sparse_attention.h #include aclnnop/aclnn_generic_block_sparse_attention_metadata.h // 1. 构造输入/输出 aclTensorquery: [T,N1,D] fp16key/value: [numBlocks,blockSize,N2,D] fp16 // sparseBlockIdx: [N2,totalQBlocks,topK] int32sparseBlockCount: [N2,totalQBlocks] int32 // cuSeqLengthsQ: [B1] int64sequsedKv: [B] int32blockTable: [B,maxBlocks] int32 // 构造方式见样例中的 CreateAclTensor 辅助函数aclrtMalloc aclrtMemcpy aclCreateTensor // 2. 先调用 Metadata 算子生成任务切分结果 uint64_t metadataWorkspaceSize 0; aclOpExecutor* metadataExecutor nullptr; ret aclnnGenericBlockSparseAttentionMetadataGetWorkspaceSize( sparseIdx, sparseCount, cuSeqQ, nullptr, nullptr, sequsedKv, S1, S2, N1, N2, D, blockShape, layoutQ, layoutKv, /*layoutSparsePattern*/4, /*maskType*/1, /*quantType*/0, /*softmaxPrecision*/1, -1, -1, 0, 0, metadata, metadataWorkspaceSize, metadataExecutor); // 申请 workspace 后执行 ret aclnnGenericBlockSparseAttentionMetadata(metadataWorkspaceAddr, metadataWorkspaceSize, metadataExecutor, stream); ret aclrtSynchronizeStream(stream); // 3. 调用主算子 GetWorkspaceSize 并执行 uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnGenericBlockSparseAttentionGetWorkspaceSize( q, k, v, sparseIdx, sparseCount, metadata, /*attenMask*/nullptr, /*qDequantScale*/nullptr, /*kDequantScale*/nullptr, /*vDequantScale*/nullptr, /*pQuantScale*/nullptr, cuSeqQ, /*cuSeqLengthsKv*/nullptr, /*sequsedQ*/nullptr, sequsedKv, blockTable, blockShape, layoutQ, layoutKv, /*layoutSparsePattern*/4, scaleValue, /*maskType*/1, /*quantType*/0, /*dstTypeMax*/0.0, /*softmaxPrecision*/1, /*winLeft*/-1, /*winRight*/-1, /*returnSoftmaxlse*/0, /*residualBlockMode*/0, /*isConsistentTopK*/0, attnOut, /*softmaxLse*/nullptr, workspaceSize, executor); // 申请 workspace 后执行 ret aclnnGenericBlockSparseAttention(workspaceAddr, workspaceSize, executor, stream); ret aclrtSynchronizeStream(stream);第一段接口完成入参校验常见返回码如下返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001query/key/value/sparseBlockIdx/sparseBlockCount/attentionOut 等必选指针为空ACLNN_ERR_PARAM_INVALID161002layout、maskType、blockShape、softmaxPrecision、quantType、returnSoftmaxlse、layoutSparsePattern、residualBlockMode、isConsistentTopK 等与约束不匹配ACLNN_ERR_INNER_NULLPTR561103metadata 为空或 Contiguous/InferShape 失败如 layout 不支持、缺少 blockTable 等PyTorch APITorchNPU 扩展接口PyTorch 侧接口定义于 torch_extension/generic_block_sparse_attention.py通过 TorchNPU 的cann_ops_transformer扩展以torch.library机制注册torch.ops.cann_ops_transformer.generic_block_sparse_attention。前置 metadata 接口cann_ops_transformer.generic_block_sparse_attention_metadata( sparse_block_idx, sparse_block_count, num_heads_q, num_heads_kv, head_dim, block_shape, *, cu_seqlens_qNone, cu_seqlens_kvNone, seqused_qNone, seqused_kvNone, max_seqlen_q-1, max_seqlen_kv-1, is_packed_gqaTrue, layout_qTND, layout_kvPA_BBND, mask_mode1, quant_mode0, softmax_precision1, win_left-1, win_right-1, ) - Tensor主算子接口cann_ops_transformer.generic_block_sparse_attention( q, k, v, sparse_block_idx, sparse_block_count, block_shape, *, metadataNone, attn_maskNone, q_dequant_scaleNone, k_dequant_scaleNone, v_dequant_scaleNone, p_quant_scaleNone, cu_seqlens_qNone, cu_seqlens_kvNone, seqused_qNone, seqused_kvNone, block_tableNone, is_packed_gqaTrue, layout_qTND, layout_kvPA_BBND, softmax_scale0.0, mask_mode1, quant_mode0, dst_type_max0.0, softmax_precision1, win_left-1, win_right-1, return_softmax_lseFalse, attention_out_dtypeNone, ) - (Tensor, Tensor)枚举说明quant_mode与mask_mode在 Python 接口中支持传入IntEnum枚举或对应 int 值枚举定义于cann_ops_transformer.ops.generic_block_sparse_attention源码见 generic_block_sparse_attention.pyquant_mode 枚举QuantMode枚举名值含义NO_QUANT0非量化默认值FP8_E4M3_STATIC_PER_GROUP1FP8_E4M3 静态 per-groupFP8_E4M3_DYNAMIC_MX2FP8_E4M3 动态 MXFP4_E2M1_DYNAMIC_OCP3FP4_E2M1 动态 OCPFP4_E2M1_DYNAMIC_CX4FP4_E2M1 动态 CXFP8_E4M3_STATIC_CAST_P5FP8_E4M3 静态 cast Pmask_mode 枚举MaskMode枚举名值含义NO_MASK0不加 maskCAUSAL1Causal 模式默认值WINDOW2Window 模式当前仅支持 mask_mode 1CAUSALquant_mode 当前仅支持 0NO_QUANT与 5FP8_E4M3_STATIC_CAST_P。Python 联合调用示例TND PA_BBND以下示例取自 torchapi_generic_block_sparse_attention.mdimport math import torch import torch_npu import cann_ops_transformer torch_npu.npu.set_device(0) B, Q_N, KV_N, Q_S, KV_S, D 1, 32, 8, 128, 256, 128 block_x, block_y 1, 128 block_size 128 top_k 16 Q_T B * Q_S num_blocks math.ceil(KV_S / block_size) total_q_blocks math.ceil(Q_S / block_x) q torch.randn(Q_T, Q_N, D, dtypetorch.float16, devicenpu) k torch.randn(num_blocks, block_size, KV_N, D, dtypetorch.float16, devicenpu) v torch.randn(num_blocks, block_size, KV_N, D, dtypetorch.float16, devicenpu) sparse_block_idx torch.randint( 0, num_blocks, (KV_N, total_q_blocks, top_k), dtypetorch.int32, devicenpu ) sparse_block_count torch.full((KV_N, total_q_blocks), top_k, dtypetorch.int32, devicenpu) cu_seqlens_q torch.tensor([0, Q_S], dtypetorch.int64, devicenpu) seqused_kv torch.tensor([KV_S], dtypetorch.int32, devicenpu) block_table torch.arange(num_blocks, dtypetorch.int32, devicenpu).view(1, -1) block_shape [block_x, block_y] metadata cann_ops_transformer.ops.generic_block_sparse_attention_metadata( sparse_block_idx, sparse_block_count, Q_N, KV_N, D, block_shape, cu_seqlens_qcu_seqlens_q, seqused_kvseqused_kv, max_seqlen_qQ_S, max_seqlen_kvKV_S, is_packed_gqaTrue, layout_qTND, layout_kvPA_BBND, mask_mode1, quant_mode0, softmax_precision1, ) attention_out, softmax_lse cann_ops_transformer.ops.generic_block_sparse_attention( q, k, v, sparse_block_idx, sparse_block_count, block_shape, metadatametadata, cu_seqlens_qcu_seqlens_q, seqused_kvseqused_kv, block_tableblock_table, is_packed_gqaTrue, layout_qTND, layout_kvPA_BBND, softmax_scale1.0 / (D ** 0.5), mask_mode1, quant_mode0, softmax_precision1, return_softmax_lseFalse, ) torch_npu.npu.synchronize() assert attention_out.shape q.shape assert attention_out.dtype q.dtype返回值说明generic_block_sparse_attention_metadata返回 shape 为(1024,)的 int32 Tensor任务切分数据generic_block_sparse_attention返回attention_outshape 与 q 一致quant_mode0 时 dtype 默认与 q 一致quant_mode≠0 时由 attention_out_dtype 指定和softmax_lsereturn_softmax_lseTrue 时 TND 布局下输出 shape 为(Q_T, Q_N, 1)的 float32 Tensor否则为空 Tensor。该逻辑在 generic_block_sparse_attention.py 的register_meta中有对应实现。Layout 与基准维度速查基准符号命名含义BBatch SizeT / totalQTokensquery 的 Total tokens所有 batch 序列长度累加和N / headNumquery 的 head 数N1numKeyValueHeadskey/value 的 head 数N2D / headDimHead Dim且满足 D H / NnumBlocksPaged KV Cache 的物理页数blockSize每一页容纳的 token 数maxNumBlocksPerBatchblockTable 第二维须 ≥ ceilDiv(maxKvSeqLength, blockSize)totalQBlocks按存储长度分块后的 Q 块总数$\sum_i \mathrm{ceilDiv}(qStorageLen_i, blockShapeX)$totalKBlocks按存储长度分块后的 KV 块总数$\sum_i \mathrm{ceilDiv}(kvStorageLen_i, blockShapeY)$maxKvBlockCount / topKsparseBlockIdx 最后一维须不小于 sparseBlockCount 中所有元素的最大值当前上限为 256blockShapeX / blockShapeY稀疏块在 Q 方向、KV 方向的块大小layoutSparsePattern 与 sparse 张量 shapelayoutSparsePattern决定 sparseBlockIdx、sparseBlockCount 的 shape当前仅支持值 4layoutSparsePatternsparseBlockIdxsparseBlockCount描述0[batch, N2, maxQBlockCount, maxKvBlockCount][batch, N2, maxQBlockCount]同 group qHead 共享 pattern表示每个 Q 块选了哪些 KV 块1[batch, N2, maxKvBlockCount, maxQBlockCount][batch, N2, maxKvBlockCount]同 group qHead 共享 pattern表示每个 KV 块选了哪些 Q 块2[batch, N1, maxQBlockCount, maxKvBlockCount][batch, N1, maxQBlockCount]同 group qHead 独立 pattern表示每个 Q 块选了哪些 KV 块3[batch, N1, maxKvBlockCount, maxQBlockCount][batch, N1, maxKvBlockCount]同 group qHead 独立 pattern表示每个 KV 块选了哪些 Q 块4[N2, totalQBlocks, maxKvBlockCount][N2, totalQBlocks]同 group qHead 共享 pattern表示每个 Q 块选了哪些 KV 块当前支持5[N2, totalKBlocks, maxQBlockCount][N2, totalKBlocks]同 group qHead 共享 pattern表示每个 KV 块选了哪些 Q 块6[N1, totalQBlocks, maxKvBlockCount][N1, totalQBlocks]同 group qHead 独立 pattern表示每个 Q 块选了哪些 KV 块7[N1, totalKBlocks, maxQBlockCount][N1, totalKBlocks]同 group qHead 独立 pattern表示每个 KV 块选了哪些 Q 块Paged Attention 相关blockTablekvLayoutKey/Value shape非空shape 为 [batch, maxNumBlocksPerBatch]代表使能 paged cachePA_BBND[numBlocks, blockSize, numKeyValueHeads, headDim]PA_BNBD[numBlocks, numKeyValueHeads, blockSize, headDim]空代表不使能 paged cache接收原始 KVTND / BSND / BNSD[totalKTokens, N2, D] / [batch, maxKvSeqLength, N2, D] / [batch, N2, maxKvSeqLength, D]当前必须传入非空 blockTable且 layoutKv 为 PA_BBNDblockTable 为 nullptr原始 KV及 PA_BNBD 当前不支持。PagedAttention 开启情况下还必须传入 sequsedKv。源码与测试佐证算子定义generic_block_sparse_attention_def.cpp 定义了全部输入、输出与属性的注册信息及默认值。InferShapegeneric_block_sparse_attention_infershape.cpp 负责输出 shape/dtype 推导。Tilinggeneric_block_sparse_attention_tiling.cpp 实现基于 metadata 的分核与任务切分单测见 test_generic_block_sparse_attention_tiling.cpp。内核入口generic_block_sparse_attention.cpp 通过 TILING_KEY 分发到 arch22/arch35 各实现。PyTorch 绑定generic_block_sparse_attention.py 与 csrc/generic_block_sparse_attention.cpp 完成 torch.library 注册与底层调用。UT 测试test_aclnn_generic_block_sparse_attention.cpp 覆盖 aclnn 接口的 op_api 级验证test_generic_block_sparse_attention_infershape.cpp 覆盖 shape 推导。总结GenericBlockSparseAttention 是 CANN ops-transformer 中面向长序列、稀疏注意力场景的高性能算子它通过sparseBlockIdx/sparseBlockCount描述每个 Q 块对 KV 块的稀疏选择借助独立的 metadata 前置算子完成负载均衡任务切分并基于 CATLASS 模板库在 arch22/arch35 两代芯片上提供 FP16/BF16/FP8 计算路径。当前版本聚焦于 TND PA_BBND 组合Packed GQA Paged KV Cache支持内置 causal mask 与 quantType 0/5 两种量化模式。开发者既可以通过 aclnn 两段式 C 接口在异构流程中精细控制 workspace 与执行器也可以通过 TorchNPU 的cann_ops_transformer.generic_block_sparse_attention接口在 PyTorch 图/命令式流程中直接使用两份调用示例均可直接在支持产品上运行验证。赞分享算子库人工智能深度学习Ascend【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-transformer点击查看免费下载相关推荐CANN ops-transformer 算子详解aclnnBlockSparseAttentionV2 块稀疏注意力接口CANN ops transformer 算子详解aclnnBlockSparseAttentionV2 块稀疏注意力接口 导读 aclnnBlockSpar算子库人工智能深度学习AscendCANN ops-transformer 中的 RainFusionAttention 算子块级稀疏注意力原理、ACLNN 两段式接口与实战调用CANN ops transformer 中的 RainFusionAttention 算子块级稀疏注意力原理、ACLNN 两段式接口与实战调用 本篇技术指南算子库人工智能深度学习AscendopenEuler-agreements社区协作指南如何在AtomGit平台高效贡献openEuler agreements社区协作指南如何在AtomGit平台高效贡献 前往项目官网免费下载 https://ar.openeuler.org算子库人工智能深度学习Ascend上一篇彻底解放Mac生产力AeroSpace中exec-and-forget命令的避坑指南下一篇2025超全Actual Budget开发环境搭建指南从源码到桌面应用全流程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考