RainFusionAttention算子

概述

RainFusionAttention是一个基于CATLASS模板库实现的高性能稀疏注意力算子,支持灵活的块级稀疏模式。该算子从catlass的32_sparse_attention_infer示例迁移而来,针对ops-transformer-dev框架进行了适配。

功能特性

灵活的稀疏块模式: 支持自定义的x*y块级稀疏模式,通过blockShape参数指定
TND格式KV Cache: 支持TND [T, N, D]格式的KV cache布局（注：BNSD格式暂不支持）
多种Q输入布局: 支持BSH、TND两种Query输入布局（注：BNSD格式暂不支持）
高性能计算: 基于CATLASS模板库,充分利用昇腾A2硬件特性
多精度支持: 支持FP16和BF16两种精度
非对齐序列: 支持序列长度不能被块大小整除的场景,自动处理边界情况

接口定义

aclnnRainFusionAttentionGetWorkspaceSize

aclnnStatus aclnnRainFusionAttentionGetWorkspaceSize(
    const aclTensor *query,                    // Query输入 [B, S, H] or [T, N, D] or [B, N, S, D]
    const aclTensor *key,                      // Key输入 (TND or BNSD格式)
    const aclTensor *value,                    // Value输入 (TND or BNSD格式)
    const aclTensor *attenMask,                // Attention mask (可选)
    const aclIntArray *actualSeqLengths,       // 实际Q序列长度 (可选)
    const aclIntArray *actualSeqLengthsKv,     // 实际KV序列长度 (可选)
    const aclTensor *blockTable,               // Block表 (可选,用于PagedAttention)
    const aclTensor *selectIdx,                // 稀疏块索引 [T, headNum, maxKvBlockNum]
    const aclTensor *selectNumIdx,             // 每个Q块的KV块数量 [T, headNum]
    const char *qInputLayout,                  // Q输入布局: "BSH", "TND", "BNSD"
    const char *kvInputLayout,                 // KV输入布局: "TND", "BNSD"
    int64_t numKeyValueHeads,                  // KV头数
    int64_t maskType,                          // Mask类型
    double scaleValue,                         // 缩放因子
    const aclIntArray *blockShape,             // 块形状 [blockShapeX, blockShapeY]
    const aclTensor *attentionOut,             // 输出tensor
    const aclTensor *softmaxLse,               // Softmax LSE输出 (可选)
    uint64_t *workspaceSize,                   // 返回workspace大小
    aclOpExecutor **executor);                 // 返回executor

aclnnRainFusionAttention

aclnnStatus aclnnRainFusionAttention(
    void *workspace,                           // workspace地址
    uint64_t workspaceSize,                    // workspace大小
    aclOpExecutor *executor,                   // executor
    const aclrtStream stream);                 // ACL stream

参数说明

输入参数

query: Query tensor,支持以下布局
- BSH: [batch, seqlen, num_heads * head_dim]
- TND: [total_tokens, num_heads, head_dim]
key/value: KV cache tensor,仅支持TND布局
- TND: [total_kv_tokens, num_kv_heads, head_dim]
- 注：BNSD格式暂不支持，如需使用请转换为TND格式
selectIdx: 稀疏块索引数组 [T, headNum, maxKvBlockNum]
- T: 所有batch中Q方向切块的总数
- headNum: 头数
- maxKvBlockNum: 最大KV块数量
- 存储每个Q块选择的KV块索引
selectNumIdx: KV块数量数组 [T, headNum]
- 存储每个Q块实际选择的KV块数量
blockShape: 稀疏块形状 [blockShapeX, blockShapeY]
- blockShapeX: Q方向块大小
- blockShapeY: KV方向块大小

输出参数

attentionOut: 注意力输出,shape与query相同
softmaxLse: Softmax log-sum-exp输出 (可选)

使用示例

// 1. 准备输入数据
int64_t batch = 1;
int64_t qSeqlen = 256;
int64_t kvSeqlen = 512;
int64_t numHeads = 8;
int64_t numKvHeads = 8;
int64_t headDim = 128;
int64_t blockShapeX = 64;
int64_t blockShapeY = 64;

// 2. 创建tensor (略)

// 3. 设置blockShape
int64_t blockShapeData[2] = {blockShapeX, blockShapeY};
aclIntArray *blockShape = aclCreateIntArray(blockShapeData, 2);

// 4. 调用算子
uint64_t workspaceSize = 0;
aclOpExecutor *executor = nullptr;

aclnnStatus ret = aclnnRainFusionAttentionGetWorkspaceSize(
    query, key, value, nullptr, nullptr, nullptr, nullptr,
    selectIdx, selectNumIdx, "BSH", "TND",
    numKvHeads, 0, scaleValue, blockShape,
    attentionOut, nullptr, &workspaceSize, &executor);

if (ret == ACLNN_SUCCESS) {
    void *workspace = nullptr;
    aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    
    ret = aclnnRainFusionAttention(workspace, workspaceSize, executor, stream);
    
    aclrtFree(workspace);
}

KV Cache布局说明

TND格式 [T, N, D] (当前支持)

T: 总token数量
N: KV头数
D: 头维度
适用于变长序列场景
这是当前唯一支持的格式

BNSD格式 [B, N, S, D] (暂不支持)

BNSD格式暂未实现，计划在未来版本中支持
如需使用BNSD格式的KV cache，请先转换为TND格式
转换方法：将BNSD [B, N, S, D] reshape为 TND [B*S, N, D]

稀疏模式说明

SelectIdx索引格式

RainFusionAttention使用新的selectIdx索引格式,不需要传统的sBlockIdx计算:

selectIdx: 形状为[T, headNum, maxKvblockNum]的稀疏块索引数组
selectNumIdx: 形状为[T, headNum]的稀疏块数量数组
T: 所有batch中Q方向切块的总数(按blockShapeX计算)
blockShapeX: Q方向基本块大小
blockShapeY: KV方向基本块大小

示例

对于qSeqlen=256, kvSeqlen=512, blockShapeX=64, blockShapeY=64:

Q方向: ceil(256/64)=4块 [0, 1, 2, 3]
KV方向: ceil(512/64)=8块 [0, 1, 2, 3, 4, 5, 6, 7]

每个Q块可以选择任意KV块组合：
- Q块0选择的KV块: [0, 2, 5]
- Q块1选择的KV块: [1, 3, 6]
- Q块2选择的KV块: [0, 4, 7]
- Q块3选择的KV块: [2, 5, 6]

性能特点

内存效率: 相比密集注意力,显著减少内存使用
计算效率: 只计算指定的稀疏块,减少计算量
灵活配置: 支持任意blockShape的稀疏块配置
硬件优化: 充分利用昇腾A2的AI Core和Vector Core

注意事项

KV Cache格式: 当前仅支持TND格式，BNSD格式暂不支持
Q输入格式: 支持BSH和TND格式，BNSD格式暂不支持
qSeqlen和kvSeqlen不需要被blockShape整除,支持非对齐场景
实际分块数通过向上取整计算: ceil(qSeqlen/blockShapeX)和ceil(kvSeqlen/blockShapeY)
稀疏块索引必须在有效范围内
最后一块可能不完整,系统会自动处理边界情况
建议稀疏度不要过低,以保持计算效率
kvInputLayout参数必须设置为"TND"

编译说明

算子编译集成在ops-transformer-dev的统一编译流程中:

bash build.sh --soc=Ascend910B3

版本历史

v1.0.0: 初始版本,从catlass 32_sparse_attention_infer迁移