Skip to main content

combined_backward

burn_mamba::mamba2::ssd::serial_recalculated::combined_backward

Function combined_backward 

Source 
pub fn combined_backward<B: Backend>(
    d_y_bnlhp: F<B, 5>,
    d_final_bhpr: F<B, 4>,
    x_bnlhp: F<B, 5>,
    dt_discretized_bhnl: F<B, 4>,
    b_bnlhr: F<B, 5>,
    c_bnlhr: F<B, 5>,
    d_h: F<B, 1>,
    initial_state_bhpr: F<B, 4>,
    a_decay_h: F<B, 1>,
) -> CombinedGrads<B>
Expand description

Memory-efficient backward for the Mamba-2 chunkwise SSD.

Recomputes the forward intermediates (K1-K4) from the saved inputs, then runs a reverse per-chunk loop that fuses the K5 (BLUE + ORANGE) backward with the K4 state-passing backward. K3/K2/K1 backwards run as single batched ops once the loop has collected all per-chunk slices.

§Arguments

  • d_y_bnlhp — upstream gradient of the SSD output
  • d_final_bhpr — upstream gradient of the final SSM state
  • x_bnlhp, dt_discretized_bhnl, b_bnlhr, c_bnlhr, d_h, initial_state_bhpr, a_decay_h — the seven saved forward inputs

§Returns

One CombinedGrads struct containing gradients for all 7 inputs.