mamba2_chunk_scan 算子说明
功能和实现说明
mamba2_chunk_scan 用于在 MambaV2 Prefill 阶段对 chunk 内状态执行 selective scan 运算：对来自前序 chunk 的传播状态、chunk 内 delta 信息以及 gating/bias 进行结合，根据时间步顺序进行递推累计，生成当前 chunk 的有效输出。 该算子包含两个matmul和两部分vector计算，因此实现为CVCV融合算子，通过VC并行提升计算性能。
计算流
Kernel输入输出（I/O）
输入
| Tensor | shape | dtype |
|---|---|---|
| ct | BCLGN | FP16 |
| bt | BCLGN | FP16 |
| xt | BCLHP | FP16 |
| dt | H | FP16 |
| inter_attn | BCHLP | FP32 |
| dacs | BCHL | FP32 |
| dtout | BCHL | FP32 |
输出
| Tensor | shape | dtype |
|---|---|---|
| final_attn | BCLHP | FP32 |
参数说明：
B: batch size
C: number of chunks
L: chunk size
H: number of head
G: ngroups
N: state size
P: head dim
其中C*L为padding后的序列长度
调用方式
import npu_ops_transformer_ext
final_attn = torch.ops.npu_ops_transformer_ext.mamba2_chunk_scan(ct, bt, xt, dt, inter_attn, dacs, dtout)
测试方法
见当前目录 tests/
python test_chunk_scan.py