核心加速API
本文档描述 mindiesd 包通过 layers 模块对外暴露的接口。所有接口均可通过 from mindiesd import <接口名> 直接导入使用。
FA 系列
FA（Flash Attention）系列接口提供昇腾亲和的注意力计算能力，涵盖标准注意力、变长序列注意力和稀疏注意力场景。
| 接口名 | 类型 | 功能描述 |
|---|---|---|
attention_forward |
函数 | 标准注意力前向计算，支持自动算子寻优 |
attention_forward_varlen |
函数 | 变长序列注意力前向计算 |
sparse_attention |
函数 | 稀疏注意力前向计算，支持 rf_v2 / ada_bsa 稀疏策略 |
attention_forward
标准注意力前向计算接口，支持多种底层算子（PFA、FASCore、LaserAttention 等）和自动寻优。
from mindiesd import attention_forward
函数签名
attention_forward(
query, key, value,
attn_mask=None,
scale=None,
fused=True,
head_first=False,
**kwargs
) -> torch.Tensor
参数说明
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
query |
torch.Tensor |
是 | - | 查询张量，4D，布局为 [B,S,N,D] 或 [B,N,S,D] |
key |
torch.Tensor |
是 | - | 键张量，4D，布局与 query 一致 |
value |
torch.Tensor |
是 | - | 值张量，4D，布局与 query 一致 |
attn_mask |
torch.Tensor |
否 | None |
注意力掩码 |
scale |
float |
否 | None |
缩放因子，为 None 时自动取 head_dim ** -0.5 |
fused |
bool |
否 | True |
是否使用融合算子，False 时回退到原生计算 |
head_first |
bool |
否 | False |
头维度是否在序列维度之前，True 表示 [B,N,S,D]，False 表示 [B,S,N,D] |
kwargs.opt_mode |
str |
否 | "runtime" |
算子调度模式，支持 "runtime"、"static"、"manual" |
kwargs.op_type |
str |
否 | "fused_attn_score" |
算子类型，仅在 opt_mode="manual" 时生效，支持 "prompt_flash_attn"、"fused_attn_score"、"ascend_laser_attention" |
kwargs.layout |
str |
否 | "BNSD" |
算子布局，仅在 opt_mode="manual" 时生效，支持 "BNSD"、"BSND"、"BSH" |
返回值
torch.Tensor：注意力计算结果，布局与输入一致。
使用示例
import torch
from mindiesd import attention_forward
query = torch.randn(2, 4096, 24, 128, device="npu", dtype=torch.float16)
key = torch.randn(2, 4096, 24, 128, device="npu", dtype=torch.float16)
value = torch.randn(2, 4096, 24, 128, device="npu", dtype=torch.float16)
out = attention_forward(query, key, value)
迁移指南
- 从
torch.nn.functional.scaled_dot_product_attention迁移时，输入布局需从[B,N,S,D]调整为[B,S,N,D]，并去掉transpose操作。 - 从
flash_attn.flash_attn_func迁移时，输入布局已为[B,S,N,D]，可直接替换。 - 本接口仅提供前向推理，不支持反向梯度计算，迁移时需去掉
dropout并将输入张量的requires_grad设为False。
attention_forward_varlen
变长序列注意力前向计算接口，适用于同一 batch 内序列长度不一致的场景。
from mindiesd import attention_forward_varlen
函数签名
attention_forward_varlen(
q, k, v,
cu_seqlens_q,
cu_seqlens_k,
max_seqlen_q=None,
max_seqlen_k=None,
dropout_p=0.0,
softmax_scale=None,
causal=False,
window_size=None,
softcap=None,
alibi_slopes=None,
deterministic=None,
return_attn_probs=None,
block_table=None
) -> torch.Tensor
参数说明
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
q |
torch.Tensor |
是 | - | 查询张量，3D，布局为 [T, N, D]（T 为所有序列 token 总数） |
k |
torch.Tensor |
是 | - | 键张量，3D，布局为 [T, N, D] |
v |
torch.Tensor |
是 | - | 值张量，3D，布局为 [T, N, D] |
cu_seqlens_q |
torch.Tensor |
是 | - | 查询序列的累积长度，1D 张量，形状为 (batch_size + 1,)，dtype 为 torch.int32 |
cu_seqlens_k |
torch.Tensor |
是 | - | 键序列的累积长度，1D 张量，形状为 (batch_size + 1,)，dtype 为 torch.int32 |
max_seqlen_q |
int |
否 | None |
预留参数 |
max_seqlen_k |
int |
否 | None |
预留参数 |
dropout_p |
float |
否 | 0.0 |
Dropout 概率，当前仅支持 0.0 |
softmax_scale |
float |
否 | None |
缩放因子，为 None 时自动取 head_dim ** -0.5 |
causal |
bool |
否 | False |
是否使用因果注意力掩码 |
window_size |
int |
否 | None |
预留参数 |
softcap |
float |
否 | None |
预留参数 |
alibi_slopes |
torch.Tensor |
否 | None |
预留参数 |
deterministic |
bool |
否 | None |
预留参数 |
return_attn_probs |
bool |
否 | None |
预留参数 |
block_table |
torch.Tensor |
否 | None |
预留参数 |
返回值
torch.Tensor：注意力计算结果，形状为 (total, nheads, headdim)。
使用示例
import torch
from mindiesd import attention_forward_varlen
q = torch.randn(8192, 24, 128, device="npu", dtype=torch.float16)
k = torch.randn(8192, 24, 128, device="npu", dtype=torch.float16)
v = torch.randn(8192, 24, 128, device="npu", dtype=torch.float16)
cu_seqlens_q = torch.tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32, device="npu")
cu_seqlens_k = torch.tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32, device="npu")
out = attention_forward_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, causal=False)
迁移指南
- 从
flash_attn.flash_attn_varlen_func迁移时，接口参数基本一致，可直接替换调用。
sparse_attention
稀疏注意力前向计算接口，支持 RainFusion（rf_v2 / rf_v3）和自适应块稀疏（ada_bsa）两种稀疏策略。
from mindiesd import sparse_attention
函数签名
sparse_attention(
q, k, v,
attn_mask=None,
scale=None,
is_causal=False,
head_num=1,
input_layout="BNSD",
inner_precise=0,
sparse_type=None,
txt_len=0,
block_size=128,
latent_shape_q=None,
latent_shape_k=None,
keep_sink=True,
keep_recent=True,
cdf_threshold=1.0,
sparsity=0.0,
**kwargs
) -> torch.Tensor
参数说明
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
q |
torch.Tensor |
是 | - | 查询张量，4D，布局由 input_layout 决定 |
k |
torch.Tensor |
是 | - | 键张量，4D，布局由 input_layout 决定 |
v |
torch.Tensor |
是 | - | 值张量，4D，布局由 input_layout 决定 |
attn_mask |
torch.Tensor |
否 | None |
注意力掩码，预留参数 |
scale |
float |
否 | None |
缩放因子，为 None 时自动取 head_dim ** -0.5 |
is_causal |
bool |
否 | False |
是否使用因果注意力掩码 |
head_num |
int |
否 | 1 |
注意力头数量 |
input_layout |
str |
否 | "BNSD" |
张量布局，支持 "BNSD" 或 "BSND" |
inner_precise |
int |
否 | 0 |
计算精度模式，0 为高精度，1 为高性能 |
sparse_type |
str |
否 | None |
稀疏类型，支持 None、"rf_v2"、"rf_v3"、"ada_bsa" |
txt_len |
int |
否 | 0 |
文本序列长度，仅在 sparse_type="rf_v2" 时生效 |
block_size |
int |
否 | 128 |
块大小，当前仅支持 128 |
latent_shape_q |
list |
否 | None |
查询的潜空间形状 [t, h, w]，t*h*w = qseqlen，仅在 sparse_type="rf_v2" 时生效 |
latent_shape_k |
list |
否 | None |
键的潜空间形状 [t, h, w]，t*h*w = kseqlen，仅在 sparse_type="rf_v2" 时生效 |
keep_sink |
bool |
否 | True |
是否保留 sink token，仅在 sparse_type="ada_bsa" 时生效 |
keep_recent |
bool |
否 | True |
是否保留 recent token，仅在 sparse_type="ada_bsa" 时生效 |
cdf_threshold |
float |
否 | 1.0 |
CDF 阈值，仅在 sparse_type="ada_bsa" 时生效 |
sparsity |
float |
否 | 0.0 |
稀疏率，取值范围 [0, 1]，0 表示不使用稀疏算法 |
返回值
torch.Tensor：注意力计算结果，布局与输入一致。
使用示例
import torch
from mindiesd import sparse_attention
q = torch.randn(2, 24, 4096, 128, device="npu", dtype=torch.float16)
k = torch.randn(2, 24, 4096, 128, device="npu", dtype=torch.float16)
v = torch.randn(2, 24, 4096, 128, device="npu", dtype=torch.float16)
out = sparse_attention(
q, k, v,
head_num=24,
input_layout="BNSD",
sparse_type="ada_bsa",
sparsity=0.5
)
融合算子
融合算子系列接口提供昇腾高性能融合算子，涵盖位置编码、归一化和激活函数等基础计算。
| 接口名 | 类型 | 功能描述 |
|---|---|---|
rotary_position_embedding |
函数 | 旋转位置编码（RoPE）融合算子 |
RMSNorm |
类 | RMS 归一化融合算子 |
fast_layernorm |
函数 | 高性能 LayerNorm 融合算子 |
layernorm_scale_shift |
函数 | 自适应 LayerNorm（AdaLayerNorm）融合算子 |
get_activation_layer |
函数 | 获取激活函数实例（含 NPU 加速版本） |
rotary_position_embedding
旋转位置编码（RoPE）融合算子，将位置信息通过旋转矩阵注入到查询和键张量中。
from mindiesd import rotary_position_embedding
函数签名
rotary_position_embedding(
x, cos, sin,
rotated_mode="rotated_half",
head_first=False,
fused=True
) -> torch.Tensor
参数说明
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
x |
torch.Tensor |
是 | - | 查询或键张量，4D，支持布局 [B,N,S,D]、[B,S,N,D]、[S,B,N,D] |
cos |
torch.Tensor |
是 | - | 预计算的余弦频率张量，2D [S,D] 或 4D [1,1,S,D]/[1,S,1,D]/[S,1,1,D] |
sin |
torch.Tensor |
是 | - | 预计算的正弦频率张量，维度与 cos 一致 |
rotated_mode |
str |
否 | "rotated_half" |
旋转模式："rotated_half" 为半旋转，"rotated_interleaved" 为交错旋转 |
head_first |
bool |
否 | False |
头维度是否在序列维度之前 |
fused |
bool |
否 | True |
是否使用融合算子 |
返回值
torch.Tensor：应用了旋转位置编码的张量，形状与输入 x 一致。
使用示例
import torch
from mindiesd import rotary_position_embedding
x = torch.randn(2, 4096, 24, 128, device="npu", dtype=torch.float16)
cos = torch.randn(1, 4096, 1, 128, device="npu", dtype=torch.float16)
sin = torch.randn(1, 4096, 1, 128, device="npu", dtype=torch.float16)
out = rotary_position_embedding(x, cos, sin, rotated_mode="rotated_half", head_first=False, fused=True)
旋转模式说明
- rotated_half：适用于 OpenSoraPlan、Stable Audio 等模型，将
x拆分为前后两半进行旋转。 - rotated_interleaved：适用于 HunyuanDiT、OpenSora、Flux、CogVideox 等模型，将
x按相邻元素交错进行旋转。
RMSNorm
RMS 归一化融合算子，等效于 T5LayerNorm，不涉及均值计算，专注于输入张量的根均方值。
from mindiesd import RMSNorm
类签名
RMSNorm(hidden_size, eps=1e-6)
构造参数
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
hidden_size |
int |
是 | - | 隐藏层维度大小 |
eps |
float |
否 | 1e-6 |
数值稳定性参数 |
forward 方法
forward(hidden_states, if_fused=True) -> torch.Tensor
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
hidden_states |
torch.Tensor |
是 | - | 输入张量，维度范围为 2~8 |
if_fused |
bool |
否 | True |
是否使用 NPU 融合算子 |
使用示例
import torch
from mindiesd import RMSNorm
norm = RMSNorm(1024, eps=1e-6)
x = torch.randn(2, 4096, 1024, device="npu", dtype=torch.float16)
out = norm(x)
fast_layernorm
高性能 LayerNorm 融合算子，支持多种计算精度模式。
from mindiesd import fast_layernorm
函数签名
fast_layernorm(
norm, x,
impl_mode=0,
fused=True
) -> torch.Tensor
参数说明
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
norm |
torch.nn.LayerNorm |
是 | - | PyTorch LayerNorm 实例 |
x |
torch.Tensor |
是 | - | 输入张量，3D，布局为 [B,S,H] |
impl_mode |
int |
否 | 0 |
计算模式：0 高精度、1 高性能、2 float16 模式（仅当所有输入均为 float16 时可用） |
fused |
bool |
否 | True |
是否使用融合算子，False 时回退到标准 torch.nn.LayerNorm 计算 |
返回值
torch.Tensor：LayerNorm 计算结果，形状与输入 x 一致。
使用示例
import torch
import torch.nn as nn
from mindiesd import fast_layernorm
norm = nn.LayerNorm(1024, eps=1e-5)
x = torch.randn(2, 4096, 1024, device="npu", dtype=torch.float16)
out = fast_layernorm(norm, x, impl_mode=0, fused=True)
layernorm_scale_shift
自适应 LayerNorm（AdaLayerNorm）融合算子，在 LayerNorm 基础上添加自适应缩放和偏移。
计算公式：out = layernorm(x) * (1 + scale) + shift
from mindiesd import layernorm_scale_shift
函数签名
layernorm_scale_shift(
layernorm, x, scale, shift,
fused=True
) -> torch.Tensor
参数说明
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
layernorm |
torch.nn.LayerNorm |
是 | - | PyTorch LayerNorm 实例 |
x |
torch.Tensor |
是 | - | 输入张量，3D，布局为 [B,S,H] |
scale |
torch.Tensor |
是 | - | 自适应缩放参数，2D [B,H] 或 3D [B,1,H] |
shift |
torch.Tensor |
是 | - | 自适应偏移参数，2D [B,H] 或 3D [B,1,H] |
fused |
bool |
否 | True |
是否使用融合算子 |
返回值
torch.Tensor：AdaLayerNorm 计算结果，形状与输入 x 一致。
使用示例
import torch
import torch.nn as nn
from mindiesd import layernorm_scale_shift
norm = nn.LayerNorm(1024, eps=1e-5)
x = torch.randn(2, 4096, 1024, device="npu", dtype=torch.float16)
scale = torch.randn(2, 1024, device="npu", dtype=torch.float16)
shift = torch.randn(2, 1024, device="npu", dtype=torch.float16)
out = layernorm_scale_shift(norm, x, scale, shift, fused=True)
约束条件
x的最后一维必须与scale、shift的最后一维相等。- 若
scale或shift为 3D 张量，则第二维必须为 1 或与x的第二维（序列长度）相等。
get_activation_layer
获取指定类型的激活函数实例，部分激活函数提供 NPU 加速版本。
from mindiesd import get_activation_layer
函数签名
get_activation_layer(act_type: str) -> nn.Module
参数说明
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
act_type |
str |
是 | - | 激活函数名称，不区分大小写 |
支持的激活函数
| 名称 | 对应实现 | 说明 |
|---|---|---|
"swish" |
nn.SiLU |
Swish 激活函数 |
"silu" |
nn.SiLU |
SiLU 激活函数（与 swish 等价） |
"mish" |
nn.Mish |
Mish 激活函数 |
"gelu" |
GELU |
标准 GELU |
"relu" |
nn.ReLU |
ReLU 激活函数 |
"gelu-tanh" |
GELU(approximate="tanh") |
tanh 近似 GELU |
"gelu-fast" |
GELU(approximate="fast") |
快速 GELU，使用 NPU 的 npu_fast_gelu 算子加速 |
返回值
nn.Module：对应激活函数的实例。
使用示例
from mindiesd import get_activation_layer
act = get_activation_layer("gelu-fast")
out = act(hidden_states)
Linear
自定义线性层，与 PyTorch 的 nn.Linear 用法一致，但新增 op_type 参数用于选择底层算子实现。
from mindiesd import Linear
构造参数
| 参数 | 类型 | 必选 | 默认值 | 说明 |
|---|---|---|---|---|
in_features |
int |
是 | - | 输入特征维度 |
out_features |
int |
是 | - | 输出特征维度 |
bias |
bool |
否 | True |
是否使用偏置 |
device |
str |
否 | None |
权重存储设备 |
dtype |
torch.dtype |
否 | None |
权重数据类型 |
op_type |
str |
否 | "matmulv2" |
算子类型，支持 "matmulv2"、"batchmatmulv2"、"batchmatmulv3" |
forward
forward(input) -> torch.Tensor
| 参数 | 类型 | 必选 | 说明 |
|---|---|---|---|
input |
torch.Tensor |
是 | 输入张量，最后一维须等于 in_features |