mamba2_chunk_state_passing 算子说明

功能和实现说明

mamba2_chunk_state_passing 用于在 MambaV2 Prefill 阶段进行跨 chunk 状态传递，将 chunk 内计算得到的状态按照时间顺序依次传递，并在各 chunk 之间执行指数衰减和新状态叠加，从而形成完整的跨 chunk 状态序列；同时返回最终的全局状态，用于下一阶段推理。该算子通常在 chunk_state 之后调用，用于将 chunk 内状态扩展为跨 chunk 的连续状态序列，并生成下一阶段使用的最终状态。同时，本算子在状态传递完成后，对重排后的状态张量与 ct 执行基于 Cube 的批量矩阵乘（states @ ct），实现类似 inter-attention 的跨 chunk 状态混合。 该算子实现为Vector+Cube融合算子，通过VC并行提升计算性能。

计算流

Kernel输入输出（I/O）

输入

Tensor	shape	dtype
dacs	BCLH	FP32
init_state	BHZ	FP32
states	BCHZ	FP32
ct	BCLGN	FP16

输出

Tensor	shape	dtype
inter_attn	BCHLP	FP32
final_state	BHNP	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

inter_attn, final_state = torch.ops.npu_ops_transformer_ext.mamba2_chunk_state_passing(dacs, init_state, states, ct)

测试方法

见当前目录 tests/

python test_chunk_state_passing.py