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 spark-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 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 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 fp8_kv_scale_count(&self) -> usize

Number of per-layer FP8 KV-cache scale tensors (*.k_scale) the checkpoint ships. >0 means the model carries calibrated KV scales, so FP8 KV needs no online calibration; 0 means the scales default to 1.0 (which clips BF16 into E4M3 range), so online calibration or a non-FP8 KV dtype is required. Used to log the right guidance at serve time.

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

Human name, for the teardown report and for attributing a failure.

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

Release everything this owns. Must be idempotent: the host calls it, and a Drop backstop may call it again.