metrale_model_weights/weights/store.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! `WeightStore` methods: lookup, deferral, derived tensors and the byte and
4//! FP8 accounting over the loaded set.
5
6use super::*;
7
8impl WeightStore {
9 /// Create an empty weight store (for testing).
10 pub fn empty() -> Self {
11 Self {
12 weights: HashMap::new(),
13 prepartitioned_tp: None,
14 deferred: HashMap::new(),
15 derived: DerivedStore::default(),
16 }
17 }
18
19 /// Record a tensor that was skipped at load, with its on-disk location.
20 pub fn defer(&mut self, name: String, t: DeferredTensor) {
21 self.deferred.insert(name, t);
22 }
23
24 /// Look up a deferred (not-uploaded) tensor's on-disk location.
25 pub fn deferred(&self, name: &str) -> Option<&DeferredTensor> {
26 self.deferred.get(name)
27 }
28
29 /// Every deferred tensor, name-sorted (NUMERIC on a trailing index, so
30 /// `embedders.10` sorts after `embedders.2` â a lexicographic sort here
31 /// silently mis-maps the n-gram tables, which cost a real debugging
32 /// session the first time).
33 pub fn deferred_sorted(&self) -> Vec<(&String, &DeferredTensor)> {
34 let mut v: Vec<_> = self.deferred.iter().collect();
35 v.sort_by_key(|(n, _)| split_trailing_index(n));
36 v
37 }
38
39 /// Wrap a pre-built map. Used by alternate loaders (e.g.
40 /// `fast_weights::FastSafetensorsLoader`, and the RDMA weight loader in
41 /// `metrale-storage`, which lives in a different crate and so needs this pub).
42 pub fn from_map(weights: HashMap<String, WeightTensor>) -> Self {
43 Self {
44 weights,
45 prepartitioned_tp: None,
46 deferred: HashMap::new(),
47 derived: DerivedStore::default(),
48 }
49 }
50
51 /// Get a weight tensor by name. Fails fast if not found.
52 pub fn get(&self, name: &str) -> Result<&WeightTensor> {
53 self.weights
54 .get(name)
55 .ok_or_else(|| anyhow::anyhow!("Weight '{name}' not found in store"))
56 }
57
58 /// Check if a weight exists.
59 pub fn contains(&self, name: &str) -> bool {
60 self.weights.contains_key(name)
61 }
62
63 /// Number of loaded weights.
64 pub fn len(&self) -> usize {
65 self.weights.len()
66 }
67
68 /// True if no weights are loaded.
69 pub fn is_empty(&self) -> bool {
70 self.weights.is_empty()
71 }
72
73 /// Device bytes the store still holds. Not the on-disk load estimate:
74 /// this shrinks as `free_matching` drops tensors the binders replaced.
75 pub fn resident_bytes(&self) -> usize {
76 self.weights.values().map(|t| t.byte_size()).sum()
77 }
78
79 /// Iterator over all weight names.
80 pub fn names(&self) -> impl Iterator<Item = &str> {
81 self.weights.keys().map(|s| s.as_str())
82 }
83
84 /// 2026-10-04: sha256 over the drafter's stored (pre-quantization, as-loaded) bytes: every
85 /// `mtp.*` tensor, plus the embedding and the shared lm_head if this checkpoint has them â
86 /// the `--spec-cost-model measured` drafter key. `None` when this checkpoint has no `mtp.*`
87 /// tensors (no drafter to key). Call before any `free_matching` of these names, or the hash
88 /// silently covers fewer tensors than the drafter actually has.
89 pub fn drafter_weights_sha256(&self, gpu: &dyn GpuBackend) -> Result<Option<String>> {
90 const LM_HEAD_KEYS: [&str; 3] = [
91 "lm_head.weight",
92 "language_model.lm_head.weight",
93 "model.lm_head.weight",
94 ];
95 const EMBED_KEYS: [&str; 3] = [
96 "model.embed_tokens.weight",
97 "language_model.model.embed_tokens.weight",
98 "model.language_model.embed_tokens.weight",
99 ];
100 let mut names: Vec<&str> = self.names().filter(|n| n.starts_with("mtp.")).collect();
101 if names.is_empty() {
102 return Ok(None);
103 }
104 names.extend(
105 LM_HEAD_KEYS
106 .into_iter()
107 .chain(EMBED_KEYS)
108 .filter(|&k| self.contains(k)),
109 );
110 names.sort_unstable();
111 names.dedup();
112 self.hash_named(gpu, &names).map(Some)
113 }
114
115 /// 2026-10-04: sha256 over `names`' device bytes, in the given order, each tensor preceded
116 /// by its name so two tensors cannot alias by concatenation. 64 MiB host chunks, the same
117 /// size `alloc_digest` uses, never one host buffer per tensor.
118 fn hash_named(&self, gpu: &dyn GpuBackend, names: &[&str]) -> Result<String> {
119 use sha2::{Digest, Sha256};
120 const CHUNK: usize = 64 << 20;
121 let mut buf = vec![0u8; CHUNK];
122 let mut h = Sha256::new();
123 for &name in names {
124 let t = self.get(name)?;
125 h.update(name.as_bytes());
126 h.update([0u8]);
127 let bytes = t.byte_size();
128 let mut off = 0usize;
129 while off < bytes {
130 let n = CHUNK.min(bytes - off);
131 gpu.copy_d2h(DevicePtr(t.ptr.0 + off as u64), &mut buf[..n])
132 .map_err(|e| e.context(format!("hashing weight {name}")))?;
133 h.update(&buf[..n]);
134 off += n;
135 }
136 }
137 Ok(h.finalize().iter().map(|b| format!("{b:02x}")).collect())
138 }
139
140 /// Free and forget every tensor whose name matches `pred`. Returns
141 /// `(tensors freed, bytes freed)`.
142 ///
143 /// For loaders that do NOT bind zero-copy from the store's device pointers:
144 /// they upload their own copy, so the original is dead weight the moment the
145 /// binder returns, and on a unified-memory GB10 that duplicate is the
146 /// difference between fitting a KV cache and not.
147 ///
148 /// ðŠĪ The caller owns the "is it dead?" question. A tensor bound zero-copy
149 /// (every routed expert, and the fused per-expert views in
150 /// `weight_loader/step3p7.rs`) is still live in a layer struct â freeing it
151 /// here is a use-after-free with no diagnostic. Match narrowly.
152 ///
153 /// Per-entry free is sound for the same reason `release` gives below: the
154 /// loaders allocate one `gpu.alloc` per tensor, and no loader inserts an
155 /// `.offset()` view of a shared block into this map.
156 pub fn free_matching(
157 &mut self,
158 gpu: &dyn GpuBackend,
159 pred: impl Fn(&str) -> bool,
160 ) -> Result<(usize, usize)> {
161 let doomed: Vec<String> = self.weights.keys().filter(|n| pred(n)).cloned().collect();
162 let (mut count, mut bytes) = (0usize, 0usize);
163 for name in doomed {
164 // `remove` before `free`: the map must never hold a pointer to
165 // memory that is gone, even if the free below fails.
166 let Some(t) = self.weights.remove(&name) else {
167 continue;
168 };
169 bytes += t.byte_size();
170 gpu.free(t.ptr)
171 .map_err(|e| e.context(format!("freeing weight {name}")))?;
172 count += 1;
173 }
174 Ok((count, bytes))
175 }
176
177 /// The owner for buffers a loader derives from these tensors.
178 ///
179 /// `&self` because `ModelWeightLoader::load_layers` takes `&WeightStore`;
180 /// the interior `Mutex` is the whole reason `DerivedStore` exists as a
181 /// type rather than a `Vec` field. See `weights/derived.rs`.
182 pub fn derived(&self) -> &DerivedStore {
183 &self.derived
184 }
185
186 /// Total bytes across all weight tensors on the GPU.
187 pub fn total_bytes(&self) -> usize {
188 self.weights.values().map(|w| w.byte_size()).sum()
189 }
190
191 /// Check if any tensor has FP8 dtype.
192 pub fn has_fp8_weights(&self) -> bool {
193 self.weights
194 .values()
195 .any(|w| matches!(w.dtype, WeightDtype::FP8E4M3))
196 }
197
198 /// 2026-09-28: Every attention layer that ships FP8 KV scales, resolved
199 /// through [`super::resolve_kv_scale_keys`], the resolver the per-layer
200 /// loader uses; an ambiguous or half pair is an error.
201 pub fn kv_scale_census(&self) -> Result<super::KvScaleCensus> {
202 super::kv_scale_census(self.names(), |n| self.contains(n))
203 }
204
205 /// 2026-09-28: The FP8 KV scale keys of the attention layer at `attn_prefix`
206 /// (`None` when it ships none), from [`super::resolve_kv_scale_keys`].
207 pub fn kv_scale_keys(&self, attn_prefix: &str) -> Result<Option<super::KvScaleKeys>> {
208 super::resolve_kv_scale_keys(|n| self.contains(n), attn_prefix)
209 }
210}