causal_conv1d 算子

概述

causal_conv1d是一个基于 Triton 实现的因果一维卷积算子，专为序列模型设计。因果卷积保证每个输出位置只依赖当前及过去的输入，不看未来，满足自回归模型的因果性要求。该算子利用昇腾 NPU 的并行计算能力实现高性能的卷积操作，同时支持可选的激活函数、残差连接及变长序列。

---

函数签名

def causal_conv1d(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
    residual: Optional[torch.Tensor] = None,
    initial_state: Optional[torch.Tensor] = None,
    activation: str = None,
    cu_seqlens: Optional[torch.Tensor] = None,
    output_final_state: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
    """执行因果一维卷积前向计算

    Args:
        x: 输入序列张量，形状 [B, T, D]
        weight: 卷积核权重，形状 [W, D]
        bias: 偏置，形状 [D]（可选）
        residual: 残差输入，形状 [B, T, D]（可选）
        initial_state: 序列头部历史状态，形状 [B, D, W]（可选，推理时使用）
        output_final_state: 是否输出序列末尾状态（推理时使用）
        activation: 激活函数类型，支持 'silu'/'swish'（可选）
        cu_seqlens: 变长序列边界索引，形状 [N+1]（可选）

    Returns:
        y: 卷积输出，形状 [B, T, D]
        final_state: 序列末尾状态，形状 [B, D, W]（output_final_state=True 时返回，否则为 None）
    """

---

参数说明

前向输入

参数	类型	形状	描述	是否必须
x	torch.Tensor	[B, T, D]	输入序列，B为batch大小，T为序列长度，D为特征维度	是
weight	torch.Tensor	[W, D]	卷积核权重，W为卷积核宽度（通常为4）	是
bias	torch.Tensor	[D]	每个特征维度的可学习偏置	否
residual	torch.Tensor	[B, T, D]	残差输入，与卷积输出相加	否
initial_state	torch.Tensor	[B, D, W]	序列头部历史状态，推理流式场景使用	否
output_final_state	bool	-	是否输出序列末尾状态	否，默认False
activation	str	-	激活函数类型，支持 'silu'/'swish'	否
cu_seqlens	torch.Tensor	[N+1]	变长序列边界，cu_seqlens[i] 到 cu_seqlens[i+1] 为第i条序列范围	否

前向输出

参数	形状	描述
y	[B, T, D]	卷积输出，与输入形状完全一致
final_state	[B, D, W]	序列末尾W个时刻的输入，供下一段推理使用；output_final_state=False 时为 None

反向输入

参数	形状	描述
x	[B, T, D]	前向保存的原始输入
dy	[B, T, D]	上层传入的输出梯度
dht	[B, D, W]	final_state 的梯度（可选）
weight	[W, D]	前向保存的卷积核权重
bias	[D]	前向保存的偏置（可选）
residual	[B, T, D]	前向保存的残差（可选）
initial_state	[B, D, W]	前向保存的初始状态（可选）
activation	str	激活函数类型（可选）
cu_seqlens	[N+1]	变长序列边界（可选）

反向输出

参数	形状	描述
dx	[B, T, D]	输入 x 的梯度，传给下层
dw	[W, D]	卷积核权重的梯度
db	[D]	偏置的梯度（可选）
dr	[B, T, D]	残差输入的梯度（可选）
dh0	[B, D, W]	initial_state 的梯度（可选）

---

实现原理

前向计算

对于每个输出位置 t，因果卷积的计算公式为：

y[t] = weight[0]*x[t] + weight[1]*x[t-1] + ... + weight[W-1]*x[t-W+1] + bias（可选）

即每个输出只依赖当前和过去共 W 个时刻的输入。算子通过以下步骤实现：

1. 头部补齐：在序列头部补 W-1 个时刻，保证输出长度与输入一致

   • 无 initial_state：补零
   • 有 initial_state：补历史真实值，其中 initial_state[b, d, k] 对应序列开始前第 k-W 步的值（k=0 对应最远，k=W-1 对应最近一步）

2. 二维分块：将工作按 [BT, BD] 的 tile 切分，时间维度块大小 BT、特征维度块大小 BD，充分利用 NPU Unified Buffer

3. 持久化 kernel：采用 grid-stride loop，每个 AI 核轮流处理多个 tile，减少调度开销

4. 滑动窗口累加：在 kernel 内部对 W 个时间偏移依次累加，中间结果保留在寄存器，最后一次性写回 HBM

5. 可选后处理：依次执行 bias 加法、SiLU 激活、残差连接

反向计算

基于链式法则，反向 kernel 同时计算四个梯度：

dx[t]      = sum_{i=0}^{W-1} weight[i] * dy[t+i]
dweight[i] = sum_{t} dy[t] * x[t-i]
dbias      = sum_{t} dy[t]
dh0        = 仅头部块，来自 initial_state 参与计算的梯度

反向 kernel 根据是否存在 initial_state/final_state 分三条路径处理，避免对不存在的状态进行无效计算。

---

使用示例

训练场景（无状态）

import torch
from mindspeed_ops.api.triton.convolution import causal_conv1d

B, T, D, W = 4, 2048, 2048, 4
x      = torch.randn(B, T, D, device='npu', dtype=torch.bfloat16)
weight = torch.randn(W, D, device='npu', dtype=torch.bfloat16)
bias   = torch.randn(D, device='npu', dtype=torch.bfloat16)

y, _ = causal_conv1d(x, weight, bias=bias, activation='silu')
assert y.shape == (B, T, D)

性能对比

方法	性能
PyTorch 实现	基准
Triton 实现	5x

注：性能数据基于昇腾 NPU 环境测试，具体数值可能因硬件配置、输入形状不同而有所差异。

---

常见训练 Shape

以下为模型训练时的典型输入形状（W 固定为 4）：

B	T	D	x / y 形状
32	2048	1536	[32, 2048, 1536]
16	2048	2048	[16, 2048, 2048]
8	2048	3072	[8, 2048, 3072]
4	2048	4096	[4, 2048, 4096]
4	8192	2048	[4, 8192, 2048]

---

注意事项

1. x 与 weight 的最后一维 D 必须一致
2. D 必须能被 BD 整除，否则会触发断言错误
3. 有 initial_state 时，NPU Unified Buffer 压力更大，kernel 会自动使用更保守的分块参数（BD=32，BT 缩小），属于正常行为
4. varlen 模式下各序列之间严格隔离，不会互相影响
5. 推理时 output_final_state=True 会额外执行一次状态提取操作，训练时无需开启
6. 激活函数反向需要重新计算前向输出 y，这是用计算换显存的设计，训练显存敏感场景下属于预期行为
7. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销可能超过计算收益