burn_mamba/mamba3/single_ssd/ssd/serial_recalculated/
backward.rs

1//! # Custom autodiff node for the Mamba-3 single-SSD recompute backward
2//!
3//! Implements [`Mamba3SingleSsdBackendExt`](crate::mamba3::single_ssd::ssd::Mamba3SingleSsdBackendExt)
4//! for `Autodiff<B>` via a single Burn [`Backward`](burn::backend::autodiff::ops::Backward) node.  The forward keeps only
5//! its leaf inputs; backprop replays the serial kernels and the gradient math in
6//! [`super::combined_backward`](crate::mamba3::single_ssd::ssd::serial_recalculated::combined_backward), so the large intermediates are never retained.
7//! The two outputs (`y`, `final_state`) are flattened into one tracked tensor
8//! (via [`burn_stack::utils::combined_grad`]) so one node covers both.
9
10#![allow(non_snake_case)]
11
12use crate::mamba3::single_ssd::ssd;
13use burn_stack::utils::fprim::F;
14use burn::backend::autodiff::{
15    Autodiff,
16    checkpoint::{base::Checkpointer, strategy::CheckpointStrategy},
17    grads::Gradients,
18    ops::{Backward, Ops, OpsKind},
19};
20use burn::backend::tensor::FloatTensor;
21use burn::backend::{Backend, BackendTypes};
22use ssd::serial_recalculated::{
23    Mamba3SingleSsdBackendExt,
24    combined_backward::{self, CombinedSingleSsdGrads},
25};
26
27impl<B: Backend + Mamba3SingleSsdBackendExt, C: CheckpointStrategy> Mamba3SingleSsdBackendExt
28    for Autodiff<B, C>
29{
30    /// Memory-efficient combined forward+backward for the Mamba-3 MIMO
31    /// Single-SSD.
32    ///
33    /// The two outputs (`y_bnlmhp`, `final_state_bhpr`) are flattened and
34    /// concatenated into a single 1-D tracked tensor so one `Backward<B, 7>`
35    /// node covers both. The seven differentiable inputs are `v, da, b, c,
36    /// gamma, scale, initial_state`.
37    fn single_ssd_serial_recalculated(
38        v_bnlmhp: FloatTensor<Self>,
39        da_bnlh: FloatTensor<Self>,
40        b_bnlmhr: FloatTensor<Self>,
41        c_bnlmhr: FloatTensor<Self>,
42        gamma_bnlh: FloatTensor<Self>,
43        scale_bnlh: FloatTensor<Self>,
44        initial_state_bhpr: FloatTensor<Self>,
45        siso_specialization: bool,
46    ) -> (FloatTensor<Self>, FloatTensor<Self>) {
47        #[derive(Debug)]
48        struct CombinedKernelsBackward;
49
50        #[derive(Clone, Debug)]
51        struct State<B: Backend> {
52            v_bnlmhp: <B as BackendTypes>::FloatTensorPrimitive,
53            da_bnlh: <B as BackendTypes>::FloatTensorPrimitive,
54            b_bnlmhr: <B as BackendTypes>::FloatTensorPrimitive,
55            c_bnlmhr: <B as BackendTypes>::FloatTensorPrimitive,
56            gamma_bnlh: <B as BackendTypes>::FloatTensorPrimitive,
57            scale_bnlh: <B as BackendTypes>::FloatTensorPrimitive,
58            initial_state_bhpr: <B as BackendTypes>::FloatTensorPrimitive,
59            siso_specialization: bool,
60            flat_len_y_BNLMHP: usize,
61            flat_len_final_state_BHPR: usize,
62            shape_v_bnlmhp: [usize; 6],
63            shape_da_bnlh: [usize; 4],
64            shape_b_bnlmhr: [usize; 6],
65            shape_c_bnlmhr: [usize; 6],
66            shape_gamma_bnlh: [usize; 4],
67            shape_scale_bnlh: [usize; 4],
68            shape_initial_state_bhpr: [usize; 4],
69            shape_y_bnlmhp: [usize; 6],
70            shape_final_state_bhpr: [usize; 4],
71        }
72
73        impl<B: Backend + Mamba3SingleSsdBackendExt> Backward<B, 7> for CombinedKernelsBackward {
74            type State = State<B>;
75
76            fn backward(
77                self,
78                ops: Ops<Self::State, 7>,
79                grads: &mut Gradients,
80                _checkpointer: &mut Checkpointer,
81            ) {
82                let [
83                    node_v_bnlmhp,
84                    node_da_bnlh,
85                    node_b_bnlmhr,
86                    node_c_bnlmhr,
87                    node_gamma_bnlh,
88                    node_scale_bnlh,
89                    node_initial_state_bhpr,
90                ] = ops.parents;
91
92                let d_combined: <B as BackendTypes>::FloatTensorPrimitive =
93                    grads.consume::<B>(&ops.node);
94
95                let State {
96                    v_bnlmhp,
97                    da_bnlh,
98                    b_bnlmhr,
99                    c_bnlmhr,
100                    gamma_bnlh,
101                    scale_bnlh,
102                    initial_state_bhpr,
103                    siso_specialization,
104                    flat_len_y_BNLMHP,
105                    flat_len_final_state_BHPR,
106                    shape_v_bnlmhp,
107                    shape_da_bnlh,
108                    shape_b_bnlmhr,
109                    shape_c_bnlmhr,
110                    shape_gamma_bnlh,
111                    shape_scale_bnlh,
112                    shape_initial_state_bhpr,
113                    shape_y_bnlmhp,
114                    shape_final_state_bhpr,
115                } = ops.state;
116
117                // ── Reconstruct saved tensors as rank-tagged primitives ──
118                let v_bnlmhp = F::<B, 6>::new(v_bnlmhp).reshape(shape_v_bnlmhp);
119                let da_bnlh = F::<B, 4>::new(da_bnlh).reshape(shape_da_bnlh);
120                let b_bnlmhr = F::<B, 6>::new(b_bnlmhr).reshape(shape_b_bnlmhr);
121                let c_bnlmhr = F::<B, 6>::new(c_bnlmhr).reshape(shape_c_bnlmhr);
122                let gamma_bnlh = F::<B, 4>::new(gamma_bnlh).reshape(shape_gamma_bnlh);
123                let scale_bnlh = F::<B, 4>::new(scale_bnlh).reshape(shape_scale_bnlh);
124                let initial_state_bhpr =
125                    F::<B, 4>::new(initial_state_bhpr).reshape(shape_initial_state_bhpr);
126
127                let (d_y_bnlmhp, d_final_state_bhpr) =
128                    burn_stack::utils::combined_grad::unflatten_pair::<B, 6, 4>(
129                        d_combined,
130                        flat_len_y_BNLMHP,
131                        flat_len_final_state_BHPR,
132                        shape_y_bnlmhp,
133                        shape_final_state_bhpr,
134                    );
135
136                let CombinedSingleSsdGrads {
137                    d_v_bnlmhp,
138                    d_da_bnlh,
139                    d_b_bnlmhr,
140                    d_c_bnlmhr,
141                    d_gamma_bnlh,
142                    d_scale_bnlh,
143                    d_initial_state_bhpr,
144                } = combined_backward::combined_backward(
145                    F::<B, 6>::new(d_y_bnlmhp),
146                    F::<B, 4>::new(d_final_state_bhpr),
147                    v_bnlmhp,
148                    da_bnlh,
149                    b_bnlmhr,
150                    c_bnlmhr,
151                    gamma_bnlh,
152                    scale_bnlh,
153                    initial_state_bhpr,
154                    siso_specialization,
155                );
156
157                if let Some(n) = node_v_bnlmhp {
158                    grads.register::<B>(n.id, d_v_bnlmhp.inner());
159                }
160                if let Some(n) = node_da_bnlh {
161                    grads.register::<B>(n.id, d_da_bnlh.inner());
162                }
163                if let Some(n) = node_b_bnlmhr {
164                    grads.register::<B>(n.id, d_b_bnlmhr.inner());
165                }
166                if let Some(n) = node_c_bnlmhr {
167                    grads.register::<B>(n.id, d_c_bnlmhr.inner());
168                }
169                if let Some(n) = node_gamma_bnlh {
170                    grads.register::<B>(n.id, d_gamma_bnlh.inner());
171                }
172                if let Some(n) = node_scale_bnlh {
173                    grads.register::<B>(n.id, d_scale_bnlh.inner());
174                }
175                if let Some(n) = node_initial_state_bhpr {
176                    grads.register::<B>(n.id, d_initial_state_bhpr.inner());
177                }
178            }
179        }
180
181        use burn::backend::TensorMetadata;
182        let [batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim] =
183            v_bnlmhp.primitive.shape().dims();
184        let [.., state_rank] = b_bnlmhr.primitive.shape().dims::<6>();
185
186        let flat_len_y_BNLMHP = batch * nchunks * chunk_len * mimo_rank * nheads * per_head_dim;
187        let flat_len_final_state_BHPR = batch * nheads * per_head_dim * state_rank;
188
189        let shape_v_bnlmhp: [usize; 6] =
190            [batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim];
191        let shape_da_bnlh: [usize; 4] = [batch, nchunks, chunk_len, nheads];
192        let shape_b_bnlmhr: [usize; 6] = [batch, nchunks, chunk_len, mimo_rank, nheads, state_rank];
193        let shape_c_bnlmhr: [usize; 6] = [batch, nchunks, chunk_len, mimo_rank, nheads, state_rank];
194        let shape_gamma_bnlh: [usize; 4] = [batch, nchunks, chunk_len, nheads];
195        let shape_scale_bnlh: [usize; 4] = [batch, nchunks, chunk_len, nheads];
196        let shape_initial_state_bhpr: [usize; 4] = [batch, nheads, per_head_dim, state_rank];
197        let shape_y_bnlmhp: [usize; 6] =
198            [batch, nchunks, chunk_len, mimo_rank, nheads, per_head_dim];
199        let shape_final_state_bhpr: [usize; 4] = [batch, nheads, per_head_dim, state_rank];
200
201        match CombinedKernelsBackward
202            .prepare::<C>([
203                v_bnlmhp.node.clone(),
204                da_bnlh.node.clone(),
205                b_bnlmhr.node.clone(),
206                c_bnlmhr.node.clone(),
207                gamma_bnlh.node.clone(),
208                scale_bnlh.node.clone(),
209                initial_state_bhpr.node.clone(),
210            ])
211            .compute_bound()
212            .stateful()
213        {
214            OpsKind::Tracked(prep) => {
215                let (prim_y_bnlmhp, prim_final_state_bhpr) = B::single_ssd_serial_recalculated(
216                    v_bnlmhp.primitive.clone(),
217                    da_bnlh.primitive.clone(),
218                    b_bnlmhr.primitive.clone(),
219                    c_bnlmhr.primitive.clone(),
220                    gamma_bnlh.primitive.clone(),
221                    scale_bnlh.primitive.clone(),
222                    initial_state_bhpr.primitive.clone(),
223                    siso_specialization,
224                );
225
226                let (prim_combined, _, _) = burn_stack::utils::combined_grad::flatten_pair::<B>(
227                    prim_y_bnlmhp,
228                    prim_final_state_bhpr,
229                );
230
231                let state = State {
232                    v_bnlmhp: v_bnlmhp.primitive.clone(),
233                    da_bnlh: da_bnlh.primitive.clone(),
234                    b_bnlmhr: b_bnlmhr.primitive.clone(),
235                    c_bnlmhr: c_bnlmhr.primitive.clone(),
236                    gamma_bnlh: gamma_bnlh.primitive.clone(),
237                    scale_bnlh: scale_bnlh.primitive.clone(),
238                    initial_state_bhpr: initial_state_bhpr.primitive.clone(),
239                    siso_specialization,
240                    flat_len_y_BNLMHP,
241                    flat_len_final_state_BHPR,
242                    shape_v_bnlmhp,
243                    shape_da_bnlh,
244                    shape_b_bnlmhr,
245                    shape_c_bnlmhr,
246                    shape_gamma_bnlh,
247                    shape_scale_bnlh,
248                    shape_initial_state_bhpr,
249                    shape_y_bnlmhp,
250                    shape_final_state_bhpr,
251                };
252                let tracked_combined: FloatTensor<Autodiff<B, C>> =
253                    prep.finish(state, prim_combined);
254
255                let (tracked_y_bnlmhp, tracked_final_state_bhpr) =
256                    burn_stack::utils::combined_grad::autodiff_unflatten_pair::<B, C, 6, 4>(
257                        tracked_combined,
258                        flat_len_y_BNLMHP,
259                        flat_len_final_state_BHPR,
260                        shape_y_bnlmhp,
261                        shape_final_state_bhpr,
262                    );
263
264                (tracked_y_bnlmhp, tracked_final_state_bhpr)
265            }
266
267            OpsKind::UnTracked(prep) => {
268                let (prim_y_bnlmhp, prim_final_state_bhpr) = B::single_ssd_serial_recalculated(
269                    v_bnlmhp.primitive,
270                    da_bnlh.primitive,
271                    b_bnlmhr.primitive,
272                    c_bnlmhr.primitive,
273                    gamma_bnlh.primitive,
274                    scale_bnlh.primitive,
275                    initial_state_bhpr.primitive,
276                    siso_specialization,
277                );
278
279                let (prim_combined, _, _) = burn_stack::utils::combined_grad::flatten_pair::<B>(
280                    prim_y_bnlmhp,
281                    prim_final_state_bhpr,
282                );
283
284                let tracked_combined: FloatTensor<Autodiff<B, C>> = prep.finish(prim_combined);
285
286                let (tracked_y_bnlmhp, tracked_final_state_bhpr) =
287                    burn_stack::utils::combined_grad::autodiff_unflatten_pair::<B, C, 6, 4>(
288                        tracked_combined,
289                        flat_len_y_BNLMHP,
290                        flat_len_final_state_BHPR,
291                        shape_y_bnlmhp,
292                        shape_final_state_bhpr,
293                    );
294
295                (tracked_y_bnlmhp, tracked_final_state_bhpr)
296            }
297        }
298    }
299}
burn_mamba/mamba3/single_ssd/ssd/serial_recalculated/backward.rs

burn_mamba/mamba3/single_ssd/ssd/serial_recalculated/
backward.rs
