Skip to main content

Module ssd_path

Module ssd_path 

Source
Expand description

The Mamba3SingleSsdInput bundle and its run() dispatch.

§Single-Pass SSD — Input Bundle

Sibling to crate::mamba3::double_ssd::ssd::ssd_path. Where the double-ssd pathway runs the standard SSD twice (γ-term and β-term), this module’s Mamba3SingleSsdInput runs one merged SSD pass that absorbs both contributions by scaling K with scaleₜ = γₜ + (1−λₜ₊₁) Δₜ₊₁. The same-step diagonal contribution differs (it must use γₜ, not scaleₜ) and is patched via an explicit correction term inside each variant. Reference kernels:

  • refs/state-spaces/mamba/mamba_ssm/ops/triton/mamba3/mamba3_siso_fwd.py
  • refs/state-spaces/mamba/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py

The interface is MIMO-first (matches the other burn-mamba SSD inputs), with mimo_rank = 1 collapsing to the SISO case. The algorithm is selected by [Mamba3SsdPath], shared with the double-ssd pathway.

Structs§

Mamba3SingleSsdInput
MIMO-first input bundle for the merged-form SSD.