burn_mamba/unified/
cache.rs

1//! Runtime-tagged cache collection, and where each family plugs into the
2//! block-generic stack: `impl Block`, `impl BlockConfig`, `impl CacheStack`.
3
4use burn::prelude::*;
5use burn_stack::modules::{Block, BlockConfig, CacheStack};
6
7
8/// Runtime-tagged caches: one variant per family, matching
9/// [`MambaLatentNet`](crate::unified::MambaLatentNet).
10///
11/// This is plain runtime state (not a `Module`): caches are threaded through
12/// `forward`/`step`, never recorded or optimised. (`Mamba3Caches` is itself a
13/// non-`Module` enum, so a `Module` derive here would not even apply.)
14#[derive(Debug, Clone)]
15pub enum MambaCaches {
16    /// Mamba-1 caches.
17    #[cfg(feature = "mamba1")]
18    Mamba1(crate::mamba1::prelude::Mamba1Caches),
19    /// Mamba-2 caches.
20    #[cfg(feature = "mamba2")]
21    Mamba2(crate::mamba2::prelude::Mamba2Caches),
22    /// Mamba-3 caches.
23    #[cfg(feature = "mamba3")]
24    Mamba3(crate::mamba3::prelude::Mamba3Caches),
25}
26
27// ===========================================================================
28// Per-family impls
29// ===========================================================================
30
31#[cfg(feature = "mamba2")]
32mod impl_mamba2 {
33    use super::*;
34    use crate::mamba2::prelude::{
35        Mamba2, Mamba2Cache, Mamba2CacheConfig, Mamba2Caches, Mamba2CachesConfig, Mamba2Config,
36        Mamba2SsdPath,
37    };
38
39    impl CacheStack for Mamba2Caches {
40        type Cache = Mamba2Cache;
41        fn slot_count(&self) -> usize {
42            self.caches.len()
43        }
44        fn into_slots(self) -> Vec<Option<Mamba2Cache>> {
45            self.caches.into_iter().map(Some).collect()
46        }
47        fn from_slots(slots: Vec<Option<Mamba2Cache>>) -> Self {
48            Self {
49                caches: slots.into_iter().map(Option::unwrap).collect(),
50            }
51        }
52        fn cache_to_inner(c: Mamba2Cache) -> Mamba2Cache {
53            Mamba2Cache {
54                conv_bvk: c.conv_bvk.inner(),
55                ssm_bhpr: c.ssm_bhpr.inner(),
56            }
57        }
58        fn cache_from_inner(c: Mamba2Cache) -> Mamba2Cache {
59            Mamba2Cache {
60                conv_bvk: Tensor::from_inner(c.conv_bvk),
61                ssm_bhpr: Tensor::from_inner(c.ssm_bhpr),
62            }
63        }
64    }
65
66    impl Block for Mamba2 {
67        type Cache = Mamba2Cache;
68        type Caches = Mamba2Caches;
69        type Options = Mamba2SsdPath;
70
71        fn block_forward(
72            &self,
73            x: Tensor<3>,
74            cache: Option<Mamba2Cache>,
75            options: Mamba2SsdPath,
76        ) -> (Tensor<3>, Mamba2Cache) {
77            self.forward(x, cache, options)
78        }
79        fn block_step(&self, x: Tensor<2>, cache: Option<Mamba2Cache>) -> (Tensor<2>, Mamba2Cache) {
80            self.step(x, cache)
81        }
82        fn zero_caches_3d(&self, x: &Tensor<3>, n_virtual: usize) -> Mamba2Caches {
83            let [batch, _seq, _d] = x.dims();
84            self.make_zero(batch, n_virtual, &x.device())
85        }
86        fn zero_caches_2d(&self, x: &Tensor<2>, n_virtual: usize) -> Mamba2Caches {
87            let [batch, _d] = x.dims();
88            self.make_zero(batch, n_virtual, &x.device())
89        }
90    }
91
92    impl Mamba2 {
93        fn make_zero(&self, batch: usize, n_virtual: usize, device: &Device) -> Mamba2Caches {
94            let [conv_dim, _, conv_kernel] = self.conv1d.weight.dims();
95            Mamba2CachesConfig::new(
96                n_virtual,
97                Mamba2CacheConfig {
98                    batch,
99                    state_rank: self.state_rank,
100                    conv_kernel,
101                    conv_dim,
102                    per_head_dim: self.per_head_dim(),
103                    nheads: self.nheads(),
104                },
105            )
106            .init(device)
107        }
108    }
109
110    impl BlockConfig for Mamba2Config {
111        type Block = Mamba2;
112        fn d_model(&self) -> usize {
113            self.d_model
114        }
115        fn init_block(&self, device: &Device) -> Mamba2 {
116            self.init(device)
117        }
118        #[cfg(feature = "optim")]
119        fn muon_projections(&self) -> Vec<burn_stack::optim::ProjSpec> {
120            self.muon_projections()
121        }
122    }
123}
124
125#[cfg(feature = "mamba3")]
126mod impl_mamba3 {
127    use super::*;
128    use crate::mamba3::double_ssd::prelude::Mamba3DoubleSsdCache;
129    use crate::mamba3::prelude::{Mamba3, Mamba3Cache, Mamba3Caches, Mamba3Config, Mamba3SsdPath};
130    use crate::mamba3::single_ssd::prelude::{
131        Mamba3SingleSsdCache, Mamba3SingleSsdCacheConfig, Mamba3SingleSsdCaches,
132        Mamba3SingleSsdCachesConfig,
133    };
134
135    /// Zero single-ssd caches sized from a `[batch, sequence, d_model]` input.
136    /// (A missing cache defaults to the single-ssd pathway — ≈½ the SSD memory
137    /// of double-ssd — for either rotation kind.)
138    fn zero_single_ssd_caches(
139        mamba_block: &Mamba3,
140        batch: usize,
141        n_virtual: usize,
142        device: &Device,
143    ) -> Mamba3SingleSsdCaches {
144        Mamba3SingleSsdCachesConfig::new(
145            n_virtual,
146            Mamba3SingleSsdCacheConfig {
147                batch,
148                state_rank: mamba_block.state_rank,
149                num_rope_angles: mamba_block.num_rope_angles,
150                per_head_dim: mamba_block.per_head_dim(),
151                nheads: mamba_block.nheads(),
152                mimo_rank: mamba_block.mimo_rank,
153                rotation: mamba_block.rotation,
154                num_quat_blocks: mamba_block.num_quat_blocks,
155            },
156        )
157        .init(device)
158    }
159
160    impl CacheStack for Mamba3Caches {
161        type Cache = Mamba3Cache;
162        fn slot_count(&self) -> usize {
163            self.caches_len()
164        }
165        fn into_slots(self) -> Vec<Option<Mamba3Cache>> {
166            self.into_options()
167        }
168        fn from_slots(slots: Vec<Option<Mamba3Cache>>) -> Self {
169            Self::from_options(slots)
170        }
171        fn cache_to_inner(c: Mamba3Cache) -> Mamba3Cache {
172            use crate::mamba3::prelude::RotationState;
173            fn rot(r: RotationState) -> RotationState {
174                match r {
175                    RotationState::Real(u) => RotationState::Real(u),
176                    RotationState::Angle(t) => RotationState::Angle(t.inner()),
177                    RotationState::Quaternion(t) => RotationState::Quaternion(t.inner()),
178                    RotationState::Rotor(t) => RotationState::Rotor(t.inner()),
179                }
180            }
181            match c {
182                Mamba3Cache::DoubleSsd(c) => Mamba3Cache::DoubleSsd(Mamba3DoubleSsdCache {
183                    ssm_bhpr: c.ssm_bhpr.inner(),
184                    k_state_bmhr: c.k_state_bmhr.inner(),
185                    v_state_bhp: c.v_state_bhp.inner(),
186                    rotation: rot(c.rotation),
187                }),
188                Mamba3Cache::SingleSsd(c) => Mamba3Cache::SingleSsd(Mamba3SingleSsdCache {
189                    ssm_bhpr: c.ssm_bhpr.inner(),
190                    k_state_bmhr: c.k_state_bmhr.inner(),
191                    v_state_bhp: c.v_state_bhp.inner(),
192                    rotation: rot(c.rotation),
193                }),
194            }
195        }
196        fn cache_from_inner(c: Mamba3Cache) -> Mamba3Cache {
197            use crate::mamba3::prelude::RotationState;
198            fn rot(r: RotationState) -> RotationState {
199                match r {
200                    RotationState::Real(u) => RotationState::Real(u),
201                    RotationState::Angle(t) => RotationState::Angle(Tensor::from_inner(t)),
202                    RotationState::Quaternion(t) => {
203                        RotationState::Quaternion(Tensor::from_inner(t))
204                    }
205                    RotationState::Rotor(t) => RotationState::Rotor(Tensor::from_inner(t)),
206                }
207            }
208            match c {
209                Mamba3Cache::DoubleSsd(c) => Mamba3Cache::DoubleSsd(Mamba3DoubleSsdCache {
210                    ssm_bhpr: Tensor::from_inner(c.ssm_bhpr),
211                    k_state_bmhr: Tensor::from_inner(c.k_state_bmhr),
212                    v_state_bhp: Tensor::from_inner(c.v_state_bhp),
213                    rotation: rot(c.rotation),
214                }),
215                Mamba3Cache::SingleSsd(c) => Mamba3Cache::SingleSsd(Mamba3SingleSsdCache {
216                    ssm_bhpr: Tensor::from_inner(c.ssm_bhpr),
217                    k_state_bmhr: Tensor::from_inner(c.k_state_bmhr),
218                    v_state_bhp: Tensor::from_inner(c.v_state_bhp),
219                    rotation: rot(c.rotation),
220                }),
221            }
222        }
223    }
224
225    impl Block for Mamba3 {
226        type Cache = Mamba3Cache;
227        type Caches = Mamba3Caches;
228        type Options = Mamba3SsdPath;
229
230        fn block_forward(
231            &self,
232            x: Tensor<3>,
233            cache: Option<Mamba3Cache>,
234            options: Mamba3SsdPath,
235        ) -> (Tensor<3>, Mamba3Cache) {
236            self.forward(x, cache, options)
237        }
238        fn block_step(&self, x: Tensor<2>, cache: Option<Mamba3Cache>) -> (Tensor<2>, Mamba3Cache) {
239            self.step(x, cache)
240        }
241        fn block_step_infinite(&self, x: Tensor<2>) -> Tensor<2> {
242            self.step_infinite(x)
243        }
244        fn zero_caches_3d(&self, x: &Tensor<3>, n_virtual: usize) -> Mamba3Caches {
245            let [batch, _seq, _d] = x.dims();
246            zero_single_ssd_caches(self, batch, n_virtual, &x.device()).into()
247        }
248        fn zero_caches_2d(&self, x: &Tensor<2>, n_virtual: usize) -> Mamba3Caches {
249            let [batch, _d] = x.dims();
250            zero_single_ssd_caches(self, batch, n_virtual, &x.device()).into()
251        }
252    }
253
254    impl BlockConfig for Mamba3Config {
255        type Block = Mamba3;
256        fn d_model(&self) -> usize {
257            self.d_model
258        }
259        fn init_block(&self, device: &Device) -> Mamba3 {
260            self.init(device)
261        }
262        #[cfg(feature = "optim")]
263        fn muon_projections(&self) -> Vec<burn_stack::optim::ProjSpec> {
264            self.muon_projections()
265        }
266    }
267}
268
269#[cfg(feature = "mamba1")]
270mod impl_mamba1 {
271    use super::*;
272    use crate::mamba1::prelude::{
273        Mamba1, Mamba1Cache, Mamba1CacheConfig, Mamba1Caches, Mamba1CachesConfig, Mamba1Config,
274    };
275
276    impl CacheStack for Mamba1Caches {
277        type Cache = Mamba1Cache;
278        fn slot_count(&self) -> usize {
279            self.caches.len()
280        }
281        fn into_slots(self) -> Vec<Option<Mamba1Cache>> {
282            self.caches.into_iter().map(Some).collect()
283        }
284        fn from_slots(slots: Vec<Option<Mamba1Cache>>) -> Self {
285            Self {
286                caches: slots.into_iter().map(Option::unwrap).collect(),
287            }
288        }
289        fn cache_to_inner(c: Mamba1Cache) -> Mamba1Cache {
290            Mamba1Cache {
291                conv_bik: c.conv_bik.inner(),
292                ssm_bir: c.ssm_bir.inner(),
293            }
294        }
295        fn cache_from_inner(c: Mamba1Cache) -> Mamba1Cache {
296            Mamba1Cache {
297                conv_bik: Tensor::from_inner(c.conv_bik),
298                ssm_bir: Tensor::from_inner(c.ssm_bir),
299            }
300        }
301    }
302
303    impl Block for Mamba1 {
304        type Cache = Mamba1Cache;
305        type Caches = Mamba1Caches;
306        /// Mamba-1 has no SSD chunking, so there is no path selector.
307        type Options = ();
308
309        fn block_forward(
310            &self,
311            x: Tensor<3>,
312            cache: Option<Mamba1Cache>,
313            _options: (),
314        ) -> (Tensor<3>, Mamba1Cache) {
315            self.forward(x, cache)
316        }
317        fn block_step(&self, x: Tensor<2>, cache: Option<Mamba1Cache>) -> (Tensor<2>, Mamba1Cache) {
318            self.step(x, cache)
319        }
320        fn zero_caches_3d(&self, x: &Tensor<3>, n_virtual: usize) -> Mamba1Caches {
321            let [batch, _seq, _d] = x.dims();
322            self.make_zero(batch, n_virtual, &x.device())
323        }
324        fn zero_caches_2d(&self, x: &Tensor<2>, n_virtual: usize) -> Mamba1Caches {
325            let [batch, _d] = x.dims();
326            self.make_zero(batch, n_virtual, &x.device())
327        }
328    }
329
330    impl Mamba1 {
331        fn cache_config(&self, batch: usize) -> Mamba1CacheConfig {
332            let [d_inner, state_rank] = self.a_log.dims();
333            let [_, _, conv_kernel] = self.conv1d.weight.dims();
334            Mamba1CacheConfig::new(batch, d_inner)
335                .with_state_rank(state_rank)
336                .with_conv_kernel(conv_kernel)
337        }
338        fn make_zero(&self, batch: usize, n_virtual: usize, device: &Device) -> Mamba1Caches {
339            Mamba1CachesConfig::new(n_virtual, self.cache_config(batch)).init(device)
340        }
341    }
342
343    impl BlockConfig for Mamba1Config {
344        type Block = Mamba1;
345        fn d_model(&self) -> usize {
346            self.d_model
347        }
348        fn init_block(&self, device: &Device) -> Mamba1 {
349            self.init(device)
350        }
351        #[cfg(feature = "optim")]
352        fn muon_projections(&self) -> Vec<burn_stack::optim::ProjSpec> {
353            self.muon_projections()
354        }
355    }
356}
burn_mamba/unified/cache.rs

burn_mamba/unified/
cache.rs
