pub fn mamba2_ssm_prefill(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
h_state: DevicePtr,
x: DevicePtr,
b_proj: DevicePtr,
c_proj: DevicePtr,
dt_raw: DevicePtr,
a_log: DevicePtr,
d_param: DevicePtr,
dt_bias: DevicePtr,
output: DevicePtr,
batch_size: u32,
seq_len: u32,
num_heads: u32,
head_dim: u32,
state_size: u32,
n_groups: u32,
dt_min: f32,
dt_max: f32,
x_stride: u32,
bc_stride: u32,
dt_stride: u32,
y_stride: u32,
stream: u64,
) -> Result<()>Expand description
2026-09-25: Mamba-2 SSM prefill (mamba2_ssm_prefill): the token-sequential
recurrence over seq_len tokens in one launch. x_stride, bc_stride,
dt_stride and y_stride are BF16 elements between consecutive tokens.