Skip to main content

Module serial

burn_mamba::mamba2::ssd

Module serial 

Source

Functionsยง

k1_ssd_chunk_cumsum
Based on the Kernel 1 Triton reference _chunk_cumsum_fwd_kernel (ssd_chunk_state.py).
k2_ssd_bmm
Based on the Kernel 2 Triton reference _bmm_chunk_fwd_kernel (ssd_bmm.py).
k3_ssd_chunk_state
Based on the Kernel 3 Triton reference _chunk_state_fwd_kernel (ssd_chunk_state.py).
k4_ssd_state_passing
Based on the Kernel 4 Triton reference _state_passing_fwd_kernel (ssd_state_passing.py).
k5_ssd_chunk_scan
Based on the Kernel 5 Triton reference _chunk_scan_fwd_kernel (ssd_chunk_scan.py).