// SPDX-License-Identifier: MIT OR Apache-2.0
//! 2026-09-26: Weight-store loading: the checkpoint, the DFlash drafter and
//! the startup LoRA adapters.
//!
//! Owner: server startup (`met serve`).
//! Invariants: none beyond the types.
use std::path::Path;
use anyhow::{Context, Result};
use metrale_config::ModelConfig;
use crate::cli;
#[cfg(test)]
#[path = "weights_allowlist_tests.rs"]
mod weights_allowlist_tests;
/// 2026-09-26: The load pre-flight's peak-memory multiplier for this model;
/// `None` leaves the loader's own (1.3x, or 1.5x with FP8 tensors).
pub(crate) fn quant_multiplier(config: &ModelConfig) -> Option<f64> {
if config.model_type == "glm5_next" {
// 2026-09-26: 1.10 is a headroom allowance, not a measured value.
Some(1.10)
} else if config.model_type == "minimax_m2" || config.model_type == "step3p7" {
Some(1.02)
} else if config
.quantization_config
.as_ref()
.is_some_and(|qc| qc.quant_method == "fp8")
{
Some(1.05)
} else {
None
}
}
pub(crate) fn load_weight_store(
args: &cli::ServeArgs,
config: &ModelConfig,
model_dir: &Path,
gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
ep_rank: usize,
ep_size: usize,
oom_reserve_bytes: usize,
) -> Result<metrale_model_weights::weights::WeightStore> {
use metrale_model_weights::weights::WeightLoader;
if matches!(
config.model_type.as_str(),
"kimi_k3" | "kimi_linear" | "Kimi-K3"
) {
anyhow::ensure!(
ep_size == 1,
"K3 rank-aware loading supports TP only; expert parallelism is not implemented"
);
anyhow::ensure!(
metrale_model_weights::weights::find_gguf(model_dir).is_none(),
"K3 rank-aware loading requires safetensors"
);
tracing::info!(
rank = config.tp_rank,
world = config.tp_world_size,
"K3 uses rank-aware safetensors loading before GPU allocation"
);
return metrale_model_weights::weights::K3SafetensorsLoader::new(config.clone())?
.load(model_dir, gpu, oom_reserve_bytes)
.context("Failed to load rank-local K3 weights");
}
let mult = quant_multiplier(config);
// 2026-09-26: Any `.gguf` file in the directory selects `GgufLoader`.
if metrale_model_weights::weights::find_gguf(model_dir).is_some() {
tracing::info!("Detected GGUF weights; using GgufLoader (GPU dequant → BF16)");
let mut loader = if ep_size > 1 {
metrale_model_weights::weights::GgufLoader::with_ep(
ep_rank,
ep_size,
config.num_experts,
)
} else {
metrale_model_weights::weights::GgufLoader::new()
};
loader.peak_memory_multiplier = mult;
let store = loader
.load(model_dir, gpu, oom_reserve_bytes)
.context("Failed to load model weights (GGUF loader)")?;
tracing::info!("Loaded {} weight tensors (GGUF)", store.len());
return Ok(store);
}
let use_fast_load =
!args.no_fast_load && std::env::var("METRALE_FAST_LOAD").ok().as_deref() != Some("0");
let store = if use_fast_load {
#[cfg(unix)]
{
tracing::info!("Using fast weight loader (O_DIRECT + pipelined read/copy)");
let mut loader = if ep_size > 1 {
metrale_model_weights::fast_weights::FastSafetensorsLoader::with_ep(
ep_rank,
ep_size,
config.num_experts,
)
} else {
metrale_model_weights::fast_weights::FastSafetensorsLoader::new()
};
loader.peak_memory_multiplier = mult;
loader.skip_activation_scales = skip_activation_scales(config);
loader.skip_mtp = skip_mtp(config);
// 2026-09-26: A loader that binds no vision encoder skips the
// tower here. `factory::build` frees an unbound tower too, but
// only after the reserve preflight has measured free memory.
loader.skip_vision = !binds_vision(config);
if loader.skip_vision {
tracing::info!(
"Vision tower: not loaded — the weight loader for model_type '{}' is a \
text-only port and binds no vision encoder.",
config.model_type,
);
}
loader.defer = defer_hook(config);
loader.defer_charge = defer_charge(config);
loader.prefetch_shards = args.fast_load_prefetch_shards
|| std::env::var("METRALE_FAST_LOAD_PREFETCH_SHARDS")
.ok()
.is_some_and(|v| v == "1" || v.eq_ignore_ascii_case("true"));
if loader.prefetch_shards {
tracing::info!("Fast weight loader shard prefetch/readahead enabled");
}
loader
.load(model_dir, gpu, oom_reserve_bytes)
.context("Failed to load model weights (fast loader)")?
}
#[cfg(not(unix))]
{
anyhow::bail!("--fast-load requires a Unix host (needs O_DIRECT / posix_fadvise)");
}
} else {
let mut loader = if ep_size > 1 {
metrale_model_weights::weights::SafetensorsLoader::with_ep(
ep_rank,
ep_size,
config.num_experts,
)
} else {
metrale_model_weights::weights::SafetensorsLoader::new()
};
loader.peak_memory_multiplier = mult;
loader.skip_activation_scales = skip_activation_scales(config);
loader.skip_mtp = skip_mtp(config);
loader.defer = defer_hook(config);
loader.defer_charge = defer_charge(config);
loader
.load(model_dir, gpu, oom_reserve_bytes)
.context("Failed to load model weights")?
};
tracing::info!("Loaded {} weight tensors", store.len());
Ok(store)
}
pub(crate) fn load_dflash_drafter(
args: &cli::ServeArgs,
ptx_set: &metrale_kernels::TargetPtxSet,
gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
) -> Result<
Option<(
metrale_model_weights::weights::WeightStore,
metrale_model_arch::weight_loader::DflashConfig,
)>,
> {
use metrale_model_weights::weights::WeightLoader;
if !args.dflash {
return Ok(None);
}
let drafter_dir = super::dflash_gamma::dflash_drafter_dir(
args,
super::dflash_gamma::model_default_drafter(ptx_set),
)?;
let drafter_config = super::dflash_gamma::read_dflash_config(&drafter_dir)?;
// 2026-09-26: The drafter allocates outside the KV budget, so its whole
// footprint is estimated from the checkpoint before anything is allocated:
// safetensors bytes on disk, the drafter KV cache, `fused_kv`, the
// DFlash2 selector's host copies, the weight copies (in FP8, half the store
// plus twice `vocab x hidden`; 2026-10-10: scaled to the
// `--draft-weight-quantization` tier, none on `declared`), and 300 MiB of
// scratch.
let store_bytes: u64 = std::fs::read_dir(&drafter_dir)
.ok()
.into_iter()
.flatten()
.flatten()
.filter(|e| e.path().extension().is_some_and(|x| x == "safetensors"))
// 2026-09-26: `std::fs::metadata` follows symlinks and
// `DirEntry::metadata` does not; a Hugging Face snapshot directory
// holds symlinks into `blobs/`.
.filter_map(|e| std::fs::metadata(e.path()).ok().map(|m| m.len()))
.sum();
let c = &drafter_config;
let kv_dim = c.num_key_value_heads * c.head_dim;
let drafter_kv =
(args.max_seq_len as u64) * (c.num_hidden_layers as u64) * 2 * (kv_dim as u64) * 2;
let fused_kv = (c.num_hidden_layers as u64) * 2 * (kv_dim as u64) * (c.hidden_size as u64) * 2;
let selector_host = c
.dflash_config
.as_ref()
.map(|d| d.selector_rank)
.filter(|r| *r > 0)
.map(|r| {
2 * (c.vocab_size as u64) * (r as u64) * 2 + (r as u64) * (c.hidden_size as u64) * 2
})
.unwrap_or(0);
let lm_head = (c.vocab_size as u64) * (c.hidden_size as u64);
let weight_copies = args
.dflash_args
.draft_weight_quantization
.copy_bytes_from_fp8(store_bytes / 2 + 2 * lm_head);
let scratch_est: u64 = 300 << 20;
let estimate =
store_bytes + drafter_kv + fused_kv + selector_host + weight_copies + scratch_est;
let free = gpu.free_memory().unwrap_or(0) as u64;
let total = gpu.total_memory().unwrap_or(0) as u64;
// 2026-09-26: `total x (1 - util)`, the memory outside the utilization
// budget, must still be free after the drafter loads.
let headroom = (total as f64 * (1.0 - args.gpu_memory_utilization)) as u64;
if free < estimate + headroom {
anyhow::bail!(
"DFlash drafter would over-commit unified memory: estimated footprint {:.2} GB (weights {:.2} + drafter-KV {:.2} + fused_kv {:.2} + selector-host {:.2} + weight-copies {:.2} + scratch {:.2}) but only {:.2} GB free with {:.2} GB headroom pledged by --gpu-memory-utilization {:.2}. On GB10 this would SWAP the host, not error. Lower --max-seq-len, lower --gpu-memory-utilization pressure elsewhere, or a --draft-weight-quantization tier with smaller copies (nvfp4, or declared for none).",
estimate as f64 / 1e9,
store_bytes as f64 / 1e9,
drafter_kv as f64 / 1e9,
fused_kv as f64 / 1e9,
selector_host as f64 / 1e9,
weight_copies as f64 / 1e9,
scratch_est as f64 / 1e9,
free as f64 / 1e9,
headroom as f64 / 1e9,
args.gpu_memory_utilization,
);
}
tracing::info!(
"DFlash footprint pre-flight: estimate {:.2} GB (weights {:.2}, drafter-KV {:.2}, weight-copies {:.2}, selector-host {:.2}) vs {:.2} GB free, {:.2} GB headroom — OK",
estimate as f64 / 1e9,
store_bytes as f64 / 1e9,
drafter_kv as f64 / 1e9,
weight_copies as f64 / 1e9,
selector_host as f64 / 1e9,
free as f64 / 1e9,
headroom as f64 / 1e9,
);
let mut loader = metrale_model_weights::weights::SafetensorsLoader::new();
loader.peak_memory_multiplier = None;
let drafter_store = loader
.load(&drafter_dir, gpu, 0)
.context("Failed to load DFlash drafter weights")?;
tracing::info!(
"DFlash drafter store: {} tensors, {} bytes",
drafter_store.len(),
drafter_store.total_bytes()
);
Ok(Some((drafter_store, drafter_config)))
}
/// 2026-09-26: Whether the target checkpoint's LM head (`lm_head.weight`,
/// with or without a `language_model.` or `model.` prefix) is FP8 E4M3. Reads
/// only the safetensors header of the shard holding it; any failure gives
/// `false`.
fn target_ships_native_fp8_lm_head(args: &cli::ServeArgs) -> bool {
fn inner(args: &cli::ServeArgs) -> Option<bool> {
let dir = if let Some(p) = &args.model_from_path {
p.clone()
} else {
crate::model_resolver::resolve_model_dir(
args.model.as_deref()?,
args.cache_dir.as_deref(),
)
.ok()?
};
const KEYS: [&str; 3] = [
"lm_head.weight",
"language_model.lm_head.weight",
"model.lm_head.weight",
];
let shard = if let Ok(idx) = std::fs::read(dir.join("model.safetensors.index.json")) {
let idx: serde_json::Value = serde_json::from_slice(&idx).ok()?;
let map = idx.get("weight_map")?;
KEYS.iter()
.find_map(|k| map.get(*k).and_then(|v| v.as_str()))
.map(|s| dir.join(s))?
} else {
dir.join("model.safetensors")
};
use std::io::Read as _;
let mut f = std::fs::File::open(shard).ok()?;
let mut len8 = [0u8; 8];
f.read_exact(&mut len8).ok()?;
let hlen = u64::from_le_bytes(len8);
if hlen > 64 << 20 {
return None;
}
let mut hdr = vec![0u8; hlen as usize];
f.read_exact(&mut hdr).ok()?;
let hdr: serde_json::Value = serde_json::from_slice(&hdr).ok()?;
let dtype = KEYS
.iter()
.find_map(|k| hdr.get(*k))
.and_then(|t| t.get("dtype"))
.and_then(|d| d.as_str())?;
Some(dtype == "F8_E4M3")
}
inner(args).unwrap_or(false)
}
/// 2026-09-26: A LoRA adapter loaded at startup: its own `WeightStore` and its
/// parsed PEFT config, one per `--lora-adapter NAME=PATH`.
pub(crate) struct LoraAdapterState {
pub name: String,
pub peft_config: metrale_config::PeftAdapterConfig,
pub store: metrale_model_weights::weights::WeightStore,
}
/// 2026-09-26: Resolve and load every `--lora-adapter` into its own device
/// `WeightStore`, in command-line order; empty when none is given. Errors on
/// more adapters than `--max-loras` before loading any, and on a repeated
/// name or a rank above `--max-lora-rank` (64 when unset) when it is reached.
pub(crate) fn load_lora_adapters(
args: &cli::ServeArgs,
gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
) -> Result<Vec<LoraAdapterState>> {
if args.lora_adapter.is_empty() {
return Ok(Vec::new());
}
if args.lora_adapter.len() > args.max_loras {
anyhow::bail!(
"--lora-adapter given {} times but --max-loras={} (pool has {} slots); \
raise --max-loras or stage the extras on an $METRALE_LORA_PEER",
args.lora_adapter.len(),
args.max_loras,
args.max_loras,
);
}
let mut states: Vec<LoraAdapterState> = Vec::with_capacity(args.lora_adapter.len());
for (name, spec) in &args.lora_adapter {
if states.iter().any(|s| &s.name == name) {
anyhow::bail!("--lora-adapter name '{name}' given twice (names must be unique)");
}
tracing::info!("LoRA: resolving adapter '{name}' from '{spec}'");
let adapter_dir =
crate::model_resolver::resolve_adapter_dir(spec, args.cache_dir.as_deref())
.context("Failed to resolve LoRA adapter")?;
let cfg_path = adapter_dir.join("adapter_config.json");
let raw = std::fs::read_to_string(&cfg_path)
.with_context(|| format!("Failed to read {}", cfg_path.display()))?;
// 2026-09-26: Scaling is per adapter: `lora_alpha / r`, or
// `lora_alpha / sqrt(r)` with `use_rslora` (`PeftAdapterConfig::scaling`).
let peft_config = metrale_config::parse_peft_adapter_config(&raw)
.with_context(|| format!("Failed to parse {}", cfg_path.display()))?;
let rank_ceiling = args.max_lora_rank.unwrap_or(64);
if peft_config.r > rank_ceiling {
anyhow::bail!(
"LoRA adapter '{}' has r={} > --max-lora-rank {} — raise the flag \
(slot pool is rank-padded to it) or use a smaller adapter",
name,
peft_config.r,
rank_ceiling,
);
}
let store =
metrale_model_weights::weights::adapter::load_adapter_safetensors(&adapter_dir, gpu, 0)
.context("Failed to load LoRA adapter weights")?;
tracing::info!(
"LoRA adapter '{}': {} tensors, {} bytes loaded; r={}, alpha={}, \
use_rslora={}, scaling={:.6}, target_modules={:?}",
name,
store.len(),
store.total_bytes(),
peft_config.r,
peft_config.lora_alpha,
peft_config.use_rslora,
peft_config.scaling(),
peft_config.target_modules,
);
states.push(LoraAdapterState {
name: name.clone(),
peft_config,
store,
});
}
Ok(states)
}
/// 2026-09-26: Whether the loaders skip the `*.input_scale` activation scales
/// for this model: `qwen4_exp`.
///
/// Both loaders skip exactly the names ending in `.input_scale`
/// (`SafetensorsLoader`, `fast_weights::skip`), so `.weight_scale` and
/// `.weight_scale_2` still load. The skip is listed per model because some
/// loaders read `input_scale`: `step3p7` on its own path, and
/// `weight_map/model_a.rs` whenever the tensor is present. 2026-10-08:
/// `glm5_next` reads them for its declared W4A4 path; its defer hook keeps
/// them off the device instead (`glm5_next_load/act_scale.rs`).
fn skip_activation_scales(config: &ModelConfig) -> bool {
matches!(config.model_type.as_str(), "qwen4_exp")
}
/// 2026-09-26: Whether `mtp.*` is left unloaded: `qwen4_exp`, whose
/// `Qwen4ExpWeightLoader::load_mtp_weights` returns `None`.
fn skip_mtp(config: &ModelConfig) -> bool {
matches!(config.model_type.as_str(), "qwen4_exp")
}
/// 2026-09-26: Whether the model's weight loader binds a vision encoder;
/// `true` when the model type does not resolve to a loader.
fn binds_vision(config: &metrale_config::ModelConfig) -> bool {
metrale_model_engine::factory::loader_for_config(config)
.map(|l| l.binds_vision_encoder())
.unwrap_or(true)
}
/// 2026-09-26: The tensors the model's weight loader reads from disk itself
/// at bind time (`defer_predicate`); they are recorded with their location,
/// not uploaded. `None` when the model type does not resolve or the loader
/// defers nothing, the trait default. Unlike `skip_activation_scales` and
/// `skip_mtp`, the loader declares this itself, so no model list is kept here.
fn defer_hook(config: &ModelConfig) -> Option<metrale_model_weights::weights::DeferHook> {
let hook = metrale_model_engine::factory::loader_for_config(config)
.ok()?
.defer_predicate(config)?;
tracing::info!(
"Weight loader for model_type '{}' defers part of the checkpoint: those tensors are \
recorded with their on-disk location and read by the model's own loader, never \
uploaded.",
config.model_type,
);
Some(hook)
}
/// 2026-10-17: The device bytes the model's weight loader allocates for the tensors it defers
/// (`defer_charge`), so the pre-flight counts them. `None` when the model type does not resolve
/// or the loader states none, the trait default.
fn defer_charge(config: &ModelConfig) -> Option<metrale_model_weights::weights::DeferCharge> {
metrale_model_engine::factory::loader_for_config(config)
.ok()?
.defer_charge(config)
}