mamba2_chunk_cumsum 算子说明
功能和实现说明
mamba2_chunk_cumsum 用于在 MambaV2 Prefill 阶段对 chunk 内部执行按时间步的累积求和操作，实现 SSM 中状态量在 chunk 维度上的递推更新。算子对输入序列在 S 维度按 chunk_size 拆分，并在每个 chunk 内按照因果顺序执行 cumulative sum，用于后续 chunk 状态更新与 selective scan 计算。本算子基于 Vector 实现累积求和计算，支持 FP16/FP32 输入输出。
计算流
Kernel输入输出（I/O）
输入
| Tensor | shape | dtype |
|---|---|---|
| at | H | FP32 |
| dt | BCLH | FP16 |
| dt_bias | H | FP16 |
| dt_mask | BCLH | FP16 |
输出
| Tensor | shape | dtype |
|---|---|---|
| dtout | BCLH | FP32 |
| dacs | BCLH | FP32 |
| dacs_chunk | BCH | FP32 |
参数说明：
B: batch size
C: number of chunks
L: chunk size
H: number of head
其中C*L为padding后的序列长度
调用方式
import npu_ops_transformer_ext
out = torch.ops.npu_ops_transformer_ext.mamba2_chunk_cumsum(at, dt, dt_bias, dt_mask)
测试方法
见当前目录 tests/
python test_chunk_cumsum.py