企业官网建设流程全解析

• 人工智能
• 大模型
• 模型推理服务
• Ascend
• CANN

导读

本文围绕 vllm-ascend（vLLM 在华为昇腾上的社区维护硬件插件）仓库中 aclnnSparseFlashAttention 算子文档 展开，系统讲解 Sparse Flash Attention（SFA）在大序列长度推理场景下的设计动机、计算原理、CANN aclnn 两段式接口原型、全部入参/出参规范、产品约束以及可复制的调用示例，并结合仓库中的 PyTorch 适配层与 vLLM 侧真实调用链给出源码级佐证。读完本文，你将掌握如何通过aclnnSparseFlashAttentionGetWorkspaceSize+aclnnSparseFlashAttention两段式接口在昇腾设备上驱动稀疏注意力计算，理解每个参数的含义、shape 与取值范围，并了解该算子如何被 vLLM 的稀疏 MLA（Multi-head Latent Attention）推理路径实际使用。

功能定位：只算关键部分的高效注意力

Sparse Flash Attention（SFA）是针对大序列长度推理场景的高效注意力计算模块。其核心思路是：在注意力计算中不再对完整的 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 主计算路径分别注册了half（FP16）与bfloat16_t（BF16）的实现（op_kernel/sparse_flash_attention.cpp），且按不同昇腾架构（arch22、arch35）拆分了 cube/vector 服务代码，与文档中「支持 FLOAT16、BFLOAT16」的数据类型约束完全一致。

产品支持情况

补充说明：README.md 中进一步列出 Atlas 200I/500 A2 推理产品、Atlas 推理系列产品、Atlas 训练系列产品均为×（不支持），因此该算子仅面向上述三档支持的硬件平台。

两段式接口：函数原型

CANN 的 aclnn 算子采用两段式接口设计：必须先调用aclnnSparseFlashAttentionGetWorkspaceSize获取计算所需 workspace 大小以及包含了算子计算流程的执行器（executor），再调用aclnnSparseFlashAttention执行计算。

第一段：aclnnSparseFlashAttentionGetWorkspaceSize

aclnnStatus 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)

第二段：aclnnSparseFlashAttention

aclnnStatus aclnnSparseFlashAttention( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)

其中第二段接口只接收 4 个参数：Device 侧 workspace 内存地址、由第一段接口返回的 workspace 大小、包含算子计算流程的 executor，以及指定执行任务的 Stream。

参数详解：维度符号与完整参数表

在逐参数展开之前，先明确文档中的维度符号约定：

• B（Batch Size）：输入样本批量大小；
• S（Sequence Length）：输入样本序列长度；
• H（Head Size）：hidden 层的大小；
• N（Head Num）：多头数；
• D（Head Dim）：hidden 层最小的单元尺寸，满足 D = H / N；
• T：所有 Batch 输入样本序列长度的累加和。

同时，Q_S 和 S1 表示 query shape 中的 S，KV_S 和 S2 表示 key shape 中的 S，Q_N 和 N1 表示 num_query_heads，KV_N 和 N2 表示 num_key_value_heads，T1 表示 query shape 中的 T，T2 表示 key shape 中的输入样本序列长度的累加和。

输入 Tensor 参数

属性（标量）参数

输出 Tensor 参数

返回值与错误码

接口返回aclnnStatus状态码。第一段接口会完成入参校验，以下场景会报错：

约束说明

• 确定性计算：aclnnSparseFlashAttention 默认确定性实现。
• 该接口支持推理场景下使用（README 补充：同时支持图模式）。
• N1支持 1~64 和 128。
• block_size为一个 block 的 token 数，取值为16 的倍数，且最大支持1024。
• query 中的D与 key、value 的 D 值相等且为512；query_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 / aclrtCreateStream）

int 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. 构造 aclTensor（H2D 拷贝 + aclCreateTensor）

示例通过模板函数CreateAclTensor将 host 数据搬移到 device，并按连续 strides 构造 ND 格式的 aclTensor：

template <typename T> int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_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::vector<int64_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 场景：

注意 softmax 输出 shape 中的最后一个维度 16 恰好对应公式 (B, N2, S1, N1/N2) 中的 N1/N2（N1=16、N2=1 时即为 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)同步，然后通过aclrtMemcpy（DEVICE_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会校验 layout（TORCH_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( query=query.contiguous(), key=cache, value=cache, sparse_indices=indices, scale_value=scale, sparse_block_size=1, # Token-wise 稀疏化 block_table=metadata.block_table, # PageAttention block 映射表 actual_seq_lengths_query=metadata.query_start_loc[1:].to(torch.int32), actual_seq_lengths_kv=metadata.seq_lens.to(torch.int32), query_rope=None, key_rope=None, layout_query="TND", # TND 布局 layout_kv="PA_BSND", # PageAttention + BSND 块布局 sparse_mode=3, # rightDownCausal mask attention_mode=2, # MLA-absorb return_softmax_lse=False, )

这段代码体现了算子文档在真实系统中的典型用法：

• layout 组合：query 用TND（T1 为所有 batch token 累加和），KV 用PA_BSND（PageAttention 的 block 存储），因此actual_seq_lengths_query传入前缀和式 token 数、actual_seq_lengths_kv传入各 batch 有效 token 数，符合文档「layoutKv 为 TND 或 PA_BSND 时必须传入」的要求；
• sparse_block_size=1：在 vLLM 的 token 级索引选择场景下启用 Token-wise 稀疏化；
• attention_mode=2：MLA-absorb 模式，配合 MLA 结构中各 head 共享 KV 的语义（代码中的query_start_loc、seq_lens均为 MLA 场景常用的 seqlen 表示）；
• sparse_indices由索引器（indexer）产生，且 vLLM 侧保证每行有效索引靠前、无效值（-1）靠后的布局，与文档「每行有效值均在前半部分，无效值均在后半部分」的约定一致。

同时，在sfa_v1.py的_execute_sparse_flash_attention_process（vllm_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 侧调用链 为二次开发、调试与性能分析提供了可直接对照的参考实现。

• 人工智能
• 大模型
• 模型推理服务
• Ascend
• CANN