人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity maintained hardware plugin for vLLM on Huawei Ascend项目地址https://gitcode.com/gh_mirrors/vl/vllm-ascend点击查看免费下载导读本文围绕 vllm-ascendvLLM 在华为昇腾上的社区维护硬件插件仓库中 aclnnSparseFlashAttention 算子文档 展开系统讲解 Sparse Flash AttentionSFA在大序列长度推理场景下的设计动机、计算原理、CANN aclnn 两段式接口原型、全部入参/出参规范、产品约束以及可复制的调用示例并结合仓库中的 PyTorch 适配层与 vLLM 侧真实调用链给出源码级佐证。读完本文你将掌握如何通过aclnnSparseFlashAttentionGetWorkspaceSizeaclnnSparseFlashAttention两段式接口在昇腾设备上驱动稀疏注意力计算理解每个参数的含义、shape 与取值范围并了解该算子如何被 vLLM 的稀疏 MLAMulti-head Latent Attention推理路径实际使用。功能定位只算关键部分的高效注意力Sparse Flash AttentionSFA是针对大序列长度推理场景的高效注意力计算模块。其核心思路是在注意力计算中不再对完整的 Key/Value 序列做全量 softmax 加权而是只计算关键部分从而大幅减少计算量。公式如下$$ \text{softmax}(\frac{Q\tilde{K}^T}{\sqrt{d_k}})\tilde{V} $$其中 $\tilde{K},\tilde{V}$ 为基于某种选择算法如lightning_indexer得到的重要性较高的 Key 和 Value一般具有稀疏或分块稀疏的特征$d_k$ 为 $Q,\tilde{K}$ 每一个头的维度。从设计动机上看稀疏化天然会引入大量的离散访存——因为被选中的 Key/Value 在缓存中往往不连续导致数据搬运时间增加进而影响整体性能。按照 README.md 的说明本次公布的sparse_flash_attention是面向 Sparse Attention 的全新算子专门针对离散访存进行了指令缩减及搬运聚合的细致优化这是它与普通 FlashAttention 的关键差异。在仓库的 kernel 实现 sparse_flash_attention.cpp 中可以看到算子通过SFA_OP_IMPL宏为 MLA 主计算路径分别注册了halfFP16与bfloat16_tBF16的实现op_kernel/sparse_flash_attention.cpp且按不同昇腾架构arch22、arch35拆分了 cube/vector 服务代码与文档中「支持 FLOAT16、BFLOAT16」的数据类型约束完全一致。产品支持情况产品是否支持Ascend 950PR 950DT 系列产品√Atlas A2 系列产品√Atlas A3 系列产品√补充说明README.md 中进一步列出 Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品均为×不支持因此该算子仅面向上述三档支持的硬件平台。两段式接口函数原型CANN 的 aclnn 算子采用两段式接口设计必须先调用aclnnSparseFlashAttentionGetWorkspaceSize获取计算所需 workspace 大小以及包含了算子计算流程的执行器executor再调用aclnnSparseFlashAttention执行计算。第一段aclnnSparseFlashAttentionGetWorkspaceSizeaclnnStatus aclnnSparseFlashAttentionGetWorkspaceSize( const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *sparseIndices, const aclTensor *blockTable, const aclTensor *actualSeqLengthsQuery, const aclTensor *actualSeqLengthsKv, const aclTensor *queryRope, const aclTensor *keyRope, double scaleValue, int64_t sparseBlockSize, char *layoutQuery, char *layoutKv, int64_t sparseMode, int64_t preTokens, int64_t nextTokens, int64_t attentionMode, bool returnSoftmaxLse, const aclTensor *attentionOutOut, const aclTensor *softmaxMaxOut, const aclTensor *softmaxSumOut, uint64_t *workspaceSize, aclOpExecutor **executor)第二段aclnnSparseFlashAttentionaclnnStatus aclnnSparseFlashAttention( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)其中第二段接口只接收 4 个参数Device 侧 workspace 内存地址、由第一段接口返回的 workspace 大小、包含算子计算流程的 executor以及指定执行任务的 Stream。参数详解维度符号与完整参数表在逐参数展开之前先明确文档中的维度符号约定BBatch Size输入样本批量大小SSequence Length输入样本序列长度HHead Sizehidden 层的大小NHead Num多头数DHead Dimhidden 层最小的单元尺寸满足 D H / NT所有 Batch 输入样本序列长度的累加和。同时Q_S 和 S1 表示 query shape 中的 SKV_S 和 S2 表示 key shape 中的 SQ_N 和 N1 表示 num_query_headsKV_N 和 N2 表示 num_key_value_headsT1 表示 query shape 中的 TT2 表示 key shape 中的输入样本序列长度的累加和。输入 Tensor 参数参数名输入/输出描述使用说明数据类型数据格式维度(shape)query输入attention 结构的 Query 输入不支持空 tensorFLOAT16、BFLOAT16NDlayoutQuery 为 BSND 时 (B,S1,N1,D)为 TND 时 (T1,N1,D)key输入attention 结构的 Key 输入不支持空 tensorblock_num 为 PageAttention 时 block 总数FLOAT16、BFLOAT16NDlayoutKv 为 PA_BSND 时 (block_num, block_size, KV_N, D)为 BSND 时 (B, S2, KV_N, D)为 TND 时 (T2, KV_N, D)value输入attention 结构的 Value 输入不支持空 tensorFLOAT16、BFLOAT16NDshape 与 key 的 shape 一致sparseIndices输入离散取 kvCache 的索引不支持空 tensorsparse_size 为一次离散选取的 block 数需保证每行有效值均在前半部分、无效值均在后半部分且 sparse_size 0INT32NDlayoutQuery 为 BSND 时 (B, Q_S, KV_N, sparse_size)为 TND 时 (Q_T, KV_N, sparse_size)blockTable输入PageAttention 中 kvCache 存储使用的 block 映射表不支持空 tensor第二维长度不小于所有 batch 中最大的 S2 对应的 block 数量即 S2_max / block_size 向上取整INT32ND(B, S2/block_size)actualSeqLengthsQuery输入不同 Batch 中 query 的有效 token 数不支持空 tensor不指定可传 None表示与 query 的 S 长度相同每个 Batch 的有效 token 数不超过 query 维度 S 且不小于 0layoutQuery 为 TND 时必须传入且以元素数量作为 B 值每个元素表示当前 batch 与之前所有 batch 的 token 数总和INT32ND(B,)actualSeqLengthsKv输入不同 Batch 中 key 和 value 的有效 token 数不支持空 tensor不指定可传 None表示与 key 的 S 长度相同每个 Batch 的有效 token 数不超过 key/value 维度 S 且不小于 0layoutKv 为 TND 或 PA_BSND 时必须传入TND 时每个元素表示当前 batch 与之前所有 batch 的 token 数总和前缀和后一个元素必须 ≥ 前一个元素INT32ND(B,)queryRope输入MLA 结构中 query 的 rope 信息不支持空 tensorFLOAT16、BFLOAT16NDlayoutQuery 为 TND 时 (B,S1,N1,Dr)为 BSND 时 (T1,N1,Dr)keyRope输入MLA 结构中 key 的 rope 信息不支持空 tensorFLOAT16、BFLOAT16NDlayoutKv 为 TND 时 (B,S1,N1,Dr)为 BSND 时 (T1,N1,Dr)为 PA_BSND 时 (block_num,block_size,N2,Dr)属性标量参数参数名输入/输出描述使用说明scaleValuedouble输入缩放系数数值类型 FLOAT16sparseBlockSizeint64_t输入sparse 阶段的 block 大小为 1 时是 Token-wise 稀疏化场景将每个 token 视为独立单元评估每个查询 token 与每个键值 token 之间的独立关联程度大于 1 且小于等于 128 时是 Block-wise 稀疏化场景将 token 序列划分为固定大小的连续块块内 token 共享相同的稀疏化决策layoutQuerychar输入query 的数据排布格式默认值 BSND支持 BSND 和 TNDlayoutKvchar输入key 的数据排布格式默认值 BSND支持 TND、BSND 和 PA_BSND其中 PA_BSND 在使能 PageAttention 时使用sparseModeint64_t输入sparse 的模式为 0 时代表全部计算为 3 时代表 rightDownCausal 模式的 mask对应以右下顶点往左上为划分线的下三角场景preTokensint64_t输入用于稀疏计算表示 attention 需要和前几个 Token 计算关联仅支持默认值 2^63-1nextTokensint64_t输入用于稀疏计算表示 attention 需要和后几个 Token 计算关联仅支持默认值 2^63-1attentionModeint64_t输入attention 的模式仅支持传入 2表示 MLA-absorb 模式returnSoftmaxLsebool输入是否返回 softmax_max 和 softmax_sumTrue 返回、False 不返回默认 False该参数仅在训练且 layoutKv 不为 PA_BSND 场景支持输出 Tensor 参数参数名输入/输出描述数据类型数据格式维度(shape)attentionOut输出公式中的输出FLOAT16、BFLOAT16NDlayoutQuery 为 BSND 时 (B,S1,N1,D)为 TND 时 (T1,N1,D)softmaxMaxOut输出query 乘 key 的结果取 max 得到 softmax_maxFLOATNDlayoutQuery 为 BSND 时 (B,N2,S1,N1/N2)为 TND 时 (N2,T1,N1/N2)softmaxSumOut输出query 乘 key 的结果减去 softmax_max 后取 exp 再求 sum 得到 softmax_sumFLOATNDlayoutQuery 为 BSND 时 (B,N2,S1,N1/N2)为 TND 时 (N2,T1,N1/N2)workspaceSizeuint64_t*输出需要在 Device 侧申请的 workspace 大小---executoraclOpExecutor输出op 执行器包含算子计算流程---返回值与错误码接口返回aclnnStatus状态码。第一段接口会完成入参校验以下场景会报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001如果传入参数是必选输入、输出或必选属性且是空指针则返回 161001ACLNN_ERR_PARAM_INVALID161002query、key、value、sparseIndices、blockTable、actualSeqLengthsQuery、actualSeqLengthsKv、queryRope、keyRope、scaleValue、sparseBlockSize、layoutQuery、layoutKv、sparseMode、attentionMode、returnSoftmaxLse、attentionOut、softmaxMaxOut、softmaxSumOut 的数据类型和数据格式不在支持的范围内约束说明确定性计算aclnnSparseFlashAttention 默认确定性实现。该接口支持推理场景下使用README 补充同时支持图模式。N1支持 1~64 和 128。block_size为一个 block 的 token 数取值为16 的倍数且最大支持1024。query 中的D与 key、value 的 D 值相等且为512query_rope 中的Dr与 key_rope 的 Dr 值相等且为64。query、key、value 的数据类型必须保持一致。支持sparse_block_size 整除 block_size。分产品约束Ascend 950PR 950DT 系列产品只支持 sparse_block_size 为 1。Atlas A3 系列产品、Atlas A2 系列产品支持 sparse_block_size 取值 [1,128]且要求是 2 的幂次方在 PageAttention 场景下要求 sparse_block_size 整除 block_size。调用示例完整可运行的 C 样例文档给出的完整 C 示例与仓库 examples/test_aclnn_sparse_flash_attention.cpp 基本一致展示了从环境初始化到结果输出的全流程。下面按阶段拆解关键代码。1. 环境初始化aclInit / aclrtSetDevice / aclrtCreateStreamint Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret; } ret aclrtSetDevice(deviceId); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret; } ret aclrtCreateStream(stream); if (!CHECK_RET(ret ACL_SUCCESS)) { LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret; } return 0; }2. 构造 aclTensorH2D 拷贝 aclCreateTensor示例通过模板函数CreateAclTensor将 host 数据搬移到 device并按连续 strides 构造 ND 格式的 aclTensortemplate typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); if (!CHECK_RET(ret ACL_SUCCESS)) return ret; ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); if (!CHECK_RET(ret ACL_SUCCESS)) return ret; std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; }3. 初始化各输入/输出 Tensor 的 shape示例使用 BSND layout、单头 MLA 场景TensorShape数据类型query / key / value{1, 2, 1, 512}B1, S2, N1, D512ACL_FLOAT16sparseIndices{1, 2, 1, 2}sparse_size2ACL_INT32attentionOut{1, 2, 1, 512}ACL_FLOAT16softmaxMax / softmaxSum{1, 2, 1, 16}ACL_FLOATqueryRope / keyRope{1, 2, 1, 64}Dr64ACL_FLOAT16注意 softmax 输出 shape 中的最后一个维度 16 恰好对应公式 (B, N2, S1, N1/N2) 中的 N1/N2N116、N21 时即为 16与文档输出 shape 规范吻合。4. 设置属性并执行两段式调用int64_t d 2; double scaleValue 1 / sqrt(d); int64_t sparseBlockSize 64; constexpr const char layerOutStr[] BSND; // layoutQuery 与 layoutKv 均为 BSND int64_t sparseMode 3; // rightDownCausal mask int64_t preTokens 9223372036854775807; // 2^63-1 int64_t nextTokens 9223372036854775807; // 2^63-1 int64_t attentionMode 2; // MLA-absorb bool returnSoftmaxLse false; aclOpExecutor* executor; // 第一段获取 workspaceSize 与 executor int ret aclnnSparseFlashAttentionGetWorkspaceSize( resources.queryTensor, resources.keyTensor, resources.valueTensor, resources.sparseIndicesTensor, nullptr, nullptr, nullptr, // blockTable/seqLens 传 None resources.queryRopeTensor, resources.keyRopeTensor, scaleValue, sparseBlockSize, layoutQuery, layoutKv, sparseMode, preTokens, nextTokens, attentionMode, returnSoftmaxLse, resources.attentionOutTensor, resources.softmaxMaxTensor, resources.softmaxSumTensor, workspaceSize, executor); // 按需申请 workspace if (*workspaceSize 0ULL) { ret aclrtMalloc(workspaceAddr, *workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 第二段执行计算 ret aclnnSparseFlashAttention(*workspaceAddr, *workspaceSize, executor, stream);5. 同步、打印输出与资源清理执行后调用aclrtSynchronizeStream(stream)同步然后通过aclrtMemcpyDEVICE_TO_HOST将attentionOut、softmaxMax、softmaxSum拷回 host 打印最后依次aclDestroyTensor、aclrtFree、aclrtDestroyStream、aclrtResetDevice、aclFinalize完成清理。完整示例中还用到securec.h的memcpy_s将字符串BSND安全拷贝到layoutQuery/layoutKv的字符数组——由于 aclnn 接口的 layout 参数以char*传入必须保证字符串以\0结尾且缓冲区足够。源码佐证PyTorch 适配层与 vLLM 侧真实调用链1. C/PyTorch 适配层npu_sparse_flash_attention仓库在 sparse_flash_attention_torch_adpt.h 中提供了npu_sparse_flash_attention适配函数将 PyTorch Tensor 封装为 aclnn 调用入口先做非空校验TORCH_CHECK(query.numel() 0 ...)对应文档「不支持空 tensor」约束construct_sparse_flash_attention_output_tensor会校验 layoutTORCH_CHECK(layout_query_str BSND || layout_query_str TND ...)并根据 layout 与 query 维度TND 时 3 维、BSND 时 4 维自动推导attention_output、softmax_max、softmax_sum的输出 shape最终通过EXEC_NPU_CMD(aclnnSparseFlashAttention, ...)宏依次传入与文档函数原型完全一致的 19 个参数并返回三元组(attention_output, softmax_max, softmax_sum)。从该适配层的输出 shape 推导逻辑可以印证文档中的规范BSND 时 softmax 输出为(B, N2, S1, N1/N2)代码中为query.size(0), key.size(2), query.size(1), query.size(2)/key.size(2)TND 时为(N2, T1, N1/N2)当returnSoftmaxLse为 False 时 softmax 输出 shape 置为{0}空 tensor。2. vLLM 侧集成sfa_v1.py 的 sparse_mla 调用在 vllm_ascend/attention/sfa_v1.py 的sparse_mla函数中vLLM 推理路径通过torch.ops._C_ascend.npu_sparse_flash_attention真实调用该算子参数设置与文档约束一一对应result torch.ops._C_ascend.npu_sparse_flash_attention( queryquery.contiguous(), keycache, valuecache, sparse_indicesindices, scale_valuescale, sparse_block_size1, # Token-wise 稀疏化 block_tablemetadata.block_table, # PageAttention block 映射表 actual_seq_lengths_querymetadata.query_start_loc[1:].to(torch.int32), actual_seq_lengths_kvmetadata.seq_lens.to(torch.int32), query_ropeNone, key_ropeNone, layout_queryTND, # TND 布局 layout_kvPA_BSND, # PageAttention BSND 块布局 sparse_mode3, # rightDownCausal mask attention_mode2, # MLA-absorb return_softmax_lseFalse, )这段代码体现了算子文档在真实系统中的典型用法layout 组合query 用TNDT1 为所有 batch token 累加和KV 用PA_BSNDPageAttention 的 block 存储因此actual_seq_lengths_query传入前缀和式 token 数、actual_seq_lengths_kv传入各 batch 有效 token 数符合文档「layoutKv 为 TND 或 PA_BSND 时必须传入」的要求sparse_block_size1在 vLLM 的 token 级索引选择场景下启用 Token-wise 稀疏化attention_mode2MLA-absorb 模式配合 MLA 结构中各 head 共享 KV 的语义代码中的query_start_loc、seq_lens均为 MLA 场景常用的 seqlen 表示sparse_indices由索引器indexer产生且 vLLM 侧保证每行有效索引靠前、无效值-1靠后的布局与文档「每行有效值均在前半部分无效值均在后半部分」的约定一致。同时在sfa_v1.py的_execute_sparse_flash_attention_processvllm_ascend/attention/sfa_v1.py中vLLM 还负责将逻辑 block 大小的 KV cache 重排为算子要求的 kernel block 大小视图并在执行后对未写入的行做掩码填充说明该算子面向的是 MLA 稀疏推理的高频执行路径。总结aclnnSparseFlashAttention 是 vllm-ascend 中支撑大序列稀疏推理的核心算子采用「GetWorkspaceSize 执行」两段式 aclnn 接口支持 BSND/TND/PA_BSND 多种布局组合、Token-wise/Block-wise 两种稀疏粒度、MLA-absorb 模式与 softmax LSE 输出。理解其参数语义尤其是 sparseIndices 的有效值前置约定、actualSeqLengths 在前缀和场景下的语义、sparse_block_size 与 block_size 的整除关系是正确调用的前提。仓库中的 PyTorch 适配层 与 vLLM 侧调用链 为二次开发、调试与性能分析提供了可直接对照的参考实现。赞分享人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity maintained hardware plugin for vLLM on Huawei Ascend项目地址https://gitcode.com/gh_mirrors/vl/vllm-ascend点击查看免费下载相关推荐vllm-ascend 稀疏注意力量化索引算子 QuantLightningIndexerV2 深度解析aclnn 两段式接口、量化模式与两级 TopK 候选块方案vllm ascend 稀疏注意力量化索引算子 QuantLightningIndexerV2 深度解析aclnn 两段式接口、量化模式与两级 TopK 候选人工智能大模型模型推理服务AscendCANNvllm-ascend SparseFlashAttention 算子深度解析稀疏注意力计算、参数说明与 aclnn 调用实战vllm ascend SparseFlashAttention 算子深度解析稀疏注意力计算、参数说明与 aclnn 调用实战 vllm ascend 是 v人工智能大模型模型推理服务AscendCANNvllm-ascend 中 aclnnSparseFlashMla 稀疏 MLA 注意力算子深度解析与调用指南vllm ascend 中 aclnnSparseFlashMla 稀疏 MLA 注意力算子深度解析与调用指南 导读 aclnnSparseFlashMla 是人工智能大模型模型推理服务AscendCANN上一篇终极指南werift-webrtc开源项目常见问题解决方案下一篇Wasmoon 常见问题终极解决方案轻松解决 Lua VM 运行难题创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考