metrale_model_layers/layers/ops/
derived_weights.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Per-model memo of GPU re-encodings of a model's weights, keyed
4//! by the source weight's device pointer.
5//!
6//! The derivations are: FP8 dequantized to BF16, block-scaled FP8
7//! re-quantized to row-wise FP8, NVFP4 bytes transposed for CUTLASS, and FP8
8//! packed to CUTLASS NVFP4. The key is a raw device pointer, which another
9//! model's allocator can hand out again; because each model owns its memo, a
10//! reused address cannot hit an entry made for a different model's weight.
11//!
12//! Owner: model-layers ops.
13//! Invariants:
14//! - Each `Derivation` has its own map, so a key never returns another kind's
15//!   value.
16//! - A key keeps the first value stored for it; a later insert is dropped.
17//! - `release` empties every map, and frees only the values, never the keys.
18
19use std::collections::HashMap;
20// 2026-09-25: parking_lot's mutex does not poison, so a panic elsewhere cannot
21// block teardown.
22use parking_lot::Mutex;
23
24/// 2026-09-25: Per-model memo of derived weight encodings, one map per
25/// [`Derivation`].
26#[derive(Default)]
27pub struct DerivedWeights {
28    /// 2026-09-25: FP8 weight ptr → `(row-wise FP8 ptr, per-row scale ptr)`.
29    rowwise_fp8: Mutex<HashMap<u64, (u64, u64)>>,
30    /// 2026-09-25: FP8 weight ptr → BF16 ptr.
31    bf16: Mutex<HashMap<u64, u64>>,
32    /// 2026-09-25: NVFP4 weight ptr → `[N,K/2]` transposed bytes ptr.
33    cutlass_nvfp4_t: Mutex<HashMap<u64, u64>>,
34    /// 2026-09-25: FP8 weight ptr → `(CUTLASS NVFP4 ptr, scale ptr)`.
35    cutlass_nvfp4_from_fp8: Mutex<HashMap<u64, (u64, u64)>>,
36}
37
38/// 2026-09-25: Which derivation a lookup is for.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum Derivation {
41    RowwiseFp8,
42    Bf16,
43    CutlassNvfp4Transposed,
44    CutlassNvfp4FromFp8,
45}
46
47impl DerivedWeights {
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    /// 2026-09-25: Memo read for a single-pointer derivation, for callers that
53    /// allocate and launch between the lookup and [`Self::insert_ptr`]. A
54    /// pair-valued `kind` returns `None`.
55    pub fn get_ptr(&self, kind: Derivation, key: u64) -> Option<u64> {
56        let map = match kind {
57            Derivation::Bf16 => &self.bf16,
58            Derivation::CutlassNvfp4Transposed => &self.cutlass_nvfp4_t,
59            _ => return None,
60        };
61        map.lock().get(&key).copied()
62    }
63
64    pub fn insert_ptr(&self, kind: Derivation, key: u64, value: u64) {
65        let map = match kind {
66            Derivation::Bf16 => &self.bf16,
67            Derivation::CutlassNvfp4Transposed => &self.cutlass_nvfp4_t,
68            _ => return,
69        };
70        map.lock().entry(key).or_insert(value);
71    }
72
73    pub fn get_pair(&self, kind: Derivation, key: u64) -> Option<(u64, u64)> {
74        let map = match kind {
75            Derivation::RowwiseFp8 => &self.rowwise_fp8,
76            Derivation::CutlassNvfp4FromFp8 => &self.cutlass_nvfp4_from_fp8,
77            _ => return None,
78        };
79        map.lock().get(&key).copied()
80    }
81
82    pub fn insert_pair(&self, kind: Derivation, key: u64, value: (u64, u64)) {
83        let map = match kind {
84            Derivation::RowwiseFp8 => &self.rowwise_fp8,
85            Derivation::CutlassNvfp4FromFp8 => &self.cutlass_nvfp4_from_fp8,
86            _ => return,
87        };
88        map.lock().entry(key).or_insert(value);
89    }
90
91    /// 2026-09-25: Look up a single-pointer derivation, building it on a miss.
92    /// A failed build stores nothing.
93    ///
94    /// `build` runs outside the lock. If two threads race on one key, both
95    /// build, the first stored value is kept and returned to both, and the
96    /// loser's value is not freed.
97    pub fn get_or_build_ptr(
98        &self,
99        kind: Derivation,
100        key: u64,
101        build: impl FnOnce() -> anyhow::Result<u64>,
102    ) -> anyhow::Result<u64> {
103        let map = match kind {
104            Derivation::Bf16 => &self.bf16,
105            Derivation::CutlassNvfp4Transposed => &self.cutlass_nvfp4_t,
106            _ => unreachable!("pair-valued derivation routed to get_or_build_ptr"),
107        };
108        if let Some(&hit) = map.lock().get(&key) {
109            return Ok(hit);
110        }
111        let built = build()?;
112        Ok(*map.lock().entry(key).or_insert(built))
113    }
114
115    /// 2026-09-25: [`Self::get_or_build_ptr`] for a pair-valued derivation.
116    pub fn get_or_build_pair(
117        &self,
118        kind: Derivation,
119        key: u64,
120        build: impl FnOnce() -> anyhow::Result<(u64, u64)>,
121    ) -> anyhow::Result<(u64, u64)> {
122        let map = match kind {
123            Derivation::RowwiseFp8 => &self.rowwise_fp8,
124            Derivation::CutlassNvfp4FromFp8 => &self.cutlass_nvfp4_from_fp8,
125            _ => unreachable!("single-valued derivation routed to get_or_build_pair"),
126        };
127        if let Some(&hit) = map.lock().get(&key) {
128            return Ok(hit);
129        }
130        let built = build()?;
131        Ok(*map.lock().entry(key).or_insert(built))
132    }
133
134    /// 2026-09-25: Total memoized entries across the four maps.
135    pub fn len(&self) -> usize {
136        self.rowwise_fp8.lock().len()
137            + self.bf16.lock().len()
138            + self.cutlass_nvfp4_t.lock().len()
139            + self.cutlass_nvfp4_from_fp8.lock().len()
140    }
141
142    pub fn is_empty(&self) -> bool {
143        self.len() == 0
144    }
145}
146
147impl std::fmt::Debug for DerivedWeights {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        f.debug_struct("DerivedWeights")
150            .field("entries", &self.len())
151            .finish()
152    }
153}
154
155/// 2026-09-25: Release every derived allocation.
156///
157/// The keys are the source weights' pointers, which the weight store owns and
158/// frees; only the values are freed here. Every free is attempted, and the
159/// first error is returned.
160impl metrale_core::scope::ModelResource<dyn metrale_gpu_runtime::gpu::GpuBackend>
161    for DerivedWeights
162{
163    fn label(&self) -> &'static str {
164        "derived weights"
165    }
166
167    fn release(&mut self, gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend) -> anyhow::Result<()> {
168        let mut owned: Vec<u64> = Vec::new();
169        // 2026-09-25: Drain, so a later lookup cannot return a freed pointer.
170        for (_, (a, b)) in self.rowwise_fp8.lock().drain() {
171            owned.push(a);
172            owned.push(b);
173        }
174        owned.extend(self.bf16.lock().drain().map(|(_, v)| v));
175        owned.extend(self.cutlass_nvfp4_t.lock().drain().map(|(_, v)| v));
176        for (_, (a, b)) in self.cutlass_nvfp4_from_fp8.lock().drain() {
177            owned.push(a);
178            owned.push(b);
179        }
180        let mut first_error = None;
181        for raw in owned {
182            if let Err(e) = gpu.free(metrale_gpu_runtime::gpu::DevicePtr(raw))
183                && first_error.is_none()
184            {
185                first_error = Some(e);
186            }
187        }
188        match first_error {
189            Some(e) => Err(e),
190            None => Ok(()),
191        }
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use metrale_core::scope::ModelResource;
199    use metrale_gpu_runtime::gpu::GpuBackend;
200    use metrale_gpu_runtime::gpu::mock::MockGpuBackend;
201    use std::sync::atomic::{AtomicUsize, Ordering};
202
203    #[test]
204    fn a_fresh_cache_is_empty() {
205        assert!(DerivedWeights::new().is_empty());
206    }
207
208    #[test]
209    fn a_derivation_is_built_once_per_key() {
210        let d = DerivedWeights::new();
211        let builds = AtomicUsize::new(0);
212        let build_for = |ptr: u64| {
213            d.get_or_build_ptr(Derivation::Bf16, ptr, || {
214                builds.fetch_add(1, Ordering::Relaxed);
215                Ok(ptr + 1000)
216            })
217            .unwrap()
218        };
219        assert_eq!(build_for(10), 1010);
220        assert_eq!(build_for(10), 1010);
221        assert_eq!(builds.load(Ordering::Relaxed), 1);
222        assert_eq!(build_for(20), 1020);
223        assert_eq!(builds.load(Ordering::Relaxed), 2);
224        assert_eq!(d.len(), 2);
225    }
226
227    #[test]
228    fn the_derivations_do_not_share_a_keyspace() {
229        // 2026-09-25: One source pointer, two encodings, two separate entries.
230        let d = DerivedWeights::new();
231        let bf16 = d
232            .get_or_build_ptr(Derivation::Bf16, 0x1000, || Ok(0xB16))
233            .unwrap();
234        let nvfp4 = d
235            .get_or_build_ptr(Derivation::CutlassNvfp4Transposed, 0x1000, || Ok(0x4444))
236            .unwrap();
237        assert_eq!(bf16, 0xB16);
238        assert_eq!(nvfp4, 0x4444);
239        assert_eq!(d.len(), 2);
240    }
241
242    #[test]
243    fn a_failed_build_is_not_memoized() {
244        let d = DerivedWeights::new();
245        assert!(
246            d.get_or_build_ptr(Derivation::Bf16, 7, || anyhow::bail!("oom"))
247                .is_err()
248        );
249        assert!(d.is_empty(), "a failure must not poison the key");
250        assert_eq!(
251            d.get_or_build_ptr(Derivation::Bf16, 7, || Ok(99)).unwrap(),
252            99
253        );
254    }
255
256    #[test]
257    fn two_models_memoize_independently_even_on_a_recycled_pointer() {
258        // 2026-09-25: Model A memoizes a derivation for an address, A is
259        // dropped, and model B sees the same address for a different weight.
260        let a = DerivedWeights::new();
261        let recycled = 0x7f00_0000u64;
262        assert_eq!(
263            a.get_or_build_ptr(Derivation::Bf16, recycled, || Ok(0xAAAA))
264                .unwrap(),
265            0xAAAA
266        );
267        drop(a);
268
269        let b = DerivedWeights::new();
270        assert_eq!(
271            b.get_or_build_ptr(Derivation::Bf16, recycled, || Ok(0xBBBB))
272                .unwrap(),
273            0xBBBB,
274            "the same address must resolve to the NEW model's derivation"
275        );
276    }
277
278    #[test]
279    fn pair_derivations_round_trip_without_sharing_a_keyspace() {
280        let d = DerivedWeights::new();
281        let got = d
282            .get_or_build_pair(Derivation::RowwiseFp8, 5, || Ok((11, 22)))
283            .unwrap();
284        assert_eq!(got, (11, 22));
285        assert_eq!(
286            d.get_or_build_pair(Derivation::CutlassNvfp4FromFp8, 5, || Ok((33, 44)))
287                .unwrap(),
288            (33, 44)
289        );
290        assert_eq!(
291            d.get_or_build_pair(Derivation::RowwiseFp8, 5, || Ok((99, 99)))
292                .unwrap(),
293            (11, 22),
294            "cached, not rebuilt"
295        );
296        assert_eq!(d.len(), 2);
297    }
298
299    #[test]
300    fn release_frees_only_derived_values_and_drains_every_map() {
301        let gpu = MockGpuBackend::new();
302        let mut d = DerivedWeights::new();
303        let keys: Vec<u64> = (0..4).map(|_| gpu.alloc(1).unwrap().0).collect();
304        let values: Vec<u64> = (0..6).map(|_| gpu.alloc(1).unwrap().0).collect();
305
306        d.insert_pair(Derivation::RowwiseFp8, keys[0], (values[0], values[1]));
307        d.insert_ptr(Derivation::Bf16, keys[1], values[2]);
308        d.insert_ptr(Derivation::CutlassNvfp4Transposed, keys[2], values[3]);
309        d.insert_pair(
310            Derivation::CutlassNvfp4FromFp8,
311            keys[3],
312            (values[4], values[5]),
313        );
314        assert_eq!(gpu.alloc_count(), 10);
315        assert_eq!(d.len(), 4);
316
317        d.release(&gpu).unwrap();
318
319        assert!(d.is_empty(), "freed pointers must not remain memoized");
320        assert_eq!(gpu.alloc_count(), 4, "source-weight keys remain GPU-owned");
321        for key in keys {
322            gpu.free(metrale_gpu_runtime::gpu::DevicePtr(key)).unwrap();
323        }
324    }
325}
metrale_model_layers/layers/ops/derived_weights.rs

metrale_model_layers/layers/ops/
derived_weights.rs
