spark_runtime/weights/
release.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3use super::WeightStore;
4use crate::gpu::GpuBackend;
5
6/// Release every weight tensor.
7///
8/// Safe to free per-entry because the loaders allocate per-tensor: the fast
9/// path calls `gpu.alloc(meta.len)` once per tensor before inserting it
10/// (`fast_weights/mod.rs:360-388`), and no loader inserts an `.offset()` view of
11/// a shared block into this map. (Fused per-expert views DO exist — see
12/// `weight_loader/step3p7.rs:93` — but they live in the layer structs that own
13/// the fused allocation, not here, so this cannot double-free them.)
14impl avarok_core::scope::ModelResource<dyn GpuBackend> for WeightStore {
15    fn label(&self) -> &'static str {
16        "weight store"
17    }
18
19    fn release(&mut self, gpu: &dyn GpuBackend) -> anyhow::Result<()> {
20        // Derived buffers FIRST: they are re-encodings of the tensors below and
21        // nothing reads one after the other is gone, but freeing the source a
22        // derivation was built from while the derivation is still listed would
23        // make a later failure here impossible to attribute.
24        let mut first_error = self.derived.release(gpu).err();
25        // `drain` rather than iterate: the map must not be left holding
26        // pointers to memory that is gone, and it makes this idempotent.
27        for (name, tensor) in self.weights.drain() {
28            if let Err(e) = gpu.free(tensor.ptr)
29                && first_error.is_none()
30            {
31                first_error = Some(e.context(format!("freeing weight {name}")));
32            }
33        }
34        match first_error {
35            Some(e) => Err(e),
36            None => Ok(()),
37        }
38    }
39}
spark_runtime/weights/release.rs

spark_runtime/weights/
release.rs
