Struct WeightStore 

pub struct WeightStore { /* private fields */ }

All model weights loaded onto the GPU, keyed by HuggingFace name.

Implementations

impl WeightStore

pub fn empty() -> Self

Create an empty weight store (for testing).

pub fn defer(&mut self, name: String, t: DeferredTensor)

Record a tensor that was skipped at load, with its on-disk location.

pub fn deferred(&self, name: &str) -> Option<&DeferredTensor>

Look up a deferred (not-uploaded) tensor’s on-disk location.

pub fn deferred_sorted(&self) -> Vec<(&String, &DeferredTensor)>

Every deferred tensor, name-sorted (NUMERIC on a trailing index, so embedders.10 sorts after embedders.2 — a lexicographic sort here silently mis-maps the n-gram tables, which cost a real debugging session the first time).

pub fn from_map(weights: HashMap<String, WeightTensor>) -> Self

Wrap a pre-built map. Used by alternate loaders (e.g. fast_weights::FastSafetensorsLoader, and the RDMA weight loader in metrale-storage, which lives in a different crate and so needs this pub).

pub fn get(&self, name: &str) -> Result<&WeightTensor>

Get a weight tensor by name. Fails fast if not found.

pub fn contains(&self, name: &str) -> bool

Check if a weight exists.

pub fn len(&self) -> usize

Number of loaded weights.

pub fn is_empty(&self) -> bool

True if no weights are loaded.

pub fn resident_bytes(&self) -> usize

Device bytes the store still holds. Not the on-disk load estimate: this shrinks as free_matching drops tensors the binders replaced.

pub fn names(&self) -> impl Iterator<Item = &str>

Iterator over all weight names.

pub fn drafter_weights_sha256( &self, gpu: &dyn GpuBackend, ) -> Result<Option<String>>

2026-10-04: sha256 over the drafter’s stored (pre-quantization, as-loaded) bytes: every mtp.* tensor, plus the embedding and the shared lm_head if this checkpoint has them — the --spec-cost-model measured drafter key. None when this checkpoint has no mtp.* tensors (no drafter to key). Call before any free_matching of these names, or the hash silently covers fewer tensors than the drafter actually has.

pub fn free_matching( &mut self, gpu: &dyn GpuBackend, pred: impl Fn(&str) -> bool, ) -> Result<(usize, usize)>

Free and forget every tensor whose name matches pred. Returns (tensors freed, bytes freed).

For loaders that do NOT bind zero-copy from the store’s device pointers: they upload their own copy, so the original is dead weight the moment the binder returns, and on a unified-memory GB10 that duplicate is the difference between fitting a KV cache and not.

🪤 The caller owns the “is it dead?” question. A tensor bound zero-copy (every routed expert, and the fused per-expert views in weight_loader/step3p7.rs) is still live in a layer struct — freeing it here is a use-after-free with no diagnostic. Match narrowly.

Per-entry free is sound for the same reason release gives below: the loaders allocate one gpu.alloc per tensor, and no loader inserts an .offset() view of a shared block into this map.

pub fn derived(&self) -> &DerivedStore

The owner for buffers a loader derives from these tensors.

&self because ModelWeightLoader::load_layers takes &WeightStore; the interior Mutex is the whole reason DerivedStore exists as a type rather than a Vec field. See weights/derived.rs.

pub fn total_bytes(&self) -> usize

Total bytes across all weight tensors on the GPU.

pub fn has_fp8_weights(&self) -> bool

Check if any tensor has FP8 dtype.

pub fn kv_scale_census(&self) -> Result<KvScaleCensus>

2026-09-28: Every attention layer that ships FP8 KV scales, resolved through super::resolve_kv_scale_keys, the resolver the per-layer loader uses; an ambiguous or half pair is an error.

pub fn kv_scale_keys(&self, attn_prefix: &str) -> Result<Option<KvScaleKeys>>

2026-09-28: The FP8 KV scale keys of the attention layer at attn_prefix (None when it ships none), from super::resolve_kv_scale_keys.

impl WeightStore

pub fn prepartitioned_tp(&self) -> Option<(usize, usize)>

Set only after the rank-aware loader has validated and sliced all tensors.

Trait Implementations

impl ModelResource<dyn GpuBackend> for WeightStore

Release every weight tensor.

Safe to free per-entry because the loaders allocate per-tensor: the fast path calls gpu.alloc(meta.len) once per tensor before inserting it (fast_weights/mod.rs:360-388), and no loader inserts an .offset() view of a shared block into this map. (Fused per-expert views DO exist — see weight_loader/step3p7.rs:93 — but they live in the layer structs that own the fused allocation, not here, so this cannot double-free them.)

fn label(&self) -> &'static str

2026-09-25: Human name, used to attribute a failure in Teardown::release_all’s error.

fn release(&mut self, gpu: &dyn GpuBackend) -> Result<()>

2026-09-25: Release everything this owns. Must be idempotent.