pub fn mamba2_ssm_prefill_persistent(
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
Persistent Mamba-2 SSM prefill: H in shared memory, reduces global traffic. Same parameters and launch config as mamba2_ssm_prefill.