mamba2_chunk_state 算子说明
功能和实现说明
mamba2_chunk_state 用于在 MambaV2 Prefill 阶段进行 chunk 内的离散时间状态更新，根据 chunk_cumsum 得到的累积量dacs/dacs_chunk和状态更新因子dtout进行状态递推，输出 chunk 内每一步的状态序列，并生成用于下一 chunk 的最终隐藏状态。本算子实现为 Vector+cube 融合算子，支持 FP16/FP32。
计算流
Kernel输入输出（I/O）
输入
| Tensor | shape | dtype |
|---|---|---|
| dtout | BCLH | FP32 |
| dacs | BCLH | FP32 |
| bt | BCLGN | FP16 |
| xt | BCLHP | FP16 |
输出
| Tensor | shape | dtype |
|---|---|---|
| states | BCHNP | 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
out = torch.ops.npu_ops_transformer_ext.mamba2_chunk_state(dtout, dacs, bt, xt)
测试方法
见当前目录 tests/
python test_chunk_state.py