mamba2_ssd_scan

Function mamba2_ssd_scan 

Source
pub fn mamba2_ssd_scan(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    h_state: DevicePtr,
    x: DevicePtr,
    b_proj: DevicePtr,
    c_proj: DevicePtr,
    d_param: DevicePtr,
    dt_f32: DevicePtr,
    da_cs: DevicePtr,
    cb: DevicePtr,
    output: DevicePtr,
    seq_len: u32,
    num_heads: u32,
    head_dim: u32,
    state_size: u32,
    n_groups: u32,
    nchunks: u32,
    batch_size: u32,
    x_stride: u32,
    bc_stride: u32,
    y_stride: u32,
    stream: u64,
) -> Result<()>
Expand description

K3: fused chunk_state + state_passing + chunk_scan (h0 stays in shared memory, so the per-chunk states tensor vLLM round-trips through DRAM never exists).