From b9786d883512ebb45ed3753a6a97d8d5954eff45 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 17:18:17 +0200 Subject: [PATCH 01/17] feat(device-mesh): add pure manifest and llama store pilot --- crates/hipfire-arch-llama/Cargo.toml | 1 + crates/hipfire-arch-llama/src/arch.rs | 154 ++- crates/hipfire-arch-llama/src/arch_model.rs | 19 +- crates/hipfire-arch-llama/src/carrier.rs | 134 ++- crates/hipfire-runtime/src/lib.rs | 2 + crates/hipfire-runtime/src/model_load.rs | 83 ++ crates/hipfire-runtime/src/weight_manifest.rs | 1004 +++++++++++++++++ crates/hipfire-runtime/src/weight_store.rs | 724 ++++++++++++ 8 files changed, 2068 insertions(+), 53 deletions(-) create mode 100644 crates/hipfire-runtime/src/weight_manifest.rs create mode 100644 crates/hipfire-runtime/src/weight_store.rs diff --git a/crates/hipfire-arch-llama/Cargo.toml b/crates/hipfire-arch-llama/Cargo.toml index f063bdf25..265db1f2a 100644 --- a/crates/hipfire-arch-llama/Cargo.toml +++ b/crates/hipfire-arch-llama/Cargo.toml @@ -11,6 +11,7 @@ lab = [] [dependencies] hipfire-runtime = { path = "../hipfire-runtime" } +hipfire-hardware = { path = "../hipfire-hardware" } hipfire-dispatch = { path = "../hipfire-dispatch", features = ["from-hip-error"] } hip-bridge = { path = "../hip-bridge" } rdna-compute = { path = "../rdna-compute" } diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index ec7322dbf..920601367 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -19,7 +19,10 @@ use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::{self, HfqFile}; use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; -use rdna_compute::Gpu; +use hipfire_runtime::weight_manifest::{ + FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, +}; +use rdna_compute::{DType, Gpu}; use hipfire_dispatch::context::DispatchCtx; use hipfire_dispatch::pipeline::{execute_steps, GemvInput, Step}; @@ -43,12 +46,8 @@ impl Architecture for Llama { fn arch_id() -> u32 { // `arch_id = 0` is the canonical LLaMA-family marker. The - // actual arch_id loaded at runtime is on `HfqFile::arch_id` - // and is either 0 (LLaMA / Mistral) or 1 (plain Qwen3 / - // Qwen2); both share this trait impl. The qwen3-norm flag - // is read off the HFQ metadata inside `config_from_hfq`, - // so the bring-up triple does not need a separate marker - // type per arch_id. + // actual id loaded at runtime is on `HfqFile::arch_id` and may + // differ for plain Qwen3/Qwen2; config parsing resolves that. 0 } @@ -57,13 +56,6 @@ impl Architecture for Llama { } fn config_from_hfq(hfq: &HfqFile) -> Result { - // `hfq::config_from_hfq` is the LLaMA-family HFQ metadata - // parser — emits a `LlamaConfig` with the appropriate - // `ModelArch` (Llama vs Qwen3) tag. It lives in the runtime - // crate because the qwen35 hybrid path's pflash drafter also - // calls it via `hfq::config_from_hfq` for its "Plain" - // variant. See arch-llama/src/lib.rs for the colocation - // rationale. hfq::config_from_hfq(hfq) } @@ -72,27 +64,141 @@ impl Architecture for Llama { cfg: &Self::Config, gpu: &mut Gpu, ) -> Result { - // `hfq::load_weights_hfq` is the LLaMA-family HFQ tensor - // loader. Same colocation reasoning as `config_from_hfq`. hfq::load_weights_hfq(hfq, cfg, gpu) .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}")) } fn new_state(gpu: &mut Gpu, cfg: &Self::Config) -> Result { - // The LLaMA-arch "state" is the `ForwardScratch` — persistent - // GPU scratch buffers reused across decode steps. There is no - // separate recurrent state (LLaMA is full-attention only). ForwardScratch::new(gpu, cfg) .map_err(|e| format!("llama: ForwardScratch::new failed: {e:?}")) } // Optional overrides: defaults from `hipfire_runtime::arch` already // assume Qwen3.5 family conventions. LLaMA / Mistral / Qwen3 don't - // emit `` blocks, but PR 11 keeps the override surface - // empty here on purpose — the daemon's existing per-`arch_id` - // policy choices stay unchanged. Future PRs that consolidate - // policy through the trait can populate these (LLaMA: no - // strip_think, no Qwen-specific blocked tokens). + // emit `` blocks, but the existing policy choices stay unchanged. +} + +impl Llama { + /// Pure dense LLaMA-family weight declaration. Source names remain + /// logical; carriers translate them to HFQ/safetensors namespaces. + pub fn weight_manifest(cfg: &LlamaConfig) -> Vec { + use ShardPolicy::*; + let (dim, hidden, head_dim) = (cfg.dim, cfg.hidden_dim, cfg.head_dim); + let (heads, kv_heads) = (cfg.n_heads, cfg.n_kv_heads); + let mut manifest = Vec::with_capacity(cfg.n_layers * 11 + 3); + manifest.push(WeightEntry::model( + "token_embd", + vec![cfg.vocab_size, dim], + DType::F16, + Pin(PinTarget::Embed), + )); + for layer in 0..cfg.n_layers { + manifest.push(WeightEntry::layer( + "wq", + layer, + vec![heads * head_dim, dim], + DType::F16, + FusedQkv { + q_heads: heads, + kv_heads, + head_dim, + layout: FusedQkvLayout::Qkv, + }, + )); + manifest.push(WeightEntry::layer( + "wk", + layer, + vec![kv_heads * head_dim, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "wv", + layer, + vec![kv_heads * head_dim, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "wo", + layer, + vec![dim, heads * head_dim], + DType::F16, + RowShard { axis: 1 }, + )); + manifest.push(WeightEntry::layer( + "ffn_gate", + layer, + vec![hidden, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "ffn_up", + layer, + vec![hidden, dim], + DType::F16, + ColumnShard { axis: 0 }, + )); + manifest.push(WeightEntry::layer( + "ffn_down", + layer, + vec![dim, hidden], + DType::F16, + RowShard { axis: 1 }, + )); + manifest.push(WeightEntry::layer( + "attn_norm", + layer, + vec![dim], + DType::F32, + Replicate, + )); + manifest.push(WeightEntry::layer( + "ffn_norm", + layer, + vec![dim], + DType::F32, + Replicate, + )); + if cfg.has_qk_norm { + manifest.push(WeightEntry::layer( + "q_norm", + layer, + vec![head_dim], + DType::F32, + Replicate, + )); + manifest.push(WeightEntry::layer( + "k_norm", + layer, + vec![head_dim], + DType::F32, + Replicate, + )); + } + } + manifest.push(WeightEntry::model( + "output_norm", + vec![dim], + DType::F32, + Replicate, + )); + manifest.push(WeightEntry::model( + "lm_head", + vec![cfg.vocab_size, dim], + DType::F16, + Pin(PinTarget::Output), + )); + manifest + } + + /// Pure state declaration for the full-attention LLaMA family. + pub fn state_manifest(cfg: &LlamaConfig) -> Vec { + (0..cfg.n_layers) + .map(|layer| StateEntry::new(StateKind::Kv { quant: String::new() }, layer)) + .collect() + } } // ── Dispatch integration ───────────────────────────────────────── diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 008b88eaf..7e7d3d7ce 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -40,19 +40,22 @@ impl ArchModel for LlamaBundle { weights, scratch, kv, + manifest_plan: _, + weight_store, + mesh, dflash_extract_layers: _, dspark_weights: _, dspark_assets: _, } = *self; - // Mirror unload_model ModelState::Llama arm exactly (lib.rs:3041): - // b.scratch.free_gpu(gpu); - // b.weights.free_gpu(gpu); - // note(b.kv.free_gpu(gpu)…) - // Ordering matters: scratch → weights → kv. dspark sidecars (when - // present) are reclaimed via the speculator/spec scratch paths, not - // here — matching the current unload_model which also does not handle - // them in this arm. + // Mirror the existing unload ordering: scratch → store/weights → kv. + // A committed store is only released here, through the ArchModel owner; + // no store destructor or independent carrier free path exists. scratch.free_gpu(gpu); + if let Some(store) = weight_store { + if let Err((_, error)) = store.release_on_owner(&mesh, gpu) { + eprintln!("llama: refusing weight-store release: {error}"); + } + } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); } diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 736eeb867..c41b80298 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -4,17 +4,30 @@ use crate::dspark_body::Qwen3DrafterAssets; use crate::Llama; +use hipfire_hardware::DeviceMesh; use hipfire_runtime::arch::Architecture; use hipfire_runtime::dspark_core::DsparkWeights; -use hipfire_runtime::llama::{ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; +use hipfire_runtime::llama::{ + ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LlamaConfig, LlamaWeights, +}; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan}; +use hipfire_runtime::weight_store::WeightStore; pub struct LlamaBundle { pub config: LlamaConfig, pub weights: LlamaWeights, pub scratch: ForwardScratch, pub kv: KvCache, + /// Pure declaration/placement plan captured at load time. The plan has no + /// GPU handles and is immutable after publication. + pub manifest_plan: ManifestPlan, + /// A pilot store is attached only after its handles are assembled under + /// this bundle. It is crate-visible so callers cannot create an independent + /// unload owner; `ArchModel::free_gpu` is the sole release path. + pub(crate) weight_store: Option, + pub(crate) mesh: DeviceMesh, /// Decoder-layer indices whose residual hidden states a hidden-conditioned /// drafter (DFlash / EAGLE) wants captured, ascending order. Empty = no /// capture (the `SpecTarget::dflash_extract_layers` default of `None`). The @@ -25,26 +38,34 @@ pub struct LlamaBundle { /// was found or speculation was disabled. Task-10 wires the speculator build. pub dspark_weights: Option, /// Loaded DSpark drafter body assets (5-layer dense-GQA transformer + - /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. + /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } /// Build the LLaMA GPU bundle from an HFQ or safetensors-directory source. /// -/// Verbatim relocation of the carrier's `(config, weights, kv, scratch)` -/// seam: HFQ via `Architecture` trait, Dir via ParoQuant loaders. Error -/// strings are byte-identical to the prior inline carrier block. +/// The source/config path remains architecture-owned. Once it resolves, the +/// carrier publishes a pure Single manifest plan. Every fallible GPU stage +/// explicitly releases earlier allocations before returning an error; no +/// implicit GPU-buffer destructor is introduced. pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { let (config, weights, kv, scratch) = match src { ModelSource::Hfq(mut hfq) => { - let config = ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; + let config = + ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; let weights = ::load_weights(&mut hfq, &config, ctx.gpu)?; hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - // Size scratch (flash-attention partials) for the runtime KV cap so the - // asym/flash attends, which index partials by ceil(physical_cap/128), don't - // overflow it (the trait `new_state` only knows the model's declared max). - let scratch = ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) - .map_err(|e| format!("llama: ForwardScratch::new_with_max_seq failed: {e:?}"))?; + // The plain LLaMA path has no independent cap resolver. PR #661's + // physical-cap behavior is owned by the existing upstream KV plan. + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + )); + } + }; let dims = KvDims { layers: KvLayers::Flat(config.n_layers), n_kv_heads: config.n_kv_heads, @@ -52,7 +73,7 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result::from_mode( + let kv = match ::from_mode( hipfire_runtime::kv_mode::resolve( ctx.kv_mode_override.unwrap_or(""), &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, @@ -61,8 +82,16 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result::from_mode failed: {e}"))?; + ) { + Ok(kv) => kv, + Err(error) => { + scratch.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ::from_mode failed: {error}" + )); + } + }; (config, weights, kv, scratch) } ModelSource::Dir(source) => { @@ -74,7 +103,6 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Result Result::from_mode( + let kv = match ::from_mode( rr.mode, KvTarget::Single(ctx.gpu), &dims, - ) - .map_err(|e| format!("KvCache: {e}"))?; - let scratch = ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) - .map_err(|e| format!("ForwardScratch::new_with_max_seq: {e:?}"))?; + ) { + Ok(kv) => kv, + Err(error) => { + weights.free_gpu(ctx.gpu); + return Err(format!("KvCache: {error}")); + } + }; + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let _ = kv.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "ForwardScratch::new_with_max_seq: {error:?}" + )); + } + }; (config, weights, kv, scratch) } }; + + // Pure plan publication happens after source/config resolution and before + // the bundle becomes visible to the loader. It performs no GPU or file IO. + let mesh = DeviceMesh::single(); + let manifest = Llama::weight_manifest(&config); + let state = Llama::state_manifest(&config); + let manifest_plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) + .map_err(|e| format!("llama: manifest planning failed: {e}"))?; + Ok(LlamaBundle { config, weights, scratch, kv, + manifest_plan, + weight_store: None, + mesh, dflash_extract_layers: Vec::new(), dspark_weights: None, dspark_assets: None, @@ -121,6 +177,42 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Result<(), (WeightStore, String)> { + if self.weight_store.is_some() { + return Err((store, "llama: weight store already attached".into())); + } + let Some(origin) = store.origin() else { + return Err((store, "llama: weight store has no origin".into())); + }; + if origin.mesh_epoch() != self.mesh.epoch() { + return Err(( + store, + format!( + "llama: weight store origin epoch {:?} does not match bundle epoch {:?}", + origin.mesh_epoch(), + self.mesh.epoch() + ), + )); + } + self.weight_store = Some(store); + Ok(()) + } + + /// The immutable mesh identity used by this bundle's manifest plan. + /// Callers that run the Single pilot must pass this exact mesh to + /// `fulfill_manifest`; constructing a fresh `DeviceMesh::single()` would + /// intentionally fail the origin check. + pub fn manifest_mesh(&self) -> &DeviceMesh { + &self.mesh + } + /// Set the decoder-layer indices whose residual hidden states the /// hidden-conditioned drafter wants captured (ascending order). The /// speculator calls this with `dflash::DflashConfig::target_layer_ids`. diff --git a/crates/hipfire-runtime/src/lib.rs b/crates/hipfire-runtime/src/lib.rs index ebf81aff3..b8c833190 100644 --- a/crates/hipfire-runtime/src/lib.rs +++ b/crates/hipfire-runtime/src/lib.rs @@ -74,5 +74,7 @@ pub mod tokenizer; pub mod calibration; pub mod tool_call; pub mod weight_backend; +pub mod weight_manifest; +pub mod weight_store; pub use crate::arch::{maybe_screen_mmq, screen_weight_tensor, MmqScreenable}; diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index dadc63f8b..46fcad664 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -6,6 +6,7 @@ //! per-tensor dequant), which `WeightSource::read_layer` calls internally. use crate::llama::{EmbeddingFormat, WeightTensor}; +use hipfire_hardware::{DeviceMesh, DimKind}; use hip_bridge::HipResult; use hipfire_hardware::Gpus; use rdna_compute::{Gpu, GpuTensor}; @@ -30,6 +31,58 @@ impl Layout { layer_to_device: (0..n_layers).map(|i| g.device_for_layer(i)).collect(), } } + + /// Build the canonical stage/rank-0 view from an admitted mesh. The + /// manifest planner owns the full stage grid; this legacy loader view + /// selects rank zero for each layer so existing orchestrators continue to + /// have one deterministic device index until their typed mesh path lands. + pub fn from_mesh(mesh: &DeviceMesh, n_layers: usize) -> Self { + let mut output_coord = mesh.coord_of(0); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + output_coord[index] = mesh.size_of(DimKind::Pp).saturating_sub(1); + } + let layer_to_device = (0..n_layers) + .map(|layer| { + let mut coord = mesh.coord_of(0); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + coord[index] = mesh.stage_for_layer(layer, n_layers); + } + mesh.device_of(&coord) + }) + .collect(); + Self { + output_device: mesh.device_of(&output_coord), + layer_to_device, + } + } + + /// Validate the pure layout before any source preparation or GPU upload. + pub fn validate(&self, n_devices: usize, n_layers: usize) -> Result<(), String> { + if self.output_device >= n_devices { + return Err(format!( + "layout output device {} outside device count {}", + self.output_device, n_devices + )); + } + if self.layer_to_device.len() != n_layers { + return Err(format!( + "layout has {} layer assignments, expected {n_layers}", + self.layer_to_device.len() + )); + } + if let Some((layer, &device)) = self + .layer_to_device + .iter() + .enumerate() + .find(|(_, &device)| device >= n_devices) + { + return Err(format!( + "layout layer {layer} device {device} outside device count {n_devices}" + )); + } + Ok(()) + } + pub fn device_for_layer(&self, i: usize) -> usize { self.layer_to_device[i] } @@ -80,6 +133,15 @@ pub fn load_weights( devices: &mut [Gpu], layout: &Layout, ) -> HipResult> { + if devices.is_empty() { + return Err(hip_bridge::HipError::new( + 0, + "load_weights: at least one device is required", + )); + } + layout + .validate(devices.len(), source.n_layers()) + .map_err(|reason| hip_bridge::HipError::new(0, &reason))?; source.prepare(devices.len())?; let out_dev = layout.output_device(); let can_alias = devices.len() == 1; @@ -114,4 +176,25 @@ mod tests { assert_eq!(l.device_for_layer(i), 0); } } + + #[test] + fn mesh_layout_selects_stage_rank_zero_without_io() { + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let layout = Layout::from_mesh(&mesh, 4); + assert_eq!(layout.output_device(), 2); + assert_eq!( + (0..4) + .map(|layer| layout.device_for_layer(layer)) + .collect::>(), + vec![0, 0, 2, 2] + ); + assert!(layout.validate(mesh.n_devices(), 4).is_ok()); + } + + #[test] + fn invalid_layout_is_rejected_before_source_work() { + let layout = Layout::single(2); + assert!(layout.validate(0, 2).is_err()); + assert!(layout.validate(1, 3).is_err()); + } } diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs new file mode 100644 index 000000000..530421db3 --- /dev/null +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -0,0 +1,1004 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Pure logical model declarations and device-mesh planning. +//! +//! A manifest describes *what* an architecture needs. [`plan_manifest`] resolves +//! those declarations against one already-admitted rectangular +//! [`hipfire_hardware::DeviceMesh`] and describes *where* each declaration and +//! synchronization point belongs. This module deliberately has no GPU, file, +//! carrier, quantizer, or allocation dependency; fulfillment is separate. +//! +//! The manifest is the single source of truth for collectives. A row-sharded +//! projection contributes one ordered tensor collective over `Tp`, an +//! expert-sharded projection contributes one over `Ep`, and pipeline boundaries +//! come from the mesh. Executors consume this schedule rather than add +//! family-local reductions. + +use crate::tp_shard::ExpertAssign; +use hipfire_hardware::{CollectiveHint, DeviceMesh, DimKind}; +use rdna_compute::DType; +use std::collections::HashSet; + +/// Derive the collective required by one weight policy. +/// +/// The returned hint is per declared operation. Two different row-sharded +/// operations in one layer are two distinct schedule entries and both execute +/// once. +#[inline] +pub fn collective_for_policy(policy: &ShardPolicy) -> Option { + match policy { + ShardPolicy::RowShard { .. } => Some(CollectiveHint::AllReduce { kind: DimKind::Tp }), + ShardPolicy::ExpertSharded { .. } => { + Some(CollectiveHint::AllReduce { kind: DimKind::Ep }) + } + ShardPolicy::ExpertTensorSharded { inner, .. } => collective_for_policy(inner), + _ => None, + } +} + +/// Non-layer placement targets resolved from mesh stage coordinates. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum PinTarget { + /// Token embedding, pinned to pipeline stage zero. + Embed, + /// Final norm/language head, pinned to the final pipeline stage. + Output, +} + +/// Optional placement override. It is separate from [`ShardPolicy`] so a tied +/// logical identity can be materialized at an output stage without changing +/// the source declaration. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +pub enum PlacementHint { + /// Resolve placement from the policy and layer scope. + #[default] + Policy, + /// Resolve placement from a mesh-derived pin target. + Pin(PinTarget), +} + +/// Source dtype acceptance. The logical manifest dtype remains an architecture +/// expectation; fulfillment preserves the source dtype and never silently +/// converts representation. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum SourceDType { + Any, + Exact(DType), + OneOf(Vec), +} + +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct DTypeConstraint { + pub source: SourceDType, +} + +impl DTypeConstraint { + pub fn any_source() -> Self { + Self { + source: SourceDType::Any, + } + } + + pub fn source_exact(dtype: DType) -> Self { + Self { + source: SourceDType::Exact(dtype), + } + } + + pub fn source_from_sources(sources: Vec) -> Self { + Self { + source: SourceDType::OneOf(sources), + } + } + + pub fn accepts(&self, dtype: DType) -> bool { + match &self.source { + SourceDType::Any => true, + SourceDType::Exact(expected) => *expected == dtype, + SourceDType::OneOf(allowed) => allowed.contains(&dtype), + } + } +} + +/// The block ordering of a fused projection. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum FusedQkvLayout { + /// `[Q | K | V]`. + Qkv, + /// `[Q | gate]`. + QGate, + /// `[Q | K | V | Z]`. + QkvZ, +} + +/// How one logical tensor is projected onto mesh devices. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum ShardPolicy { + /// A complete tensor on every device in the owning compute grid. + Replicate, + /// Split the output dimension `axis` across `Tp`. + ColumnShard { axis: usize }, + /// Split the input dimension `axis` across `Tp`; the consumer reduces. + RowShard { axis: usize }, + /// Assign complete expert tensors across `Ep` ranks. + ExpertSharded { + n_experts: usize, + assign: ExpertAssign, + }, + /// Fused QKV projection with head-aware block boundaries. + FusedQkv { + q_heads: usize, + kv_heads: usize, + head_dim: usize, + layout: FusedQkvLayout, + }, + /// Per-head projection (DeltaNet state/projections). + HeadSharded { n_heads: usize, head_dim: usize }, + /// Alias another logical source in the same manifest scope. + Tied { source: String }, + /// Pin to a mesh-derived non-layer stage. + Pin(PinTarget), + /// Split vocabulary rows across `Tp`. + VocabShard { axis: usize }, + /// Split each expert tensor across `Tp`. The inner policy is normally + /// `ColumnShard { axis: 1 }` for gate/up or `RowShard { axis: 2 }` for down. + ExpertTensorSharded { + n_experts: usize, + inner: Box, + }, +} + +/// A logical weight declaration. No source filename or GPU handle belongs +/// here; architecture carriers resolve those at fulfillment time. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct WeightEntry { + pub name: String, + pub layer: Option, + pub logical_shape: Vec, + pub dtype: DType, + pub dtype_constraint: DTypeConstraint, + pub placement: PlacementHint, + pub policy: ShardPolicy, +} + +impl WeightEntry { + pub fn model( + name: impl Into, + logical_shape: Vec, + dtype: DType, + policy: ShardPolicy, + ) -> Self { + Self::model_with_dtype_constraint( + name, + logical_shape, + dtype, + DTypeConstraint::any_source(), + policy, + ) + } + + pub fn model_with_dtype_constraint( + name: impl Into, + logical_shape: Vec, + dtype: DType, + dtype_constraint: DTypeConstraint, + policy: ShardPolicy, + ) -> Self { + Self { + name: name.into(), + layer: None, + logical_shape, + dtype, + dtype_constraint, + placement: PlacementHint::Policy, + policy, + } + } + + pub fn layer( + name: impl Into, + layer: usize, + logical_shape: Vec, + dtype: DType, + policy: ShardPolicy, + ) -> Self { + Self::layer_with_dtype_constraint( + name, + layer, + logical_shape, + dtype, + DTypeConstraint::any_source(), + policy, + ) + } + + pub fn layer_with_dtype_constraint( + name: impl Into, + layer: usize, + logical_shape: Vec, + dtype: DType, + dtype_constraint: DTypeConstraint, + policy: ShardPolicy, + ) -> Self { + Self { + name: name.into(), + layer: Some(layer), + logical_shape, + dtype, + dtype_constraint, + placement: PlacementHint::Policy, + policy, + } + } + + pub fn with_placement(mut self, placement: PlacementHint) -> Self { + self.placement = placement; + self + } + + /// Stable identity used by source resolvers and store keys. + pub fn identity(&self) -> (&str, Option) { + (&self.name, self.layer) + } +} + +/// Per-layer state declaration. Actual cache representation remains in the +/// architecture/model owner; this records logical placement scope only. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +pub enum StateKind { + Kv { quant: String }, + Recurrent, + Conv, +} + +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct StateEntry { + pub kind: StateKind, + pub layer: usize, +} + +impl StateEntry { + pub fn new(kind: StateKind, layer: usize) -> Self { + Self { kind, layer } + } +} + +/// One fully resolved weight placement. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct WeightPlacement { + pub name: String, + pub layer: Option, + pub devices: Vec, +} + +/// One ordered collective implied by one manifest operation. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct CollectiveScheduleEntry { + pub name: String, + pub layer: usize, + pub hint: CollectiveHint, +} + +/// Complete pure compilation of declarations against a mesh. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct ManifestPlan { + pub weights: Vec, + /// State and the global devices on which that state is resident. + pub state: Vec<(StateEntry, Vec)>, + /// Ordered `(layer, hint)` schedule retained for executor integration. + pub layer_collectives: Vec<(usize, CollectiveHint)>, + /// Named schedule entries, allowing an executor to prove no operation was + /// silently omitted or scheduled twice. + pub collective_schedule: Vec, + /// PP boundary hints in ascending after-layer order. + pub band_xfers: Vec<(usize, CollectiveHint)>, +} + +fn base_coord_for(entry: &WeightEntry, mesh: &DeviceMesh, n_layers: usize) -> Vec { + let stage = match (entry.placement, &entry.policy, entry.layer) { + (PlacementHint::Pin(PinTarget::Embed), _, _) + | (PlacementHint::Policy, ShardPolicy::Pin(PinTarget::Embed), _) => 0, + (PlacementHint::Pin(PinTarget::Output), _, _) + | (PlacementHint::Policy, ShardPolicy::Pin(PinTarget::Output), _) => { + mesh.size_of(DimKind::Pp).saturating_sub(1) + } + (PlacementHint::Policy, _, Some(layer)) => mesh.stage_for_layer(layer, n_layers), + (PlacementHint::Policy, _, None) => 0, + }; + let mut coord = mesh.coord_of(0); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + coord[index] = stage; + } + coord +} + +/// Compute global placement without touching a source, GPU, or allocator. +pub fn placement_devices(entry: &WeightEntry, mesh: &DeviceMesh, n_layers: usize) -> Vec { + let coord = base_coord_for(entry, mesh, n_layers); + match &entry.policy { + ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } => vec![mesh.device_of(&coord)], + ShardPolicy::ExpertSharded { .. } => mesh.group_along(DimKind::Ep, &coord), + ShardPolicy::ExpertTensorSharded { .. } => mesh.group_along(DimKind::Tp, &coord), + _ => mesh.stage_devices(&coord), + } +} + +/// Ordered per-operation collective schedule. This deliberately does not +/// deduplicate by `(layer, kind)`: two distinct row-sharded projections in one +/// layer represent two distinct output points and each must reduce once. +pub fn collective_schedule(manifest: &[WeightEntry]) -> Vec { + manifest + .iter() + .filter_map(|entry| { + let layer = entry.layer?; + let hint = collective_for_policy(&entry.policy)?; + Some(CollectiveScheduleEntry { + name: entry.name.clone(), + layer, + hint, + }) + }) + .collect() +} + +/// Compact schedule view consumed by executor adapters. +pub fn layer_collectives(manifest: &[WeightEntry]) -> Vec<(usize, CollectiveHint)> { + collective_schedule(manifest) + .into_iter() + .map(|entry| (entry.layer, entry.hint)) + .collect() +} + +fn validate_shape(entry: &WeightEntry) -> Result<(), String> { + if entry.name.is_empty() { + return Err("manifest entry has an empty name".to_string()); + } + if entry.logical_shape.is_empty() || entry.logical_shape.iter().any(|&d| d == 0) { + return Err(format!( + "{}[layer {:?}]: logical_shape {:?} must be non-empty", + entry.name, entry.layer, entry.logical_shape + )); + } + Ok(()) +} + +/// Validate logical shard math and tied source identity before fulfillment. +pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result<(), String> { + let mut identities = HashSet::new(); + for entry in manifest { + validate_shape(entry)?; + if !identities.insert(entry.identity()) { + return Err(format!( + "duplicate manifest identity ('{}', {:?})", + entry.name, entry.layer + )); + } + } + + let tp = mesh.size_of(DimKind::Tp); + for entry in manifest { + let context = format!("{}[layer {:?}]", entry.name, entry.layer); + match &entry.policy { + ShardPolicy::ColumnShard { axis } + | ShardPolicy::RowShard { axis } + | ShardPolicy::VocabShard { axis } => { + let dim = entry.logical_shape.get(*axis).ok_or_else(|| { + format!("{context}: shard axis {axis} outside logical shape") + })?; + if tp > 1 && dim % tp != 0 { + return Err(format!( + "{context}: shard dim {dim} (axis {axis}) not divisible by Tp={tp}" + )); + } + } + ShardPolicy::FusedQkv { + q_heads, + kv_heads, + head_dim, + .. + } => { + if *q_heads == 0 || *kv_heads == 0 || *head_dim == 0 { + return Err(format!("{context}: fused QKV geometry must be non-zero")); + } + if tp > 1 && (q_heads % tp != 0 || kv_heads % tp != 0) { + return Err(format!( + "{context}: q_heads={q_heads}/kv_heads={kv_heads} not divisible by Tp={tp}" + )); + } + } + ShardPolicy::HeadSharded { n_heads, head_dim } => { + if *n_heads == 0 || *head_dim == 0 { + return Err(format!("{context}: head geometry must be non-zero")); + } + if tp > 1 && n_heads % tp != 0 { + return Err(format!( + "{context}: n_heads={n_heads} not divisible by Tp={tp}" + )); + } + } + ShardPolicy::Tied { source } => { + if source.is_empty() { + return Err(format!("{context}: tied source is empty")); + } + let source_entry = manifest + .iter() + .find(|candidate| candidate.name == *source && candidate.layer == entry.layer) + .ok_or_else(|| { + format!("{context}: Tied source '{source}' has no manifest entry in scope") + })?; + if source_entry.identity() == entry.identity() { + return Err(format!("{context}: an entry cannot tie to itself")); + } + } + ShardPolicy::ExpertSharded { n_experts, .. } => { + if *n_experts == 0 || entry.logical_shape.first() != Some(n_experts) { + return Err(format!( + "{context}: logical_shape {:?} first dimension must equal n_experts={n_experts}", + entry.logical_shape + )); + } + } + ShardPolicy::ExpertTensorSharded { n_experts, inner } => { + if *n_experts == 0 || entry.logical_shape.first() != Some(n_experts) { + return Err(format!( + "{context}: ExpertTensorSharded shape {:?} must start with n_experts={n_experts}", + entry.logical_shape + )); + } + let axis = match inner.as_ref() { + ShardPolicy::ColumnShard { axis: 1 } + | ShardPolicy::RowShard { axis: 2 } => match inner.as_ref() { + ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => *axis, + _ => unreachable!(), + }, + ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => { + return Err(format!( + "{context}: ExpertTensorSharded inner axis {axis} is incompatible with [expert, projection, hidden]" + )); + } + other => { + return Err(format!( + "{context}: ExpertTensorSharded inner policy {other:?} is unsupported" + )); + } + }; + let dim = entry.logical_shape.get(axis).copied().ok_or_else(|| { + format!("{context}: ExpertTensorSharded axis {axis} outside shape") + })?; + if tp > 1 && dim % tp != 0 { + return Err(format!( + "{context}: ExpertTensorSharded dim {dim} (axis {axis}) not divisible by Tp={tp}" + )); + } + } + ShardPolicy::Replicate | ShardPolicy::Pin(_) => {} + } + } + Ok(()) +} + +/// Compile declarations against a mesh. +pub fn plan_manifest( + weights: &[WeightEntry], + state: &[StateEntry], + mesh: &DeviceMesh, + n_layers: usize, +) -> Result { + validate_manifest(weights, mesh)?; + let mut state_ids = HashSet::new(); + for entry in state { + if entry.layer >= n_layers { + return Err(format!( + "state {:?} layer {} outside n_layers={n_layers}", + entry.kind, entry.layer + )); + } + if !state_ids.insert((&entry.kind, entry.layer)) { + return Err(format!( + "duplicate state declaration {:?}[layer {}]", + entry.kind, entry.layer + )); + } + } + let schedule = collective_schedule(weights); + let layer_collectives = schedule + .iter() + .map(|entry| (entry.layer, entry.hint)) + .collect(); + let weight_placements = weights + .iter() + .map(|entry| WeightPlacement { + name: entry.name.clone(), + layer: entry.layer, + devices: placement_devices(entry, mesh, n_layers), + }) + .collect(); + let state_placements = state + .iter() + .map(|entry| { + let mut coord = mesh.coord_of(0); + let stage = mesh.stage_for_layer(entry.layer, n_layers); + if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { + coord[index] = stage; + } + (entry.clone(), mesh.stage_devices(&coord)) + }) + .collect(); + let band_xfers = (0..n_layers) + .filter_map(|layer| mesh.band_xfer_after(layer, n_layers).map(|hint| (layer, hint))) + .collect(); + Ok(ManifestPlan { + weights: weight_placements, + state: state_placements, + layer_collectives, + collective_schedule: schedule, + band_xfers, + }) +} + +// ── Logical expert source identity ───────────────────────────────────────── + +/// How one logical expert group is distributed. This declaration is consumed +/// by the G5 executor-owned sealed plan; no rank assignment is resolved here. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum ExpertParallelism { + Single, + TensorParallel, + ExpertParallel, +} + +/// Stable source identities for expert projections. These names are manifest +/// references, not on-disk paths; the carrier/source resolver owns translation +/// to an artifact namespace. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum ExpertSourceLayout { + PackedFused { + gate_up: String, + down: String, + sidecars: Vec, + }, + PackedSeparate { + gate: String, + up: String, + down: String, + sidecars: Vec, + }, + PerExpertFused { + gate_up: Vec, + down: Vec, + sidecars: Vec, + }, + PerExpertSeparate { + gate: Vec, + up: Vec, + down: Vec, + sidecars: Vec, + }, +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub struct ExpertResourceRequirements { + pub bytes_per_expert: usize, + pub alignment: usize, +} + +/// Architecture-declared identity and source description of one expert group. +/// G5 derives rank ownership and seals the executor plan from this value. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct ExpertGroupSpec { + pub group: String, + pub layer: Option, + pub n_experts: usize, + pub parallelism: ExpertParallelism, + pub assignment: ExpertAssign, + pub source_layout: ExpertSourceLayout, + pub resources: ExpertResourceRequirements, + pub router: String, + pub execution: String, +} + +fn expert_context(spec: &ExpertGroupSpec) -> String { + format!("expert group '{}' layer {:?}", spec.group, spec.layer) +} + +fn source_names(layout: &ExpertSourceLayout) -> Vec<(&'static str, Vec)> { + match layout { + ExpertSourceLayout::PackedFused { + gate_up, + down, + sidecars, + } => vec![ + ("gate_up", vec![gate_up.clone()]), + ("down", vec![down.clone()]), + ("sidecar", sidecars.clone()), + ], + ExpertSourceLayout::PackedSeparate { + gate, + up, + down, + sidecars, + } => vec![ + ("gate", vec![gate.clone()]), + ("up", vec![up.clone()]), + ("down", vec![down.clone()]), + ("sidecar", sidecars.clone()), + ], + ExpertSourceLayout::PerExpertFused { + gate_up, + down, + sidecars, + } => vec![ + ("gate_up", gate_up.clone()), + ("down", down.clone()), + ("sidecar", sidecars.clone()), + ], + ExpertSourceLayout::PerExpertSeparate { + gate, + up, + down, + sidecars, + } => vec![ + ("gate", gate.clone()), + ("up", up.clone()), + ("down", down.clone()), + ("sidecar", sidecars.clone()), + ], + } +} + +fn manifest_entry<'a>( + spec: &ExpertGroupSpec, + manifest: &'a [WeightEntry], + label: &str, + name: &str, +) -> Result<&'a WeightEntry, String> { + let context = expert_context(spec); + if name.is_empty() { + return Err(format!("{context}: {label} reference is empty")); + } + manifest + .iter() + .find(|entry| entry.name == name && entry.layer == spec.layer) + .ok_or_else(|| format!("{context}: {label} reference '{name}' not found")) +} + +fn source_policy_matches( + spec: &ExpertGroupSpec, + label: &str, + policy: &ShardPolicy, +) -> bool { + match spec.parallelism { + ExpertParallelism::Single => matches!( + policy, + ShardPolicy::Replicate + | ShardPolicy::Pin(_) + | ShardPolicy::Tied { .. } + ), + ExpertParallelism::TensorParallel => match (label, policy) { + ( + "gate_up" | "gate" | "up", + ShardPolicy::ExpertTensorSharded { n_experts, inner }, + ) => { + *n_experts == spec.n_experts + && matches!(inner.as_ref(), ShardPolicy::ColumnShard { axis: 1 }) + } + ("down", ShardPolicy::ExpertTensorSharded { n_experts, inner }) => { + *n_experts == spec.n_experts + && matches!(inner.as_ref(), ShardPolicy::RowShard { axis: 2 }) + } + ("sidecar", ShardPolicy::Replicate | ShardPolicy::Tied { .. }) => true, + _ => false, + }, + ExpertParallelism::ExpertParallel => match (label, policy) { + ( + "gate_up" | "gate" | "up" | "down", + ShardPolicy::ExpertSharded { n_experts, assign }, + ) => *n_experts == spec.n_experts && *assign == spec.assignment, + ("sidecar", ShardPolicy::Replicate | ShardPolicy::Tied { .. }) => true, + _ => false, + }, + } +} + +fn source_shape_matches( + spec: &ExpertGroupSpec, + label: &str, + per_expert: bool, + entry: &WeightEntry, +) -> Result<(), String> { + let context = expert_context(spec); + if !source_policy_matches(spec, label, &entry.policy) { + return Err(format!( + "{context}: {label} source '{}' has incompatible policy {:?}", + entry.name, entry.policy + )); + } + if entry.logical_shape.len() < 2 { + return Err(format!( + "{context}: {label} source '{}' shape {:?} is too short", + entry.name, entry.logical_shape + )); + } + if !per_expert && entry.logical_shape.first() != Some(&spec.n_experts) { + return Err(format!( + "{context}: {label} source '{}' shape {:?} must start in n_experts={}", + entry.name, entry.logical_shape, spec.n_experts + )); + } + Ok(()) +} + +fn validate_expert_sources( + spec: &ExpertGroupSpec, + manifest: &[WeightEntry], +) -> Result<(), String> { + let context = expert_context(spec); + let router = manifest_entry(spec, manifest, "router", &spec.router)?; + if !matches!(router.logical_shape.len(), 1 | 2) + || router.logical_shape.last() != Some(&spec.n_experts) + { + return Err(format!( + "{context}: router '{}' shape {:?} must end in n_experts={}", + router.name, router.logical_shape, spec.n_experts + )); + } + if !matches!( + router.policy, + ShardPolicy::Replicate | ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } + ) { + return Err(format!( + "{context}: router '{}' has incompatible policy {:?}", + router.name, router.policy + )); + } + let per_expert = matches!( + spec.source_layout, + ExpertSourceLayout::PerExpertFused { .. } | ExpertSourceLayout::PerExpertSeparate { .. } + ); + if per_expert && spec.parallelism != ExpertParallelism::Single { + return Err(format!( + "{context}: per-expert source layout is only admitted for Single" + )); + } + + for (label, names) in source_names(&spec.source_layout) { + if names.is_empty() { + continue; + } + if per_expert && label != "sidecar" && names.len() != spec.n_experts { + return Err(format!( + "{context}: {label} source count={} != n_experts={}", + names.len(), + spec.n_experts + )); + } + let mut seen = HashSet::new(); + let mut shape: Option> = None; + for (index, name) in names.iter().enumerate() { + if !seen.insert(name.as_str()) { + return Err(format!( + "{context}: duplicate {label} source '{name}' at index {index}" + )); + } + source_shape_matches(spec, label, per_expert, entry)?; + if per_expert { + if let Some(previous) = &shape { + if previous != &entry.logical_shape { + return Err(format!( + "{context}: {label}[{index}] shape {:?} differs from {:?}", + entry.logical_shape, previous + )); + } + } else { + shape = Some(entry.logical_shape.clone()); + } + } + } + } + Ok(()) +} + +/// Validate logical expert source identities. Rank assignment remains owned by +/// G5; this function only proves source names, shapes, and scope are coherent. +pub fn validate_expert_group_specs( + specs: &[ExpertGroupSpec], + manifest: &[WeightEntry], +) -> Result<(), String> { + let mut identities = HashSet::new(); + for entry in manifest { + if !identities.insert(entry.identity()) { + return Err(format!( + "duplicate manifest identity ('{}', {:?})", + entry.name, entry.layer + )); + } + } + let mut groups = HashSet::new(); + for spec in specs { + let context = expert_context(spec); + if spec.group.is_empty() || spec.router.is_empty() || spec.execution.is_empty() { + return Err(format!("{context}: group/router/execution identities must be non-empty")); + } + if spec.n_experts == 0 || spec.resources.bytes_per_expert == 0 { + return Err(format!("{context}: n_experts and bytes_per_expert must be non-zero")); + } + if spec.resources.alignment == 0 || !spec.resources.alignment.is_power_of_two() { + return Err(format!("{context}: alignment must be a non-zero power of two")); + } + if !groups.insert((&spec.group, spec.layer)) { + return Err(format!("{context}: duplicate group/layer identity")); + } + validate_expert_sources(spec, manifest)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn layer_entry(name: &str, layer: usize, policy: ShardPolicy) -> WeightEntry { + WeightEntry::layer(name, layer, vec![8, 8], DType::F16, policy) + } + + #[test] + fn placement_and_boundaries_use_named_mesh() { + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let embed = WeightEntry::model( + "token_embd", + vec![32, 8], + DType::F16, + ShardPolicy::Pin(PinTarget::Embed), + ); + let row = layer_entry("wo", 1, ShardPolicy::RowShard { axis: 1 }); + assert_eq!(placement_devices(&embed, &mesh, 4), vec![0]); + assert_eq!(placement_devices(&row, &mesh, 4), vec![2, 3]); + let plan = plan_manifest( + &[layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 }), row], + &[], + &mesh, + 4, + ) + .unwrap(); + assert_eq!(plan.layer_collectives.len(), 2); + assert_eq!( + plan.band_xfers, + vec![(1, CollectiveHint::BandXfer { src: 0, dst: 1 })] + ); + } + + #[test] + fn schedule_is_ordered_per_operation_not_deduped() { + let manifest = vec![ + layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 }), + layer_entry("down", 0, ShardPolicy::RowShard { axis: 1 }), + ]; + assert_eq!( + layer_collectives(&manifest), + vec![ + (0, CollectiveHint::AllReduce { kind: DimKind::Tp }), + (0, CollectiveHint::AllReduce { kind: DimKind::Tp }), + ] + ); + assert_eq!(collective_schedule(&manifest)[0].name, "wo"); + assert_eq!(collective_schedule(&manifest)[1].name, "down"); + } + + #[test] + fn validation_covers_divisibility_ties_and_expert_shape() { + let tp3 = DeviceMesh::rect(&[(DimKind::Tp, 3)]); + assert!(validate_manifest( + &[layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 })], + &tp3 + ) + .is_err()); + let tied = vec![ + WeightEntry::model( + "embed", + vec![8, 8], + DType::F16, + ShardPolicy::Pin(PinTarget::Embed), + ), + WeightEntry::model( + "lm_head", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "embed".into(), + }, + ), + ]; + assert!(validate_manifest(&tied, &DeviceMesh::single()).is_ok()); + let bad_expert = WeightEntry::layer( + "experts", + 0, + vec![3, 8, 8], + DType::F16, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ); + assert!(validate_manifest(&[bad_expert], &DeviceMesh::single()).is_err()); + } + + #[test] + fn expert_source_identity_and_shape_are_checked() { + let manifest = vec![ + WeightEntry::layer( + "router", + 0, + vec![8, 4], + DType::F16, + ShardPolicy::Replicate, + ), + WeightEntry::layer( + "gate_up", + 0, + vec![4, 8, 8], + DType::F16, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ), + WeightEntry::layer( + "down", + 0, + vec![4, 8, 8], + DType::F16, + ShardPolicy::ExpertSharded { + n_experts: 4, + assign: ExpertAssign::Stride, + }, + ), + ]; + let spec = ExpertGroupSpec { + group: "ffn".into(), + layer: Some(0), + n_experts: 4, + parallelism: ExpertParallelism::ExpertParallel, + assignment: ExpertAssign::Stride, + source_layout: ExpertSourceLayout::PackedFused { + gate_up: "gate_up".into(), + down: "down".into(), + sidecars: Vec::new(), + }, + resources: ExpertResourceRequirements { + bytes_per_expert: 1024, + alignment: 256, + }, + router: "router".into(), + execution: "moe.ffn".into(), + }; + assert!(validate_expert_group_specs(&[spec], &manifest).is_ok()); + let bad = ExpertGroupSpec { + source_layout: ExpertSourceLayout::PackedFused { + gate_up: "missing".into(), + down: "down".into(), + sidecars: Vec::new(), + }, + ..ExpertGroupSpec { + group: "ffn2".into(), + layer: Some(0), + n_experts: 4, + parallelism: ExpertParallelism::ExpertParallel, + assignment: ExpertAssign::Stride, + source_layout: ExpertSourceLayout::PackedFused { + gate_up: "gate_up".into(), + down: "down".into(), + sidecars: Vec::new(), + }, + resources: ExpertResourceRequirements { + bytes_per_expert: 1024, + alignment: 256, + }, + router: "router".into(), + execution: "moe.ffn".into(), + } + }; + assert!(validate_expert_group_specs(&[bad], &manifest).is_err()); + } +} diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs new file mode 100644 index 000000000..a66336985 --- /dev/null +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -0,0 +1,724 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 Kaden Schutt +// hipfire — see LICENSE and NOTICE in the project root. + +//! Transactional fulfillment for the pure weight manifest. +//! +//! [`crate::weight_manifest::plan_manifest`] owns the CPU-only "where". This +//! module owns the narrow "how" pilot for a plain LLaMA Single target: a +//! source callback supplies already-resolved bytes and dtype, the store uploads +//! them, and the first failure explicitly rolls back every resident buffer. +//! +//! The store is not a model owner. It has no `Drop` implementation and never +//! frees GPU buffers implicitly. A carrier may move a committed store into its +//! existing `ArchModel` owner; that owner must call [`WeightStore::release_on_owner`] +//! during its existing teardown path. `take` transfers a resident handle to the +//! owner that is assembling typed weights, and therefore removes the cell from +//! the store's cleanup set. + +use crate::weight_manifest::{placement_devices, ShardPolicy, WeightEntry}; +use hipfire_hardware::{DeviceMesh, MeshEpoch}; +use rdna_compute::{DType, Gpu, GpuTensor}; +use std::collections::HashMap; + +/// Stable logical placement identity. Layer is part of the key because a +/// per-layer name such as `wq` appears once for every decoder block. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +pub struct WeightPlacementKey { + pub name: String, + pub layer: Option, + pub device: usize, +} + +impl WeightPlacementKey { + pub fn new(name: impl Into, layer: Option, device: usize) -> Self { + Self { + name: name.into(), + layer, + device, + } + } +} + +/// The immutable projection applied to one logical source before upload. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum WeightProjectionKind { + Static, + ColumnShard, + RowShard, + FusedQkv, + HeadSharded, + VocabShard, + ExpertCompact, + ExpertTensor, +} + +/// Value-owned placement metadata. It contains no GPU or source-file +/// representation and remains stable after a handle is taken from the store. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct WeightProjection { + pub kind: WeightProjectionKind, + pub axis: Option, + pub rank: usize, + pub world_size: usize, + pub logical_shape: Vec, + pub dtype: DType, +} + +fn projection_for(entry: &WeightEntry, rank: usize, world_size: usize, dtype: DType) -> WeightProjection { + let (kind, axis) = match &entry.policy { + ShardPolicy::ColumnShard { axis } => (WeightProjectionKind::ColumnShard, Some(*axis)), + ShardPolicy::RowShard { axis } => (WeightProjectionKind::RowShard, Some(*axis)), + ShardPolicy::FusedQkv { .. } => (WeightProjectionKind::FusedQkv, None), + ShardPolicy::HeadSharded { .. } => (WeightProjectionKind::HeadSharded, None), + ShardPolicy::VocabShard { axis } => (WeightProjectionKind::VocabShard, Some(*axis)), + ShardPolicy::ExpertSharded { .. } => (WeightProjectionKind::ExpertCompact, None), + ShardPolicy::ExpertTensorSharded { .. } => (WeightProjectionKind::ExpertTensor, None), + ShardPolicy::Replicate | ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } => { + (WeightProjectionKind::Static, None) + } + }; + WeightProjection { + kind, + axis, + rank, + world_size, + logical_shape: entry.logical_shape.clone(), + dtype, + } +} + +/// A resident GPU tensor or a symbolic alias to another logical source. +/// +/// Aliases own no buffer. Resident buffers have no implicit destructor; the +/// current model owner explicitly consumes them through its teardown method. +pub enum WeightHandle { + Resident(GpuTensor), + Alias(String), +} + +/// Identity captured at the start of a load. It is deliberately immutable and +/// contains only mesh generation, logical rank, and physical device identity. +/// No policy or source representation is smuggled into the origin. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +pub struct WeightOrigin { + mesh_epoch: MeshEpoch, + logical_rank: usize, + physical_device: i32, +} + +impl WeightOrigin { + pub fn from_parts(mesh_epoch: MeshEpoch, logical_rank: usize, physical_device: i32) -> Self { + Self { + mesh_epoch, + logical_rank, + physical_device, + } + } + + pub fn for_single(mesh: &DeviceMesh, gpu: &Gpu) -> Self { + Self::from_parts(mesh.epoch(), 0, gpu.device_id) + } + + pub fn mesh_epoch(self) -> MeshEpoch { + self.mesh_epoch + } + + pub fn logical_rank(self) -> usize { + self.logical_rank + } + + pub fn physical_device(self) -> i32 { + self.physical_device + } +} + +/// Errors that are detected before a store is allowed to release a resident +/// buffer. Origin mismatch always returns the store to the caller unchanged. +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum WeightStoreError { + OriginMismatch { + expected: WeightOrigin, + actual: WeightOrigin, + }, + UnboundOrigin, + DuplicatePlacement(WeightPlacementKey), + MissingPlacement(WeightPlacementKey), + InvalidTarget(String), +} + +impl std::fmt::Display for WeightStoreError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::OriginMismatch { expected, actual } => write!( + f, + "weight store origin mismatch: expected {:?}, got {:?}", + expected, actual + ), + Self::UnboundOrigin => write!(f, "weight store has no target origin"), + Self::DuplicatePlacement(key) => write!( + f, + "duplicate weight placement {}[layer {:?}] on device {}", + key.name, key.layer, key.device + ), + Self::MissingPlacement(key) => write!( + f, + "missing weight placement {}[layer {:?}] on device {}", + key.name, key.layer, key.device + ), + Self::InvalidTarget(message) => write!(f, "invalid weight store target: {message}"), + } + } +} + +impl std::error::Error for WeightStoreError {} + +/// Error identifying the first failed manifest cell. The store has already +/// been rolled back before this value is returned by [`fulfill_manifest`]. +#[derive(Debug)] +pub struct FulfillError { + pub name: String, + pub layer: Option, + pub device: usize, + pub reason: String, +} + +impl std::fmt::Display for FulfillError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "fulfill_manifest: {}[layer {:?}] on device {}: {}", + self.name, self.layer, self.device, self.reason + ) + } +} + +impl std::error::Error for FulfillError {} + +/// Load-side placement container. It records one immutable projection per +/// `(name, layer, device)` and captures the target origin once. +/// +/// There is intentionally no `Drop` implementation. A `WeightStore` that is +/// abandoned without explicit rollback/release leaks rather than guessing a +/// GPU owner; production callers keep it beneath `ArchModel`. +#[derive(Default)] +pub struct WeightStore { + placements: HashMap, + projections: HashMap, + origin: Option, +} + +impl WeightStore { + pub fn new() -> Self { + Self::default() + } + + pub fn with_origin(origin: WeightOrigin) -> Self { + Self { + placements: HashMap::new(), + projections: HashMap::new(), + origin: Some(origin), + } + } + + pub fn origin(&self) -> Option { + self.origin + } + + pub fn len(&self) -> usize { + self.placements.len() + } + + pub fn is_empty(&self) -> bool { + self.placements.is_empty() + } + + pub fn contains(&self, name: &str, layer: Option, device: usize) -> bool { + self.placements + .contains_key(&WeightPlacementKey::new(name, layer, device)) + } + + pub fn get(&self, name: &str, layer: Option, device: usize) -> Option<&WeightHandle> { + self.placements + .get(&WeightPlacementKey::new(name, layer, device)) + } + + pub fn projection( + &self, + name: &str, + layer: Option, + device: usize, + ) -> Option<&WeightProjection> { + self.projections + .get(&WeightPlacementKey::new(name, layer, device)) + } + + pub fn devices_for(&self, name: &str, layer: Option) -> Vec { + let mut devices: Vec<_> = self + .placements + .keys() + .filter(|key| key.name == name && key.layer == layer) + .map(|key| key.device) + .collect(); + devices.sort_unstable(); + devices + } + + fn insert( + &mut self, + key: WeightPlacementKey, + handle: WeightHandle, + projection: WeightProjection, + ) -> Result<(), WeightStoreError> { + if self.placements.contains_key(&key) { + return Err(WeightStoreError::DuplicatePlacement(key)); + } + self.placements.insert(key.clone(), handle); + self.projections.insert(key, projection); + Ok(()) + } + + /// Stage a symbolic alias without GPU work. Used for tied declarations and + /// CPU ownership tests; aliases never participate in release. + pub fn stage_alias( + &mut self, + name: impl Into, + layer: Option, + device: usize, + source: impl Into, + projection: WeightProjection, + ) -> Result<(), WeightStoreError> { + self.insert( + WeightPlacementKey::new(name, layer, device), + WeightHandle::Alias(source.into()), + projection, + ) + } + + /// Move a handle out of the store. The projection is removed with it so no + /// stale metadata can describe a cell that the store no longer owns. + pub fn take( + &mut self, + name: &str, + layer: Option, + device: usize, + ) -> Option { + let key = WeightPlacementKey::new(name, layer, device); + self.projections.remove(&key); + self.placements.remove(&key) + } + + pub fn take_with_projection( + &mut self, + name: &str, + layer: Option, + device: usize, + ) -> Option<(WeightHandle, WeightProjection)> { + let key = WeightPlacementKey::new(name, layer, device); + let handle = self.placements.remove(&key)?; + let projection = self.projections.remove(&key)?; + Some((handle, projection)) + } + + /// Start a typed assembly transaction. Handles moved through the + /// transaction are restored to this store if the transaction is dropped + /// before `finalize`; no GPU free or second owner is introduced. + pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { + WeightStoreAssembly { + store: self, + taken: Vec::new(), + committed: false, + } + } + + /// Compare a store's captured origin with an already-resolved target + /// identity. This pure seam is used by fault-path tests and by owner + /// teardown after target resolution. + pub fn validate_origin_value( + &self, + expected: WeightOrigin, + ) -> Result<(), WeightStoreError> { + let actual = self.origin.ok_or(WeightStoreError::UnboundOrigin)?; + if actual != expected { + return Err(WeightStoreError::OriginMismatch { expected, actual }); + } + Ok(()) + } + + /// Verify that this store is still being handled by the same mesh/device + /// target. No GPU calls occur on mismatch. + pub fn validate_origin( + &self, + mesh: &DeviceMesh, + gpu: &Gpu, + ) -> Result<(), WeightStoreError> { + self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) + } + + /// Explicit owner teardown. This method is intentionally consuming and has + /// no implicit/drop fallback; `ArchModel::free_gpu` is the production call + /// site. Callers must validate the target first with [`validate_origin`]. + /// On mismatch the original store is returned unchanged for retry by the + /// owner; no GPU call occurs. + pub fn release_on_owner( + self, + mesh: &DeviceMesh, + gpu: &mut Gpu, + ) -> Result<(), (Self, WeightStoreError)> { + if let Err(error) = self.validate_origin(mesh, gpu) { + return Err((self, error)); + } + self.release_unchecked(gpu); + Ok(()) + } + + /// Explicit rollback for a failed transaction. It consumes the partial + /// store and frees every resident buffer on the single owning GPU. + fn rollback(self, gpu: &Gpu) { + self.release_unchecked(gpu); + } + + fn release_unchecked(self, gpu: &Gpu) { + for handle in self.placements.into_values() { + if let WeightHandle::Resident(tensor) = handle { + // Rollback is deliberately direct and best-effort, matching + // the existing loader's explicit owner teardown. The store + // never relies on a destructor to release GPU memory. + let _ = gpu.hip.free(tensor.buf); + } + } + } +} + +/// One resident/alias handle temporarily moved during typed assembly. +pub struct TakenWeight { + pub key: WeightPlacementKey, + pub handle: WeightHandle, + pub projection: WeightProjection, +} + +/// Rollback-owning assembly transaction. Dropping it restores every taken cell +/// to the parent store; it never frees a GPU buffer implicitly. +pub struct WeightStoreAssembly<'a> { + store: &'a mut WeightStore, + taken: Vec, + committed: bool, +} + +impl<'a> WeightStoreAssembly<'a> { + pub fn take( + &mut self, + name: &str, + layer: Option, + device: usize, + ) -> Option { + let key = WeightPlacementKey::new(name, layer, device); + let (handle, projection) = self.store.take_with_projection(name, layer, device)?; + let slot = self.taken.len(); + self.taken.push(TakenWeight { + key, + handle, + projection, + }); + Some(slot) + } + + pub fn commit(self) -> WeightStoreAssemblyGuard<'a> { + WeightStoreAssemblyGuard { inner: self } + } +} + +impl Drop for WeightStoreAssembly<'_> { + fn drop(&mut self) { + if self.committed { + return; + } + for taken in self.taken.drain(..) { + let _ = self + .store + .insert(taken.key, taken.handle, taken.projection); + } + } +} + +/// Guard retained while the typed architecture object is being built. If it +/// is dropped before `finalize`, all handles return to the parent store. +pub struct WeightStoreAssemblyGuard<'a> { + inner: WeightStoreAssembly<'a>, +} + +impl WeightStoreAssemblyGuard<'_> { + pub fn get(&self, slot: usize) -> Option<&WeightHandle> { + self.inner.taken.get(slot).map(|taken| &taken.handle) + } + + pub fn projection(&self, slot: usize) -> Option<&WeightProjection> { + self.inner.taken.get(slot).map(|taken| &taken.projection) + } + + /// Transfer the taken handles to the existing ArchModel-owned typed + /// weights. This is the sole operation that removes them from rollback + /// ownership. + pub fn finalize(mut self) -> Vec { + self.inner.committed = true; + std::mem::take(&mut self.inner.taken) + } +} + +fn target_error(mesh: &DeviceMesh) -> Option { + (mesh.n_devices() != 1).then(|| FulfillError { + name: "".to_string(), + layer: None, + device: 0, + reason: format!( + "plain LLaMA Single fulfillment requires one logical device, got {}", + mesh.n_devices() + ), + }) +} + +/// Fulfill a manifest for a plain LLaMA Single target. +/// +/// The source callback is the architecture-owned namespace seam and returns +/// raw bytes plus the actual source dtype. No file/GGUF/HFQ type crosses this +/// API. On the first source, dtype, or upload failure every earlier resident is +/// explicitly released before the error is returned. +pub fn fulfill_manifest_single( + weights: &[WeightEntry], + mesh: &DeviceMesh, + n_layers: usize, + gpu: &Gpu, + source: F, +) -> Result +where + F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, +{ + if let Some(error) = target_error(mesh) { + return Err(error); + } + if let Err(reason) = crate::weight_manifest::validate_manifest(weights, mesh) { + return Err(FulfillError { + name: "".to_string(), + layer: None, + device: 0, + reason, + }); + } + + let origin = WeightOrigin::for_single(mesh, gpu); + let mut store = WeightStore::with_origin(origin); + for entry in weights { + let devices = placement_devices(entry, mesh, n_layers); + if devices.as_slice() != [0] { + let error = FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: devices.first().copied().unwrap_or(0), + reason: format!( + "Single placement resolved to {:?}, expected [0]", + devices + ), + }; + store.rollback(gpu); + return Err(error); + } + let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); + if let ShardPolicy::Tied { source: source_name } = &entry.policy { + let projection = projection_for(entry, 0, 1, entry.dtype); + if let Err(reason) = store.insert( + key, + WeightHandle::Alias(source_name.clone()), + projection, + ) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: reason.to_string(), + }); + } + continue; + } + + let (bytes, dtype) = match source(entry) { + Ok(value) => value, + Err(reason) => { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!("source read failed: {reason}"), + }); + } + }; + if !entry.dtype_constraint.accepts(dtype) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "source dtype {dtype:?} violates constraint {:?}", + entry.dtype_constraint + ), + }); + } + let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { + Ok(tensor) => tensor, + Err(error) => { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!("upload_raw failed: {error}"), + }); + } + }; + tensor.dtype = dtype; + let projection = projection_for(entry, 0, 1, dtype); + if let Err(reason) = store.insert(key, WeightHandle::Resident(tensor), projection) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: reason.to_string(), + }); + } + } + Ok(store) +} + +/// Canonical name used by the manifest fulfillment seam. The target is +/// deliberately Single-only in this pilot; multi-device fulfillment belongs to +/// the admitted mesh/G5 integration and must not grow a second owner here. +pub fn fulfill_manifest( + weights: &[WeightEntry], + mesh: &DeviceMesh, + n_layers: usize, + gpu: &Gpu, + source: F, +) -> Result +where + F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, +{ + fulfill_manifest_single(weights, mesh, n_layers, gpu, source) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::weight_manifest::{PinTarget, ShardPolicy}; + use hipfire_hardware::DimKind; + + fn projection(dtype: DType) -> WeightProjection { + WeightProjection { + kind: WeightProjectionKind::Static, + axis: None, + rank: 0, + world_size: 1, + logical_shape: vec![1], + dtype, + } + } + + #[test] + fn origin_mismatch_is_detected_before_gpu_release() { + let first = DeviceMesh::single(); + let second = DeviceMesh::single(); + let actual = WeightOrigin::from_parts(first.epoch(), 0, 0); + let expected = WeightOrigin::from_parts(second.epoch(), 0, 0); + let store = WeightStore::with_origin(actual); + let error = store.validate_origin_value(expected).unwrap_err(); + assert!(matches!( + error, + WeightStoreError::OriginMismatch { + expected: got_expected, + actual: got_actual + } if got_expected == expected && got_actual == actual + )); + } + + #[test] + fn staged_rollback_removes_handles_and_projection_together() { + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); + let mut store = WeightStore::with_origin(origin); + store + .stage_alias("first", None, 0, "source", projection(DType::F16)) + .unwrap(); + store + .stage_alias("second", Some(2), 0, "source", projection(DType::F16)) + .unwrap(); + assert_eq!(store.len(), 2); + let first = store.take_with_projection("first", None, 0).unwrap(); + assert!(matches!(first.0, WeightHandle::Alias(_))); + assert!(store.projection("first", None, 0).is_none()); + assert_eq!(store.len(), 1); + let second = store.take("second", Some(2), 0).unwrap(); + assert!(matches!(second, WeightHandle::Alias(_))); + assert!(store.is_empty()); + } + + #[test] + fn assembly_drop_restores_staged_handles() { + let mesh = DeviceMesh::single(); + let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); + store + .stage_alias("x", None, 0, "source", projection(DType::F16)) + .unwrap(); + { + let mut assembly = store.begin_assembly(); + assert_eq!(assembly.take("x", None, 0), Some(0)); + assert!(assembly.get(0).is_some()); + } + assert!(store.contains("x", None, 0)); + assert!(store.projection("x", None, 0).is_some()); + } + + #[test] + fn repeated_unload_lookup_cannot_reclaim_a_transferred_cell() { + let mesh = DeviceMesh::single(); + let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); + store + .stage_alias("x", None, 0, "source", projection(DType::F16)) + .unwrap(); + let _owned = store.take("x", None, 0).unwrap(); + assert!(store.take("x", None, 0).is_none()); + assert!(store.projection("x", None, 0).is_none()); + assert!(store.is_empty()); + } + + #[test] + fn duplicate_projection_is_rejected_without_replacing_identity() { + let mesh = DeviceMesh::single(); + let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); + store + .stage_alias("x", None, 0, "source-a", projection(DType::F16)) + .unwrap(); + let error = store + .stage_alias("x", None, 0, "source-b", projection(DType::F32)) + .unwrap_err(); + assert!(matches!(error, WeightStoreError::DuplicatePlacement(_))); + assert!(matches!(store.get("x", None, 0), Some(WeightHandle::Alias(source)) if source == "source-a")); + assert_eq!(store.projection("x", None, 0).unwrap().dtype, DType::F16); + } + + #[test] + fn single_target_refuses_multi_device_before_source_or_gpu_work() { + let mesh = DeviceMesh::rect(&[(DimKind::Tp, 2)]); + let entry = WeightEntry::model( + "embed", + vec![2, 2], + DType::F16, + ShardPolicy::Pin(PinTarget::Embed), + ); + // The target guard is pure and can be checked without constructing a + // Gpu; the closure would be unreachable on this path. + assert!(target_error(&mesh).is_some()); + assert_eq!(placement_devices(&entry, &mesh, 1), vec![0]); + } +} From f907ef7501ccaa401a50d00e8e9d6f42b67676df Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 17:49:21 +0200 Subject: [PATCH 02/17] fix(device-mesh): bind expert manifest source entries --- crates/hipfire-runtime/src/weight_manifest.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 530421db3..26c64ad2e 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -782,6 +782,7 @@ fn validate_expert_sources( "{context}: duplicate {label} source '{name}' at index {index}" )); } + let entry = manifest_entry(spec, manifest, &format!("{label}[{index}]"), name)?; source_shape_matches(spec, label, per_expert, entry)?; if per_expert { if let Some(previous) = &shape { From ccd50edc8995e220dfedfa122d7d08f7c7d1625d Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 18:55:36 +0200 Subject: [PATCH 03/17] fix(device-mesh): complete llama weight store pilot --- crates/hipfire-arch-llama/src/arch.rs | 99 ++- crates/hipfire-arch-llama/src/arch_model.rs | 61 +- crates/hipfire-arch-llama/src/carrier.rs | 758 +++++++++++++++--- crates/hipfire-runtime/src/weight_backend.rs | 9 + crates/hipfire-runtime/src/weight_manifest.rs | 129 +++ crates/hipfire-runtime/src/weight_store.rs | 254 +++++- 6 files changed, 1146 insertions(+), 164 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index 920601367..f154f0c1b 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -20,7 +20,7 @@ use hipfire_runtime::hfq::{self, HfqFile}; use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::weight_manifest::{ - FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, + DTypeConstraint, FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, }; use rdna_compute::{DType, Gpu}; @@ -39,6 +39,58 @@ use hipfire_runtime::llama::{attention_family, AttnParams, KvTierInputs, KvTierP /// see [`hipfire_arch_qwen35::Qwen35`] for those. pub struct Llama; +fn linear_source_constraint() -> DTypeConstraint { + DTypeConstraint::source_from_sources(vec![ + DType::F32, + DType::Q4F16G64, + DType::Q8_0, + DType::Q4K, + DType::Q8HFQ, + DType::HFQ4G256, + DType::HFQ4G128, + DType::HFQ6G256, + DType::HFQ2G256, + DType::HFQ2G128, + DType::HFQ3G256, + DType::HFQ3G128, + DType::MQ4G256, + DType::MQ8G256, + DType::MQ6G256, + DType::MQ3G256, + DType::MQ2G256, + DType::MQ2G256Lloyd, + DType::MQ2G256LloydU, + DType::MQ3G256Lloyd, + DType::HFP4G32, + DType::MFP4G32, + DType::MQ4G256Lloyd, + DType::MQ2G256GL, + DType::MQ3G256GL, + DType::TQ2G128, + DType::BQ1G128, + DType::MQ4G256V2, + DType::MQ4CG256, + DType::MQ6G256V2, + DType::MQ5G256V2, + DType::MQ3G256V2, + DType::MQ2G256V2, + ]) +} + +fn embedding_source_constraint() -> DTypeConstraint { + DTypeConstraint::source_from_sources(vec![ + DType::F32, + DType::Q8_0, + DType::Q4K, + DType::HFQ4G256, + DType::HFQ4G128, + ]) +} + +fn norm_source_constraint() -> DTypeConstraint { + DTypeConstraint::source_exact(DType::F32) +} + impl Architecture for Llama { type Weights = LlamaWeights; type State = ForwardScratch; @@ -85,19 +137,24 @@ impl Llama { use ShardPolicy::*; let (dim, hidden, head_dim) = (cfg.dim, cfg.hidden_dim, cfg.head_dim); let (heads, kv_heads) = (cfg.n_heads, cfg.n_kv_heads); + let linear = linear_source_constraint(); + let embedding = embedding_source_constraint(); + let norm = norm_source_constraint(); let mut manifest = Vec::with_capacity(cfg.n_layers * 11 + 3); - manifest.push(WeightEntry::model( + manifest.push(WeightEntry::model_with_dtype_constraint( "token_embd", vec![cfg.vocab_size, dim], DType::F16, + embedding, Pin(PinTarget::Embed), )); for layer in 0..cfg.n_layers { - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wq", layer, vec![heads * head_dim, dim], DType::F16, + linear.clone(), FusedQkv { q_heads: heads, kv_heads, @@ -105,89 +162,101 @@ impl Llama { layout: FusedQkvLayout::Qkv, }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wk", layer, vec![kv_heads * head_dim, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wv", layer, vec![kv_heads * head_dim, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "wo", layer, vec![dim, heads * head_dim], DType::F16, + linear.clone(), RowShard { axis: 1 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_gate", layer, vec![hidden, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_up", layer, vec![hidden, dim], DType::F16, + linear.clone(), ColumnShard { axis: 0 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_down", layer, vec![dim, hidden], DType::F16, + linear.clone(), RowShard { axis: 1 }, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "attn_norm", layer, vec![dim], DType::F32, + norm.clone(), Replicate, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "ffn_norm", layer, vec![dim], DType::F32, + norm.clone(), Replicate, )); if cfg.has_qk_norm { - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "q_norm", layer, vec![head_dim], DType::F32, + norm.clone(), Replicate, )); - manifest.push(WeightEntry::layer( + manifest.push(WeightEntry::layer_with_dtype_constraint( "k_norm", layer, vec![head_dim], DType::F32, + norm.clone(), Replicate, )); } } - manifest.push(WeightEntry::model( + manifest.push(WeightEntry::model_with_dtype_constraint( "output_norm", vec![dim], DType::F32, + norm, Replicate, )); - manifest.push(WeightEntry::model( + manifest.push(WeightEntry::model_with_dtype_constraint( "lm_head", vec![cfg.vocab_size, dim], DType::F16, + linear, Pin(PinTarget::Output), )); manifest diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 7e7d3d7ce..9d1297df9 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -4,10 +4,26 @@ use hipfire_runtime::arch_model::ArchModel; use hipfire_runtime::llama::KvCache; +use hipfire_runtime::weight_store::{ + WeightHandle, WeightOrigin, WeightStore, WeightStoreError, +}; use rdna_compute::Gpu; use crate::carrier::LlamaBundle; +fn drain_weight_store( + store: WeightStore, + origin: WeightOrigin, + gpu: &mut Gpu, +) -> Result<(), (WeightStore, WeightStoreError)> { + for handle in store.take_all(origin)? { + if let WeightHandle::Resident(tensor) = handle { + let _ = gpu.free_tensor(tensor); + } + } + Ok(()) +} + impl ArchModel for LlamaBundle { fn dim(&self) -> usize { self.config.dim @@ -35,6 +51,17 @@ impl ArchModel for LlamaBundle { } fn free_gpu(self: Box, gpu: &mut Gpu) { + // Validate before destructuring the consuming owner. A mismatch must + // leave the resident store attached to an owner that can be retried; + // leaking the boxed owner is safer than dropping the only cleanup + // authority. + if let Some(store) = self.weight_store.as_ref() { + if let Err(error) = store.validate_origin_value(self.weight_origin) { + eprintln!("llama: refusing weight-store release: {error}"); + let _ = Box::into_raw(self); + return; + } + } let LlamaBundle { config: _, weights, @@ -42,21 +69,47 @@ impl ArchModel for LlamaBundle { kv, manifest_plan: _, weight_store, - mesh, + weight_origin, + mesh: _, dflash_extract_layers: _, dspark_weights: _, dspark_assets: _, } = *self; // Mirror the existing unload ordering: scratch → store/weights → kv. - // A committed store is only released here, through the ArchModel owner; - // no store destructor or independent carrier free path exists. scratch.free_gpu(gpu); if let Some(store) = weight_store { - if let Err((_, error)) = store.release_on_owner(&mesh, gpu) { + if let Err((store, error)) = drain_weight_store(store, weight_origin, gpu) { + // This is defensive because the pre-check above used the + // same immutable origin. Never discard a rejected store. eprintln!("llama: refusing weight-store release: {error}"); + std::mem::forget(store); } } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); } } + +#[cfg(test)] +mod tests { + use super::*; + use hipfire_hardware::DeviceMesh; + use hipfire_runtime::weight_manifest::{ShardPolicy, WeightEntry}; + use hipfire_runtime::weight_store::fulfill_manifest_single; + + #[test] + fn arch_owner_unload_drains_residents_and_repeated_empty_unload_is_safe() { + let Ok(mut gpu) = Gpu::init() else { + return; + }; + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::for_single(&mesh, &gpu); + let entry = WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); + let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], rdna_compute::DType::F32)) + }) + .unwrap(); + assert!(drain_weight_store(store, origin, &mut gpu).is_ok()); + assert!(drain_weight_store(WeightStore::with_origin(origin), origin, &mut gpu).is_ok()); + } +} diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index c41b80298..86b4dcfe5 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -7,13 +7,21 @@ use crate::Llama; use hipfire_hardware::DeviceMesh; use hipfire_runtime::arch::Architecture; use hipfire_runtime::dspark_core::DsparkWeights; +use hipfire_runtime::hfq::HfqFile; use hipfire_runtime::llama::KvCacheExt; use hipfire_runtime::llama::{ - ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LlamaConfig, LlamaWeights, + EmbeddingFormat, ForwardScratch, KvCache, KvDims, KvLayers, KvTarget, LayerWeights, + LlamaConfig, LlamaWeights, WeightTensor, }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; -use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan}; -use hipfire_runtime::weight_store::WeightStore; +use hipfire_runtime::weight_backend::hfq_weight_dtype; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; +use hipfire_runtime::weight_store::{ + TakenWeight, WeightHandle, WeightStore, WeightStoreAssembly, WeightStoreAssemblyGuard, + WeightOrigin, +}; +use rdna_compute::{DType, GpuTensor}; +use std::collections::HashMap; pub struct LlamaBundle { pub config: LlamaConfig, @@ -27,6 +35,9 @@ pub struct LlamaBundle { /// this bundle. It is crate-visible so callers cannot create an independent /// unload owner; `ArchModel::free_gpu` is the sole release path. pub(crate) weight_store: Option, + /// Exact target identity captured when this bundle was admitted. The + /// owner uses it to validate every store-origin component before teardown. + pub(crate) weight_origin: WeightOrigin, pub(crate) mesh: DeviceMesh, /// Decoder-layer indices whose residual hidden states a hidden-conditioned /// drafter (DFlash / EAGLE) wants captured, ascending order. Empty = no @@ -38,139 +49,539 @@ pub struct LlamaBundle { /// was found or speculation was disabled. Task-10 wires the speculator build. pub dspark_weights: Option, /// Loaded DSpark drafter body assets (5-layer dense-GQA transformer + + /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } +fn plan_single(config: &LlamaConfig) -> Result<(DeviceMesh, ManifestPlan), String> { + let mesh = DeviceMesh::single(); + let manifest = Llama::weight_manifest(config); + let state = Llama::state_manifest(config); + let plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) + .map_err(|e| format!("llama: manifest planning failed: {e}"))?; + Ok((mesh, plan)) +} -/// Build the LLaMA GPU bundle from an HFQ or safetensors-directory source. -/// -/// The source/config path remains architecture-owned. Once it resolves, the -/// carrier publishes a pure Single manifest plan. Every fallible GPU stage -/// explicitly releases earlier allocations before returning an error; no -/// implicit GPU-buffer destructor is introduced. -pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { - let (config, weights, kv, scratch) = match src { - ModelSource::Hfq(mut hfq) => { - let config = - ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; - let weights = ::load_weights(&mut hfq, &config, ctx.gpu)?; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - // The plain LLaMA path has no independent cap resolver. PR #661's - // physical-cap behavior is owned by the existing upstream KV plan. - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" - )); - } - }; - let dims = KvDims { - layers: KvLayers::Flat(config.n_layers), - n_kv_heads: config.n_kv_heads, - head_dim: config.head_dim, - max_seq: ctx.max_seq, - physical_cap: None, - }; - let kv = match ::from_mode( - hipfire_runtime::kv_mode::resolve( - ctx.kv_mode_override.unwrap_or(""), - &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, - config.head_dim, - ) - .mode, - KvTarget::Single(ctx.gpu), - &dims, +fn llama_kv_dims(config: &LlamaConfig, max_seq: usize, physical_cap: Option) -> KvDims { + KvDims { + layers: KvLayers::Flat(config.n_layers), + n_kv_heads: config.n_kv_heads, + head_dim: config.head_dim, + max_seq, + physical_cap, + } +} + +fn hfq_layer_names(layer: usize, relative: &str) -> Vec { + vec![ + format!("model.layers.{layer}.{relative}.weight"), + format!("layers.{layer}.{relative}.weight"), + ] +} + +fn hfq_entry_names(entry: &WeightEntry) -> Result, String> { + let names = match (entry.name.as_str(), entry.layer) { + ("token_embd", None) => vec!["model.embed_tokens.weight".to_string()], + ("output_norm", None) => vec!["model.norm.weight".to_string()], + ("lm_head", None) => vec![ + "lm_head.weight".to_string(), + "model.lm_head.weight".to_string(), + "model.language_model.lm_head.weight".to_string(), + ], + ("wq", Some(layer)) => hfq_layer_names(layer, "self_attn.q_proj"), + ("wk", Some(layer)) => hfq_layer_names(layer, "self_attn.k_proj"), + ("wv", Some(layer)) => hfq_layer_names(layer, "self_attn.v_proj"), + ("wo", Some(layer)) => hfq_layer_names(layer, "self_attn.o_proj"), + ("ffn_gate", Some(layer)) => hfq_layer_names(layer, "mlp.gate_proj"), + ("ffn_up", Some(layer)) => hfq_layer_names(layer, "mlp.up_proj"), + ("ffn_down", Some(layer)) => hfq_layer_names(layer, "mlp.down_proj"), + ("attn_norm", Some(layer)) => hfq_layer_names(layer, "input_layernorm"), + ("ffn_norm", Some(layer)) => hfq_layer_names(layer, "post_attention_layernorm"), + ("q_norm", Some(layer)) => hfq_layer_names(layer, "self_attn.q_norm"), + ("k_norm", Some(layer)) => hfq_layer_names(layer, "self_attn.k_norm"), + (name, layer) => { + return Err(format!( + "llama: manifest entry {name}[layer {layer:?}] has no HFQ source mapping" + )); + } + }; + Ok(names) +} + +fn hfq_entry_data(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, u8), String> { + for name in hfq_entry_names(entry)? { + if let Some((info, data)) = hfq.tensor_data_vec(&name) { + if !matches!( + entry.name.as_str(), + "token_embd" | "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm" ) { - Ok(kv) => kv, - Err(error) => { - scratch.free_gpu(ctx.gpu); - weights.free_gpu(ctx.gpu); + let sidecar = match name.strip_suffix(".weight") { + Some(stem) => format!("{stem}.awq_scale.weight"), + None => format!("{name}.awq_scale.weight"), + }; + if hfq.find_tensor_info(&sidecar).is_some() { return Err(format!( - "llama: ::from_mode failed: {error}" + "llama: AWQ sidecar {sidecar} is not represented by the manifest pilot" )); } - }; - (config, weights, kv, scratch) + } + return Ok((data, info.quant_type)); + } + } + if entry.name == "lm_head" && entry.layer.is_none() { + if let Some((info, data)) = hfq.tensor_data_vec("model.embed_tokens.weight") { + return Ok((data, info.quant_type)); + } + } + Err(format!( + "llama: source tensor for {}[layer {:?}] is missing", + entry.name, entry.layer + )) +} + +fn f32_bytes_from_hfq(quant_type: u8, data: &[u8], name: &str) -> Result, String> { + let mut bytes = Vec::with_capacity(match quant_type { + 1 | 16 => data.len() * 2, + 2 => data.len(), + _ => 0, + }); + match quant_type { + 1 => { + let chunks = data.chunks_exact(2); + if !chunks.remainder().is_empty() { + return Err(format!("{name}: truncated F16 payload")); + } + for chunk in chunks { + bytes.extend_from_slice( + &hipfire_runtime::llama::f16_to_f32(u16::from_le_bytes([chunk[0], chunk[1]])) + .to_le_bytes(), + ); + } + } + 2 => { + if !data.len().is_multiple_of(4) { + return Err(format!("{name}: truncated F32 payload")); + } + bytes.extend_from_slice(data); } - ModelSource::Dir(source) => { - let config = - hipfire_runtime::hfq::config_from_safetensors_llama(&source).map_err(|e| { - format!("failed to parse LLaMA/Qwen3 config from config.json: {e}") - })?; - let weights = - hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) - .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - let kv_mode_str = ctx - .kv_mode_override - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); - let rr = hipfire_runtime::kv_mode::resolve( - &kv_mode_str, - &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, - config.head_dim, - ); - if let Some(w) = rr.warning { - eprintln!( - " KV cache: {w} (site {})", - hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site + 16 => { + let chunks = data.chunks_exact(2); + if !chunks.remainder().is_empty() { + return Err(format!("{name}: truncated BF16 payload")); + } + for chunk in chunks { + bytes.extend_from_slice( + &f32::from_bits( + u16::from_le_bytes([chunk[0], chunk[1]]) as u32 * (1 << 16), + ) + .to_le_bytes(), ); } - let dims = KvDims { - layers: KvLayers::Flat(config.n_layers), - n_kv_heads: config.n_kv_heads, - head_dim: config.head_dim, - max_seq: ctx.max_seq, - physical_cap: Some(ctx.max_seq), - }; - let kv = match ::from_mode( - rr.mode, - KvTarget::Single(ctx.gpu), - &dims, - ) { - Ok(kv) => kv, - Err(error) => { - weights.free_gpu(ctx.gpu); - return Err(format!("KvCache: {error}")); - } - }; - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - let _ = kv.free_gpu(ctx.gpu); - weights.free_gpu(ctx.gpu); - return Err(format!( - "ForwardScratch::new_with_max_seq: {error:?}" - )); - } - }; - (config, weights, kv, scratch) } - }; + other => { + return Err(format!( + "{name}: quant_type={other} is not a host float payload" + )); + } + } + Ok(bytes) +} - // Pure plan publication happens after source/config resolution and before - // the bundle becomes visible to the loader. It performs no GPU or file IO. - let mesh = DeviceMesh::single(); - let manifest = Llama::weight_manifest(&config); - let state = Llama::state_manifest(&config); - let manifest_plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) - .map_err(|e| format!("llama: manifest planning failed: {e}"))?; +fn hfq_source(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, DType), String> { + let (data, quant_type) = hfq_entry_data(hfq, entry)?; + let name = format!("{}[layer {:?}]", entry.name, entry.layer); + if entry.name == "token_embd" { + return match quant_type { + 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), + 3 => Ok((data, DType::Q8_0)), + 4 => Ok((data, DType::Q4K)), + 6 => Ok((data, DType::HFQ4G256)), + 7 => Ok((data, DType::HFQ4G128)), + other => Err(format!( + "{name}: quant_type={other} is unsupported for a LLaMA embedding" + )), + }; + } + if matches!(entry.name.as_str(), "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm") + { + return Ok(( + f32_bytes_from_hfq(quant_type, &data, &name)?, + DType::F32, + )); + } + match quant_type { + 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), + other => hfq_weight_dtype(other) + .map(|dtype| (data, dtype)) + .ok_or_else(|| format!("{name}: unsupported HFQ quant_type={other}")), + } +} - Ok(LlamaBundle { +fn take_slot( + assembly: &mut WeightStoreAssembly<'_>, + slots: &mut HashMap<(String, Option), usize>, + name: &str, + layer: Option, +) -> Result<(), String> { + let slot = assembly + .take(name, layer, 0) + .ok_or_else(|| format!("llama: fulfilled store is missing {name}[layer {layer:?}]"))?; + slots.insert((name.to_string(), layer), slot); + Ok(()) +} + +fn require_resident( + assembly: &WeightStoreAssemblyGuard<'_>, + name: &str, + layer: Option, + slot: usize, +) -> Result<(), String> { + if matches!(assembly.get(slot), Some(WeightHandle::Resident(_))) { + Ok(()) + } else { + Err(format!( + "llama: {name}[layer {layer:?}] is an alias; typed LLaMA assembly requires a resident handle" + )) + } +} + +fn resident_cell( + cells: &mut HashMap<(String, Option), TakenWeight>, + name: &str, + layer: Option, +) -> Result { + let taken = cells + .remove(&(name.to_string(), layer)) + .ok_or_else(|| format!("llama: assembled store is missing {name}[layer {layer:?}]"))?; + match taken.handle { + WeightHandle::Resident(tensor) => Ok(tensor), + WeightHandle::Alias(source) => Err(format!( + "llama: {name}[layer {layer:?}] aliases {source}, expected resident handle" + )), + } +} + +fn resident_weight( + cells: &mut HashMap<(String, Option), TakenWeight>, + name: &str, + layer: Option, + m: usize, + k: usize, +) -> Result { + let tensor = resident_cell(cells, name, layer)?; + let dtype = tensor.dtype; + Ok(WeightTensor { + buf: tensor, + gpu_dtype: dtype, + m, + k, + row_stride: dtype.row_stride(k), + paro: None, + awq_scale: None, + }) +} + +fn embedding_format(dtype: DType) -> Result { + match dtype { + DType::F32 => Ok(EmbeddingFormat::F32), + DType::Q4K => Ok(EmbeddingFormat::Q4K), + DType::HFQ4G256 => Ok(EmbeddingFormat::HFQ4G256), + DType::HFQ4G128 => Ok(EmbeddingFormat::HFQ4G128), + DType::Q8_0 => Ok(EmbeddingFormat::Q8_0), + other => Err(format!( + "llama: unsupported assembled embedding dtype {other:?}" + )), + } +} + +fn assemble_llama_weights( + config: &LlamaConfig, + store: &mut WeightStore, +) -> Result { + let mut assembly = store.begin_assembly(); + let mut slots = HashMap::new(); + let mut take = |name: &str, layer: Option| take_slot(&mut assembly, &mut slots, name, layer); + + take("token_embd", None)?; + take("output_norm", None)?; + take("lm_head", None)?; + for layer in 0..config.n_layers { + for name in ["wq", "wk", "wv", "wo", "ffn_gate", "ffn_up", "ffn_down", "attn_norm", "ffn_norm"] { + take(name, Some(layer))?; + } + if config.has_qk_norm { + take("q_norm", Some(layer))?; + take("k_norm", Some(layer))?; + } + } + + drop(take); + let guard = assembly.commit(); + for ((name, layer), slot) in &slots { + require_resident(&guard, name, *layer, *slot)?; + } + let cells: HashMap<_, _> = guard + .finalize() + .into_iter() + .map(|taken| ((taken.key.name.clone(), taken.key.layer), taken)) + .collect(); + let mut cells = cells; + let token_embd = resident_cell(&mut cells, "token_embd", None)?; + let embd_format = embedding_format(token_embd.dtype)?; + let output_norm = resident_cell(&mut cells, "output_norm", None)?; + let output = resident_weight( + &mut cells, + "lm_head", + None, + config.vocab_size, + config.dim, + )?; + let mut layers = Vec::with_capacity(config.n_layers); + for layer in 0..config.n_layers { + let q_norm = if config.has_qk_norm { + Some(resident_cell(&mut cells, "q_norm", Some(layer))?) + } else { + None + }; + let k_norm = if config.has_qk_norm { + Some(resident_cell(&mut cells, "k_norm", Some(layer))?) + } else { + None + }; + layers.push(LayerWeights { + attn_norm: resident_cell(&mut cells, "attn_norm", Some(layer))?, + wq: resident_weight( + &mut cells, + "wq", + Some(layer), + config.n_heads * config.head_dim, + config.dim, + )?, + wk: resident_weight( + &mut cells, + "wk", + Some(layer), + config.n_kv_heads * config.head_dim, + config.dim, + )?, + wv: resident_weight( + &mut cells, + "wv", + Some(layer), + config.n_kv_heads * config.head_dim, + config.dim, + )?, + wo: resident_weight( + &mut cells, + "wo", + Some(layer), + config.dim, + config.n_heads * config.head_dim, + )?, + q_norm, + k_norm, + ffn_norm: resident_cell(&mut cells, "ffn_norm", Some(layer))?, + w_gate: resident_weight( + &mut cells, + "ffn_gate", + Some(layer), + config.hidden_dim, + config.dim, + )?, + w_up: resident_weight( + &mut cells, + "ffn_up", + Some(layer), + config.hidden_dim, + config.dim, + )?, + w_down: resident_weight( + &mut cells, + "ffn_down", + Some(layer), + config.dim, + config.hidden_dim, + )?, + }); + } + Ok(LlamaWeights { + token_embd, + embd_format, + output_norm, + output, + layers, + lm_head_aliases_embd: false, + }) +} + +/// Build the LLaMA GPU bundle from an HFQ or safetensors-directory source. +/// +/// The HFQ plain-LLaMA Single path is the production manifest pilot: planning +/// and source admission happen first, fulfillment uploads transactionally, and +/// typed handles are moved into `LlamaWeights` before the committed remainder +/// is published beneath this bundle's owner. The directory path remains on its +/// existing ParoQuant loader until that source has an equivalent representation +/// resolver. +pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { + let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = + match src { + ModelSource::Hfq(hfq) => { + let config = + ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; + let (mesh, manifest_plan) = plan_single(&config)?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let mut store = hipfire_runtime::weight_store::fulfill_manifest( + &Llama::weight_manifest(&config), + &mesh, + config.n_layers, + ctx.gpu, + |entry| hfq_source(&hfq, entry), + ) + .map_err(|e| format!("llama: {e}"))?; + let weights = match assemble_llama_weights(&config, &mut store) { + Ok(weights) => weights, + Err(error) => { + store.rollback_unpublished(ctx.gpu); + return Err(error); + } + }; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + // The plain LLaMA path has no independent cap resolver. PR + // #661's physical-cap behavior is owned by the existing + // upstream KV plan. + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + store.rollback_unpublished(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + )); + } + }; + let dims = llama_kv_dims(&config, ctx.max_seq, None); + let kv = match ::from_mode( + hipfire_runtime::kv_mode::resolve( + ctx.kv_mode_override.unwrap_or(""), + &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, + config.head_dim, + ) + .mode, + KvTarget::Single(ctx.gpu), + &dims, + ) { + Ok(kv) => kv, + Err(error) => { + scratch.free_gpu(ctx.gpu); + store.rollback_unpublished(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ::from_mode failed: {error}" + )); + } + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + Some(store), + mesh, + weight_origin, + ) + } + ModelSource::Dir(source) => { + let config = + hipfire_runtime::hfq::config_from_safetensors_llama(&source).map_err(|e| { + format!("failed to parse LLaMA/Qwen3 config from config.json: {e}") + })?; + let (mesh, manifest_plan) = plan_single(&config)?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let weights = + hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) + .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + let kv_mode_str = ctx + .kv_mode_override + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); + let rr = hipfire_runtime::kv_mode::resolve( + &kv_mode_str, + &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, + config.head_dim, + ); + if let Some(w) = rr.warning { + eprintln!( + " KV cache: {w} (site {})", + hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site + ); + } + let dims = llama_kv_dims(&config, ctx.max_seq, Some(ctx.max_seq)); + let kv = match ::from_mode( + rr.mode, + KvTarget::Single(ctx.gpu), + &dims, + ) { + Ok(kv) => kv, + Err(error) => { + weights.free_gpu(ctx.gpu); + return Err(format!("KvCache: {error}")); + } + }; + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let _ = kv.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!( + "ForwardScratch::new_with_max_seq: {error:?}" + )); + } + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + None, + mesh, + weight_origin, + ) + } + }; + + let mut bundle = LlamaBundle { config, weights, scratch, kv, manifest_plan, weight_store: None, + weight_origin, mesh, dflash_extract_layers: Vec::new(), dspark_weights: None, dspark_assets: None, - }) + }; + if let Some(store) = weight_store { + if let Err((store, error)) = bundle.attach_weight_store(store) { + let LlamaBundle { + weights, + scratch, + kv, + .. + } = bundle; + store.rollback_unpublished(ctx.gpu); + scratch.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + let _ = kv.free_gpu(ctx.gpu); + return Err(error); + } + } + Ok(bundle) } /// Alias matching the `load__bundle` naming convention in the task. @@ -178,8 +589,9 @@ pub use load_bundle as load_llama_bundle; impl LlamaBundle { /// Attach a store whose resident handles have been assembled for this - /// bundle. The target mesh epoch is checked before publication; teardown - /// remains exclusively in `ArchModel::free_gpu`. On rejection the store is + /// bundle. The complete target origin (mesh epoch, logical rank, and + /// physical device) is checked before publication; teardown remains + /// exclusively in `ArchModel::free_gpu`. On rejection the store is /// returned unchanged so the caller can retry against the right owner. pub fn attach_weight_store( &mut self, @@ -188,18 +600,8 @@ impl LlamaBundle { if self.weight_store.is_some() { return Err((store, "llama: weight store already attached".into())); } - let Some(origin) = store.origin() else { - return Err((store, "llama: weight store has no origin".into())); - }; - if origin.mesh_epoch() != self.mesh.epoch() { - return Err(( - store, - format!( - "llama: weight store origin epoch {:?} does not match bundle epoch {:?}", - origin.mesh_epoch(), - self.mesh.epoch() - ), - )); + if let Err(error) = store.validate_origin_value(self.weight_origin) { + return Err((store, format!("llama: weight store origin rejected: {error}"))); } self.weight_store = Some(store); Ok(()) @@ -224,3 +626,107 @@ impl LlamaBundle { self.dflash_extract_layers = layers; } } + +#[cfg(test)] +mod tests { + use super::*; + use hipfire_runtime::llama::ModelArch; + use hipfire_runtime::weight_store::{WeightProjection, WeightProjectionKind}; + + fn config() -> LlamaConfig { + LlamaConfig { + arch: ModelArch::Llama, + dim: 4, + hidden_dim: 8, + n_layers: 1, + n_heads: 1, + n_kv_heads: 1, + vocab_size: 8, + head_dim: 4, + norm_eps: 1e-5, + max_seq_len: 32, + rope_freq_base: 10_000.0, + bos_token: 1, + eos_token: 2, + has_qk_norm: false, + } + } + + fn alias_projection() -> WeightProjection { + WeightProjection { + kind: WeightProjectionKind::Static, + axis: None, + rank: 0, + world_size: 1, + logical_shape: vec![1], + dtype: DType::F32, + } + } + + #[test] + fn single_plan_covers_every_typed_llama_handle() { + let (mesh, plan) = plan_single(&config()).unwrap(); + let manifest = Llama::weight_manifest(&config()); + assert_eq!(mesh.n_devices(), 1); + assert_eq!(plan.weights.len(), 12); + assert_eq!(plan.state.len(), 1); + assert!(plan.collective_schedule.iter().any(|entry| entry.name == "wo")); + assert!(manifest[0].dtype_constraint.accepts(DType::HFQ4G256)); + assert!(manifest[1].dtype_constraint.accepts(DType::MQ4G256)); + assert!(manifest[9].dtype_constraint.accepts(DType::F32)); + assert!(!manifest[9].dtype_constraint.accepts(DType::F16)); + } + + #[test] + fn typed_assembly_rolls_back_when_a_cell_is_not_resident() { + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); + let mut store = WeightStore::with_origin(origin); + for name in ["token_embd", "output_norm", "lm_head"] { + store + .stage_alias(name, None, 0, "source", alias_projection()) + .unwrap(); + } + let error = match assemble_llama_weights( + &LlamaConfig { + n_layers: 0, + ..config() + }, + &mut store, + ) { + Ok(_) => panic!("alias unexpectedly assembled as typed weights"), + Err(error) => error, + }; + assert!(error.contains("alias")); + assert_eq!(store.len(), 3); + assert!(store.contains("token_embd", None, 0)); + assert!(store.projection("lm_head", None, 0).is_some()); + } + + + #[test] + fn hfq_float_widening_matches_legacy_f32_representation() { + let f16_one = [0x00, 0x3c, 0x00, 0xc0]; + let actual = f32_bytes_from_hfq(1, &f16_one, "test").unwrap(); + let expected = [1.0f32, -2.0f32] + .into_iter() + .flat_map(f32::to_le_bytes) + .collect::>(); + assert_eq!(actual, expected); + } + + #[test] + fn manifest_constraints_admit_every_pilot_representation() { + let manifest = Llama::weight_manifest(&config()); + assert!(manifest[0].dtype_constraint.accepts(DType::HFQ4G256)); + assert!(manifest[1].dtype_constraint.accepts(DType::MQ4G256)); + assert!(manifest[9].dtype_constraint.accepts(DType::F32)); + assert!(!manifest[9].dtype_constraint.accepts(DType::F16)); + } + #[test] + fn physical_cap_remains_separate_from_configured_max_seq() { + let dims = llama_kv_dims(&config(), 32_768, Some(4_096)); + assert_eq!(dims.max_seq, 32_768); + assert_eq!(dims.physical_cap, Some(4_096)); + } +} diff --git a/crates/hipfire-runtime/src/weight_backend.rs b/crates/hipfire-runtime/src/weight_backend.rs index 43a7d09dc..b06a6076c 100644 --- a/crates/hipfire-runtime/src/weight_backend.rs +++ b/crates/hipfire-runtime/src/weight_backend.rs @@ -470,6 +470,15 @@ pub(crate) fn raw_codec(quant_type: u8) -> Option<&'static RawCodec> { RAW_CODECS.iter().find(|c| c.quant_type == quant_type) } +/// Return the compute representation used for a raw HFQ weight payload. +/// +/// Host-decoded source types (F16/F32/BF16) intentionally return `None`; +/// callers must widen those payloads before upload so the result matches the +/// established LLaMA loader semantics. +pub fn hfq_weight_dtype(quant_type: u8) -> Option { + raw_codec(quant_type).map(|codec| codec.dtype) +} + /// Decode a passthrough quant format: enforce the K%256 guard (via DType), /// upload bytes verbatim, build the `WeightTensor` with the dtype + its /// DType-derived row_stride. `name` is the caller context for the guard panic. diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 26c64ad2e..4e9f91023 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -364,6 +364,23 @@ fn validate_shape(entry: &WeightEntry) -> Result<(), String> { Ok(()) } +pub(crate) fn validate_weight_layers( + manifest: &[WeightEntry], + n_layers: usize, +) -> Result<(), String> { + for entry in manifest { + if let Some(layer) = entry.layer { + if layer >= n_layers { + return Err(format!( + "{} layer {} outside n_layers={n_layers}", + entry.name, layer + )); + } + } + } + Ok(()) +} + /// Validate logical shard math and tied source identity before fulfillment. pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result<(), String> { let mut identities = HashSet::new(); @@ -431,6 +448,30 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< if source_entry.identity() == entry.identity() { return Err(format!("{context}: an entry cannot tie to itself")); } + if source_entry.logical_shape != entry.logical_shape { + return Err(format!( + "{context}: tied source '{source}' shape {:?} does not match {:?}", + source_entry.logical_shape, entry.logical_shape + )); + } + if source_entry.dtype != entry.dtype { + return Err(format!( + "{context}: tied source '{source}' dtype {:?} does not match {:?}", + source_entry.dtype, entry.dtype + )); + } + if !entry.dtype_constraint.accepts(source_entry.dtype) + || !source_entry.dtype_constraint.accepts(entry.dtype) + { + return Err(format!( + "{context}: tied source '{source}' violates the source dtype contract" + )); + } + if matches!(&source_entry.policy, ShardPolicy::Tied { .. }) { + return Err(format!( + "{context}: tied source '{source}' is itself tied; chains and cycles are unsupported" + )); + } } ShardPolicy::ExpertSharded { n_experts, .. } => { if *n_experts == 0 || entry.logical_shape.first() != Some(n_experts) { @@ -486,6 +527,7 @@ pub fn plan_manifest( mesh: &DeviceMesh, n_layers: usize, ) -> Result { + validate_weight_layers(weights, n_layers)?; validate_manifest(weights, mesh)?; let mut state_ids = HashSet::new(); for entry in state { @@ -1002,4 +1044,91 @@ mod tests { }; assert!(validate_expert_group_specs(&[bad], &manifest).is_err()); } + + #[test] + fn planning_rejects_weight_layer_at_n_layers_and_accepts_last_layer() { + let mesh = DeviceMesh::single(); + let valid = layer_entry("w", 2, ShardPolicy::Replicate); + assert!(plan_manifest(&[valid], &[], &mesh, 3).is_ok()); + let out_of_range = layer_entry("w", 3, ShardPolicy::Replicate); + let error = plan_manifest(&[out_of_range], &[], &mesh, 3).unwrap_err(); + assert!(error.contains("outside n_layers=3")); + } + + #[test] + fn tied_entries_require_matching_representation_and_no_tied_chain() { + let source = WeightEntry::model( + "source", + vec![8, 8], + DType::F16, + ShardPolicy::Replicate, + ); + let shape_mismatch = WeightEntry::model( + "shape_mismatch", + vec![8, 4], + DType::F16, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + assert!(validate_manifest( + &[source.clone(), shape_mismatch], + &DeviceMesh::single() + ) + .is_err()); + + let dtype_mismatch = WeightEntry::model( + "dtype_mismatch", + vec![8, 8], + DType::F32, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + assert!(validate_manifest( + &[source.clone(), dtype_mismatch], + &DeviceMesh::single() + ) + .is_err()); + + let chained_source = WeightEntry::model( + "chained_source", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + let chain = WeightEntry::model( + "chain", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "chained_source".into(), + }, + ); + assert!(validate_manifest( + &[source, chained_source, chain], + &DeviceMesh::single() + ) + .is_err()); + + let cycle_a = WeightEntry::model( + "cycle_a", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "cycle_b".into(), + }, + ); + let cycle_b = WeightEntry::model( + "cycle_b", + vec![8, 8], + DType::F16, + ShardPolicy::Tied { + source: "cycle_a".into(), + }, + ); + assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single()).is_err()); + } } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index a66336985..80b8b9bd8 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -11,16 +11,21 @@ //! //! The store is not a model owner. It has no `Drop` implementation and never //! frees GPU buffers implicitly. A carrier may move a committed store into its -//! existing `ArchModel` owner; that owner must call [`WeightStore::release_on_owner`] -//! during its existing teardown path. `take` transfers a resident handle to the -//! owner that is assembling typed weights, and therefore removes the cell from -//! the store's cleanup set. +//! existing `ArchModel` owner; that owner must transfer its resident handles +//! through [`WeightStore::take_all`] during the existing teardown path. +//! `take` transfers a resident handle to the owner that is assembling typed +//! weights, and therefore removes the cell from the store's cleanup set. use crate::weight_manifest::{placement_devices, ShardPolicy, WeightEntry}; use hipfire_hardware::{DeviceMesh, MeshEpoch}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::collections::HashMap; +#[cfg(test)] +thread_local! { + static RESIDENT_RELEASES: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + /// Stable logical placement identity. Layer is part of the key because a /// per-layer name such as `wq` appears once for every decoder block. #[derive(Clone, PartialEq, Eq, Hash, Debug)] @@ -199,8 +204,8 @@ impl std::error::Error for FulfillError {} /// `(name, layer, device)` and captures the target origin once. /// /// There is intentionally no `Drop` implementation. A `WeightStore` that is -/// abandoned without explicit rollback/release leaks rather than guessing a -/// GPU owner; production callers keep it beneath `ArchModel`. +/// abandoned without explicit rollback or owner transfer leaks rather than +/// guessing a GPU owner; production callers keep it beneath `ArchModel`. #[derive(Default)] pub struct WeightStore { placements: HashMap, @@ -355,21 +360,25 @@ impl WeightStore { self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) } - /// Explicit owner teardown. This method is intentionally consuming and has - /// no implicit/drop fallback; `ArchModel::free_gpu` is the production call - /// site. Callers must validate the target first with [`validate_origin`]. - /// On mismatch the original store is returned unchanged for retry by the - /// owner; no GPU call occurs. - pub fn release_on_owner( + /// Roll back a fulfilled store before it is published beneath a model + /// owner. This is the only public consuming GPU-free operation: callers + /// may use it while a load transaction is still unpublished, but an + /// attached store can only be drained by the model owner via `take_all`. + pub fn rollback_unpublished(self, gpu: &Gpu) { + self.release_unchecked(gpu); + } + + /// Transfer every resident/alias handle to the model owner after checking + /// the complete captured origin. On mismatch, the original store is + /// returned unchanged so the owner can retry against the correct target. + pub fn take_all( self, - mesh: &DeviceMesh, - gpu: &mut Gpu, - ) -> Result<(), (Self, WeightStoreError)> { - if let Err(error) = self.validate_origin(mesh, gpu) { + expected: WeightOrigin, + ) -> Result, (Self, WeightStoreError)> { + if let Err(error) = self.validate_origin_value(expected) { return Err((self, error)); } - self.release_unchecked(gpu); - Ok(()) + Ok(self.placements.into_values().collect()) } /// Explicit rollback for a failed transaction. It consumes the partial @@ -385,6 +394,8 @@ impl WeightStore { // the existing loader's explicit owner teardown. The store // never relies on a destructor to release GPU memory. let _ = gpu.hip.free(tensor.buf); + #[cfg(test)] + RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); } } } @@ -496,7 +507,9 @@ where if let Some(error) = target_error(mesh) { return Err(error); } - if let Err(reason) = crate::weight_manifest::validate_manifest(weights, mesh) { + if let Err(reason) = crate::weight_manifest::validate_weight_layers(weights, n_layers) + .and_then(|_| crate::weight_manifest::validate_manifest(weights, mesh)) + { return Err(FulfillError { name: "".to_string(), layer: None, @@ -565,6 +578,27 @@ where ), }); } + if matches!(dtype, DType::F32 | DType::F16 | DType::BF16) { + let expected_bytes = entry + .logical_shape + .iter() + .try_fold(1usize, |count, &dim| count.checked_mul(dim)) + .and_then(|elements| elements.checked_mul(dtype.size())); + if expected_bytes != Some(bytes.len()) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "source payload has {} bytes, expected {:?} for {dtype:?} {:?}", + bytes.len(), + expected_bytes, + entry.logical_shape + ), + }); + } + } let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { Ok(tensor) => tensor, Err(error) => { @@ -721,4 +755,186 @@ mod tests { assert!(target_error(&mesh).is_some()); assert_eq!(placement_devices(&entry, &mesh, 1), vec![0]); } + + #[test] + fn successful_single_fulfillment_commits_resident_projection() { + let Ok(gpu) = Gpu::init() else { + return; + }; + let mesh = DeviceMesh::single(); + let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); + let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], DType::F32)) + }) + .unwrap(); + assert_eq!(store.len(), 1); + assert!(matches!( + store.get("resident", None, 0), + Some(WeightHandle::Resident(tensor)) if tensor.dtype == DType::F32 + )); + assert_eq!( + store.projection("resident", None, 0).unwrap().dtype, + DType::F32 + ); + store.rollback_unpublished(&gpu); + } + + #[test] + fn full_origin_mismatch_returns_resident_store_unchanged() { + let first = DeviceMesh::single(); + let second = DeviceMesh::single(); + let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); + let expected = WeightOrigin::from_parts(second.epoch(), 4, 12); + let mut store = WeightStore::with_origin(actual); + store + .stage_alias("resident", None, 0, "source", projection(DType::F16)) + .unwrap(); + let (store, error) = match store.take_all(expected) { + Ok(_) => panic!("origin mismatch unexpectedly succeeded"), + Err(value) => value, + }; + assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); + assert_eq!(store.origin(), Some(actual)); + assert!(store.contains("resident", None, 0)); + assert!(store.projection("resident", None, 0).is_some()); + } + + #[test] + fn full_origin_mismatch_does_not_free_a_resident_store() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); + let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], DType::F32)) + }) + .unwrap(); + let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); + let (store, error) = match store.take_all(expected) { + Ok(_) => panic!("origin mismatch unexpectedly succeeded"), + Err(value) => value, + }; + assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); + assert_eq!(store.len(), 1); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + 0, + "origin rejection must not free resident buffers" + ); + store.rollback_unpublished(&gpu); + assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); + } + + #[test] + fn owner_transfer_is_consuming_and_empty_transfer_is_idempotent() { + let mesh = DeviceMesh::single(); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); + let mut store = WeightStore::with_origin(origin); + store + .stage_alias("owned", None, 0, "source", projection(DType::F16)) + .unwrap(); + let handles = store.take_all(origin).unwrap(); + assert_eq!(handles.len(), 1); + assert!(matches!( + handles.into_iter().next(), + Some(WeightHandle::Alias(_)) + )); + let second = WeightStore::with_origin(origin).take_all(origin).unwrap(); + assert!(second.is_empty()); + } + + #[test] + fn source_failure_after_resident_upload_rolls_back_everything() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let entries = vec![ + WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), + WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), + ]; + let error = fulfill_manifest_single(&entries, &mesh, 1, &gpu, |entry| { + if entry.name == "first" { + Ok((vec![0; 4], DType::F32)) + } else { + Err("injected source failure".into()) + } + }) + .unwrap_err(); + assert_eq!(error.name, "second"); + assert!(error.reason.contains("source read failed")); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + 1, + "the first resident allocation must be explicitly freed" + ); + } + + #[test] + fn dtype_failure_after_resident_upload_rolls_back_everything() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let constraint = DTypeConstraint::source_exact(DType::F32); + let entries = vec![ + WeightEntry::model_with_dtype_constraint( + "first", + vec![1], + DType::F32, + constraint.clone(), + ShardPolicy::Replicate, + ), + WeightEntry::model_with_dtype_constraint( + "second", + vec![1], + DType::F32, + constraint, + ShardPolicy::Replicate, + ), + ]; + let error = fulfill_manifest_single(&entries, &mesh, 1, &gpu, |entry| { + if entry.name == "first" { + Ok((vec![0; 4], DType::F32)) + } else { + Ok((vec![0; 2], DType::F16)) + } + }) + .unwrap_err(); + assert_eq!(error.name, "second"); + assert!(error.reason.contains("violates constraint")); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + 1, + "the first resident allocation must be explicitly freed" + ); + } + + #[test] + fn malformed_upload_payload_after_resident_allocation_rolls_back() { + let Ok(gpu) = Gpu::init() else { + return; + }; + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single(); + let entries = vec![ + WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), + WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), + ]; + let error = fulfill_manifest_single(&entries, &mesh, 1, &gpu, |entry| { + if entry.name == "first" { + Ok((vec![0; 4], DType::F32)) + } else { + Ok((vec![0; 1], DType::F32)) + } + }) + .unwrap_err(); + assert_eq!(error.name, "second"); + assert!(error.reason.contains("payload")); + assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); + } } From 6cee389b882a0ff3536a75fd8ca8c6fe86df9032 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 19:36:49 +0200 Subject: [PATCH 04/17] fix(device-mesh): close llama manifest pilot review gaps --- crates/hipfire-arch-llama/src/arch.rs | 28 + crates/hipfire-arch-llama/src/arch_model.rs | 59 +- crates/hipfire-arch-llama/src/carrier.rs | 582 ++++++++++++++---- crates/hipfire-runtime/src/hfq.rs | 53 +- crates/hipfire-runtime/src/model_load.rs | 14 +- crates/hipfire-runtime/src/weight_manifest.rs | 46 +- crates/hipfire-runtime/src/weight_store.rs | 254 +++++--- 7 files changed, 790 insertions(+), 246 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index f154f0c1b..27c60aa9e 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -262,6 +262,34 @@ impl Llama { manifest } + /// Build the manifest for an HFQ source after source classification. + /// + /// A separate `lm_head.weight` is a resident output projection. When the + /// source omits it, the declaration is a true tie to `token_embd`; the + /// output placement remains pinned to the final stage while the source + /// representation contract is copied from the embedding entry. + pub fn weight_manifest_for_hfq( + cfg: &LlamaConfig, + has_separate_lm_head: bool, + ) -> Vec { + let mut manifest = Self::weight_manifest(cfg); + if !has_separate_lm_head { + let embedding_constraint = manifest + .first() + .expect("LLaMA manifest always contains token_embd") + .dtype_constraint + .clone(); + let output = manifest + .last_mut() + .expect("LLaMA manifest always contains lm_head"); + output.dtype_constraint = embedding_constraint; + output.policy = ShardPolicy::Tied { + source: "token_embd".into(), + }; + } + manifest + } + /// Pure state declaration for the full-attention LLaMA family. pub fn state_manifest(cfg: &LlamaConfig) -> Vec { (0..cfg.n_layers) diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 9d1297df9..763dbbb29 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -4,26 +4,10 @@ use hipfire_runtime::arch_model::ArchModel; use hipfire_runtime::llama::KvCache; -use hipfire_runtime::weight_store::{ - WeightHandle, WeightOrigin, WeightStore, WeightStoreError, -}; use rdna_compute::Gpu; use crate::carrier::LlamaBundle; -fn drain_weight_store( - store: WeightStore, - origin: WeightOrigin, - gpu: &mut Gpu, -) -> Result<(), (WeightStore, WeightStoreError)> { - for handle in store.take_all(origin)? { - if let WeightHandle::Resident(tensor) = handle { - let _ = gpu.free_tensor(tensor); - } - } - Ok(()) -} - impl ArchModel for LlamaBundle { fn dim(&self) -> usize { self.config.dim @@ -51,17 +35,6 @@ impl ArchModel for LlamaBundle { } fn free_gpu(self: Box, gpu: &mut Gpu) { - // Validate before destructuring the consuming owner. A mismatch must - // leave the resident store attached to an owner that can be retried; - // leaking the boxed owner is safer than dropping the only cleanup - // authority. - if let Some(store) = self.weight_store.as_ref() { - if let Err(error) = store.validate_origin_value(self.weight_origin) { - eprintln!("llama: refusing weight-store release: {error}"); - let _ = Box::into_raw(self); - return; - } - } let LlamaBundle { config: _, weights, @@ -69,7 +42,7 @@ impl ArchModel for LlamaBundle { kv, manifest_plan: _, weight_store, - weight_origin, + weight_origin: _, mesh: _, dflash_extract_layers: _, dspark_weights: _, @@ -78,12 +51,10 @@ impl ArchModel for LlamaBundle { // Mirror the existing unload ordering: scratch → store/weights → kv. scratch.free_gpu(gpu); if let Some(store) = weight_store { - if let Err((store, error)) = drain_weight_store(store, weight_origin, gpu) { - // This is defensive because the pre-check above used the - // same immutable origin. Never discard a rejected store. - eprintln!("llama: refusing weight-store release: {error}"); - std::mem::forget(store); - } + // Attachment already checked the complete origin and created this + // owner capability. There is no mismatch branch to leak the model: + // an attached store can only be drained by this consuming owner. + store.drain(gpu); } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); @@ -95,21 +66,27 @@ mod tests { use super::*; use hipfire_hardware::DeviceMesh; use hipfire_runtime::weight_manifest::{ShardPolicy, WeightEntry}; - use hipfire_runtime::weight_store::fulfill_manifest_single; + use hipfire_runtime::weight_store::{ + fulfill_manifest_single, WeightLoadTransaction, WeightOrigin, WeightStore, + }; #[test] - fn arch_owner_unload_drains_residents_and_repeated_empty_unload_is_safe() { - let Ok(mut gpu) = Gpu::init() else { + fn attached_owner_drain_is_consuming_and_empty_drain_is_safe() { + let Ok(gpu) = Gpu::init() else { return; }; let mesh = DeviceMesh::single(); let origin = WeightOrigin::for_single(&mesh, &gpu); - let entry = WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); - let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + let entry = + WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); + let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], rdna_compute::DType::F32)) }) .unwrap(); - assert!(drain_weight_store(store, origin, &mut gpu).is_ok()); - assert!(drain_weight_store(WeightStore::with_origin(origin), origin, &mut gpu).is_ok()); + transaction.publish(origin).unwrap().drain(&gpu); + WeightLoadTransaction::new(WeightStore::with_origin(origin)) + .publish(origin) + .unwrap() + .drain(&gpu); } } diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 86b4dcfe5..5a8b5ebe1 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -15,10 +15,9 @@ use hipfire_runtime::llama::{ }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; use hipfire_runtime::weight_backend::hfq_weight_dtype; -use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_store::{ - TakenWeight, WeightHandle, WeightStore, WeightStoreAssembly, WeightStoreAssemblyGuard, - WeightOrigin, + AttachedWeightStore, TakenWeight, WeightHandle, WeightLoadTransaction, + WeightStoreAssembly, WeightStoreAssemblyGuard, WeightOrigin, }; use rdna_compute::{DType, GpuTensor}; use std::collections::HashMap; @@ -34,11 +33,11 @@ pub struct LlamaBundle { /// A pilot store is attached only after its handles are assembled under /// this bundle. It is crate-visible so callers cannot create an independent /// unload owner; `ArchModel::free_gpu` is the sole release path. - pub(crate) weight_store: Option, - /// Exact target identity captured when this bundle was admitted. The - /// owner uses it to validate every store-origin component before teardown. + pub(crate) weight_store: Option, + /// Exact target identity captured before publication. The attached store + /// binds this identity into its private drain capability, so teardown + /// cannot encounter an origin mismatch. pub(crate) weight_origin: WeightOrigin, - pub(crate) mesh: DeviceMesh, /// Decoder-layer indices whose residual hidden states a hidden-conditioned /// drafter (DFlash / EAGLE) wants captured, ascending order. Empty = no /// capture (the `SpecTarget::dflash_extract_layers` default of `None`). The @@ -53,9 +52,12 @@ pub struct LlamaBundle { /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } -fn plan_single(config: &LlamaConfig) -> Result<(DeviceMesh, ManifestPlan), String> { +fn plan_single( + config: &LlamaConfig, + has_separate_lm_head: bool, +) -> Result<(DeviceMesh, ManifestPlan), String> { let mesh = DeviceMesh::single(); - let manifest = Llama::weight_manifest(config); + let manifest = Llama::weight_manifest_for_hfq(config, has_separate_lm_head); let state = Llama::state_manifest(config); let plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) .map_err(|e| format!("llama: manifest planning failed: {e}"))?; @@ -230,18 +232,25 @@ fn take_slot( Ok(()) } -fn require_resident( +fn require_materialized( assembly: &WeightStoreAssemblyGuard<'_>, name: &str, layer: Option, slot: usize, ) -> Result<(), String> { - if matches!(assembly.get(slot), Some(WeightHandle::Resident(_))) { - Ok(()) - } else { - Err(format!( - "llama: {name}[layer {layer:?}] is an alias; typed LLaMA assembly requires a resident handle" - )) + match assembly.get(slot) { + Some(WeightHandle::Resident(_)) => Ok(()), + Some(WeightHandle::Alias(source)) + if name == "lm_head" && layer.is_none() && source == "token_embd" => + { + Ok(()) + } + Some(WeightHandle::Alias(source)) => Err(format!( + "llama: {name}[layer {layer:?}] aliases {source}; only lm_head may tie token_embd" + )), + None => Err(format!( + "llama: {name}[layer {layer:?}] assembly slot {slot} is missing" + )), } } @@ -249,15 +258,13 @@ fn resident_cell( cells: &mut HashMap<(String, Option), TakenWeight>, name: &str, layer: Option, -) -> Result { - let taken = cells - .remove(&(name.to_string(), layer)) - .ok_or_else(|| format!("llama: assembled store is missing {name}[layer {layer:?}]"))?; - match taken.handle { - WeightHandle::Resident(tensor) => Ok(tensor), - WeightHandle::Alias(source) => Err(format!( - "llama: {name}[layer {layer:?}] aliases {source}, expected resident handle" - )), +) -> GpuTensor { + match cells.remove(&(name.to_string(), layer)) { + Some(TakenWeight { + handle: WeightHandle::Resident(tensor), + .. + }) => tensor, + _ => unreachable!("validated LLaMA assembly lost resident {name}[layer {layer:?}]"), } } @@ -267,10 +274,10 @@ fn resident_weight( layer: Option, m: usize, k: usize, -) -> Result { - let tensor = resident_cell(cells, name, layer)?; +) -> WeightTensor { + let tensor = resident_cell(cells, name, layer); let dtype = tensor.dtype; - Ok(WeightTensor { + WeightTensor { buf: tensor, gpu_dtype: dtype, m, @@ -278,7 +285,27 @@ fn resident_weight( row_stride: dtype.row_stride(k), paro: None, awq_scale: None, - }) + } +} + +fn tied_weight( + cells: &mut HashMap<(String, Option), TakenWeight>, + token_embd: &GpuTensor, + embd_format: EmbeddingFormat, + name: &str, + layer: Option, + m: usize, + k: usize, +) -> WeightTensor { + match cells.remove(&(name.to_string(), layer)) { + Some(TakenWeight { + handle: WeightHandle::Alias(source), + .. + }) if source == "token_embd" => { + hipfire_runtime::weight_backend::tied_lm_head_alias(token_embd, embd_format, m, k) + } + _ => unreachable!("validated LLaMA assembly lost tied {name}[layer {layer:?}]"), + } } fn embedding_format(dtype: DType) -> Result { @@ -296,17 +323,29 @@ fn embedding_format(dtype: DType) -> Result { fn assemble_llama_weights( config: &LlamaConfig, - store: &mut WeightStore, + transaction: &mut WeightLoadTransaction, ) -> Result { - let mut assembly = store.begin_assembly(); + let mut assembly = transaction.begin_assembly(); let mut slots = HashMap::new(); - let mut take = |name: &str, layer: Option| take_slot(&mut assembly, &mut slots, name, layer); + let mut take = |name: &str, layer: Option| { + take_slot(&mut assembly, &mut slots, name, layer) + }; take("token_embd", None)?; take("output_norm", None)?; take("lm_head", None)?; for layer in 0..config.n_layers { - for name in ["wq", "wk", "wv", "wo", "ffn_gate", "ffn_up", "ffn_down", "attn_norm", "ffn_norm"] { + for name in [ + "wq", + "wk", + "wv", + "wo", + "ffn_gate", + "ffn_up", + "ffn_down", + "attn_norm", + "ffn_norm", + ] { take(name, Some(layer))?; } if config.has_qk_norm { @@ -318,99 +357,124 @@ fn assemble_llama_weights( drop(take); let guard = assembly.commit(); for ((name, layer), slot) in &slots { - require_resident(&guard, name, *layer, *slot)?; + require_materialized(&guard, name, *layer, *slot)?; } + let token_slot = slots[&("token_embd".to_string(), None)]; + let token_dtype = match guard.get(token_slot) { + Some(WeightHandle::Resident(tensor)) => tensor.dtype, + _ => unreachable!("validated token_embd is not resident"), + }; + let embd_format = embedding_format(token_dtype)?; let cells: HashMap<_, _> = guard .finalize() .into_iter() .map(|taken| ((taken.key.name.clone(), taken.key.layer), taken)) .collect(); let mut cells = cells; - let token_embd = resident_cell(&mut cells, "token_embd", None)?; - let embd_format = embedding_format(token_embd.dtype)?; - let output_norm = resident_cell(&mut cells, "output_norm", None)?; - let output = resident_weight( - &mut cells, - "lm_head", - None, - config.vocab_size, - config.dim, - )?; + let token_embd = resident_cell(&mut cells, "token_embd", None); + let output_norm = resident_cell(&mut cells, "output_norm", None); + let lm_head_aliases_embd = matches!( + cells.get(&("lm_head".to_string(), None)), + Some(TakenWeight { + handle: WeightHandle::Alias(_), + .. + }) + ); + let output = if lm_head_aliases_embd { + tied_weight( + &mut cells, + &token_embd, + embd_format, + "lm_head", + None, + config.vocab_size, + config.dim, + ) + } else { + resident_weight( + &mut cells, + "lm_head", + None, + config.vocab_size, + config.dim, + ) + }; let mut layers = Vec::with_capacity(config.n_layers); for layer in 0..config.n_layers { let q_norm = if config.has_qk_norm { - Some(resident_cell(&mut cells, "q_norm", Some(layer))?) + Some(resident_cell(&mut cells, "q_norm", Some(layer))) } else { None }; let k_norm = if config.has_qk_norm { - Some(resident_cell(&mut cells, "k_norm", Some(layer))?) + Some(resident_cell(&mut cells, "k_norm", Some(layer))) } else { None }; layers.push(LayerWeights { - attn_norm: resident_cell(&mut cells, "attn_norm", Some(layer))?, + attn_norm: resident_cell(&mut cells, "attn_norm", Some(layer)), wq: resident_weight( &mut cells, "wq", Some(layer), config.n_heads * config.head_dim, config.dim, - )?, + ), wk: resident_weight( &mut cells, "wk", Some(layer), config.n_kv_heads * config.head_dim, config.dim, - )?, + ), wv: resident_weight( &mut cells, "wv", Some(layer), config.n_kv_heads * config.head_dim, config.dim, - )?, + ), wo: resident_weight( &mut cells, "wo", Some(layer), config.dim, config.n_heads * config.head_dim, - )?, + ), q_norm, k_norm, - ffn_norm: resident_cell(&mut cells, "ffn_norm", Some(layer))?, + ffn_norm: resident_cell(&mut cells, "ffn_norm", Some(layer)), w_gate: resident_weight( &mut cells, "ffn_gate", Some(layer), config.hidden_dim, config.dim, - )?, + ), w_up: resident_weight( &mut cells, "ffn_up", Some(layer), config.hidden_dim, config.dim, - )?, + ), w_down: resident_weight( &mut cells, "ffn_down", Some(layer), config.dim, config.hidden_dim, - )?, + ), }); } + debug_assert!(cells.is_empty(), "validated LLaMA assembly left cells"); Ok(LlamaWeights { token_embd, embd_format, output_norm, output, layers, - lm_head_aliases_embd: false, + lm_head_aliases_embd, }) } @@ -422,43 +486,84 @@ fn assemble_llama_weights( /// is published beneath this bundle's owner. The directory path remains on its /// existing ParoQuant loader until that source has an equivalent representation /// resolver. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum HfqLoadRoute { + /// Plain, non-AWQ files admitted to the manifest/typed-assembly pilot. + ManifestPlainLlama, + /// Files carrying AWQ scale sidecars retain the established loader until + /// sidecar ownership is represented by the manifest transaction. + LegacyAwq, +} + +fn classify_hfq_route(hfq: &HfqFile) -> HfqLoadRoute { + if hfq.has_awq_sidecars() { + HfqLoadRoute::LegacyAwq + } else { + HfqLoadRoute::ManifestPlainLlama + } +} + pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = match src { ModelSource::Hfq(hfq) => { let config = ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; - let (mesh, manifest_plan) = plan_single(&config)?; + // Admission and route classification are pure source checks. + // They must run before any manifest fulfillment or GPU upload. + hipfire_runtime::hfq::validate_llama_hfq_admission(&hfq) + .map_err(|e| e.to_string())?; + let has_separate_lm_head = + hfq.find_tensor_info("lm_head.weight").is_some(); + let route = classify_hfq_route(&hfq); + eprintln!("llama: HFQ source route = {route:?}"); + let (mesh, manifest_plan) = plan_single(&config, has_separate_lm_head)?; let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); - let mut store = hipfire_runtime::weight_store::fulfill_manifest( - &Llama::weight_manifest(&config), - &mesh, - config.n_layers, - ctx.gpu, - |entry| hfq_source(&hfq, entry), - ) - .map_err(|e| format!("llama: {e}"))?; - let weights = match assemble_llama_weights(&config, &mut store) { - Ok(weights) => weights, - Err(error) => { - store.rollback_unpublished(ctx.gpu); - return Err(error); + let (weights, mut weight_store) = match route { + HfqLoadRoute::LegacyAwq => { + let weights = + hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, ctx.gpu) + .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}"))?; + (weights, None) + } + HfqLoadRoute::ManifestPlainLlama => { + let manifest = + Llama::weight_manifest_for_hfq(&config, has_separate_lm_head); + let mut transaction = hipfire_runtime::weight_store::fulfill_manifest( + &manifest, + &mesh, + config.n_layers, + ctx.gpu, + |entry| hfq_source(&hfq, entry), + ) + .map_err(|e| format!("llama: {e}"))?; + let weights = match assemble_llama_weights(&config, &mut transaction) { + Ok(weights) => weights, + Err(error) => { + transaction.rollback(ctx.gpu); + return Err(error); + } + }; + (weights, Some(transaction)) } }; hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); // The plain LLaMA path has no independent cap resolver. PR // #661's physical-cap behavior is owned by the existing // upstream KV plan. - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - store.rollback_unpublished(ctx.gpu); - weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" - )); - } - }; + let scratch = + match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu); + } + weights.free_gpu(ctx.gpu); + return Err(format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + )); + } + }; let dims = llama_kv_dims(&config, ctx.max_seq, None); let kv = match ::from_mode( hipfire_runtime::kv_mode::resolve( @@ -473,7 +578,9 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result kv, Err(error) => { scratch.free_gpu(ctx.gpu); - store.rollback_unpublished(ctx.gpu); + if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu); + } weights.free_gpu(ctx.gpu); return Err(format!( "llama: ::from_mode failed: {error}" @@ -486,7 +593,7 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Result Result Result Result<(), (WeightStore, String)> { + transaction: WeightLoadTransaction, + ) -> Result<(), (WeightLoadTransaction, String)> { if self.weight_store.is_some() { - return Err((store, "llama: weight store already attached".into())); - } - if let Err(error) = store.validate_origin_value(self.weight_origin) { - return Err((store, format!("llama: weight store origin rejected: {error}"))); + return Err(( + transaction, + "llama: weight store already attached".into(), + )); } - self.weight_store = Some(store); + let attached = match transaction.publish(self.weight_origin) { + Ok(attached) => attached, + Err((transaction, error)) => { + return Err((transaction, format!("llama: weight store origin rejected: {error}"))); + } + }; + self.weight_store = Some(attached); Ok(()) } @@ -631,7 +743,135 @@ impl LlamaBundle { mod tests { use super::*; use hipfire_runtime::llama::ModelArch; - use hipfire_runtime::weight_store::{WeightProjection, WeightProjectionKind}; + use hipfire_runtime::arch_model::ArchModel; + use hipfire_runtime::hfq::{ + write_hfqm_package_mem, HfqFile, HfqMemTensor, + }; + use hipfire_runtime::kv_backend::KvBackend; + use hipfire_runtime::kv_mode::KvMode; + use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; + use hipfire_runtime::llama::{KvCache, KvCacheExt, KvDims, KvLayers, KvTarget}; + use hipfire_runtime::weight_store::{ + WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, + }; + use std::path::{Path, PathBuf}; + use std::time::{SystemTime, UNIX_EPOCH}; + + fn hfq_tensor(name: &str, shape: &[u32], quant_type: u8, bytes: usize) -> HfqMemTensor { + HfqMemTensor { + name: name.into(), + quant_type, + shape: shape.to_vec(), + group_size: 0, + data: vec![0; bytes], + } + } + + fn f32_hfq_tensor(name: &str, shape: &[u32], malformed: bool) -> HfqMemTensor { + let elements = shape.iter().map(|&dim| dim as usize).product::(); + hfq_tensor(name, shape, 2, if malformed { 4 } else { elements * 4 }) + } + + fn fixture_hfq( + with_awq_sidecar: bool, + with_q_proj_bias: bool, + malformed_output_norm: bool, + separate_lm_head: bool, + ) -> (PathBuf, HfqFile) { + let mut tensors = vec![ + f32_hfq_tensor("model.embed_tokens.weight", &[2, 32], false), + f32_hfq_tensor("model.norm.weight", &[32], false), + f32_hfq_tensor("model.layers.0.self_attn.q_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.self_attn.k_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.self_attn.v_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.self_attn.o_proj.weight", &[32, 32], false), + f32_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32], false), + f32_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32], false), + f32_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64], false), + f32_hfq_tensor( + "model.layers.0.input_layernorm.weight", + &[32], + false, + ), + f32_hfq_tensor( + "model.layers.0.post_attention_layernorm.weight", + &[32], + false, + ), + ]; + if malformed_output_norm { + tensors[1] = f32_hfq_tensor("model.norm.weight", &[32], true); + } + if with_awq_sidecar { + tensors.push(hfq_tensor( + "model.layers.0.self_attn.q_proj.awq_scale.weight", + &[32], + 1, + 32 * 2, + )); + } + if with_q_proj_bias { + tensors.push(hfq_tensor( + "model.layers.0.self_attn.q_proj.bias", + &[32], + 1, + 32 * 2, + )); + } + if separate_lm_head { + tensors.push(f32_hfq_tensor("lm_head.weight", &[2, 32], false)); + } + let metadata = r#"{ + "config": { + "model_type": "llama", + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "intermediate_size": 64, + "vocab_size": 2, + "head_dim": 32, + "rms_norm_eps": 0.00001, + "max_position_embeddings": 8, + "rope_theta": 10000.0 + } + }"#; + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before epoch") + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "hipfire-g3-{}-{nonce}.hfq", + std::process::id() + )); + write_hfqm_package_mem(&path, 0, metadata, &tensors).expect("write HFQ fixture"); + let hfq = HfqFile::open(&path).expect("open HFQ fixture"); + (path, hfq) + } + + fn load_ctx<'a>( + path: &'a Path, + gpu: &'a mut rdna_compute::Gpu, + cask: &'a CaskConfig, + ) -> LoadCtx<'a> { + LoadCtx { + path: path.to_str().expect("fixture path is UTF-8"), + max_seq: 8, + deepseek4_compute_placement: Default::default(), + deepseek4_experts_per_token: None, + draft_path: None, + kv_mode_override: Some("q8"), + kv_backend: KvBackend::Contiguous, + kv_adaptive_override: None, + state_quant_override: None, + cask, + pp: 1, + spec: SpecLoadCfg::default(), + gpu, + gemma4_drafter_path: None, + gemma4_draft_len: 3, + } + } fn config() -> LlamaConfig { LlamaConfig { @@ -662,10 +902,9 @@ mod tests { dtype: DType::F32, } } - #[test] fn single_plan_covers_every_typed_llama_handle() { - let (mesh, plan) = plan_single(&config()).unwrap(); + let (mesh, plan) = plan_single(&config(), true).unwrap(); let manifest = Llama::weight_manifest(&config()); assert_eq!(mesh.n_devices(), 1); assert_eq!(plan.weights.len(), 12); @@ -687,20 +926,21 @@ mod tests { .stage_alias(name, None, 0, "source", alias_projection()) .unwrap(); } + let mut transaction = WeightLoadTransaction::new(store); let error = match assemble_llama_weights( &LlamaConfig { n_layers: 0, ..config() }, - &mut store, + &mut transaction, ) { Ok(_) => panic!("alias unexpectedly assembled as typed weights"), Err(error) => error, }; assert!(error.contains("alias")); - assert_eq!(store.len(), 3); - assert!(store.contains("token_embd", None, 0)); - assert!(store.projection("lm_head", None, 0).is_some()); + assert_eq!(transaction.len(), 3); + assert!(transaction.contains("token_embd", None, 0)); + assert!(transaction.projection("lm_head", None, 0).is_some()); } @@ -729,4 +969,136 @@ mod tests { assert_eq!(dims.max_seq, 32_768); assert_eq!(dims.physical_cap, Some(4_096)); } + + #[test] + fn missing_lm_head_manifest_declares_a_tied_embedding_alias() { + let manifest = Llama::weight_manifest_for_hfq(&config(), false); + let token = &manifest[0]; + let output = manifest.last().expect("manifest has lm_head"); + assert!(matches!( + output.policy, + ShardPolicy::Tied { ref source } if source == "token_embd" + )); + assert!(token + .dtype_constraint + .same_source_set(&output.dtype_constraint)); + } + + #[test] + fn production_hfq_single_route_aliases_missing_lm_head_without_second_allocation() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, false, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); + drop(ctx); + assert!(bundle.weights.lm_head_aliases_embd); + assert_eq!( + bundle.weights.output.buf.buf.as_ptr(), + bundle.weights.token_embd.buf.as_ptr() + ); + assert!(bundle.weight_store.is_some()); + Box::new(bundle).free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_awq_route_preserves_legacy_loader() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(true, false, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) + .expect("AWQ fixture must use the legacy HFQ loader"); + drop(ctx); + assert!(bundle.weight_store.is_none()); + assert!(bundle.weights.lm_head_aliases_embd); + Box::new(bundle).free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_biased_hfq_is_rejected_before_manifest_upload() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, true, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let error = match load_bundle(ModelSource::Hfq(hfq), &mut ctx) { + Ok(_) => panic!("biased HFQ unexpectedly loaded"), + Err(error) => error, + }; + drop(ctx); + assert!(error.contains("q_proj.bias")); + assert!(error.contains("refusing to load Qwen2")); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_post_resident_failure_returns_clean_load_error() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, false, true, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let error = match load_bundle(ModelSource::Hfq(hfq), &mut ctx) { + Ok(_) => panic!("malformed output norm unexpectedly loaded"), + Err(error) => error, + }; + drop(ctx); + assert!(error.contains("source payload") || error.contains("output_norm")); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn production_manifest_matches_legacy_alias_contract() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let (path, hfq) = fixture_hfq(false, false, false, false); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); + let manifest_alias = bundle.weights.lm_head_aliases_embd; + drop(ctx); + Box::new(bundle).free_gpu(&mut gpu); + + let hfq = HfqFile::open(&path).expect("reopen HFQ fixture"); + let config = ::config_from_hfq(&hfq).expect("fixture config"); + let legacy = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, &mut gpu) + .expect("legacy loader fixture"); + assert_eq!(manifest_alias, legacy.lm_head_aliases_embd); + assert_eq!(legacy.embd_format, EmbeddingFormat::F32); + legacy.free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + + #[test] + fn physical_cap_is_honored_by_upstream_kv_constructor() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + let dims = KvDims { + layers: KvLayers::Flat(1), + n_kv_heads: 1, + head_dim: 32, + max_seq: 8, + physical_cap: Some(4), + }; + let cache = ::from_mode( + KvMode::Q8, + KvTarget::Single(&mut gpu), + &dims, + ) + .expect("upstream Q8 constructor"); + assert_eq!(cache.max_seq, 8); + assert_eq!(cache.physical_cap, 4); + let _ = cache.free_gpu(&mut gpu); + } } diff --git a/crates/hipfire-runtime/src/hfq.rs b/crates/hipfire-runtime/src/hfq.rs index c15678e42..54a68a0b6 100644 --- a/crates/hipfire-runtime/src/hfq.rs +++ b/crates/hipfire-runtime/src/hfq.rs @@ -432,6 +432,19 @@ impl HfqFile { self.overlay.is_some() } + /// Whether this source carries any AWQ scale sidecar. The carrier uses + /// this classification before allocation so supported sidecars stay on the + /// legacy loader until the manifest resolver can represent them. + pub fn has_awq_sidecars(&self) -> bool { + self.tensors + .iter() + .any(|tensor| tensor.name.ends_with(".awq_scale.weight")) + || self + .overlay + .as_ref() + .is_some_and(|overlay| overlay.has_awq_sidecars()) + } + /// Open an HFQM container that lives inside a larger file, starting at /// `base_offset`. Used by the bundled `.mq4-mtp` loader to parse the /// MTP section embedded after the trunk's tensor data. @@ -1638,6 +1651,7 @@ fn load_embedding_llama( .expect("embed_tokens not found"); // Q4K embeddings are llama-family-only (GGUF-derived). qwen2/qwen35 have no // Q4K embedding-lookup kernel — that is why the shared `load_embedding` / + // `embed_classify` deliberately rejects qt 4 (rejecting at load gives a clean // error instead of an "unsupported embedding format" panic deep in the qwen // forward pass). So Q4K stays an explicit llama-only branch here; everything @@ -1651,28 +1665,13 @@ fn load_embedding_llama( load_embedding(gpu, info.quant_type, data, config.vocab_size, config.dim) } -/// Load LLaMA weights from an HFQ file onto GPU. -pub fn load_weights_hfq( - hfq: &HfqFile, - config: &LlamaConfig, - gpu: &mut Gpu, -) -> HipResult { - // R2 guard: the LLaMA-family loader does NOT read Q/K/V proj bias — - // `LayerWeights` has no `wq_bias` / `wk_bias` / `wv_bias` fields and - // the per-layer load below only names `*.q_proj.weight`. Qwen2 - // requires those biases (`attention_bias=true` is the modeling - // default). The quantiser used to auto-tag every Qwen2 model as - // `arch_id=1`, which the daemon dispatches to this loader; the - // result was silently-wrong outputs with no warning. As of the - // `--arch-id` flag (see `hipfire-quantize`), Qwen2 models should be - // tagged `arch_id=7` and dispatched to `hipfire-arch-qwen2`. - // - // If we see `q_proj.bias` while loading as the LLaMA family, the - // input is a mis-tagged Qwen2 HFQ. Refuse hard with a pointer at - // the correct path. (Detection by manifest is robust to either the - // model_type tag or the model family — both LLaMA and Qwen3 lack - // these bias tensors, so any HFQ with `model.layers.0.self_attn.q_proj.bias` - // is by definition a Qwen2-family input.) +/// Reject a mis-tagged Qwen2 HFQ before any model allocation. +/// +/// The LLaMA-family `LayerWeights` type has no attention-bias tensors. A +/// Qwen2 file carrying `q_proj.bias` would therefore load and produce +/// silently-wrong output unless this admission check runs before every loader +/// route, including the manifest pilot. +pub fn validate_llama_hfq_admission(hfq: &HfqFile) -> HipResult<()> { if hfq .find_tensor_info("model.layers.0.self_attn.q_proj.bias") .is_some() @@ -1696,6 +1695,16 @@ pub fn load_weights_hfq( ), )); } + Ok(()) +} + +/// Load LLaMA weights from an HFQ file onto GPU. +pub fn load_weights_hfq( + hfq: &HfqFile, + config: &LlamaConfig, + gpu: &mut Gpu, +) -> HipResult { + validate_llama_hfq_admission(hfq)?; let mut source = LlamaHfqSource { hfq, cfg: config }; let layout = crate::model_load::Layout::single(config.n_layers); diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index 46fcad664..d233630a6 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -92,15 +92,15 @@ impl Layout { } /// Neutral result of the orchestrator. Each arch assembles its own weights -/// struct from this (qwen35 adds `pager`; llama drops `lm_head_aliases_embd`). +/// struct from this (qwen35 adds `pager`). pub struct LoadedWeights { pub token_embd: GpuTensor, pub embd_format: EmbeddingFormat, pub output_norm: GpuTensor, pub output: WeightTensor, pub layers: Vec, - /// True iff the tied lm_head aliases the embedding buffer (qwen35 single-GPU); - /// llama always returns `false` (it reuploads). + /// True iff the tied lm_head aliases the embedding buffer on this + /// single-device route; false means a separate output allocation exists. pub lm_head_aliases_embd: bool, } @@ -112,12 +112,10 @@ pub trait WeightSource { fn n_layers(&self) -> usize; /// Pre-load hook. HFQ drops the mmap when n==1; PaRo rejects n>1; llama no-op. fn prepare(&mut self, n_devices: usize) -> HipResult<()>; - fn read_embed(&mut self, gpu: &mut Gpu) -> HipResult<(GpuTensor, EmbeddingFormat)>; - fn read_final_norm(&mut self, gpu: &mut Gpu) -> HipResult; - /// `can_alias` is true iff embed and output share a device (n==1); the impl - /// decides whether to use it (qwen35 aliases; llama ignores it and reuploads). + /// `can_alias` is true iff embed and output share a device (n==1); the + /// implementation decides whether to use it (single-device LLaMA and + /// qwen35 alias tied embeddings; multi-device routes re-materialize). fn read_output( - &mut self, gpu: &mut Gpu, embd: &GpuTensor, embd_fmt: EmbeddingFormat, diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 4e9f91023..044bc62c6 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -100,6 +100,27 @@ impl DTypeConstraint { SourceDType::OneOf(allowed) => allowed.contains(&dtype), } } + + /// Whether two source constraints admit exactly the same representation + /// set. Variant spelling is not part of the contract: `Exact(F16)` and + /// `OneOf([F16])` are equivalent, while `Any` is never equivalent to a + /// finite list. + pub fn same_source_set(&self, other: &Self) -> bool { + fn finite_equal(left: &[DType], right: &[DType]) -> bool { + left.iter().all(|dtype| right.contains(dtype)) + && right.iter().all(|dtype| left.contains(dtype)) + } + match (&self.source, &other.source) { + (SourceDType::Any, SourceDType::Any) => true, + (SourceDType::Any, _) | (_, SourceDType::Any) => false, + (SourceDType::Exact(left), SourceDType::Exact(right)) => left == right, + (SourceDType::Exact(dtype), SourceDType::OneOf(values)) + | (SourceDType::OneOf(values), SourceDType::Exact(dtype)) => { + values.iter().all(|value| value == dtype) + } + (SourceDType::OneOf(left), SourceDType::OneOf(right)) => finite_equal(left, right), + } + } } /// The block ordering of a fused projection. @@ -460,7 +481,10 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< source_entry.dtype, entry.dtype )); } - if !entry.dtype_constraint.accepts(source_entry.dtype) + if !source_entry + .dtype_constraint + .same_source_set(&entry.dtype_constraint) + || !entry.dtype_constraint.accepts(source_entry.dtype) || !source_entry.dtype_constraint.accepts(entry.dtype) { return Err(format!( @@ -1131,4 +1155,24 @@ mod tests { ); assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single()).is_err()); } + #[test] + fn tied_entries_reject_different_source_sets_with_equal_logical_dtype() { + let source = WeightEntry::model( + "source", + vec![8, 8], + DType::F16, + ShardPolicy::Replicate, + ); + let tied = WeightEntry::model_with_dtype_constraint( + "tied", + vec![8, 8], + DType::F16, + DTypeConstraint::source_exact(DType::F16), + ShardPolicy::Tied { + source: "source".into(), + }, + ); + let error = validate_manifest(&[source, tied], &DeviceMesh::single()).unwrap_err(); + assert!(error.contains("source dtype contract")); + } } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index 80b8b9bd8..99ac6cc88 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -10,12 +10,12 @@ //! them, and the first failure explicitly rolls back every resident buffer. //! //! The store is not a model owner. It has no `Drop` implementation and never -//! frees GPU buffers implicitly. A carrier may move a committed store into its -//! existing `ArchModel` owner; that owner must transfer its resident handles -//! through [`WeightStore::take_all`] during the existing teardown path. -//! `take` transfers a resident handle to the owner that is assembling typed -//! weights, and therefore removes the cell from the store's cleanup set. - +//! frees GPU buffers implicitly. A carrier moves a committed transaction into +//! its existing `ArchModel` owner; that owner consumes the private drain +//! capability during the existing teardown path. +//! `WeightStoreAssembly::take` transfers a resident handle to the owner that is +//! assembling typed weights, and therefore removes the cell from the store's +//! cleanup set. use crate::weight_manifest::{placement_devices, ShardPolicy, WeightEntry}; use hipfire_hardware::{DeviceMesh, MeshEpoch}; use rdna_compute::{DType, Gpu, GpuTensor}; @@ -201,11 +201,9 @@ impl std::fmt::Display for FulfillError { impl std::error::Error for FulfillError {} /// Load-side placement container. It records one immutable projection per -/// `(name, layer, device)` and captures the target origin once. -/// -/// There is intentionally no `Drop` implementation. A `WeightStore` that is -/// abandoned without explicit rollback or owner transfer leaks rather than -/// guessing a GPU owner; production callers keep it beneath `ArchModel`. +/// `(name, layer, device)` and captures the target origin once. The container +/// itself has no consuming teardown API: lifecycle transitions are represented +/// by [`WeightLoadTransaction`] and [`AttachedWeightStore`]. #[derive(Default)] pub struct WeightStore { placements: HashMap, @@ -213,6 +211,145 @@ pub struct WeightStore { origin: Option, } +/// The only owner that may roll back resident allocations before publication. +/// +/// A transaction owns the store until [`Self::publish`] transfers it into the +/// attached owner held by the architecture bundle. It deliberately has no +/// implicit `Drop` cleanup because the GPU is not available to a destructor. +pub struct WeightLoadTransaction { + store: Option, +} + +impl std::fmt::Debug for WeightLoadTransaction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WeightLoadTransaction") + .field("origin", &self.origin()) + .field("len", &self.len()) + .finish() + } +} + +/// The resident-store capability returned by a committed load transaction. +/// +/// The backing store and its drain capability are private. Architecture +/// owners receive this value during attachment and consume it exactly once +/// during unload; no public `WeightStore` method can drain an attached store. +pub struct AttachedWeightStore { + store: WeightStore, + capability: WeightStoreDrainCapability, +} + +struct WeightStoreDrainCapability { + origin: WeightOrigin, +} + +impl WeightLoadTransaction { + pub fn new(store: WeightStore) -> Self { + Self { store: Some(store) } + } + + pub fn origin(&self) -> Option { + self.store.as_ref().and_then(WeightStore::origin) + } + + pub fn len(&self) -> usize { + self.store.as_ref().map_or(0, WeightStore::len) + } + + pub fn is_empty(&self) -> bool { + self.store.as_ref().map_or(true, WeightStore::is_empty) + } + + pub fn contains(&self, name: &str, layer: Option, device: usize) -> bool { + self.store + .as_ref() + .is_some_and(|store| store.contains(name, layer, device)) + } + + pub fn get( + &self, + name: &str, + layer: Option, + device: usize, + ) -> Option<&WeightHandle> { + self.store + .as_ref() + .and_then(|store| store.get(name, layer, device)) + } + + pub fn projection( + &self, + name: &str, + layer: Option, + device: usize, + ) -> Option<&WeightProjection> { + self.store + .as_ref() + .and_then(|store| store.projection(name, layer, device)) + } + + pub fn devices_for(&self, name: &str, layer: Option) -> Vec { + self.store + .as_ref() + .map_or_else(Vec::new, |store| store.devices_for(name, layer)) + } + + /// Start typed assembly while this load is still unpublished. + pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { + self.store + .as_mut() + .expect("weight load transaction was already consumed") + .begin_assembly() + } + + /// Consume this transaction and release every resident handle it owns. + /// This is intentionally the only rollback operation exposed by the + /// lifecycle API. + pub fn rollback(mut self, gpu: &Gpu) { + if let Some(store) = self.store.take() { + store.rollback(gpu); + } + } + + /// Publish the store beneath an architecture owner after checking the + /// complete immutable target identity. A mismatch returns this + /// transaction unchanged so the caller can retry or roll it back. + pub fn publish( + mut self, + expected: WeightOrigin, + ) -> Result { + let store = self + .store + .take() + .expect("weight load transaction was already consumed"); + if let Err(error) = store.validate_origin_value(expected) { + self.store = Some(store); + return Err((self, error)); + } + Ok(AttachedWeightStore { + store, + capability: WeightStoreDrainCapability { origin: expected }, + }) + } +} + +impl AttachedWeightStore { + /// Drain resident handles through the private owner capability. The + /// capability is established only by `WeightLoadTransaction::publish`, so + /// origin mismatch is impossible after attachment. + pub fn drain(self, gpu: &Gpu) { + let Self { store, capability } = self; + capability.drain(store, gpu); + } +} + +impl WeightStoreDrainCapability { + fn drain(self, store: WeightStore, gpu: &Gpu) { + debug_assert_eq!(store.origin, Some(self.origin)); + store.release_unchecked(gpu); + } +} + impl WeightStore { pub fn new() -> Self { Self::default() @@ -300,9 +437,10 @@ impl WeightStore { ) } - /// Move a handle out of the store. The projection is removed with it so no - /// stale metadata can describe a cell that the store no longer owns. - pub fn take( + /// Move a handle out of the store. This is private to the assembly + /// capability so arbitrary store holders cannot independently tear down a + /// resident allocation. + fn take( &mut self, name: &str, layer: Option, @@ -313,7 +451,7 @@ impl WeightStore { self.placements.remove(&key) } - pub fn take_with_projection( + fn take_with_projection( &mut self, name: &str, layer: Option, @@ -325,10 +463,7 @@ impl WeightStore { Some((handle, projection)) } - /// Start a typed assembly transaction. Handles moved through the - /// transaction are restored to this store if the transaction is dropped - /// before `finalize`; no GPU free or second owner is introduced. - pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { + fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { WeightStoreAssembly { store: self, taken: Vec::new(), @@ -337,8 +472,7 @@ impl WeightStore { } /// Compare a store's captured origin with an already-resolved target - /// identity. This pure seam is used by fault-path tests and by owner - /// teardown after target resolution. + /// identity. This read-only seam cannot release or extract any handle. pub fn validate_origin_value( &self, expected: WeightOrigin, @@ -360,27 +494,6 @@ impl WeightStore { self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) } - /// Roll back a fulfilled store before it is published beneath a model - /// owner. This is the only public consuming GPU-free operation: callers - /// may use it while a load transaction is still unpublished, but an - /// attached store can only be drained by the model owner via `take_all`. - pub fn rollback_unpublished(self, gpu: &Gpu) { - self.release_unchecked(gpu); - } - - /// Transfer every resident/alias handle to the model owner after checking - /// the complete captured origin. On mismatch, the original store is - /// returned unchanged so the owner can retry against the correct target. - pub fn take_all( - self, - expected: WeightOrigin, - ) -> Result, (Self, WeightStoreError)> { - if let Err(error) = self.validate_origin_value(expected) { - return Err((self, error)); - } - Ok(self.placements.into_values().collect()) - } - /// Explicit rollback for a failed transaction. It consumes the partial /// store and frees every resident buffer on the single owning GPU. fn rollback(self, gpu: &Gpu) { @@ -500,7 +613,7 @@ pub fn fulfill_manifest_single( n_layers: usize, gpu: &Gpu, source: F, -) -> Result +) -> Result where F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, { @@ -623,7 +736,7 @@ where }); } } - Ok(store) + Ok(WeightLoadTransaction::new(store)) } /// Canonical name used by the manifest fulfillment seam. The target is @@ -635,7 +748,7 @@ pub fn fulfill_manifest( n_layers: usize, gpu: &Gpu, source: F, -) -> Result +) -> Result where F: Fn(&WeightEntry) -> Result<(Vec, DType), String>, { @@ -763,24 +876,24 @@ mod tests { }; let mesh = DeviceMesh::single(); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) }) .unwrap(); - assert_eq!(store.len(), 1); + assert_eq!(transaction.len(), 1); assert!(matches!( - store.get("resident", None, 0), + transaction.get("resident", None, 0), Some(WeightHandle::Resident(tensor)) if tensor.dtype == DType::F32 )); assert_eq!( - store.projection("resident", None, 0).unwrap().dtype, + transaction.projection("resident", None, 0).unwrap().dtype, DType::F32 ); - store.rollback_unpublished(&gpu); + transaction.rollback(&gpu); } #[test] - fn full_origin_mismatch_returns_resident_store_unchanged() { + fn full_origin_mismatch_returns_unpublished_transaction_unchanged() { let first = DeviceMesh::single(); let second = DeviceMesh::single(); let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); @@ -789,60 +902,63 @@ mod tests { store .stage_alias("resident", None, 0, "source", projection(DType::F16)) .unwrap(); - let (store, error) = match store.take_all(expected) { + let transaction = WeightLoadTransaction::new(store); + let (transaction, error) = match transaction.publish(expected) { Ok(_) => panic!("origin mismatch unexpectedly succeeded"), Err(value) => value, }; assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); - assert_eq!(store.origin(), Some(actual)); - assert!(store.contains("resident", None, 0)); - assert!(store.projection("resident", None, 0).is_some()); + assert_eq!(transaction.origin(), Some(actual)); + assert!(transaction.contains("resident", None, 0)); + assert!(transaction.projection("resident", None, 0).is_some()); } #[test] - fn full_origin_mismatch_does_not_free_a_resident_store() { + fn full_origin_mismatch_does_not_free_a_resident_transaction() { let Ok(gpu) = Gpu::init() else { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); let mesh = DeviceMesh::single(); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let store = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { + let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) }) .unwrap(); let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); - let (store, error) = match store.take_all(expected) { + let (transaction, error) = match transaction.publish(expected) { Ok(_) => panic!("origin mismatch unexpectedly succeeded"), Err(value) => value, }; assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); - assert_eq!(store.len(), 1); + assert_eq!(transaction.len(), 1); assert_eq!( RESIDENT_RELEASES.with(std::cell::Cell::get), 0, "origin rejection must not free resident buffers" ); - store.rollback_unpublished(&gpu); + transaction.rollback(&gpu); assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } #[test] - fn owner_transfer_is_consuming_and_empty_transfer_is_idempotent() { + fn attached_owner_transfer_is_consuming_and_empty_transfer_is_safe() { let mesh = DeviceMesh::single(); let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); let mut store = WeightStore::with_origin(origin); store .stage_alias("owned", None, 0, "source", projection(DType::F16)) .unwrap(); - let handles = store.take_all(origin).unwrap(); - assert_eq!(handles.len(), 1); - assert!(matches!( - handles.into_iter().next(), - Some(WeightHandle::Alias(_)) - )); - let second = WeightStore::with_origin(origin).take_all(origin).unwrap(); - assert!(second.is_empty()); + let transaction = WeightLoadTransaction::new(store); + let attached = transaction.publish(origin).unwrap(); + let Ok(gpu) = Gpu::init() else { + return; + }; + attached.drain(&gpu); + let empty = WeightLoadTransaction::new(WeightStore::with_origin(origin)) + .publish(origin) + .unwrap(); + empty.drain(&gpu); } #[test] From 5173c5e37cce4ffd89434e365b0d749109659fe1 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 20:03:16 +0200 Subject: [PATCH 05/17] fix(device-mesh): close final G3 manifest review gaps --- crates/hipfire-arch-llama/src/arch_model.rs | 29 --- crates/hipfire-arch-llama/src/carrier.rs | 127 ++++++++-- crates/hipfire-runtime/src/weight_store.rs | 250 +++++++++++++------- 3 files changed, 268 insertions(+), 138 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 763dbbb29..f21e75b72 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -61,32 +61,3 @@ impl ArchModel for LlamaBundle { } } -#[cfg(test)] -mod tests { - use super::*; - use hipfire_hardware::DeviceMesh; - use hipfire_runtime::weight_manifest::{ShardPolicy, WeightEntry}; - use hipfire_runtime::weight_store::{ - fulfill_manifest_single, WeightLoadTransaction, WeightOrigin, WeightStore, - }; - - #[test] - fn attached_owner_drain_is_consuming_and_empty_drain_is_safe() { - let Ok(gpu) = Gpu::init() else { - return; - }; - let mesh = DeviceMesh::single(); - let origin = WeightOrigin::for_single(&mesh, &gpu); - let entry = - WeightEntry::model("owned", vec![1], rdna_compute::DType::F32, ShardPolicy::Replicate); - let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { - Ok((vec![0; 4], rdna_compute::DType::F32)) - }) - .unwrap(); - transaction.publish(origin).unwrap().drain(&gpu); - WeightLoadTransaction::new(WeightStore::with_origin(origin)) - .publish(origin) - .unwrap() - .drain(&gpu); - } -} diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 5a8b5ebe1..f80caf8b5 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -16,8 +16,8 @@ use hipfire_runtime::llama::{ use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; use hipfire_runtime::weight_backend::hfq_weight_dtype; use hipfire_runtime::weight_store::{ - AttachedWeightStore, TakenWeight, WeightHandle, WeightLoadTransaction, - WeightStoreAssembly, WeightStoreAssemblyGuard, WeightOrigin, + TakenWeight, WeightHandle, WeightLoadTransaction, WeightOrigin, WeightStoreAssembly, + WeightStoreAssemblyGuard, WeightStoreError, }; use rdna_compute::{DType, GpuTensor}; use std::collections::HashMap; @@ -52,6 +52,32 @@ pub struct LlamaBundle { /// block-only KvCache/scratch). `None` when `dspark_weights` is `None`. pub dspark_assets: Option, } + +/// Crate-private attached owner for the manifest transaction. +/// +/// The runtime transaction stays public only long enough for the load carrier +/// to assemble or roll it back. Once wrapped here, the only consuming path is +/// the crate's [`hipfire_runtime::arch_model::ArchModel::free_gpu`] implementation. +pub(crate) struct AttachedWeightStore { + transaction: WeightLoadTransaction, +} + +impl AttachedWeightStore { + fn from_transaction( + transaction: WeightLoadTransaction, + expected: WeightOrigin, + ) -> Result { + if let Err(error) = transaction.validate_origin_value(expected) { + return Err((transaction, error)); + } + Ok(Self { transaction }) + } + + pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) { + self.transaction.rollback(gpu); + } +} + fn plan_single( config: &LlamaConfig, has_separate_lm_head: bool, @@ -697,9 +723,9 @@ pub use load_bundle as load_llama_bundle; impl LlamaBundle { /// Attach an unpublished load transaction after validating the complete - /// target identity. Publication creates the sole resident-store drain - /// capability; a rejected transaction is returned unchanged. - pub fn attach_weight_store( + /// target identity. The resulting owner is crate-private and can only be + /// consumed by `ArchModel::free_gpu`. + fn attach_weight_store( &mut self, transaction: WeightLoadTransaction, ) -> Result<(), (WeightLoadTransaction, String)> { @@ -709,10 +735,16 @@ impl LlamaBundle { "llama: weight store already attached".into(), )); } - let attached = match transaction.publish(self.weight_origin) { + let attached = match AttachedWeightStore::from_transaction( + transaction, + self.weight_origin, + ) { Ok(attached) => attached, Err((transaction, error)) => { - return Err((transaction, format!("llama: weight store origin rejected: {error}"))); + return Err(( + transaction, + format!("llama: weight store origin rejected: {error}"), + )); } }; self.weight_store = Some(attached); @@ -750,7 +782,10 @@ mod tests { use hipfire_runtime::kv_backend::KvBackend; use hipfire_runtime::kv_mode::KvMode; use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; - use hipfire_runtime::llama::{KvCache, KvCacheExt, KvDims, KvLayers, KvTarget}; + use hipfire_runtime::llama::{ + weight_gemv, KvCache, KvCacheExt, KvDims, KvLayers, KvTarget, + }; + use hipfire_runtime::weight_store::test_support; use hipfire_runtime::weight_store::{ WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, }; @@ -769,7 +804,20 @@ mod tests { fn f32_hfq_tensor(name: &str, shape: &[u32], malformed: bool) -> HfqMemTensor { let elements = shape.iter().map(|&dim| dim as usize).product::(); - hfq_tensor(name, shape, 2, if malformed { 4 } else { elements * 4 }) + let data = if malformed { + vec![0; 4] + } else { + (0..elements) + .flat_map(|value| ((value as f32) + 1.0).to_le_bytes()) + .collect() + }; + HfqMemTensor { + name: name.into(), + quant_type: 2, + shape: shape.to_vec(), + group_size: 0, + data, + } } fn fixture_hfq( @@ -902,6 +950,24 @@ mod tests { dtype: DType::F32, } } + + fn output_for( + gpu: &mut rdna_compute::Gpu, + weights: &LlamaWeights, + hidden: &[f32], + ) -> Vec { + let input = gpu + .upload_f32(hidden, &[hidden.len()]) + .expect("upload deterministic output input"); + let output = gpu + .alloc_tensor(&[weights.output.m], DType::F32) + .expect("allocate deterministic output"); + weight_gemv(gpu, &weights.output, &input, &output).expect("run output projection"); + let values = gpu.download_f32(&output).expect("download output projection"); + let _ = gpu.free_tensor(output); + let _ = gpu.free_tensor(input); + values + } #[test] fn single_plan_covers_every_typed_llama_handle() { let (mesh, plan) = plan_single(&config(), true).unwrap(); @@ -1040,42 +1106,67 @@ mod tests { } #[test] - fn production_post_resident_failure_returns_clean_load_error() { + fn production_post_resident_failure_reclaims_every_uploaded_allocation() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { return; }; - let (path, hfq) = fixture_hfq(false, false, true, false); + let (path, hfq) = fixture_hfq(false, false, false, false); + test_support::reset(); + test_support::arm_fail_after_upload(1); let cask = CaskConfig::default(); let mut ctx = load_ctx(&path, &mut gpu, &cask); let error = match load_bundle(ModelSource::Hfq(hfq), &mut ctx) { - Ok(_) => panic!("malformed output norm unexpectedly loaded"), + Ok(_) => panic!("post-upload fault unexpectedly succeeded"), Err(error) => error, }; drop(ctx); - assert!(error.contains("source payload") || error.contains("output_norm")); + test_support::clear_faults(); + assert!(error.contains("test fault injected after resident upload")); + let allocations = test_support::resident_allocations(); + assert!(allocations > 0, "fault must follow a resident upload"); + assert_eq!( + allocations, + test_support::resident_releases(), + "every resident allocation must be reclaimed on load failure" + ); std::fs::remove_file(path).expect("remove HFQ fixture"); } #[test] - fn production_manifest_matches_legacy_alias_contract() { + fn production_manifest_matches_legacy_numerical_output() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { return; }; let (path, hfq) = fixture_hfq(false, false, false, false); let cask = CaskConfig::default(); let mut ctx = load_ctx(&path, &mut gpu, &cask); - let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); - let manifest_alias = bundle.weights.lm_head_aliases_embd; + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) + .expect("load plain HFQ fixture through manifest path"); drop(ctx); + assert!(bundle.weights.lm_head_aliases_embd); + let hidden: Vec = (0..bundle.config.dim) + .map(|index| (index as f32 + 1.0) / 17.0) + .collect(); + let manifest_output = output_for(&mut gpu, &bundle.weights, &hidden); Box::new(bundle).free_gpu(&mut gpu); let hfq = HfqFile::open(&path).expect("reopen HFQ fixture"); let config = ::config_from_hfq(&hfq).expect("fixture config"); let legacy = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, &mut gpu) - .expect("legacy loader fixture"); - assert_eq!(manifest_alias, legacy.lm_head_aliases_embd); + .expect("load legacy HFQ fixture"); + assert!(legacy.lm_head_aliases_embd); assert_eq!(legacy.embd_format, EmbeddingFormat::F32); + let legacy_output = output_for(&mut gpu, &legacy, &hidden); legacy.free_gpu(&mut gpu); + + assert_eq!( + manifest_output, legacy_output, + "manifest and legacy output projections must agree for the same input" + ); + assert!( + manifest_output.iter().any(|value| *value != 0.0), + "parity assertion must observe a non-zero numerical output" + ); std::fs::remove_file(path).expect("remove HFQ fixture"); } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index 99ac6cc88..76b919b08 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -11,8 +11,8 @@ //! //! The store is not a model owner. It has no `Drop` implementation and never //! frees GPU buffers implicitly. A carrier moves a committed transaction into -//! its existing `ArchModel` owner; that owner consumes the private drain -//! capability during the existing teardown path. +//! its existing `ArchModel` owner; that owner consumes the architecture-private +//! attached owner during the existing teardown path. //! `WeightStoreAssembly::take` transfers a resident handle to the owner that is //! assembling typed weights, and therefore removes the cell from the store's //! cleanup set. @@ -21,9 +21,64 @@ use hipfire_hardware::{DeviceMesh, MeshEpoch}; use rdna_compute::{DType, Gpu, GpuTensor}; use std::collections::HashMap; -#[cfg(test)] thread_local! { + static RESIDENT_ALLOCATIONS: std::cell::Cell = + const { std::cell::Cell::new(0) }; static RESIDENT_RELEASES: std::cell::Cell = const { std::cell::Cell::new(0) }; + static FAIL_AFTER_UPLOAD: std::cell::Cell> = + const { std::cell::Cell::new(None) }; +} + +/// Test-only allocation accounting and deterministic post-upload fault seam. +/// +/// The production loader calls the same release path regardless of whether +/// this seam is armed. Callers should use [`reset`] before a scenario and +/// [`clear_faults`] after it so a failed test cannot poison a later one. +#[doc(hidden)] +pub mod test_support { + use super::{ + FAIL_AFTER_UPLOAD, RESIDENT_ALLOCATIONS, RESIDENT_RELEASES, + }; + + pub fn reset() { + RESIDENT_ALLOCATIONS.with(|count| count.set(0)); + RESIDENT_RELEASES.with(|count| count.set(0)); + clear_faults(); + } + + pub fn arm_fail_after_upload(upload_number: usize) { + assert!(upload_number > 0, "upload fault threshold must be non-zero"); + FAIL_AFTER_UPLOAD.with(|fault| fault.set(Some(upload_number))); + } + + pub fn clear_faults() { + FAIL_AFTER_UPLOAD.with(|fault| fault.set(None)); + } + + pub fn resident_allocations() -> usize { + RESIDENT_ALLOCATIONS.with(std::cell::Cell::get) + } + + pub fn resident_releases() -> usize { + RESIDENT_RELEASES.with(std::cell::Cell::get) + } + + pub(super) fn record_resident_upload() -> bool { + let allocation = RESIDENT_ALLOCATIONS.with(|count| { + let next = count.get() + 1; + count.set(next); + next + }); + FAIL_AFTER_UPLOAD.with(|fault| { + let should_fail = fault + .get() + .is_some_and(|upload_number| allocation >= upload_number); + if should_fail { + fault.set(None); + } + should_fail + }) + } } /// Stable logical placement identity. Layer is part of the key because a @@ -202,8 +257,8 @@ impl std::error::Error for FulfillError {} /// Load-side placement container. It records one immutable projection per /// `(name, layer, device)` and captures the target origin once. The container -/// itself has no consuming teardown API: lifecycle transitions are represented -/// by [`WeightLoadTransaction`] and [`AttachedWeightStore`]. +/// itself has no consuming teardown API; lifecycle transitions are represented +/// by [`WeightLoadTransaction`] and the architecture-private attached owner. #[derive(Default)] pub struct WeightStore { placements: HashMap, @@ -213,9 +268,9 @@ pub struct WeightStore { /// The only owner that may roll back resident allocations before publication. /// -/// A transaction owns the store until [`Self::publish`] transfers it into the -/// attached owner held by the architecture bundle. It deliberately has no -/// implicit `Drop` cleanup because the GPU is not available to a destructor. +/// A transaction owns the store until the architecture carrier consumes it +/// into its crate-private attached owner. It deliberately has no implicit +/// `Drop` cleanup because the GPU is not available to a destructor. pub struct WeightLoadTransaction { store: Option, } @@ -229,20 +284,6 @@ impl std::fmt::Debug for WeightLoadTransaction { } } -/// The resident-store capability returned by a committed load transaction. -/// -/// The backing store and its drain capability are private. Architecture -/// owners receive this value during attachment and consume it exactly once -/// during unload; no public `WeightStore` method can drain an attached store. -pub struct AttachedWeightStore { - store: WeightStore, - capability: WeightStoreDrainCapability, -} - -struct WeightStoreDrainCapability { - origin: WeightOrigin, -} - impl WeightLoadTransaction { pub fn new(store: WeightStore) -> Self { Self { store: Some(store) } @@ -294,6 +335,19 @@ impl WeightLoadTransaction { .map_or_else(Vec::new, |store| store.devices_for(name, layer)) } + /// Compare the unpublished transaction's captured target with an admitted + /// owner identity. This read-only check is used before the carrier wraps + /// the transaction in its private attached owner. + pub fn validate_origin_value( + &self, + expected: WeightOrigin, + ) -> Result<(), WeightStoreError> { + self.store.as_ref().map_or( + Err(WeightStoreError::UnboundOrigin), + |store| store.validate_origin_value(expected), + ) + } + /// Start typed assembly while this load is still unpublished. pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { self.store @@ -310,44 +364,6 @@ impl WeightLoadTransaction { store.rollback(gpu); } } - - /// Publish the store beneath an architecture owner after checking the - /// complete immutable target identity. A mismatch returns this - /// transaction unchanged so the caller can retry or roll it back. - pub fn publish( - mut self, - expected: WeightOrigin, - ) -> Result { - let store = self - .store - .take() - .expect("weight load transaction was already consumed"); - if let Err(error) = store.validate_origin_value(expected) { - self.store = Some(store); - return Err((self, error)); - } - Ok(AttachedWeightStore { - store, - capability: WeightStoreDrainCapability { origin: expected }, - }) - } -} - -impl AttachedWeightStore { - /// Drain resident handles through the private owner capability. The - /// capability is established only by `WeightLoadTransaction::publish`, so - /// origin mismatch is impossible after attachment. - pub fn drain(self, gpu: &Gpu) { - let Self { store, capability } = self; - capability.drain(store, gpu); - } -} - -impl WeightStoreDrainCapability { - fn drain(self, store: WeightStore, gpu: &Gpu) { - debug_assert_eq!(store.origin, Some(self.origin)); - store.release_unchecked(gpu); - } } impl WeightStore { @@ -507,7 +523,6 @@ impl WeightStore { // the existing loader's explicit owner teardown. The store // never relies on a destructor to release GPU memory. let _ = gpu.hip.free(tensor.buf); - #[cfg(test)] RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); } } @@ -650,7 +665,34 @@ where } let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); if let ShardPolicy::Tied { source: source_name } = &entry.policy { - let projection = projection_for(entry, 0, 1, entry.dtype); + let source_dtype = match store.get(source_name, entry.layer, 0) { + Some(WeightHandle::Resident(tensor)) => Some(tensor.dtype), + Some(WeightHandle::Alias(_)) | None => None, + }; + let Some(actual_dtype) = source_dtype else { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "tied source '{source_name}' is unresolved or has no actual resident dtype" + ), + }); + }; + if !entry.dtype_constraint.accepts(actual_dtype) { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!( + "tied source '{source_name}' actual dtype {actual_dtype:?} is excluded by constraint {:?}", + entry.dtype_constraint + ), + }); + } + let projection = projection_for(entry, 0, 1, actual_dtype); if let Err(reason) = store.insert( key, WeightHandle::Alias(source_name.clone()), @@ -735,6 +777,15 @@ where reason: reason.to_string(), }); } + if test_support::record_resident_upload() { + store.rollback(gpu); + return Err(FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: "test fault injected after resident upload".into(), + }); + } } Ok(WeightLoadTransaction::new(store)) } @@ -758,7 +809,7 @@ where #[cfg(test)] mod tests { use super::*; - use crate::weight_manifest::{PinTarget, ShardPolicy}; + use crate::weight_manifest::{DTypeConstraint, PinTarget, ShardPolicy}; use hipfire_hardware::DimKind; fn projection(dtype: DType) -> WeightProjection { @@ -869,6 +920,48 @@ mod tests { assert_eq!(placement_devices(&entry, &mesh, 1), vec![0]); } + #[test] + fn tied_projection_preserves_fulfilled_source_dtype() { + let Ok(gpu) = Gpu::init() else { + return; + }; + let mesh = DeviceMesh::single(); + let constraint = DTypeConstraint::source_from_sources(vec![DType::F16, DType::F32]); + let source = WeightEntry::model_with_dtype_constraint( + "source", + vec![1], + DType::F16, + constraint.clone(), + ShardPolicy::Replicate, + ); + let alias = WeightEntry::model_with_dtype_constraint( + "alias", + vec![1], + DType::F16, + constraint, + ShardPolicy::Tied { + source: "source".into(), + }, + ); + let transaction = fulfill_manifest_single( + &[source, alias], + &mesh, + 1, + &gpu, + |_| Ok((vec![0; 4], DType::F32)), + ) + .unwrap(); + assert_eq!( + transaction.projection("alias", None, 0).unwrap().dtype, + DType::F32 + ); + assert!(matches!( + transaction.get("alias", None, 0), + Some(WeightHandle::Alias(source)) if source == "source" + )); + transaction.rollback(&gpu); + } + #[test] fn successful_single_fulfillment_commits_resident_projection() { let Ok(gpu) = Gpu::init() else { @@ -893,7 +986,7 @@ mod tests { } #[test] - fn full_origin_mismatch_returns_unpublished_transaction_unchanged() { + fn full_origin_mismatch_leaves_unpublished_transaction_unchanged() { let first = DeviceMesh::single(); let second = DeviceMesh::single(); let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); @@ -903,10 +996,7 @@ mod tests { .stage_alias("resident", None, 0, "source", projection(DType::F16)) .unwrap(); let transaction = WeightLoadTransaction::new(store); - let (transaction, error) = match transaction.publish(expected) { - Ok(_) => panic!("origin mismatch unexpectedly succeeded"), - Err(value) => value, - }; + let error = transaction.validate_origin_value(expected).unwrap_err(); assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); assert_eq!(transaction.origin(), Some(actual)); assert!(transaction.contains("resident", None, 0)); @@ -926,10 +1016,7 @@ mod tests { }) .unwrap(); let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); - let (transaction, error) = match transaction.publish(expected) { - Ok(_) => panic!("origin mismatch unexpectedly succeeded"), - Err(value) => value, - }; + let error = transaction.validate_origin_value(expected).unwrap_err(); assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); assert_eq!(transaction.len(), 1); assert_eq!( @@ -941,25 +1028,6 @@ mod tests { assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } - #[test] - fn attached_owner_transfer_is_consuming_and_empty_transfer_is_safe() { - let mesh = DeviceMesh::single(); - let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); - let mut store = WeightStore::with_origin(origin); - store - .stage_alias("owned", None, 0, "source", projection(DType::F16)) - .unwrap(); - let transaction = WeightLoadTransaction::new(store); - let attached = transaction.publish(origin).unwrap(); - let Ok(gpu) = Gpu::init() else { - return; - }; - attached.drain(&gpu); - let empty = WeightLoadTransaction::new(WeightStore::with_origin(origin)) - .publish(origin) - .unwrap(); - empty.drain(&gpu); - } #[test] fn source_failure_after_resident_upload_rolls_back_everything() { From f9cb8ec1de06089d5c57ba788a6e3dcfe406312f Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Mon, 31 Aug 2026 20:30:48 +0200 Subject: [PATCH 06/17] fix(device-mesh): repair llama manifest APIs --- crates/hipfire-arch-llama/src/arch_model.rs | 4 +- crates/hipfire-arch-llama/src/carrier.rs | 151 ++++++++------ crates/hipfire-runtime/src/model_load.rs | 6 +- crates/hipfire-runtime/src/weight_manifest.rs | 34 +++- crates/hipfire-runtime/src/weight_store.rs | 186 ++++++++++++------ 5 files changed, 255 insertions(+), 126 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index f21e75b72..1de6953da 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -54,7 +54,9 @@ impl ArchModel for LlamaBundle { // Attachment already checked the complete origin and created this // owner capability. There is no mismatch branch to leak the model: // an attached store can only be drained by this consuming owner. - store.drain(gpu); + if let Err(error) = store.drain(gpu) { + eprintln!("llama: failed to release attached weight store: {error}"); + } } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index f80caf8b5..a47bd13b9 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -73,8 +73,18 @@ impl AttachedWeightStore { Ok(Self { transaction }) } - pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) { - self.transaction.rollback(gpu); + pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) -> hip_bridge::HipResult<()> { + self.transaction.rollback(gpu) + } +} + +fn with_weight_rollback_error( + reason: String, + rollback: hip_bridge::HipResult<()>, +) -> String { + match rollback { + Ok(()) => reason, + Err(error) => format!("{reason}; resident rollback failed: {error}"), } } @@ -82,7 +92,7 @@ fn plan_single( config: &LlamaConfig, has_separate_lm_head: bool, ) -> Result<(DeviceMesh, ManifestPlan), String> { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().map_err(|error| format!("llama: device mesh: {error}"))?; let manifest = Llama::weight_manifest_for_hfq(config, has_separate_lm_head); let state = Llama::state_manifest(config); let plan = plan_manifest(&manifest, &state, &mesh, config.n_layers) @@ -566,8 +576,10 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result weights, Err(error) => { - transaction.rollback(ctx.gpu); - return Err(error); + return Err(with_weight_rollback_error( + error, + transaction.rollback(ctx.gpu), + )); } }; (weights, Some(transaction)) @@ -581,12 +593,17 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result scratch, Err(error) => { - if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu); - } + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + return Err(with_weight_rollback_error( + format!( + "llama: ForwardScratch::new_with_max_seq failed: {error:?}" + ), + rollback, )); } }; @@ -604,12 +621,17 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result kv, Err(error) => { scratch.free_gpu(ctx.gpu); - if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu); - } + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; weights.free_gpu(ctx.gpu); - return Err(format!( - "llama: ::from_mode failed: {error}" + return Err(with_weight_rollback_error( + format!( + "llama: ::from_mode failed: {error}" + ), + rollback, )); } }; @@ -708,11 +730,11 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result Vec { - let input = gpu - .upload_f32(hidden, &[hidden.len()]) - .expect("upload deterministic output input"); - let output = gpu - .alloc_tensor(&[weights.output.m], DType::F32) - .expect("allocate deterministic output"); - weight_gemv(gpu, &weights.output, &input, &output).expect("run output projection"); - let values = gpu.download_f32(&output).expect("download output projection"); - let _ = gpu.free_tensor(output); - let _ = gpu.free_tensor(input); - values - } #[test] fn single_plan_covers_every_typed_llama_handle() { let (mesh, plan) = plan_single(&config(), true).unwrap(); @@ -984,7 +990,7 @@ mod tests { #[test] fn typed_assembly_rolls_back_when_a_cell_is_not_resident() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); let mut store = WeightStore::with_origin(origin); for name in ["token_embd", "output_norm", "lm_head"] { @@ -1133,43 +1139,78 @@ mod tests { } #[test] - fn production_manifest_matches_legacy_numerical_output() { + fn production_manifest_matches_legacy_forward_logits() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { return; }; let (path, hfq) = fixture_hfq(false, false, false, false); let cask = CaskConfig::default(); let mut ctx = load_ctx(&path, &mut gpu, &cask); - let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) - .expect("load plain HFQ fixture through manifest path"); + let mut bundle = + load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load plain HFQ fixture"); drop(ctx); - assert!(bundle.weights.lm_head_aliases_embd); - let hidden: Vec = (0..bundle.config.dim) - .map(|index| (index as f32 + 1.0) / 17.0) - .collect(); - let manifest_output = output_for(&mut gpu, &bundle.weights, &hidden); + + let manifest_logits = { + forward_scratch_embed( + &mut gpu, + &bundle.weights, + &bundle.config, + 1, + 0, + &bundle.scratch, + ) + .expect("manifest embedding forward"); + forward_scratch_compute( + &mut gpu, + &bundle.weights, + &bundle.config, + 0, + &mut bundle.kv, + &bundle.scratch, + ) + .expect("manifest model forward"); + gpu.download_f32(&bundle.scratch.logits) + .expect("download manifest logits") + }; Box::new(bundle).free_gpu(&mut gpu); let hfq = HfqFile::open(&path).expect("reopen HFQ fixture"); let config = ::config_from_hfq(&hfq).expect("fixture config"); let legacy = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, &mut gpu) .expect("load legacy HFQ fixture"); - assert!(legacy.lm_head_aliases_embd); - assert_eq!(legacy.embd_format, EmbeddingFormat::F32); - let legacy_output = output_for(&mut gpu, &legacy, &hidden); + let scratch = ForwardScratch::new_with_max_seq(&mut gpu, &config, 8) + .expect("allocate legacy forward scratch"); + let dims = llama_kv_dims(&config, 8, None); + let mut kv = ::from_mode( + KvMode::Q8, + KvTarget::Single(&mut gpu), + &dims, + ) + .expect("allocate legacy KV cache"); + forward_scratch_embed(&mut gpu, &legacy, &config, 1, 0, &scratch) + .expect("legacy embedding forward"); + forward_scratch_compute(&mut gpu, &legacy, &config, 0, &mut kv, &scratch) + .expect("legacy model forward"); + let legacy_logits = gpu + .download_f32(&scratch.logits) + .expect("download legacy logits"); + scratch.free_gpu(&mut gpu); + let _ = kv.free_gpu(&mut gpu); legacy.free_gpu(&mut gpu); - assert_eq!( - manifest_output, legacy_output, - "manifest and legacy output projections must agree for the same input" - ); - assert!( - manifest_output.iter().any(|value| *value != 0.0), - "parity assertion must observe a non-zero numerical output" - ); + assert_eq!(manifest_logits.len(), legacy_logits.len()); + for (index, (manifest, legacy)) in + manifest_logits.iter().zip(&legacy_logits).enumerate() + { + assert!( + (manifest - legacy).abs() <= 1e-5, + "logit mismatch at index {index}: manifest={manifest} legacy={legacy}" + ); + } std::fs::remove_file(path).expect("remove HFQ fixture"); } + #[test] fn physical_cap_is_honored_by_upstream_kv_constructor() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index d233630a6..26178a122 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -112,10 +112,13 @@ pub trait WeightSource { fn n_layers(&self) -> usize; /// Pre-load hook. HFQ drops the mmap when n==1; PaRo rejects n>1; llama no-op. fn prepare(&mut self, n_devices: usize) -> HipResult<()>; + fn read_embed(&mut self, gpu: &mut Gpu) -> HipResult<(GpuTensor, EmbeddingFormat)>; + fn read_final_norm(&mut self, gpu: &mut Gpu) -> HipResult; /// `can_alias` is true iff embed and output share a device (n==1); the /// implementation decides whether to use it (single-device LLaMA and /// qwen35 alias tied embeddings; multi-device routes re-materialize). fn read_output( + &mut self, gpu: &mut Gpu, embd: &GpuTensor, embd_fmt: EmbeddingFormat, @@ -177,7 +180,8 @@ mod tests { #[test] fn mesh_layout_selects_stage_rank_zero_without_io() { - let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]) + .expect("small test mesh construction cannot overflow"); let layout = Layout::from_mesh(&mesh, 4); assert_eq!(layout.output_device(), 2); assert_eq!( diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 044bc62c6..1a4f62a7d 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -912,7 +912,8 @@ mod tests { #[test] fn placement_and_boundaries_use_named_mesh() { - let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]); + let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]) + .expect("small test mesh construction cannot overflow"); let embed = WeightEntry::model( "token_embd", vec![32, 8], @@ -955,7 +956,8 @@ mod tests { #[test] fn validation_covers_divisibility_ties_and_expert_shape() { - let tp3 = DeviceMesh::rect(&[(DimKind::Tp, 3)]); + let tp3 = DeviceMesh::rect(&[(DimKind::Tp, 3)]) + .expect("small test mesh construction cannot overflow"); assert!(validate_manifest( &[layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 })], &tp3 @@ -977,7 +979,11 @@ mod tests { }, ), ]; - assert!(validate_manifest(&tied, &DeviceMesh::single()).is_ok()); + assert!(validate_manifest( + &tied, + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .is_ok()); let bad_expert = WeightEntry::layer( "experts", 0, @@ -988,7 +994,11 @@ mod tests { assign: ExpertAssign::Stride, }, ); - assert!(validate_manifest(&[bad_expert], &DeviceMesh::single()).is_err()); + assert!(validate_manifest( + &[bad_expert], + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .is_err()); } #[test] @@ -1071,7 +1081,7 @@ mod tests { #[test] fn planning_rejects_weight_layer_at_n_layers_and_accepts_last_layer() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let valid = layer_entry("w", 2, ShardPolicy::Replicate); assert!(plan_manifest(&[valid], &[], &mesh, 3).is_ok()); let out_of_range = layer_entry("w", 3, ShardPolicy::Replicate); @@ -1097,7 +1107,7 @@ mod tests { ); assert!(validate_manifest( &[source.clone(), shape_mismatch], - &DeviceMesh::single() + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") ) .is_err()); @@ -1111,7 +1121,7 @@ mod tests { ); assert!(validate_manifest( &[source.clone(), dtype_mismatch], - &DeviceMesh::single() + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") ) .is_err()); @@ -1133,7 +1143,7 @@ mod tests { ); assert!(validate_manifest( &[source, chained_source, chain], - &DeviceMesh::single() + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") ) .is_err()); @@ -1153,7 +1163,7 @@ mod tests { source: "cycle_a".into(), }, ); - assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single()).is_err()); + assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single().expect("single-device mesh construction cannot overflow")).is_err()); } #[test] fn tied_entries_reject_different_source_sets_with_equal_logical_dtype() { @@ -1172,7 +1182,11 @@ mod tests { source: "source".into(), }, ); - let error = validate_manifest(&[source, tied], &DeviceMesh::single()).unwrap_err(); + let error = validate_manifest( + &[source, tied], + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .unwrap_err(); assert!(error.contains("source dtype contract")); } } diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index 76b919b08..ce0019441 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -358,10 +358,13 @@ impl WeightLoadTransaction { /// Consume this transaction and release every resident handle it owns. /// This is intentionally the only rollback operation exposed by the - /// lifecycle API. - pub fn rollback(mut self, gpu: &Gpu) { + /// lifecycle API. Successful frees are reflected in the resident-release + /// accounting; any failed HIP free is returned to the caller. + pub fn rollback(mut self, gpu: &Gpu) -> hip_bridge::HipResult<()> { if let Some(store) = self.store.take() { - store.rollback(gpu); + store.rollback(gpu) + } else { + Ok(()) } } } @@ -512,20 +515,30 @@ impl WeightStore { /// Explicit rollback for a failed transaction. It consumes the partial /// store and frees every resident buffer on the single owning GPU. - fn rollback(self, gpu: &Gpu) { - self.release_unchecked(gpu); + fn rollback(self, gpu: &Gpu) -> hip_bridge::HipResult<()> { + self.release_unchecked(gpu) } - fn release_unchecked(self, gpu: &Gpu) { + fn release_unchecked(self, gpu: &Gpu) -> hip_bridge::HipResult<()> { + let mut first_error = None; for handle in self.placements.into_values() { if let WeightHandle::Resident(tensor) = handle { - // Rollback is deliberately direct and best-effort, matching - // the existing loader's explicit owner teardown. The store - // never relies on a destructor to release GPU memory. - let _ = gpu.hip.free(tensor.buf); - RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); + match gpu.hip.free(tensor.buf) { + Ok(()) => { + RESIDENT_RELEASES.with(|count| count.set(count.get() + 1)); + } + Err(error) => { + if first_error.is_none() { + first_error = Some(error); + } + } + } } } + match first_error { + Some(error) => Err(error), + None => Ok(()), + } } } @@ -615,6 +628,18 @@ fn target_error(mesh: &DeviceMesh) -> Option { ), }) } +fn rollback_fulfill_error( + store: WeightStore, + gpu: &Gpu, + mut error: FulfillError, +) -> FulfillError { + if let Err(release_error) = store.rollback(gpu) { + error + .reason + .push_str(&format!("; resident rollback failed: {release_error}")); + } + error +} /// Fulfill a manifest for a plain LLaMA Single target. /// @@ -660,8 +685,7 @@ where devices ), }; - store.rollback(gpu); - return Err(error); + return Err(rollback_fulfill_error(store, gpu, error)); } let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); if let ShardPolicy::Tied { source: source_name } = &entry.policy { @@ -670,19 +694,18 @@ where Some(WeightHandle::Alias(_)) | None => None, }; let Some(actual_dtype) = source_dtype else { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: format!( "tied source '{source_name}' is unresolved or has no actual resident dtype" ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); }; if !entry.dtype_constraint.accepts(actual_dtype) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, @@ -690,7 +713,8 @@ where "tied source '{source_name}' actual dtype {actual_dtype:?} is excluded by constraint {:?}", entry.dtype_constraint ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } let projection = projection_for(entry, 0, 1, actual_dtype); if let Err(reason) = store.insert( @@ -698,13 +722,13 @@ where WeightHandle::Alias(source_name.clone()), projection, ) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: reason.to_string(), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } continue; } @@ -712,18 +736,17 @@ where let (bytes, dtype) = match source(entry) { Ok(value) => value, Err(reason) => { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: format!("source read failed: {reason}"), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } }; if !entry.dtype_constraint.accepts(dtype) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, @@ -731,7 +754,8 @@ where "source dtype {dtype:?} violates constraint {:?}", entry.dtype_constraint ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } if matches!(dtype, DType::F32 | DType::F16 | DType::BF16) { let expected_bytes = entry @@ -740,8 +764,7 @@ where .try_fold(1usize, |count, &dim| count.checked_mul(dim)) .and_then(|elements| elements.checked_mul(dtype.size())); if expected_bytes != Some(bytes.len()) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, @@ -751,40 +774,41 @@ where expected_bytes, entry.logical_shape ), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } } let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { Ok(tensor) => tensor, Err(error) => { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: format!("upload_raw failed: {error}"), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } }; tensor.dtype = dtype; let projection = projection_for(entry, 0, 1, dtype); if let Err(reason) = store.insert(key, WeightHandle::Resident(tensor), projection) { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: reason.to_string(), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } if test_support::record_resident_upload() { - store.rollback(gpu); - return Err(FulfillError { + let error = FulfillError { name: entry.name.clone(), layer: entry.layer, device: 0, reason: "test fault injected after resident upload".into(), - }); + }; + return Err(rollback_fulfill_error(store, gpu, error)); } } Ok(WeightLoadTransaction::new(store)) @@ -825,8 +849,8 @@ mod tests { #[test] fn origin_mismatch_is_detected_before_gpu_release() { - let first = DeviceMesh::single(); - let second = DeviceMesh::single(); + let first = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let second = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let actual = WeightOrigin::from_parts(first.epoch(), 0, 0); let expected = WeightOrigin::from_parts(second.epoch(), 0, 0); let store = WeightStore::with_origin(actual); @@ -842,7 +866,7 @@ mod tests { #[test] fn staged_rollback_removes_handles_and_projection_together() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let origin = WeightOrigin::from_parts(mesh.epoch(), 0, 0); let mut store = WeightStore::with_origin(origin); store @@ -863,7 +887,7 @@ mod tests { #[test] fn assembly_drop_restores_staged_handles() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); store .stage_alias("x", None, 0, "source", projection(DType::F16)) @@ -871,7 +895,8 @@ mod tests { { let mut assembly = store.begin_assembly(); assert_eq!(assembly.take("x", None, 0), Some(0)); - assert!(assembly.get(0).is_some()); + let guard = assembly.commit(); + assert!(guard.get(0).is_some()); } assert!(store.contains("x", None, 0)); assert!(store.projection("x", None, 0).is_some()); @@ -879,7 +904,7 @@ mod tests { #[test] fn repeated_unload_lookup_cannot_reclaim_a_transferred_cell() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); store .stage_alias("x", None, 0, "source", projection(DType::F16)) @@ -892,7 +917,7 @@ mod tests { #[test] fn duplicate_projection_is_rejected_without_replacing_identity() { - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let mut store = WeightStore::with_origin(WeightOrigin::from_parts(mesh.epoch(), 0, 0)); store .stage_alias("x", None, 0, "source-a", projection(DType::F16)) @@ -907,7 +932,8 @@ mod tests { #[test] fn single_target_refuses_multi_device_before_source_or_gpu_work() { - let mesh = DeviceMesh::rect(&[(DimKind::Tp, 2)]); + let mesh = DeviceMesh::rect(&[(DimKind::Tp, 2)]) + .expect("small test mesh construction cannot overflow"); let entry = WeightEntry::model( "embed", vec![2, 2], @@ -925,7 +951,7 @@ mod tests { let Ok(gpu) = Gpu::init() else { return; }; - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let constraint = DTypeConstraint::source_from_sources(vec![DType::F16, DType::F32]); let source = WeightEntry::model_with_dtype_constraint( "source", @@ -959,7 +985,9 @@ mod tests { transaction.get("alias", None, 0), Some(WeightHandle::Alias(source)) if source == "source" )); - transaction.rollback(&gpu); + transaction + .rollback(&gpu) + .expect("resident transaction rollback must free its HIP allocation"); } #[test] @@ -967,7 +995,7 @@ mod tests { let Ok(gpu) = Gpu::init() else { return; }; - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) @@ -982,13 +1010,15 @@ mod tests { transaction.projection("resident", None, 0).unwrap().dtype, DType::F32 ); - transaction.rollback(&gpu); + transaction + .rollback(&gpu) + .expect("resident transaction rollback must free its HIP allocation"); } #[test] fn full_origin_mismatch_leaves_unpublished_transaction_unchanged() { - let first = DeviceMesh::single(); - let second = DeviceMesh::single(); + let first = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let second = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let actual = WeightOrigin::from_parts(first.epoch(), 3, 11); let expected = WeightOrigin::from_parts(second.epoch(), 4, 12); let mut store = WeightStore::with_origin(actual); @@ -1009,7 +1039,7 @@ mod tests { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { Ok((vec![0; 4], DType::F32)) @@ -1024,18 +1054,56 @@ mod tests { 0, "origin rejection must not free resident buffers" ); - transaction.rollback(&gpu); + transaction + .rollback(&gpu) + .expect("resident transaction rollback must free its HIP allocation"); assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } + #[test] + fn rollback_reports_free_failure_without_counting_release() { + let Ok(gpu) = Gpu::init() else { + return; + }; + test_support::reset(); + RESIDENT_ALLOCATIONS.with(|count| count.set(1)); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let origin = WeightOrigin::from_parts(mesh.epoch(), 0, gpu.device_id); + let mut store = WeightStore::with_origin(origin); + let borrowed = GpuTensor { + buf: unsafe { + hip_bridge::DeviceBuffer::from_raw( + std::ptr::null_mut::(), + 0, + ) + }, + shape: vec![0], + dtype: DType::F32, + }; + store + .insert( + WeightPlacementKey::new("borrowed", None, 0), + WeightHandle::Resident(borrowed), + projection(DType::F32), + ) + .expect("insert borrowed resident test handle"); + let error = WeightLoadTransaction::new(store) + .rollback(&gpu) + .expect_err("rollback must surface a failed HIP free"); + assert!(error.message.contains("borrowed")); + assert_eq!(test_support::resident_allocations(), 1); + assert_eq!(test_support::resident_releases(), 0); + test_support::reset(); + } + #[test] fn source_failure_after_resident_upload_rolls_back_everything() { let Ok(gpu) = Gpu::init() else { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entries = vec![ WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), @@ -1063,7 +1131,7 @@ mod tests { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let constraint = DTypeConstraint::source_exact(DType::F32); let entries = vec![ WeightEntry::model_with_dtype_constraint( @@ -1104,7 +1172,7 @@ mod tests { return; }; RESIDENT_RELEASES.with(|count| count.set(0)); - let mesh = DeviceMesh::single(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entries = vec![ WeightEntry::model("first", vec![1], DType::F32, ShardPolicy::Replicate), WeightEntry::model("second", vec![1], DType::F32, ShardPolicy::Replicate), From df5b8f89953ff97e1d66f647b13c4ba927907c1d Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 00:11:27 +0200 Subject: [PATCH 07/17] fix(device-mesh): restore llama manifest integration --- Cargo.lock | 1 + crates/hipfire-arch-llama/src/carrier.rs | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 0f90443e0..9692607f0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1196,6 +1196,7 @@ version = "0.3.0" dependencies = [ "hip-bridge", "hipfire-dispatch", + "hipfire-hardware", "hipfire-runtime", "rdna-compute", ] diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index a47bd13b9..7725ab82f 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -14,6 +14,8 @@ use hipfire_runtime::llama::{ LlamaConfig, LlamaWeights, WeightTensor, }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; +use hipfire_runtime::model_source::ModelSource as ModelSourceTrait; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_backend::hfq_weight_dtype; use hipfire_runtime::weight_store::{ TakenWeight, WeightHandle, WeightLoadTransaction, WeightOrigin, WeightStoreAssembly, @@ -27,6 +29,8 @@ pub struct LlamaBundle { pub weights: LlamaWeights, pub scratch: ForwardScratch, pub kv: KvCache, + /// The admitted mesh that owns this plan and the attached store origin. + pub(crate) mesh: DeviceMesh, /// Pure declaration/placement plan captured at load time. The plan has no /// GPU handles and is immutable after publication. pub manifest_plan: ManifestPlan, From d7d18cac0c0ab4d8e5cd6a7177ac665b2272cc79 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 00:29:14 +0200 Subject: [PATCH 08/17] fix(device-mesh): correct manifest pilot contracts --- crates/hipfire-arch-llama/src/carrier.rs | 45 +++++++++++-------- crates/hipfire-runtime/src/weight_manifest.rs | 4 +- 2 files changed, 27 insertions(+), 22 deletions(-) diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index 7725ab82f..b85a71a3d 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -813,6 +813,7 @@ mod tests { KvTarget, }; use hipfire_runtime::weight_store::test_support; + use hipfire_runtime::weight_manifest::ShardPolicy; use hipfire_runtime::weight_store::{ WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, }; @@ -846,6 +847,22 @@ mod tests { data, } } + fn f16_hfq_tensor(name: &str, shape: &[u32]) -> HfqMemTensor { + let elements = shape.iter().map(|&dim| dim as usize).product::(); + HfqMemTensor { + name: name.into(), + quant_type: 1, + shape: shape.to_vec(), + group_size: 0, + data: (0..elements) + .flat_map(|index| { + let bits = if index % 2 == 0 { 0x3c00u16 } else { 0x3800u16 }; + bits.to_le_bytes() + }) + .collect(), + } + } + fn fixture_hfq( with_awq_sidecar: bool, @@ -856,13 +873,13 @@ mod tests { let mut tensors = vec![ f32_hfq_tensor("model.embed_tokens.weight", &[2, 32], false), f32_hfq_tensor("model.norm.weight", &[32], false), - f32_hfq_tensor("model.layers.0.self_attn.q_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.self_attn.k_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.self_attn.v_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.self_attn.o_proj.weight", &[32, 32], false), - f32_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32], false), - f32_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32], false), - f32_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64], false), + f16_hfq_tensor("model.layers.0.self_attn.q_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.k_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.v_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.o_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32]), + f16_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32]), + f16_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64]), f32_hfq_tensor( "model.layers.0.input_layernorm.weight", &[32], @@ -1081,19 +1098,9 @@ mod tests { } #[test] - fn production_awq_route_preserves_legacy_loader() { - let Ok(mut gpu) = rdna_compute::Gpu::init() else { - return; - }; + fn production_awq_sidecar_selects_legacy_loader() { let (path, hfq) = fixture_hfq(true, false, false, false); - let cask = CaskConfig::default(); - let mut ctx = load_ctx(&path, &mut gpu, &cask); - let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) - .expect("AWQ fixture must use the legacy HFQ loader"); - drop(ctx); - assert!(bundle.weight_store.is_none()); - assert!(bundle.weights.lm_head_aliases_embd); - Box::new(bundle).free_gpu(&mut gpu); + assert_eq!(classify_hfq_route(&hfq), HfqLoadRoute::LegacyAwq); std::fs::remove_file(path).expect("remove HFQ fixture"); } diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 1a4f62a7d..65a9172f4 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -484,8 +484,6 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< if !source_entry .dtype_constraint .same_source_set(&entry.dtype_constraint) - || !entry.dtype_constraint.accepts(source_entry.dtype) - || !source_entry.dtype_constraint.accepts(entry.dtype) { return Err(format!( "{context}: tied source '{source}' violates the source dtype contract" @@ -920,7 +918,7 @@ mod tests { DType::F16, ShardPolicy::Pin(PinTarget::Embed), ); - let row = layer_entry("wo", 1, ShardPolicy::RowShard { axis: 1 }); + let row = layer_entry("wo", 2, ShardPolicy::RowShard { axis: 1 }); assert_eq!(placement_devices(&embed, &mesh, 4), vec![0]); assert_eq!(placement_devices(&row, &mesh, 4), vec![2, 3]); let plan = plan_manifest( From 3c29ecf75cb44715b6263b9ad908862614c637f7 Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 00:32:58 +0200 Subject: [PATCH 09/17] style(device-mesh): format manifest pilot --- crates/hipfire-arch-llama/src/arch.rs | 11 +- crates/hipfire-arch-llama/src/arch_model.rs | 1 - crates/hipfire-arch-llama/src/carrier.rs | 395 ++++++++---------- crates/hipfire-runtime/src/model_load.rs | 3 +- crates/hipfire-runtime/src/weight_manifest.rs | 91 ++-- crates/hipfire-runtime/src/weight_store.rs | 119 ++---- 6 files changed, 263 insertions(+), 357 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch.rs b/crates/hipfire-arch-llama/src/arch.rs index 27c60aa9e..d02289ef0 100644 --- a/crates/hipfire-arch-llama/src/arch.rs +++ b/crates/hipfire-arch-llama/src/arch.rs @@ -17,8 +17,8 @@ use hip_bridge::HipResult; use hipfire_runtime::arch::Architecture; use hipfire_runtime::hfq::{self, HfqFile}; -use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::llama::KvCacheExt; +use hipfire_runtime::llama::{ForwardScratch, KvCache, LlamaConfig, LlamaWeights}; use hipfire_runtime::weight_manifest::{ DTypeConstraint, FusedQkvLayout, PinTarget, ShardPolicy, StateEntry, StateKind, WeightEntry, }; @@ -293,7 +293,14 @@ impl Llama { /// Pure state declaration for the full-attention LLaMA family. pub fn state_manifest(cfg: &LlamaConfig) -> Vec { (0..cfg.n_layers) - .map(|layer| StateEntry::new(StateKind::Kv { quant: String::new() }, layer)) + .map(|layer| { + StateEntry::new( + StateKind::Kv { + quant: String::new(), + }, + layer, + ) + }) .collect() } } diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 1de6953da..7f374f20c 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -62,4 +62,3 @@ impl ArchModel for LlamaBundle { let _ = kv.free_gpu(gpu); } } - diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index b85a71a3d..a6bd292c0 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -15,8 +15,8 @@ use hipfire_runtime::llama::{ }; use hipfire_runtime::loader_api::{LoadCtx, ModelSource}; use hipfire_runtime::model_source::ModelSource as ModelSourceTrait; -use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_backend::hfq_weight_dtype; +use hipfire_runtime::weight_manifest::{plan_manifest, ManifestPlan, WeightEntry}; use hipfire_runtime::weight_store::{ TakenWeight, WeightHandle, WeightLoadTransaction, WeightOrigin, WeightStoreAssembly, WeightStoreAssemblyGuard, WeightStoreError, @@ -82,10 +82,7 @@ impl AttachedWeightStore { } } -fn with_weight_rollback_error( - reason: String, - rollback: hip_bridge::HipResult<()>, -) -> String { +fn with_weight_rollback_error(reason: String, rollback: hip_bridge::HipResult<()>) -> String { match rollback { Ok(()) => reason, Err(error) => format!("{reason}; resident rollback failed: {error}"), @@ -213,10 +210,8 @@ fn f32_bytes_from_hfq(quant_type: u8, data: &[u8], name: &str) -> Result } for chunk in chunks { bytes.extend_from_slice( - &f32::from_bits( - u16::from_le_bytes([chunk[0], chunk[1]]) as u32 * (1 << 16), - ) - .to_le_bytes(), + &f32::from_bits(u16::from_le_bytes([chunk[0], chunk[1]]) as u32 * (1 << 16)) + .to_le_bytes(), ); } } @@ -244,12 +239,11 @@ fn hfq_source(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, DType), St )), }; } - if matches!(entry.name.as_str(), "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm") - { - return Ok(( - f32_bytes_from_hfq(quant_type, &data, &name)?, - DType::F32, - )); + if matches!( + entry.name.as_str(), + "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm" + ) { + return Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)); } match quant_type { 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), @@ -367,9 +361,8 @@ fn assemble_llama_weights( ) -> Result { let mut assembly = transaction.begin_assembly(); let mut slots = HashMap::new(); - let mut take = |name: &str, layer: Option| { - take_slot(&mut assembly, &mut slots, name, layer) - }; + let mut take = + |name: &str, layer: Option| take_slot(&mut assembly, &mut slots, name, layer); take("token_embd", None)?; take("output_norm", None)?; @@ -431,13 +424,7 @@ fn assemble_llama_weights( config.dim, ) } else { - resident_weight( - &mut cells, - "lm_head", - None, - config.vocab_size, - config.dim, - ) + resident_weight(&mut cells, "lm_head", None, config.vocab_size, config.dim) }; let mut layers = Vec::with_capacity(config.n_layers); for layer in 0..config.n_layers { @@ -544,174 +531,159 @@ fn classify_hfq_route(hfq: &HfqFile) -> HfqLoadRoute { } pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result { - let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = - match src { - ModelSource::Hfq(hfq) => { - let config = - ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; - // Admission and route classification are pure source checks. - // They must run before any manifest fulfillment or GPU upload. - hipfire_runtime::hfq::validate_llama_hfq_admission(&hfq) - .map_err(|e| e.to_string())?; - let has_separate_lm_head = - hfq.find_tensor_info("lm_head.weight").is_some(); - let route = classify_hfq_route(&hfq); - eprintln!("llama: HFQ source route = {route:?}"); - let (mesh, manifest_plan) = plan_single(&config, has_separate_lm_head)?; - let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); - let (weights, mut weight_store) = match route { - HfqLoadRoute::LegacyAwq => { - let weights = - hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, ctx.gpu) - .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}"))?; - (weights, None) - } - HfqLoadRoute::ManifestPlainLlama => { - let manifest = - Llama::weight_manifest_for_hfq(&config, has_separate_lm_head); - let mut transaction = hipfire_runtime::weight_store::fulfill_manifest( - &manifest, - &mesh, - config.n_layers, - ctx.gpu, - |entry| hfq_source(&hfq, entry), - ) - .map_err(|e| format!("llama: {e}"))?; - let weights = match assemble_llama_weights(&config, &mut transaction) { - Ok(weights) => weights, - Err(error) => { - return Err(with_weight_rollback_error( - error, - transaction.rollback(ctx.gpu), - )); - } - }; - (weights, Some(transaction)) - } - }; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - // The plain LLaMA path has no independent cap resolver. PR - // #661's physical-cap behavior is owned by the existing - // upstream KV plan. - let scratch = - match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, + let (config, weights, kv, scratch, manifest_plan, weight_store, mesh, weight_origin) = match src + { + ModelSource::Hfq(hfq) => { + let config = + ::config_from_hfq(&hfq).map_err(|e| e.to_string())?; + // Admission and route classification are pure source checks. + // They must run before any manifest fulfillment or GPU upload. + hipfire_runtime::hfq::validate_llama_hfq_admission(&hfq).map_err(|e| e.to_string())?; + let has_separate_lm_head = hfq.find_tensor_info("lm_head.weight").is_some(); + let route = classify_hfq_route(&hfq); + eprintln!("llama: HFQ source route = {route:?}"); + let (mesh, manifest_plan) = plan_single(&config, has_separate_lm_head)?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let (weights, mut weight_store) = match route { + HfqLoadRoute::LegacyAwq => { + let weights = hipfire_runtime::hfq::load_weights_hfq(&hfq, &config, ctx.gpu) + .map_err(|e| format!("llama: load_weights_hfq failed: {e:?}"))?; + (weights, None) + } + HfqLoadRoute::ManifestPlainLlama => { + let manifest = Llama::weight_manifest_for_hfq(&config, has_separate_lm_head); + let mut transaction = hipfire_runtime::weight_store::fulfill_manifest( + &manifest, + &mesh, + config.n_layers, + ctx.gpu, + |entry| hfq_source(&hfq, entry), + ) + .map_err(|e| format!("llama: {e}"))?; + let weights = match assemble_llama_weights(&config, &mut transaction) { + Ok(weights) => weights, Err(error) => { - let rollback = if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu) - } else { - Ok(()) - }; - weights.free_gpu(ctx.gpu); return Err(with_weight_rollback_error( - format!( - "llama: ForwardScratch::new_with_max_seq failed: {error:?}" - ), - rollback, + error, + transaction.rollback(ctx.gpu), )); } }; - let dims = llama_kv_dims(&config, ctx.max_seq, None); - let kv = match ::from_mode( - hipfire_runtime::kv_mode::resolve( - ctx.kv_mode_override.unwrap_or(""), - &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, - config.head_dim, - ) - .mode, - KvTarget::Single(ctx.gpu), - &dims, - ) { - Ok(kv) => kv, - Err(error) => { - scratch.free_gpu(ctx.gpu); - let rollback = if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu) - } else { - Ok(()) - }; - weights.free_gpu(ctx.gpu); - return Err(with_weight_rollback_error( - format!( - "llama: ::from_mode failed: {error}" - ), - rollback, - )); - } - }; - ( - config, - weights, - kv, - scratch, - manifest_plan, - weight_store, - mesh, - weight_origin, - ) - } - ModelSource::Dir(source) => { - let config = - hipfire_runtime::hfq::config_from_safetensors_llama(&source).map_err(|e| { - format!("failed to parse LLaMA/Qwen3 config from config.json: {e}") - })?; - let (mesh, manifest_plan) = - plan_single(&config, source.tensor_info("lm_head.weight").is_some())?; - let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); - let weights = - hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) - .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; - hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); - let kv_mode_str = ctx - .kv_mode_override - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); - let rr = hipfire_runtime::kv_mode::resolve( - &kv_mode_str, - &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, + (weights, Some(transaction)) + } + }; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + // The plain LLaMA path has no independent cap resolver. PR + // #661's physical-cap behavior is owned by the existing + // upstream KV plan. + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; + weights.free_gpu(ctx.gpu); + return Err(with_weight_rollback_error( + format!("llama: ForwardScratch::new_with_max_seq failed: {error:?}"), + rollback, + )); + } + }; + let dims = llama_kv_dims(&config, ctx.max_seq, None); + let kv = match ::from_mode( + hipfire_runtime::kv_mode::resolve( + ctx.kv_mode_override.unwrap_or(""), + &hipfire_runtime::kv_mode::LLAMA_HFQ_POLICY, config.head_dim, - ); - if let Some(w) = rr.warning { - eprintln!( - " KV cache: {w} (site {})", - hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site - ); + ) + .mode, + KvTarget::Single(ctx.gpu), + &dims, + ) { + Ok(kv) => kv, + Err(error) => { + scratch.free_gpu(ctx.gpu); + let rollback = if let Some(transaction) = weight_store.take() { + transaction.rollback(ctx.gpu) + } else { + Ok(()) + }; + weights.free_gpu(ctx.gpu); + return Err(with_weight_rollback_error( + format!("llama: ::from_mode failed: {error}"), + rollback, + )); } - let dims = llama_kv_dims(&config, ctx.max_seq, Some(ctx.max_seq)); - let kv = match ::from_mode( - rr.mode, - KvTarget::Single(ctx.gpu), - &dims, - ) { + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + weight_store, + mesh, + weight_origin, + ) + } + ModelSource::Dir(source) => { + let config = hipfire_runtime::hfq::config_from_safetensors_llama(&source) + .map_err(|e| format!("failed to parse LLaMA/Qwen3 config from config.json: {e}"))?; + let (mesh, manifest_plan) = + plan_single(&config, source.tensor_info("lm_head.weight").is_some())?; + let weight_origin = WeightOrigin::for_single(&mesh, ctx.gpu); + let weights = + hipfire_runtime::hfq::load_weights_paroquant_llama(&source, &config, ctx.gpu) + .map_err(|e| format!("load_weights_paroquant_llama: {e:?}"))?; + hipfire_runtime::maybe_screen_mmq(&weights, ctx.gpu); + let kv_mode_str = ctx + .kv_mode_override + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .unwrap_or_else(|| hipfire_runtime::config::get().kv_mode.clone()); + let rr = hipfire_runtime::kv_mode::resolve( + &kv_mode_str, + &hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY, + config.head_dim, + ); + if let Some(w) = rr.warning { + eprintln!( + " KV cache: {w} (site {})", + hipfire_runtime::kv_mode::DIR_SAFETENSORS_POLICY.site + ); + } + let dims = llama_kv_dims(&config, ctx.max_seq, Some(ctx.max_seq)); + let kv = + match ::from_mode(rr.mode, KvTarget::Single(ctx.gpu), &dims) + { Ok(kv) => kv, Err(error) => { weights.free_gpu(ctx.gpu); return Err(format!("KvCache: {error}")); } }; - let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { - Ok(scratch) => scratch, - Err(error) => { - let _ = kv.free_gpu(ctx.gpu); - weights.free_gpu(ctx.gpu); - return Err(format!( - "ForwardScratch::new_with_max_seq: {error:?}" - )); - } - }; - ( - config, - weights, - kv, - scratch, - manifest_plan, - None, - mesh, - weight_origin, - ) - } - }; + let scratch = match ForwardScratch::new_with_max_seq(ctx.gpu, &config, ctx.max_seq) { + Ok(scratch) => scratch, + Err(error) => { + let _ = kv.free_gpu(ctx.gpu); + weights.free_gpu(ctx.gpu); + return Err(format!("ForwardScratch::new_with_max_seq: {error:?}")); + } + }; + ( + config, + weights, + kv, + scratch, + manifest_plan, + None, + mesh, + weight_origin, + ) + } + }; let mut bundle = LlamaBundle { config, @@ -756,15 +728,10 @@ impl LlamaBundle { transaction: WeightLoadTransaction, ) -> Result<(), (WeightLoadTransaction, String)> { if self.weight_store.is_some() { - return Err(( - transaction, - "llama: weight store already attached".into(), - )); + return Err((transaction, "llama: weight store already attached".into())); } - let attached = match AttachedWeightStore::from_transaction( - transaction, - self.weight_origin, - ) { + let attached = match AttachedWeightStore::from_transaction(transaction, self.weight_origin) + { Ok(attached) => attached, Err((transaction, error)) => { return Err(( @@ -800,20 +767,18 @@ impl LlamaBundle { #[cfg(test)] mod tests { use super::*; - use hipfire_runtime::llama::ModelArch; use hipfire_runtime::arch_model::ArchModel; - use hipfire_runtime::hfq::{ - write_hfqm_package_mem, HfqFile, HfqMemTensor, - }; + use hipfire_runtime::hfq::{write_hfqm_package_mem, HfqFile, HfqMemTensor}; use hipfire_runtime::kv_backend::KvBackend; use hipfire_runtime::kv_mode::KvMode; - use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; + use hipfire_runtime::llama::ModelArch; use hipfire_runtime::llama::{ forward_scratch_compute, forward_scratch_embed, KvCache, KvCacheExt, KvDims, KvLayers, KvTarget, }; - use hipfire_runtime::weight_store::test_support; + use hipfire_runtime::loader_api::{CaskConfig, LoadCtx, ModelSource, SpecLoadCfg}; use hipfire_runtime::weight_manifest::ShardPolicy; + use hipfire_runtime::weight_store::test_support; use hipfire_runtime::weight_store::{ WeightLoadTransaction, WeightOrigin, WeightProjection, WeightProjectionKind, WeightStore, }; @@ -863,7 +828,6 @@ mod tests { } } - fn fixture_hfq( with_awq_sidecar: bool, with_q_proj_bias: bool, @@ -880,11 +844,7 @@ mod tests { f16_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32]), f16_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32]), f16_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64]), - f32_hfq_tensor( - "model.layers.0.input_layernorm.weight", - &[32], - false, - ), + f32_hfq_tensor("model.layers.0.input_layernorm.weight", &[32], false), f32_hfq_tensor( "model.layers.0.post_attention_layernorm.weight", &[32], @@ -932,10 +892,8 @@ mod tests { .duration_since(UNIX_EPOCH) .expect("system clock before epoch") .as_nanos(); - let path = std::env::temp_dir().join(format!( - "hipfire-g3-{}-{nonce}.hfq", - std::process::id() - )); + let path = + std::env::temp_dir().join(format!("hipfire-g3-{}-{nonce}.hfq", std::process::id())); write_hfqm_package_mem(&path, 0, metadata, &tensors).expect("write HFQ fixture"); let hfq = HfqFile::open(&path).expect("open HFQ fixture"); (path, hfq) @@ -1002,7 +960,10 @@ mod tests { assert_eq!(mesh.n_devices(), 1); assert_eq!(plan.weights.len(), 12); assert_eq!(plan.state.len(), 1); - assert!(plan.collective_schedule.iter().any(|entry| entry.name == "wo")); + assert!(plan + .collective_schedule + .iter() + .any(|entry| entry.name == "wo")); assert!(manifest[0].dtype_constraint.accepts(DType::HFQ4G256)); assert!(manifest[1].dtype_constraint.accepts(DType::MQ4G256)); assert!(manifest[9].dtype_constraint.accepts(DType::F32)); @@ -1036,7 +997,6 @@ mod tests { assert!(transaction.projection("lm_head", None, 0).is_some()); } - #[test] fn hfq_float_widening_matches_legacy_f32_representation() { let f16_one = [0x00, 0x3c, 0x00, 0xc0]; @@ -1192,12 +1152,9 @@ mod tests { let scratch = ForwardScratch::new_with_max_seq(&mut gpu, &config, 8) .expect("allocate legacy forward scratch"); let dims = llama_kv_dims(&config, 8, None); - let mut kv = ::from_mode( - KvMode::Q8, - KvTarget::Single(&mut gpu), - &dims, - ) - .expect("allocate legacy KV cache"); + let mut kv = + ::from_mode(KvMode::Q8, KvTarget::Single(&mut gpu), &dims) + .expect("allocate legacy KV cache"); forward_scratch_embed(&mut gpu, &legacy, &config, 1, 0, &scratch) .expect("legacy embedding forward"); forward_scratch_compute(&mut gpu, &legacy, &config, 0, &mut kv, &scratch) @@ -1210,9 +1167,7 @@ mod tests { legacy.free_gpu(&mut gpu); assert_eq!(manifest_logits.len(), legacy_logits.len()); - for (index, (manifest, legacy)) in - manifest_logits.iter().zip(&legacy_logits).enumerate() - { + for (index, (manifest, legacy)) in manifest_logits.iter().zip(&legacy_logits).enumerate() { assert!( (manifest - legacy).abs() <= 1e-5, "logit mismatch at index {index}: manifest={manifest} legacy={legacy}" @@ -1221,7 +1176,6 @@ mod tests { std::fs::remove_file(path).expect("remove HFQ fixture"); } - #[test] fn physical_cap_is_honored_by_upstream_kv_constructor() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { @@ -1234,12 +1188,9 @@ mod tests { max_seq: 8, physical_cap: Some(4), }; - let cache = ::from_mode( - KvMode::Q8, - KvTarget::Single(&mut gpu), - &dims, - ) - .expect("upstream Q8 constructor"); + let cache = + ::from_mode(KvMode::Q8, KvTarget::Single(&mut gpu), &dims) + .expect("upstream Q8 constructor"); assert_eq!(cache.max_seq, 8); assert_eq!(cache.physical_cap, 4); let _ = cache.free_gpu(&mut gpu); diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index 26178a122..90b38444e 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -6,9 +6,8 @@ //! per-tensor dequant), which `WeightSource::read_layer` calls internally. use crate::llama::{EmbeddingFormat, WeightTensor}; -use hipfire_hardware::{DeviceMesh, DimKind}; use hip_bridge::HipResult; -use hipfire_hardware::Gpus; +use hipfire_hardware::{DeviceMesh, DimKind, Gpus}; use rdna_compute::{Gpu, GpuTensor}; /// Where each piece of the model lands across a device slice. `single` = the diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 65a9172f4..2f2923194 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -30,9 +30,7 @@ use std::collections::HashSet; pub fn collective_for_policy(policy: &ShardPolicy) -> Option { match policy { ShardPolicy::RowShard { .. } => Some(CollectiveHint::AllReduce { kind: DimKind::Tp }), - ShardPolicy::ExpertSharded { .. } => { - Some(CollectiveHint::AllReduce { kind: DimKind::Ep }) - } + ShardPolicy::ExpertSharded { .. } => Some(CollectiveHint::AllReduce { kind: DimKind::Ep }), ShardPolicy::ExpertTensorSharded { inner, .. } => collective_for_policy(inner), _ => None, } @@ -422,9 +420,10 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } | ShardPolicy::VocabShard { axis } => { - let dim = entry.logical_shape.get(*axis).ok_or_else(|| { - format!("{context}: shard axis {axis} outside logical shape") - })?; + let dim = entry + .logical_shape + .get(*axis) + .ok_or_else(|| format!("{context}: shard axis {axis} outside logical shape"))?; if tp > 1 && dim % tp != 0 { return Err(format!( "{context}: shard dim {dim} (axis {axis}) not divisible by Tp={tp}" @@ -511,11 +510,14 @@ pub fn validate_manifest(manifest: &[WeightEntry], mesh: &DeviceMesh) -> Result< )); } let axis = match inner.as_ref() { - ShardPolicy::ColumnShard { axis: 1 } - | ShardPolicy::RowShard { axis: 2 } => match inner.as_ref() { - ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => *axis, - _ => unreachable!(), - }, + ShardPolicy::ColumnShard { axis: 1 } | ShardPolicy::RowShard { axis: 2 } => { + match inner.as_ref() { + ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => { + *axis + } + _ => unreachable!(), + } + } ShardPolicy::ColumnShard { axis } | ShardPolicy::RowShard { axis } => { return Err(format!( "{context}: ExpertTensorSharded inner axis {axis} is incompatible with [expert, projection, hidden]" @@ -591,7 +593,10 @@ pub fn plan_manifest( }) .collect(); let band_xfers = (0..n_layers) - .filter_map(|layer| mesh.band_xfer_after(layer, n_layers).map(|hint| (layer, hint))) + .filter_map(|layer| { + mesh.band_xfer_after(layer, n_layers) + .map(|hint| (layer, hint)) + }) .collect(); Ok(ManifestPlan { weights: weight_placements, @@ -728,23 +733,14 @@ fn manifest_entry<'a>( .ok_or_else(|| format!("{context}: {label} reference '{name}' not found")) } -fn source_policy_matches( - spec: &ExpertGroupSpec, - label: &str, - policy: &ShardPolicy, -) -> bool { +fn source_policy_matches(spec: &ExpertGroupSpec, label: &str, policy: &ShardPolicy) -> bool { match spec.parallelism { ExpertParallelism::Single => matches!( policy, - ShardPolicy::Replicate - | ShardPolicy::Pin(_) - | ShardPolicy::Tied { .. } + ShardPolicy::Replicate | ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } ), ExpertParallelism::TensorParallel => match (label, policy) { - ( - "gate_up" | "gate" | "up", - ShardPolicy::ExpertTensorSharded { n_experts, inner }, - ) => { + ("gate_up" | "gate" | "up", ShardPolicy::ExpertTensorSharded { n_experts, inner }) => { *n_experts == spec.n_experts && matches!(inner.as_ref(), ShardPolicy::ColumnShard { axis: 1 }) } @@ -794,10 +790,7 @@ fn source_shape_matches( Ok(()) } -fn validate_expert_sources( - spec: &ExpertGroupSpec, - manifest: &[WeightEntry], -) -> Result<(), String> { +fn validate_expert_sources(spec: &ExpertGroupSpec, manifest: &[WeightEntry]) -> Result<(), String> { let context = expert_context(spec); let router = manifest_entry(spec, manifest, "router", &spec.router)?; if !matches!(router.logical_shape.len(), 1 | 2) @@ -884,13 +877,19 @@ pub fn validate_expert_group_specs( for spec in specs { let context = expert_context(spec); if spec.group.is_empty() || spec.router.is_empty() || spec.execution.is_empty() { - return Err(format!("{context}: group/router/execution identities must be non-empty")); + return Err(format!( + "{context}: group/router/execution identities must be non-empty" + )); } if spec.n_experts == 0 || spec.resources.bytes_per_expert == 0 { - return Err(format!("{context}: n_experts and bytes_per_expert must be non-zero")); + return Err(format!( + "{context}: n_experts and bytes_per_expert must be non-zero" + )); } if spec.resources.alignment == 0 || !spec.resources.alignment.is_power_of_two() { - return Err(format!("{context}: alignment must be a non-zero power of two")); + return Err(format!( + "{context}: alignment must be a non-zero power of two" + )); } if !groups.insert((&spec.group, spec.layer)) { return Err(format!("{context}: duplicate group/layer identity")); @@ -1002,13 +1001,7 @@ mod tests { #[test] fn expert_source_identity_and_shape_are_checked() { let manifest = vec![ - WeightEntry::layer( - "router", - 0, - vec![8, 4], - DType::F16, - ShardPolicy::Replicate, - ), + WeightEntry::layer("router", 0, vec![8, 4], DType::F16, ShardPolicy::Replicate), WeightEntry::layer( "gate_up", 0, @@ -1089,12 +1082,7 @@ mod tests { #[test] fn tied_entries_require_matching_representation_and_no_tied_chain() { - let source = WeightEntry::model( - "source", - vec![8, 8], - DType::F16, - ShardPolicy::Replicate, - ); + let source = WeightEntry::model("source", vec![8, 8], DType::F16, ShardPolicy::Replicate); let shape_mismatch = WeightEntry::model( "shape_mismatch", vec![8, 4], @@ -1161,16 +1149,15 @@ mod tests { source: "cycle_a".into(), }, ); - assert!(validate_manifest(&[cycle_a, cycle_b], &DeviceMesh::single().expect("single-device mesh construction cannot overflow")).is_err()); + assert!(validate_manifest( + &[cycle_a, cycle_b], + &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + ) + .is_err()); } #[test] fn tied_entries_reject_different_source_sets_with_equal_logical_dtype() { - let source = WeightEntry::model( - "source", - vec![8, 8], - DType::F16, - ShardPolicy::Replicate, - ); + let source = WeightEntry::model("source", vec![8, 8], DType::F16, ShardPolicy::Replicate); let tied = WeightEntry::model_with_dtype_constraint( "tied", vec![8, 8], @@ -1182,7 +1169,7 @@ mod tests { ); let error = validate_manifest( &[source, tied], - &DeviceMesh::single().expect("single-device mesh construction cannot overflow") + &DeviceMesh::single().expect("single-device mesh construction cannot overflow"), ) .unwrap_err(); assert!(error.contains("source dtype contract")); diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index ce0019441..d67bf06c0 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -36,9 +36,7 @@ thread_local! { /// [`clear_faults`] after it so a failed test cannot poison a later one. #[doc(hidden)] pub mod test_support { - use super::{ - FAIL_AFTER_UPLOAD, RESIDENT_ALLOCATIONS, RESIDENT_RELEASES, - }; + use super::{FAIL_AFTER_UPLOAD, RESIDENT_ALLOCATIONS, RESIDENT_RELEASES}; pub fn reset() { RESIDENT_ALLOCATIONS.with(|count| count.set(0)); @@ -125,7 +123,12 @@ pub struct WeightProjection { pub dtype: DType, } -fn projection_for(entry: &WeightEntry, rank: usize, world_size: usize, dtype: DType) -> WeightProjection { +fn projection_for( + entry: &WeightEntry, + rank: usize, + world_size: usize, + dtype: DType, +) -> WeightProjection { let (kind, axis) = match &entry.policy { ShardPolicy::ColumnShard { axis } => (WeightProjectionKind::ColumnShard, Some(*axis)), ShardPolicy::RowShard { axis } => (WeightProjectionKind::RowShard, Some(*axis)), @@ -307,12 +310,7 @@ impl WeightLoadTransaction { .is_some_and(|store| store.contains(name, layer, device)) } - pub fn get( - &self, - name: &str, - layer: Option, - device: usize, - ) -> Option<&WeightHandle> { + pub fn get(&self, name: &str, layer: Option, device: usize) -> Option<&WeightHandle> { self.store .as_ref() .and_then(|store| store.get(name, layer, device)) @@ -338,14 +336,12 @@ impl WeightLoadTransaction { /// Compare the unpublished transaction's captured target with an admitted /// owner identity. This read-only check is used before the carrier wraps /// the transaction in its private attached owner. - pub fn validate_origin_value( - &self, - expected: WeightOrigin, - ) -> Result<(), WeightStoreError> { - self.store.as_ref().map_or( - Err(WeightStoreError::UnboundOrigin), - |store| store.validate_origin_value(expected), - ) + pub fn validate_origin_value(&self, expected: WeightOrigin) -> Result<(), WeightStoreError> { + self.store + .as_ref() + .map_or(Err(WeightStoreError::UnboundOrigin), |store| { + store.validate_origin_value(expected) + }) } /// Start typed assembly while this load is still unpublished. @@ -459,12 +455,7 @@ impl WeightStore { /// Move a handle out of the store. This is private to the assembly /// capability so arbitrary store holders cannot independently tear down a /// resident allocation. - fn take( - &mut self, - name: &str, - layer: Option, - device: usize, - ) -> Option { + fn take(&mut self, name: &str, layer: Option, device: usize) -> Option { let key = WeightPlacementKey::new(name, layer, device); self.projections.remove(&key); self.placements.remove(&key) @@ -492,10 +483,7 @@ impl WeightStore { /// Compare a store's captured origin with an already-resolved target /// identity. This read-only seam cannot release or extract any handle. - pub fn validate_origin_value( - &self, - expected: WeightOrigin, - ) -> Result<(), WeightStoreError> { + pub fn validate_origin_value(&self, expected: WeightOrigin) -> Result<(), WeightStoreError> { let actual = self.origin.ok_or(WeightStoreError::UnboundOrigin)?; if actual != expected { return Err(WeightStoreError::OriginMismatch { expected, actual }); @@ -505,11 +493,7 @@ impl WeightStore { /// Verify that this store is still being handled by the same mesh/device /// target. No GPU calls occur on mismatch. - pub fn validate_origin( - &self, - mesh: &DeviceMesh, - gpu: &Gpu, - ) -> Result<(), WeightStoreError> { + pub fn validate_origin(&self, mesh: &DeviceMesh, gpu: &Gpu) -> Result<(), WeightStoreError> { self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) } @@ -558,12 +542,7 @@ pub struct WeightStoreAssembly<'a> { } impl<'a> WeightStoreAssembly<'a> { - pub fn take( - &mut self, - name: &str, - layer: Option, - device: usize, - ) -> Option { + pub fn take(&mut self, name: &str, layer: Option, device: usize) -> Option { let key = WeightPlacementKey::new(name, layer, device); let (handle, projection) = self.store.take_with_projection(name, layer, device)?; let slot = self.taken.len(); @@ -586,9 +565,7 @@ impl Drop for WeightStoreAssembly<'_> { return; } for taken in self.taken.drain(..) { - let _ = self - .store - .insert(taken.key, taken.handle, taken.projection); + let _ = self.store.insert(taken.key, taken.handle, taken.projection); } } } @@ -628,11 +605,7 @@ fn target_error(mesh: &DeviceMesh) -> Option { ), }) } -fn rollback_fulfill_error( - store: WeightStore, - gpu: &Gpu, - mut error: FulfillError, -) -> FulfillError { +fn rollback_fulfill_error(store: WeightStore, gpu: &Gpu, mut error: FulfillError) -> FulfillError { if let Err(release_error) = store.rollback(gpu) { error .reason @@ -680,15 +653,15 @@ where name: entry.name.clone(), layer: entry.layer, device: devices.first().copied().unwrap_or(0), - reason: format!( - "Single placement resolved to {:?}, expected [0]", - devices - ), + reason: format!("Single placement resolved to {:?}, expected [0]", devices), }; return Err(rollback_fulfill_error(store, gpu, error)); } let key = WeightPlacementKey::new(&entry.name, entry.layer, 0); - if let ShardPolicy::Tied { source: source_name } = &entry.policy { + if let ShardPolicy::Tied { + source: source_name, + } = &entry.policy + { let source_dtype = match store.get(source_name, entry.layer, 0) { Some(WeightHandle::Resident(tensor)) => Some(tensor.dtype), Some(WeightHandle::Alias(_)) | None => None, @@ -717,11 +690,9 @@ where return Err(rollback_fulfill_error(store, gpu, error)); } let projection = projection_for(entry, 0, 1, actual_dtype); - if let Err(reason) = store.insert( - key, - WeightHandle::Alias(source_name.clone()), - projection, - ) { + if let Err(reason) = + store.insert(key, WeightHandle::Alias(source_name.clone()), projection) + { let error = FulfillError { name: entry.name.clone(), layer: entry.layer, @@ -926,7 +897,9 @@ mod tests { .stage_alias("x", None, 0, "source-b", projection(DType::F32)) .unwrap_err(); assert!(matches!(error, WeightStoreError::DuplicatePlacement(_))); - assert!(matches!(store.get("x", None, 0), Some(WeightHandle::Alias(source)) if source == "source-a")); + assert!( + matches!(store.get("x", None, 0), Some(WeightHandle::Alias(source)) if source == "source-a") + ); assert_eq!(store.projection("x", None, 0).unwrap().dtype, DType::F16); } @@ -969,13 +942,9 @@ mod tests { source: "source".into(), }, ); - let transaction = fulfill_manifest_single( - &[source, alias], - &mesh, - 1, - &gpu, - |_| Ok((vec![0; 4], DType::F32)), - ) + let transaction = fulfill_manifest_single(&[source, alias], &mesh, 1, &gpu, |_| { + Ok((vec![0; 4], DType::F32)) + }) .unwrap(); assert_eq!( transaction.projection("alias", None, 0).unwrap().dtype, @@ -997,10 +966,9 @@ mod tests { }; let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { - Ok((vec![0; 4], DType::F32)) - }) - .unwrap(); + let transaction = + fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| Ok((vec![0; 4], DType::F32))) + .unwrap(); assert_eq!(transaction.len(), 1); assert!(matches!( transaction.get("resident", None, 0), @@ -1041,10 +1009,9 @@ mod tests { RESIDENT_RELEASES.with(|count| count.set(0)); let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); let entry = WeightEntry::model("resident", vec![1], DType::F32, ShardPolicy::Replicate); - let transaction = fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| { - Ok((vec![0; 4], DType::F32)) - }) - .unwrap(); + let transaction = + fulfill_manifest_single(&[entry], &mesh, 1, &gpu, |_| Ok((vec![0; 4], DType::F32))) + .unwrap(); let expected = WeightOrigin::from_parts(mesh.epoch(), 1, gpu.device_id); let error = transaction.validate_origin_value(expected).unwrap_err(); assert!(matches!(error, WeightStoreError::OriginMismatch { .. })); @@ -1060,7 +1027,6 @@ mod tests { assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } - #[test] fn rollback_reports_free_failure_without_counting_release() { let Ok(gpu) = Gpu::init() else { @@ -1073,10 +1039,7 @@ mod tests { let mut store = WeightStore::with_origin(origin); let borrowed = GpuTensor { buf: unsafe { - hip_bridge::DeviceBuffer::from_raw( - std::ptr::null_mut::(), - 0, - ) + hip_bridge::DeviceBuffer::from_raw(std::ptr::null_mut::(), 0) }, shape: vec![0], dtype: DType::F32, From 04de36cf8249797cb34d2aeb2b303e669342fd8c Mon Sep 17 00:00:00 2001 From: Bjoern Agent Date: Tue, 1 Sep 2026 08:56:53 +0200 Subject: [PATCH 10/17] fix(device-mesh): recognize alternate lm head names --- crates/hipfire-arch-llama/src/carrier.rs | 59 ++++++++++++++++++++---- 1 file changed, 51 insertions(+), 8 deletions(-) diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index a6bd292c0..c40e55ad5 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -117,16 +117,26 @@ fn hfq_layer_names(layer: usize, relative: &str) -> Vec { format!("layers.{layer}.{relative}.weight"), ] } +const HFQ_LM_HEAD_NAMES: &[&str] = &[ + "lm_head.weight", + "model.lm_head.weight", + "model.language_model.lm_head.weight", +]; + +fn hfq_has_separate_lm_head(hfq: &HfqFile) -> bool { + HFQ_LM_HEAD_NAMES + .iter() + .any(|name| hfq.find_tensor_info(name).is_some()) +} fn hfq_entry_names(entry: &WeightEntry) -> Result, String> { let names = match (entry.name.as_str(), entry.layer) { ("token_embd", None) => vec!["model.embed_tokens.weight".to_string()], ("output_norm", None) => vec!["model.norm.weight".to_string()], - ("lm_head", None) => vec![ - "lm_head.weight".to_string(), - "model.lm_head.weight".to_string(), - "model.language_model.lm_head.weight".to_string(), - ], + ("lm_head", None) => HFQ_LM_HEAD_NAMES + .iter() + .map(|name| (*name).to_string()) + .collect(), ("wq", Some(layer)) => hfq_layer_names(layer, "self_attn.q_proj"), ("wk", Some(layer)) => hfq_layer_names(layer, "self_attn.k_proj"), ("wv", Some(layer)) => hfq_layer_names(layer, "self_attn.v_proj"), @@ -539,7 +549,7 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result (PathBuf, HfqFile) { + fixture_hfq_with_lm_head( + with_awq_sidecar, + with_q_proj_bias, + malformed_output_norm, + separate_lm_head.then_some("lm_head.weight"), + ) + } + + fn fixture_hfq_with_lm_head( + with_awq_sidecar: bool, + with_q_proj_bias: bool, + malformed_output_norm: bool, + lm_head_name: Option<&str>, ) -> (PathBuf, HfqFile) { let mut tensors = vec![ f32_hfq_tensor("model.embed_tokens.weight", &[2, 32], false), @@ -870,8 +894,8 @@ mod tests { 32 * 2, )); } - if separate_lm_head { - tensors.push(f32_hfq_tensor("lm_head.weight", &[2, 32], false)); + if let Some(lm_head_name) = lm_head_name { + tensors.push(f32_hfq_tensor(lm_head_name, &[2, 32], false)); } let metadata = r#"{ "config": { @@ -1064,6 +1088,25 @@ mod tests { std::fs::remove_file(path).expect("remove HFQ fixture"); } + #[test] + fn alternate_explicit_lm_head_names_are_not_tied() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + for name in &HFQ_LM_HEAD_NAMES[1..] { + let (path, hfq) = fixture_hfq_with_lm_head(false, false, false, Some(name)); + assert!(hfq_has_separate_lm_head(&hfq)); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = + load_bundle(ModelSource::Hfq(hfq), &mut ctx).expect("load explicit lm_head"); + drop(ctx); + assert!(!bundle.weights.lm_head_aliases_embd); + Box::new(bundle).free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + } + #[test] fn production_biased_hfq_is_rejected_before_manifest_upload() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { From 5329cd49a6784c7d2194bbff2dce1c638c818adb Mon Sep 17 00:00:00 2001 From: Kaden Schutt <151092359+Kaden-Schutt@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:11:08 -0700 Subject: [PATCH 11/17] fix(device-mesh): seal manifest ownership and teardown --- crates/hipfire-arch-llama/src/arch_model.rs | 32 +- crates/hipfire-arch-llama/src/carrier.rs | 339 ++++++++++++++++--- crates/hipfire-daemon/src/main.rs | 108 ++++-- crates/hipfire-loader/src/lib.rs | 178 +++++++++- crates/hipfire-runtime/src/arch_model.rs | 8 + crates/hipfire-runtime/src/hfq.rs | 25 +- crates/hipfire-runtime/src/weight_backend.rs | 102 ++++++ crates/hipfire-runtime/src/weight_store.rs | 210 ++++++++++-- 8 files changed, 869 insertions(+), 133 deletions(-) diff --git a/crates/hipfire-arch-llama/src/arch_model.rs b/crates/hipfire-arch-llama/src/arch_model.rs index 7f374f20c..7c16b1ca2 100644 --- a/crates/hipfire-arch-llama/src/arch_model.rs +++ b/crates/hipfire-arch-llama/src/arch_model.rs @@ -34,7 +34,27 @@ impl ArchModel for LlamaBundle { Ok(()) } + fn validate_teardown_device(&self, device_id: i32) -> Result<(), String> { + if let Some(store) = self.weight_store.as_ref() { + store + .validate_device(device_id) + .map_err(|e| e.to_string())?; + } + Ok(()) + } + fn free_gpu(self: Box, gpu: &mut Gpu) { + let mut bundle = *self; + // Preflight in `unload_model` already validated the device. If drain + // still fails (e.g., hip free error), keep fail-closed quarantine + // semantics but do not claim the failure is retryable. + if let Some(store) = bundle.weight_store.take() { + if let Err((_, error)) = store.drain(gpu) { + eprintln!("llama: failed to release attached weight store: {error}"); + std::mem::forget(bundle); + return; + } + } let LlamaBundle { config: _, weights, @@ -47,17 +67,9 @@ impl ArchModel for LlamaBundle { dflash_extract_layers: _, dspark_weights: _, dspark_assets: _, - } = *self; - // Mirror the existing unload ordering: scratch → store/weights → kv. + } = bundle; + drop(weight_store); scratch.free_gpu(gpu); - if let Some(store) = weight_store { - // Attachment already checked the complete origin and created this - // owner capability. There is no mismatch branch to leak the model: - // an attached store can only be drained by this consuming owner. - if let Err(error) = store.drain(gpu) { - eprintln!("llama: failed to release attached weight store: {error}"); - } - } weights.free_gpu(gpu); let _ = kv.free_gpu(gpu); } diff --git a/crates/hipfire-arch-llama/src/carrier.rs b/crates/hipfire-arch-llama/src/carrier.rs index c40e55ad5..f053c3344 100644 --- a/crates/hipfire-arch-llama/src/carrier.rs +++ b/crates/hipfire-arch-llama/src/carrier.rs @@ -77,8 +77,25 @@ impl AttachedWeightStore { Ok(Self { transaction }) } - pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) -> hip_bridge::HipResult<()> { - self.transaction.rollback(gpu) + /// Read-only device gate for preflight. Performs zero GPU calls. + pub(crate) fn validate_device(&self, device_id: i32) -> Result<(), WeightStoreError> { + self.transaction.validate_device(device_id) + } + + /// Gated drain that validates physical device before any GPU call. + pub(crate) fn try_drain( + &mut self, + gpu: &mut rdna_compute::Gpu, + ) -> Result<(), WeightStoreError> { + self.transaction.try_rollback(gpu) + } + + pub(crate) fn drain(self, gpu: &mut rdna_compute::Gpu) -> Result<(), (Self, WeightStoreError)> { + let mut me = self; + match me.try_drain(gpu) { + Ok(()) => Ok(()), + Err(error) => Err((me, error)), + } } } @@ -89,6 +106,26 @@ fn with_weight_rollback_error(reason: String, rollback: hip_bridge::HipResult<() } } +fn with_gated_rollback_error( + reason: String, + rollback: Result<(), (WeightLoadTransaction, WeightStoreError)>, +) -> String { + match rollback { + Ok(()) => reason, + Err((_, error)) => format!("{reason}; resident rollback failed: {error}"), + } +} + +fn with_store_drain_error( + reason: String, + drain: Result<(), (AttachedWeightStore, WeightStoreError)>, +) -> String { + match drain { + Ok(()) => reason, + Err((_, error)) => format!("{reason}; weight store drain failed: {error}"), + } +} + fn plan_single( config: &LlamaConfig, has_separate_lm_head: bool, @@ -157,7 +194,10 @@ fn hfq_entry_names(entry: &WeightEntry) -> Result, String> { Ok(names) } -fn hfq_entry_data(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, u8), String> { +fn hfq_entry_data( + hfq: &HfqFile, + entry: &WeightEntry, +) -> Result<(Vec, hipfire_runtime::hfq::HfqTensorInfo), String> { for name in hfq_entry_names(entry)? { if let Some((info, data)) = hfq.tensor_data_vec(&name) { if !matches!( @@ -174,12 +214,29 @@ fn hfq_entry_data(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, u8), S )); } } - return Ok((data, info.quant_type)); + // Authority: HFQ tensor shape must exactly match the manifest logical shape. + let expected_shape: Vec = entry.logical_shape.iter().map(|&d| d as u32).collect(); + if info.shape != expected_shape { + return Err(format!( + "llama: shape mismatch for {}[layer {:?}] (candidate {name}): HFQ shape {:?} != manifest logical_shape {:?}", + entry.name, entry.layer, info.shape, entry.logical_shape + )); + } + return Ok((data, info.clone())); } } if entry.name == "lm_head" && entry.layer.is_none() { if let Some((info, data)) = hfq.tensor_data_vec("model.embed_tokens.weight") { - return Ok((data, info.quant_type)); + // Tied fallback: still validate that the source shape matches the declared lm_head shape + // (manifest declares lm_head logical_shape as [vocab, dim], same as embed). + let expected_shape: Vec = entry.logical_shape.iter().map(|&d| d as u32).collect(); + if info.shape != expected_shape { + return Err(format!( + "llama: tied lm_head shape mismatch: HFQ embed shape {:?} != manifest logical_shape {:?}", + info.shape, entry.logical_shape + )); + } + return Ok((data, info.clone())); } } Err(format!( @@ -235,32 +292,48 @@ fn f32_bytes_from_hfq(quant_type: u8, data: &[u8], name: &str) -> Result } fn hfq_source(hfq: &HfqFile, entry: &WeightEntry) -> Result<(Vec, DType), String> { - let (data, quant_type) = hfq_entry_data(hfq, entry)?; + let (data, info) = hfq_entry_data(hfq, entry)?; + let quant_type = info.quant_type; let name = format!("{}[layer {:?}]", entry.name, entry.layer); - if entry.name == "token_embd" { - return match quant_type { - 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), - 3 => Ok((data, DType::Q8_0)), - 4 => Ok((data, DType::Q4K)), - 6 => Ok((data, DType::HFQ4G256)), - 7 => Ok((data, DType::HFQ4G128)), - other => Err(format!( - "{name}: quant_type={other} is unsupported for a LLaMA embedding" - )), - }; - } - if matches!( + let (bytes, dtype) = if entry.name == "token_embd" { + match quant_type { + 1 | 2 | 16 => (f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32), + 3 => (data, DType::Q8_0), + 4 => (data, DType::Q4K), + 6 => (data, DType::HFQ4G256), + 7 => (data, DType::HFQ4G128), + other => { + return Err(format!( + "{name}: quant_type={other} is unsupported for a LLaMA embedding" + )) + } + } + } else if matches!( entry.name.as_str(), "output_norm" | "attn_norm" | "ffn_norm" | "q_norm" | "k_norm" ) { - return Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)); - } - match quant_type { - 1 | 2 | 16 => Ok((f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32)), - other => hfq_weight_dtype(other) - .map(|dtype| (data, dtype)) - .ok_or_else(|| format!("{name}: unsupported HFQ quant_type={other}")), - } + (f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32) + } else { + match quant_type { + 1 | 2 | 16 => (f32_bytes_from_hfq(quant_type, &data, &name)?, DType::F32), + other => { + let dtype = hfq_weight_dtype(other) + .ok_or_else(|| format!("{name}: unsupported HFQ quant_type={other}"))?; + (data, dtype) + } + } + }; + // Route every dtype (quantized or raw host float widened to F32) through the canonical + // layout/K-divisibility/exact-byte validator before any GPU upload. This reuses the single + // source of truth in weight_backend and fails malformed payloads before upload_raw. + hipfire_runtime::weight_backend::validate_weight_payload( + dtype, + bytes.len(), + &entry.logical_shape, + &name, + ) + .map_err(|e| format!("llama: payload validation for {name} failed: {e}"))?; + Ok((bytes, dtype)) } fn take_slot( @@ -573,10 +646,13 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result weights, Err(error) => { - return Err(with_weight_rollback_error( - error, - transaction.rollback(ctx.gpu), - )); + let rollback = transaction.try_rollback(ctx.gpu); + return Err(match rollback { + Ok(()) => error, + Err(rb_err) => { + format!("{error}; resident rollback failed: {rb_err}") + } + }); } }; (weights, Some(transaction)) @@ -589,16 +665,17 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result scratch, Err(error) => { - let rollback = if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu) - } else { - Ok(()) - }; + let rollback: Result<(), String> = + if let Some(mut transaction) = weight_store.take() { + transaction.try_rollback(ctx.gpu).map_err(|e| e.to_string()) + } else { + Ok(()) + }; weights.free_gpu(ctx.gpu); - return Err(with_weight_rollback_error( - format!("llama: ForwardScratch::new_with_max_seq failed: {error:?}"), - rollback, - )); + return Err(match rollback { + Ok(()) => format!("llama: ForwardScratch::new_with_max_seq failed: {error:?}"), + Err(rb_err) => format!("llama: ForwardScratch::new_with_max_seq failed: {error:?}; resident rollback failed: {rb_err}"), + }); } }; let dims = llama_kv_dims(&config, ctx.max_seq, None); @@ -615,16 +692,17 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result kv, Err(error) => { scratch.free_gpu(ctx.gpu); - let rollback = if let Some(transaction) = weight_store.take() { - transaction.rollback(ctx.gpu) - } else { - Ok(()) - }; + let rollback: Result<(), String> = + if let Some(mut transaction) = weight_store.take() { + transaction.try_rollback(ctx.gpu).map_err(|e| e.to_string()) + } else { + Ok(()) + }; weights.free_gpu(ctx.gpu); - return Err(with_weight_rollback_error( - format!("llama: ::from_mode failed: {error}"), - rollback, - )); + return Err(match rollback { + Ok(()) => format!("llama: ::from_mode failed: {error}"), + Err(rb_err) => format!("llama: ::from_mode failed: {error}; resident rollback failed: {rb_err}"), + }); } }; ( @@ -709,18 +787,21 @@ pub fn load_bundle(src: ModelSource, ctx: &mut LoadCtx) -> Result error, + Err(rb_err) => format!("{error}; resident rollback failed: {rb_err}"), + }); } } Ok(bundle) @@ -895,7 +976,13 @@ mod tests { )); } if let Some(lm_head_name) = lm_head_name { - tensors.push(f32_hfq_tensor(lm_head_name, &[2, 32], false)); + // AWQ sidecar path uses F16 for lm_head in legacy loader tests; + // use F16 when a sidecar is present so the legacy load_weight_tensor succeeds. + if with_awq_sidecar { + tensors.push(f16_hfq_tensor(lm_head_name, &[2, 32])); + } else { + tensors.push(f32_hfq_tensor(lm_head_name, &[2, 32], false)); + } } let metadata = r#"{ "config": { @@ -1107,6 +1194,152 @@ mod tests { } } + #[test] + fn alternate_explicit_lm_head_with_awq_sidecar_is_not_tied() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + // AWQ sidecar forces the legacy route; alternate names must still be + // recognized as explicit heads and not silently tied. + for name in &HFQ_LM_HEAD_NAMES[1..] { + let (path, hfq) = fixture_hfq_with_lm_head(true, false, false, Some(name)); + assert!(hfq_has_separate_lm_head(&hfq)); + assert!(hfq.has_awq_sidecars()); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let bundle = load_bundle(ModelSource::Hfq(hfq), &mut ctx) + .expect("load alternate lm_head with AWQ sidecar"); + drop(ctx); + assert!( + !bundle.weights.lm_head_aliases_embd, + "alternate {name} with AWQ sidecar must not be tied" + ); + Box::new(bundle).free_gpu(&mut gpu); + std::fs::remove_file(path).expect("remove HFQ fixture"); + } + } + + #[test] + fn manifest_shape_mismatch_fails_before_upload() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + // Build a tiny HFQ where embed has wrong shape vs manifest. + let mut tensors = vec![ + // manifest expects [2,32] for token_embd, give [2,16] instead + f32_hfq_tensor("model.embed_tokens.weight", &[2, 16], false), + f32_hfq_tensor("model.norm.weight", &[32], false), + f16_hfq_tensor("model.layers.0.self_attn.q_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.k_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.v_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.self_attn.o_proj.weight", &[32, 32]), + f16_hfq_tensor("model.layers.0.mlp.gate_proj.weight", &[64, 32]), + f16_hfq_tensor("model.layers.0.mlp.up_proj.weight", &[64, 32]), + f16_hfq_tensor("model.layers.0.mlp.down_proj.weight", &[32, 64]), + f32_hfq_tensor("model.layers.0.input_layernorm.weight", &[32], false), + f32_hfq_tensor( + "model.layers.0.post_attention_layernorm.weight", + &[32], + false, + ), + f32_hfq_tensor("lm_head.weight", &[2, 32], false), + ]; + let metadata = r#"{ + "config": { + "model_type": "llama", + "hidden_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "intermediate_size": 64, + "vocab_size": 2, + "head_dim": 32, + "rms_norm_eps": 0.00001, + "max_position_embeddings": 8, + "rope_theta": 10000.0 + } + }"#; + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock") + .as_nanos(); + let path = std::env::temp_dir().join(format!("hipfire-shape-mismatch-{nonce}.hfq")); + write_hfqm_package_mem(&path, 0, metadata, &tensors).expect("write"); + let hfq = HfqFile::open(&path).expect("open"); + let cask = CaskConfig::default(); + let mut ctx = load_ctx(&path, &mut gpu, &cask); + let err = match load_bundle(ModelSource::Hfq(hfq), &mut ctx) { + Ok(_) => panic!("expected load_bundle to fail on shape mismatch"), + Err(e) => e, + }; + assert!( + err.contains("shape mismatch") || err.contains("logical_shape"), + "wrong shape must fail before upload, got: {err}" + ); + drop(ctx); + std::fs::remove_file(path).expect("remove"); + } + + #[test] + fn malformed_quant_payload_fails_before_upload() { + let Ok(mut gpu) = rdna_compute::Gpu::init() else { + return; + }; + // Direct validator check: a truncated TQ2 payload must be rejected + // before any GPU upload via the canonical helper. + let dtype = rdna_compute::DType::TQ2G128; + let logical_shape = vec![32, 128]; + let bad_len = 1; + let err = hipfire_runtime::weight_backend::validate_weight_payload( + dtype, + bad_len, + &logical_shape, + "test.malformed", + ) + .unwrap_err(); + assert!( + err.contains("mismatch") + || err.contains("expects") + || err.contains("blob length") + || err.contains("payload") + ); + // Manifest path: a source returning a truncated quant payload must fail + // with a payload reason before upload_raw, and prior residents must be freed. + let mesh = hipfire_hardware::DeviceMesh::single().expect("mesh"); + let manifest = vec![ + hipfire_runtime::weight_manifest::WeightEntry::model( + "first", + vec![1], + rdna_compute::DType::F32, + hipfire_runtime::weight_manifest::ShardPolicy::Replicate, + ), + hipfire_runtime::weight_manifest::WeightEntry::model( + "bad_quant", + vec![32, 128], + rdna_compute::DType::TQ2G128, + hipfire_runtime::weight_manifest::ShardPolicy::Replicate, + ), + ]; + let fulfill_err = + hipfire_runtime::weight_store::fulfill_manifest(&manifest, &mesh, 1, &gpu, |entry| { + if entry.name == "bad_quant" { + Ok((vec![0u8; 1], rdna_compute::DType::TQ2G128)) + } else { + Ok((vec![0u8; 4], rdna_compute::DType::F32)) + } + }) + .unwrap_err(); + assert!( + fulfill_err.reason.contains("payload") + || fulfill_err.reason.contains("mismatch") + || fulfill_err.reason.contains("expects") + || fulfill_err.reason.contains("blob"), + "malformed quant must fail, got: {}", + fulfill_err.reason + ); + drop(gpu); + } + #[test] fn production_biased_hfq_is_rejected_before_manifest_upload() { let Ok(mut gpu) = rdna_compute::Gpu::init() else { diff --git a/crates/hipfire-daemon/src/main.rs b/crates/hipfire-daemon/src/main.rs index 23d1dc386..7e3dd3a00 100644 --- a/crates/hipfire-daemon/src/main.rs +++ b/crates/hipfire-daemon/src/main.rs @@ -422,7 +422,7 @@ fn write_typed_error( /// model may be published only when that prior unload succeeds (or there was /// no prior model — caller passes `Ok(())` in that case). A failed prior /// unload must never install/emit `loaded` for the new model. -fn ep_deferred_may_publish(prior_unload: &Result<(), String>) -> bool { +fn ep_deferred_may_publish(prior_unload: &Result<(), hipfire_loader::UnloadError>) -> bool { prior_unload.is_ok() } @@ -948,10 +948,14 @@ fn main() { pflash_cfg = None; if let Some(m) = model.take() { if let Err(err) = hipfire_loader::unload_model(m, &mut gpu) { + let reason = err.reason().to_owned(); + if let Some(restored) = err.into_model() { + model.replace(restored); + } emit_uncorrelated_error( &mut stdout, None, - &format!("prior unload failed: {err}"), + &format!("prior unload failed: {reason}"), "internal", false, false, @@ -1123,10 +1127,14 @@ fn main() { pflash_cfg = None; if let Some(m) = model.take() { if let Err(err) = hipfire_loader::unload_model(m, &mut gpu) { + let reason = err.reason().to_owned(); + if let Some(restored) = err.into_model() { + model.replace(restored); + } emit_uncorrelated_error( &mut stdout, None, - &format!("prior unload failed: {err}"), + &format!("prior unload failed: {reason}"), "internal", false, false, @@ -1697,20 +1705,36 @@ fn main() { Ok(()) }; if !ep_deferred_may_publish(&prior_unload) { - let prior_err = prior_unload - .err() - .unwrap_or_else(|| "prior unload failed".to_string()); + let prior_err_owned = match prior_unload { + Ok(()) => "prior unload failed".to_string(), + Err(e) => { + let reason = e.reason().to_owned(); + if let Some(restored) = e.into_model() { + model.replace(restored); + } + reason + } + }; + let prior_err = prior_err_owned.as_str(); // Roll back the newly built EP model — GpuTensor // has no Drop; must free explicitly. - let rollback_err = match hipfire_loader::unload_model(m, &mut gpu) { - Ok(()) => None, - Err(e) => Some(e), - }; - // model stays None; pflash already cleared above. - let msg = ep_deferred_handoff_error_message( - &prior_err, - rollback_err.as_deref(), - ); + let rollback_err_owned: Option = + match hipfire_loader::unload_model(m, &mut gpu) { + Ok(()) => None, + Err(e) => { + let reason = e.reason().to_owned(); + if let Some(restored) = e.into_model() { + // Only restore the new model if no prior retryable model already occupies the slot. + if model.is_none() { + model.replace(restored); + } + } + Some(reason) + } + }; + let rollback_err = rollback_err_owned.as_deref(); + let msg = + ep_deferred_handoff_error_message(prior_err, rollback_err); write_error(&mut stdout, "", &msg); continue; } @@ -3446,25 +3470,43 @@ fn main() { } } pflash_cfg = None; - let unload_result = if let Some(m) = model.take() { - hipfire_loader::unload_model(m, &mut gpu) - } else { - // No model: still retry any process-global pending VMM arenas. - hipfire_loader::ensure_vmm_ready_for_load(&mut gpu) - }; - match unload_result { - Ok(()) => { - let _ = writeln!(stdout, r#"{{"type":"unloaded"}}"#); + if let Some(m) = model.take() { + match hipfire_loader::unload_model(m, &mut gpu) { + Ok(()) => { + let _ = writeln!(stdout, r#"{{"type":"unloaded"}}"#); + } + Err(err) => { + let reason = err.reason().to_owned(); + if let Some(restored) = err.into_model() { + model.replace(restored); + } + emit_uncorrelated_error( + &mut stdout, + None, + &format!( + "unload incomplete: {reason}; VMM arenas retained for retry" + ), + "internal", + false, + false, + ); + } } - Err(err) => { - emit_uncorrelated_error( - &mut stdout, - None, - &format!("unload incomplete: {err}; VMM arenas retained for retry"), - "internal", - false, - false, - ); + } else { + match hipfire_loader::ensure_vmm_ready_for_load(&mut gpu) { + Ok(()) => { + let _ = writeln!(stdout, r#"{{"type":"unloaded"}}"#); + } + Err(err) => { + emit_uncorrelated_error( + &mut stdout, + None, + &format!("unload incomplete: {err}; VMM arenas retained for retry"), + "internal", + false, + false, + ); + } } } batch_scheduler = None; diff --git a/crates/hipfire-loader/src/lib.rs b/crates/hipfire-loader/src/lib.rs index c97fd9d3a..d46ff4f2d 100644 --- a/crates/hipfire-loader/src/lib.rs +++ b/crates/hipfire-loader/src/lib.rs @@ -3621,7 +3621,70 @@ pub fn ensure_vmm_ready_for_load(gpu: &mut rdna_compute::Gpu) -> Result<(), Stri // ─── Unload ─────────────────────────────────────────────────────────── -pub fn unload_model(mut m: LoadedModel, gpu: &mut rdna_compute::Gpu) -> Result<(), String> { +/// Typed unload error that distinguishes a retryable wrong-device preflight +/// failure (carrying the intact `LoadedModel`) from a terminal failure after +/// teardown has started (reason only). +pub enum UnloadError { + /// Preflight device mismatch — no GPU work has occurred, model is intact. + Retryable { model: LoadedModel, reason: String }, + /// Failure after teardown began — model resources may be partially freed. + Terminal(String), +} + +impl std::fmt::Debug for UnloadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UnloadError") + .field("retryable", &self.is_retryable()) + .field("reason", &self.reason()) + .finish_non_exhaustive() + } +} + +impl std::fmt::Display for UnloadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Retryable { reason, .. } => write!(f, "{reason}"), + Self::Terminal(reason) => write!(f, "{reason}"), + } + } +} + +impl std::error::Error for UnloadError {} + +impl UnloadError { + /// Human-readable reason for the failure. + pub fn reason(&self) -> &str { + match self { + Self::Retryable { reason, .. } => reason, + Self::Terminal(reason) => reason, + } + } + + /// Whether this is a retryable preflight failure carrying the model. + pub fn is_retryable(&self) -> bool { + matches!(self, Self::Retryable { .. }) + } + + /// Consume the error and return the intact model if retryable. + pub fn into_model(self) -> Option { + match self { + Self::Retryable { model, .. } => Some(model), + Self::Terminal(_) => None, + } + } +} + +/// Read-only preflight check for single-GPU teardown. Takes only the +/// copy-identity device id so tests never need a fake Gpu. Returns the +/// mismatch reason without performing any GPU work. +fn validate_unload_device(model: &LoadedModel, device_id: i32) -> Result<(), String> { + if let Some(state) = model.state.as_ref() { + state.validate_teardown_device(device_id)?; + } + Ok(()) +} + +pub fn unload_model(mut m: LoadedModel, gpu: &mut rdna_compute::Gpu) -> Result<(), UnloadError> { // EP unload-free. An EP model owns its own `Gpus` (the daemon's single `gpu` // is unused for tp>1). Without this branch a SUCCESSFUL EP unload leaked every // per-rank weight / state / partial. Free per-rank weights → state → partials @@ -3772,7 +3835,7 @@ pub fn unload_model(mut m: LoadedModel, gpu: &mut rdna_compute::Gpu) -> Result<( } let _ = gpu; if let Some(err) = ep_first_err { - return Err(err); + return Err(UnloadError::Terminal(err)); } return Ok(()); // `gpus` drops here, tearing down comms + devices. @@ -3811,16 +3874,21 @@ pub fn unload_model(mut m: LoadedModel, gpu: &mut rdna_compute::Gpu) -> Result<( let _ = gpu; return Ok(()); } + // Preflight: validate teardown device before any GPU call. + // Wrong-device must return the intact model without touching the GPU. + if let Err(reason) = validate_unload_device(&m, gpu.device_id) { + return Err(UnloadError::Retryable { model: m, reason }); + } // Quiesce retained-PM4 (if any) before freeing any captured owner. Unknown // quiescence → keep model quarantined; daemon restart is containment. if let Some(spec) = &mut m.speculator { if let Err(reason) = spec.quiesce(gpu) { eprintln!("dflash verify PM4: unload refused — unknown quiescence: {reason}"); std::mem::forget(m); - return Err(format!( + return Err(UnloadError::Terminal(format!( "dflash verify PM4: unload refused after unknown quiescence ({reason}); \ model remains quarantined until process restart" - )); + ))); } } if let Some(spec) = m.speculator.take() { @@ -3870,7 +3938,7 @@ pub fn unload_model(mut m: LoadedModel, gpu: &mut rdna_compute::Gpu) -> Result<( // failed free_tensor. Success is reported only when none remain. note(gpu.ensure_vmm_cleaned().map_err(|e| e.to_string())); match first_err { - Some(err) => Err(err), + Some(err) => Err(UnloadError::Terminal(err)), None => Ok(()), } } @@ -4842,3 +4910,103 @@ mod registry_tests { ); } } + +#[cfg(test)] +mod unload_preflight_tests { + use super::*; + use hipfire_runtime::arch_model::ArchModel; + use hipfire_runtime::llama::KvCache; + use rdna_compute::Gpu; + use std::sync::atomic::{AtomicUsize, Ordering}; + + static FREE_GPU_CALLS: AtomicUsize = AtomicUsize::new(0); + + struct RejectingBundle { + id: usize, + } + + impl ArchModel for RejectingBundle { + fn dim(&self) -> usize { + 8 + } + fn n_layers(&self) -> usize { + 2 + } + fn vocab_size(&self) -> usize { + 32 + } + fn arch_key(&self) -> &'static str { + "rejecting" + } + fn kv_cache_mut(&mut self) -> Option<&mut KvCache> { + None + } + fn validate_teardown_device(&self, _device_id: i32) -> Result<(), String> { + Err("wrong device: expected 0 got 1".to_string()) + } + fn free_gpu(self: Box, _gpu: &mut Gpu) { + FREE_GPU_CALLS.fetch_add(1, Ordering::SeqCst); + } + } + + #[test] + fn rejecting_preflight_returns_intact_model_without_free_gpu() { + FREE_GPU_CALLS.store(0, Ordering::SeqCst); + let bundle = RejectingBundle { id: 0xCAFE }; + let model = LoadedModel { + arch_id: 99, + pp: 1, + pp_gpus: None, + pp_dn_la_to_device: None, + ep: None, + state: Some(Box::new(bundle) as Box), + deepseek4_eos_tok: 0, + minimax_eos_tok: 0, + qwen35_eos_tok: 0, + mtp_mode: "auto".to_string(), + mtp_k: 3, + mtp_weights_present: false, + tokenizer: None, + seq_pos: 0, + max_seq: 2048, + physical_cap: 0, + eviction: None, + kv_adaptive: None, + conversation_tokens: Vec::new(), + asst_turn_cache: AsstTurnCache::new_from_env(), + prefill_checkpoints: Vec::new(), + dflash_checkpoints: Vec::new(), + decoded_vocab: None, + model_path: "test-mock".to_string(), + speculator: None, + chat_template: None, + rec_temperature: None, + rec_top_p: None, + rec_top_k: None, + rec_min_p: None, + rec_presence_penalty: None, + }; + // Exercise the copy-identity preflight helper directly — no fake Gpu. + let reason = super::validate_unload_device(&model, 1).expect_err("preflight should reject"); + assert!(reason.contains("wrong device")); + // Simulate the production retryable path: helper fails => caller would + // return Err(Retryable { model, reason }) without invoking free_gpu. + // Verify free_gpu was never invoked and the model remains intact. + assert_eq!( + FREE_GPU_CALLS.load(Ordering::SeqCst), + 0, + "free_gpu must not be invoked on preflight failure" + ); + assert_eq!(model.arch_id, 99); + assert_eq!(model.model_path, "test-mock"); + assert!(model.state.is_some()); + // Also verify the full unload_model retryable path still preserves the model + // when driven through the helper (indirectly). The helper is the sole + // device check, so no GPU object is needed for this assertion. + let err = UnloadError::Retryable { model, reason }; + assert!(err.is_retryable()); + let restored = err.into_model().expect("retryable must carry model"); + assert_eq!(restored.arch_id, 99); + assert!(restored.state.is_some()); + } +} diff --git a/crates/hipfire-runtime/src/arch_model.rs b/crates/hipfire-runtime/src/arch_model.rs index dac1b0e9e..9661ba467 100644 --- a/crates/hipfire-runtime/src/arch_model.rs +++ b/crates/hipfire-runtime/src/arch_model.rs @@ -99,6 +99,14 @@ pub trait ArchModel: Send + std::any::Any { /// experiment converted 15 sites of 154 and this hatch is expected to do /// better. + /// Validate that teardown on `gpu` would target the correct physical device. + /// + /// Read-only, performs zero GPU calls. Returning `Err` signals a wrong-device + /// mismatch and allows the caller to preserve the model for retry on the + /// correct device. The default is success so non-pilot arches need no change. + fn validate_teardown_device(&self, _device_id: i32) -> Result<(), String> { + Ok(()) + } /// Return every GPU buffer this model owns. /// diff --git a/crates/hipfire-runtime/src/hfq.rs b/crates/hipfire-runtime/src/hfq.rs index 54a68a0b6..b877ad9c4 100644 --- a/crates/hipfire-runtime/src/hfq.rs +++ b/crates/hipfire-runtime/src/hfq.rs @@ -1595,7 +1595,14 @@ impl WeightSource for LlamaHfqSource<'_> { ) -> HipResult<(WeightTensor, bool)> { let cfg = self.cfg; let hfq = self.hfq; - let has_separate = hfq.find_tensor("lm_head.weight").is_some(); + const LM_HEAD_CANDIDATES: &[&str] = &[ + "lm_head.weight", + "model.lm_head.weight", + "model.language_model.lm_head.weight", + ]; + let has_separate = LM_HEAD_CANDIDATES + .iter() + .any(|name| hfq.find_tensor(name).is_some()); resolve_lm_head( gpu, has_separate, @@ -1605,6 +1612,22 @@ impl WeightSource for LlamaHfqSource<'_> { cfg.vocab_size, cfg.dim, |gpu| { + // Try each explicit LM-head candidate; the first that exists is loaded + // with a direct single-name resolver so the sidecar name is derived + // from the actual on-disk name. + for cand in LM_HEAD_CANDIDATES { + if hfq.find_tensor(cand).is_some() { + if let Ok(w) = + load_weight_tensor(hfq, gpu, cand, cfg.vocab_size, cfg.dim, |n| { + vec![n.to_string()] + }) + { + return Ok(w); + } + } + } + // Fallback to the flat resolver for backward compatibility (covers + // any future naming not in the explicit set). load_weight_tensor( hfq, gpu, diff --git a/crates/hipfire-runtime/src/weight_backend.rs b/crates/hipfire-runtime/src/weight_backend.rs index b06a6076c..67e1e428c 100644 --- a/crates/hipfire-runtime/src/weight_backend.rs +++ b/crates/hipfire-runtime/src/weight_backend.rs @@ -670,6 +670,108 @@ fn validate_lowbit_layout( Ok(()) } +/// Validate that `data_len` exactly matches the published layout for `dtype` and `logical_shape`. +/// +/// This is the canonical codec/layout/K-divisibility/exact-byte validator used by both the legacy +/// HFQ loader and the manifest pilot before any GPU upload. It mirrors `decode_raw_codec`'s checks +/// but is GPU-free. Host-decoded types (F32/F16/BF16) are validated by element count; quantized +/// types are validated by K-divisibility and exact packed length for every layout that has a +/// published byte formula. A future codec need only be added here and in `decode_raw_codec`. +pub fn validate_weight_payload( + dtype: DType, + data_len: usize, + logical_shape: &[usize], + name: &str, +) -> Result<(), String> { + if matches!(dtype, DType::F32 | DType::F16 | DType::BF16) { + let expected = logical_shape + .iter() + .try_fold(1usize, |count, &dim| count.checked_mul(dim)) + .and_then(|elements| elements.checked_mul(dtype.size())); + if expected != Some(data_len) { + return Err(format!( + "source payload for {name} has {data_len} bytes, expected {expected:?} for {dtype:?} {logical_shape:?}" + )); + } + return Ok(()); + } + if logical_shape.len() != 2 { + return Err(format!( + "quant payload {name} has {dtype:?} but logical_shape {logical_shape:?} is not 2D [m,k]" + )); + } + let m = logical_shape[0]; + let k = logical_shape[1]; + if let Some(block_bytes) = lowbit_block_bytes(dtype) { + validate_lowbit_layout(dtype, data_len, m, k, name, block_bytes) + .map_err(|e| e.to_string())?; + return Ok(()); + } + if dtype.requires_k_mod_256() && k % 256 != 0 { + return Err(format!( + "{dtype:?} tensor has K={k} but kernel requires K%256==0 (caller: {name})" + )); + } + match dtype { + DType::MQ4G256V2 => { + let gpr = k / 256; + let expected = m * gpr * 136; + if data_len != expected { + return Err(format!( + "MQ4G256V2 blob length mismatch: expected {expected}, got {data_len} (M={m} K={k} caller: {name})" + )); + } + } + DType::MQ4CG256 => { + let gpr = k / 256; + let expected = m * gpr * MQ4C_GROUP_BYTES; + if data_len != expected { + return Err(format!( + "MQ4CG256 blob length mismatch: expected {expected}, got {data_len} (M={m} K={k} caller: {name})" + )); + } + } + DType::MQ6G256V2 => { + let gpr = k / 256; + let expected = m * gpr * MQ6G256V2_GROUP_BYTES; + if data_len != expected { + return Err(format!( + "MQ6G256V2 blob length mismatch: expected {expected}, got {data_len} (M={m} K={k} caller: {name})" + )); + } + } + DType::MQ5G256V2 => { + let gpr = k / 256; + let expected = m * gpr * MQ5G256V2_GROUP_BYTES; + if data_len != expected { + return Err(format!( + "MQ5G256V2 blob length mismatch: expected {expected}, got {data_len} (M={m} K={k} caller: {name})" + )); + } + } + DType::MQ3G256V2 => { + let gpr = k / 256; + let expected = m * gpr * MQ3G256V2_GROUP_BYTES; + if data_len != expected { + return Err(format!( + "MQ3G256V2 blob length mismatch: expected {expected}, got {data_len} (M={m} K={k} caller: {name})" + )); + } + } + DType::MQ2G256V2 => { + let gpr = k / 256; + let expected = m * gpr * MQ2G256V2_GROUP_BYTES; + if data_len != expected { + return Err(format!( + "MQ2G256V2 blob length mismatch: expected {expected}, got {data_len} (M={m} K={k} caller: {name})" + )); + } + } + _ => {} + } + Ok(()) +} + /// Quant `data` → device `WeightTensor [m, k]`. Moved from /// `hipfire-arch-qwen35::qwen35::load_weight_tensor_raw` (Task 2). pub fn dequant_weight_raw( diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index d67bf06c0..58bd4512f 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -343,6 +343,11 @@ impl WeightLoadTransaction { store.validate_origin_value(expected) }) } + /// Read-only physical-device gate. Performs zero GPU calls. + pub fn validate_device(&self, device_id: i32) -> Result<(), WeightStoreError> { + let store = self.store.as_ref().ok_or(WeightStoreError::UnboundOrigin)?; + store.validate_device(device_id) + } /// Start typed assembly while this load is still unpublished. pub fn begin_assembly(&mut self) -> WeightStoreAssembly<'_> { @@ -352,15 +357,47 @@ impl WeightLoadTransaction { .begin_assembly() } - /// Consume this transaction and release every resident handle it owns. - /// This is intentionally the only rollback operation exposed by the - /// lifecycle API. Successful frees are reflected in the resident-release - /// accounting; any failed HIP free is returned to the caller. - pub fn rollback(mut self, gpu: &Gpu) -> hip_bridge::HipResult<()> { - if let Some(store) = self.store.take() { - store.rollback(gpu) - } else { - Ok(()) + /// Gate on physical device before any GPU call. On mismatch returns + /// an error without consuming the allocation, preserving retryability. + pub fn try_rollback(&mut self, gpu: &Gpu) -> Result<(), WeightStoreError> { + let store = self.store.as_ref().ok_or(WeightStoreError::UnboundOrigin)?; + store.validate_device(gpu.device_id)?; + // Device matches — now consume the store and free. + let store = self.store.take().expect("store was Some after validation"); + store + .release_unchecked(gpu) + .map_err(|e| WeightStoreError::InvalidTarget(format!("rollback hip free failed: {e}"))) + } + + /// Consuming rollback gated on physical device. On mismatch the transaction + /// is returned alongside the error without any GPU call, preserving the + /// allocation for retry on the correct device. + pub fn rollback(mut self, gpu: &Gpu) -> Result<(), (Self, WeightStoreError)> { + match self.try_rollback(gpu) { + Ok(()) => Ok(()), + Err(error) => Err((self, error)), + } + } + + /// Legacy HipResult rollback for internal fulfillment paths where the device + /// is known to be correct (same gpu that created the transaction). It still + /// gates on device but maps the origin error into a HipError for backward + /// compatibility with existing `with_weight_rollback_error` callers. + pub(crate) fn rollback_hip(mut self, gpu: &Gpu) -> hip_bridge::HipResult<()> { + match self.try_rollback(gpu) { + Ok(()) => Ok(()), + Err(WeightStoreError::OriginMismatch { expected, actual }) => { + Err(hip_bridge::HipError::new( + 0, + &format!("weight store origin mismatch: expected {expected:?}, got {actual:?}"), + )) + } + Err(WeightStoreError::UnboundOrigin) => Err(hip_bridge::HipError::new( + 0, + "weight store has no target origin", + )), + Err(WeightStoreError::InvalidTarget(msg)) => Err(hip_bridge::HipError::new(0, &msg)), + Err(e) => Err(hip_bridge::HipError::new(0, &e.to_string())), } } } @@ -497,6 +534,27 @@ impl WeightStore { self.validate_origin_value(WeightOrigin::for_single(mesh, gpu)) } + /// Gate on physical device before any GPU call for this store. + pub fn validate_device(&self, device_id: i32) -> Result<(), WeightStoreError> { + let actual = self.origin.ok_or(WeightStoreError::UnboundOrigin)?; + if actual.physical_device() != device_id { + let expected = + WeightOrigin::from_parts(actual.mesh_epoch(), actual.logical_rank(), device_id); + return Err(WeightStoreError::OriginMismatch { expected, actual }); + } + Ok(()) + } + + /// Gated explicit rollback: validates device before any GPU call. + /// On mismatch the error is returned without any GPU call; the store + /// itself is consumed (its buffers remain allocated) — the transaction + /// layer above preserves retryability by not consuming on mismatch. + pub fn try_rollback(self, gpu: &Gpu) -> Result<(), WeightStoreError> { + self.validate_device(gpu.device_id)?; + self.release_unchecked(gpu) + .map_err(|e| WeightStoreError::InvalidTarget(format!("hip free failed: {e}"))) + } + /// Explicit rollback for a failed transaction. It consumes the partial /// store and frees every resident buffer on the single owning GPU. fn rollback(self, gpu: &Gpu) -> hip_bridge::HipResult<()> { @@ -728,26 +786,19 @@ where }; return Err(rollback_fulfill_error(store, gpu, error)); } - if matches!(dtype, DType::F32 | DType::F16 | DType::BF16) { - let expected_bytes = entry - .logical_shape - .iter() - .try_fold(1usize, |count, &dim| count.checked_mul(dim)) - .and_then(|elements| elements.checked_mul(dtype.size())); - if expected_bytes != Some(bytes.len()) { - let error = FulfillError { - name: entry.name.clone(), - layer: entry.layer, - device: 0, - reason: format!( - "source payload has {} bytes, expected {:?} for {dtype:?} {:?}", - bytes.len(), - expected_bytes, - entry.logical_shape - ), - }; - return Err(rollback_fulfill_error(store, gpu, error)); - } + if let Err(reason) = crate::weight_backend::validate_weight_payload( + dtype, + bytes.len(), + &entry.logical_shape, + &entry.name, + ) { + let error = FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!("payload validation failed for {dtype:?}: {reason}"), + }; + return Err(rollback_fulfill_error(store, gpu, error)); } let mut tensor = match gpu.upload_raw(&bytes, &entry.logical_shape) { Ok(tensor) => tensor, @@ -1051,10 +1102,10 @@ mod tests { projection(DType::F32), ) .expect("insert borrowed resident test handle"); - let error = WeightLoadTransaction::new(store) + let (_tx, error) = WeightLoadTransaction::new(store) .rollback(&gpu) .expect_err("rollback must surface a failed HIP free"); - assert!(error.message.contains("borrowed")); + assert!(error.to_string().contains("borrowed")); assert_eq!(test_support::resident_allocations(), 1); assert_eq!(test_support::resident_releases(), 0); test_support::reset(); @@ -1152,4 +1203,101 @@ mod tests { assert!(error.reason.contains("payload")); assert_eq!(RESIDENT_RELEASES.with(std::cell::Cell::get), 1); } + + #[test] + fn wrong_device_try_rollback_preserves_allocation_and_performs_zero_gpu_calls() { + let Ok(gpu) = Gpu::init() else { + return; + }; + test_support::reset(); + RESIDENT_RELEASES.with(|count| count.set(0)); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + // Create a store whose origin claims a different physical device than `gpu`. + let wrong_origin = WeightOrigin::from_parts(mesh.epoch(), 0, gpu.device_id + 99); + let mut wrong_store = WeightStore::with_origin(wrong_origin); + // Use stage_alias for simplicity — alias rollback does zero GPU calls but still gates. + wrong_store + .stage_alias("x", None, 0, "source", projection(DType::F32)) + .unwrap(); + let mut wrong_tx = WeightLoadTransaction::new(wrong_store); + let releases_before = RESIDENT_RELEASES.with(std::cell::Cell::get); + let err = wrong_tx.try_rollback(&gpu).unwrap_err(); + assert!(matches!(err, WeightStoreError::OriginMismatch { .. })); + // No hip::free must have been called. + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + releases_before + ); + // Allocation/ownership is preserved — transaction still has the entry. + assert!(wrong_tx.contains("x", None, 0)); + assert!(wrong_tx.projection("x", None, 0).is_some()); + // Correct-device retry: a properly-originated transaction succeeds. + let correct_origin = WeightOrigin::for_single(&mesh, &gpu); + let mut correct_store = WeightStore::with_origin(correct_origin); + correct_store + .stage_alias("y", None, 0, "source", projection(DType::F32)) + .unwrap(); + let mut correct_tx = WeightLoadTransaction::new(correct_store); + correct_tx + .try_rollback(&gpu) + .expect("correct device try_rollback must succeed"); + assert!(correct_tx.is_empty()); + test_support::reset(); + } + + #[test] + fn wrong_device_consuming_rollback_preserves_transaction_for_retry() { + let Ok(gpu) = Gpu::init() else { + return; + }; + test_support::reset(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let wrong_origin = WeightOrigin::from_parts(mesh.epoch(), 0, gpu.device_id + 77); + let mut store = WeightStore::with_origin(wrong_origin); + store + .stage_alias("preserved", None, 0, "source", projection(DType::F32)) + .unwrap(); + let tx = WeightLoadTransaction::new(store); + let releases_before = RESIDENT_RELEASES.with(std::cell::Cell::get); + let (returned_tx, err) = tx.rollback(&gpu).unwrap_err(); + assert!(matches!(err, WeightStoreError::OriginMismatch { .. })); + assert_eq!( + RESIDENT_RELEASES.with(std::cell::Cell::get), + releases_before + ); + // Returned transaction still owns the allocation — caller can retry. + assert!(returned_tx.contains("preserved", None, 0)); + // Retry with correct device after fixing origin (simulates retry on correct GPU). + // We prove the allocation is still there by validating it. + assert!(returned_tx + .validate_origin_value(WeightOrigin::from_parts( + mesh.epoch(), + 0, + gpu.device_id + 77 + )) + .is_ok()); + test_support::reset(); + } + + #[test] + fn attached_store_wrong_device_try_drain_preserves() { + // This test exercises the attached-store path indirectly via the + // transaction gate; the carrier's AttachedWeightStore::try_drain + // delegates to the same transaction gate. + let Ok(gpu) = Gpu::init() else { + return; + }; + test_support::reset(); + let mesh = DeviceMesh::single().expect("single-device mesh construction cannot overflow"); + let wrong_origin = WeightOrigin::from_parts(mesh.epoch(), 0, gpu.device_id + 55); + let mut store = WeightStore::with_origin(wrong_origin); + store + .stage_alias("attached", None, 0, "source", projection(DType::F32)) + .unwrap(); + let mut tx = WeightLoadTransaction::new(store); + let err = tx.try_rollback(&gpu).unwrap_err(); + assert!(matches!(err, WeightStoreError::OriginMismatch { .. })); + assert!(tx.contains("attached", None, 0)); + test_support::reset(); + } } From 2458ef5df258a4ed190ec3b5462ce4caa53df0d6 Mon Sep 17 00:00:00 2001 From: Kaden Schutt <151092359+Kaden-Schutt@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:18:06 -0700 Subject: [PATCH 12/17] fix(device-mesh): integrate fail-closed G3 manifest planning Propagate G1 mesh errors through layout and placement planning, preserve retryable teardown ownership in the daemon slot, and refresh generated crate maps. The daemon grows by 39 lines for typed retry restoration and rollback ownership; the ratchet records that explicit lifecycle trade. --- crates/hipfire-arch-llama/map.md | 14 +++--- crates/hipfire-daemon/map.md | 4 +- crates/hipfire-loader/map.md | 6 +-- crates/hipfire-runtime/map.md | 24 ++++++---- crates/hipfire-runtime/src/model_load.rs | 19 ++++---- crates/hipfire-runtime/src/weight_manifest.rs | 46 ++++++++++++------- crates/hipfire-runtime/src/weight_store.rs | 15 +++++- scripts/leanup-thresholds.txt | 9 ++-- 8 files changed, 82 insertions(+), 55 deletions(-) diff --git a/crates/hipfire-arch-llama/map.md b/crates/hipfire-arch-llama/map.md index 75651d8d4..c986fa413 100644 --- a/crates/hipfire-arch-llama/map.md +++ b/crates/hipfire-arch-llama/map.md @@ -22,25 +22,25 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| -| [`src/arch.rs`](src/arch.rs) | 379 | 2 | 0 | -| [`src/arch_model.rs`](src/arch_model.rs) | 59 | 0 | 0 | -| [`src/carrier.rs`](src/carrier.rs) | 134 | 3 | 0 | +| [`src/arch.rs`](src/arch.rs) | 589 | 5 | 0 | +| [`src/arch_model.rs`](src/arch_model.rs) | 76 | 0 | 0 | +| [`src/carrier.rs`](src/carrier.rs) | 1,474 | 4 | 16 | | [`src/dspark_body.rs`](src/dspark_body.rs) | 1,294 | 8 | 0 | | [`src/lib.rs`](src/lib.rs) | 67 | 6 | 0 | | [`src/spec_impl.rs`](src/spec_impl.rs) | 450 | 1 | 0 | ### Public API surface -- [`src/arch.rs`](src/arch.rs): `Llama`, `forward_scratch_layers` +- [`src/arch.rs`](src/arch.rs): `Llama`, `weight_manifest`, `weight_manifest_for_hfq`, `state_manifest`, `forward_scratch_layers` - [`src/arch_model.rs`](src/arch_model.rs): — -- [`src/carrier.rs`](src/carrier.rs): `LlamaBundle`, `load_bundle`, `set_dflash_extract_layers` +- [`src/carrier.rs`](src/carrier.rs): `LlamaBundle`, `load_bundle`, `manifest_mesh`, `set_dflash_extract_layers` - [`src/dspark_body.rs`](src/dspark_body.rs): `Qwen3DrafterAssets`, `load_qwen3_dspark`, `Qwen3DsparkScratch`, `new`, `free_gpu`, `dspark_qwen3_block_forward`, `Qwen3DsparkBody`, `build_qwen3_dspark_body` - [`src/lib.rs`](src/lib.rs): `arch`, `arch_model`, `carrier`, `dspark_body`, `spec_impl`, `hipfire_runtime` - [`src/spec_impl.rs`](src/spec_impl.rs): `LlamaSpecScratch` ### Dependencies (from `Cargo.toml`) -- path: `hip-bridge`, `hipfire-dispatch`, `hipfire-runtime`, `rdna-compute` +- path: `hip-bridge`, `hipfire-dispatch`, `hipfire-hardware`, `hipfire-runtime`, `rdna-compute` - external: — - dev: — - build: — @@ -51,6 +51,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 6 modules · 2,383 lines · 20 public items · 0 tests · 4 examples +- 6 modules · 3,950 lines · 24 public items · 16 tests · 4 examples diff --git a/crates/hipfire-daemon/map.md b/crates/hipfire-daemon/map.md index 442999268..c7f17c85c 100644 --- a/crates/hipfire-daemon/map.md +++ b/crates/hipfire-daemon/map.md @@ -23,7 +23,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| -| [`src/main.rs`](src/main.rs) | 4,152 | 1 | 0 | +| [`src/main.rs`](src/main.rs) | 4,194 | 1 | 0 | | [`src/slots.rs`](src/slots.rs) | 1,526 | 22 | 14 | ### Public API surface @@ -44,6 +44,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 2 modules · 5,678 lines · 23 public items · 14 tests · 0 examples +- 2 modules · 5,720 lines · 23 public items · 14 tests · 0 examples diff --git a/crates/hipfire-loader/map.md b/crates/hipfire-loader/map.md index 86e348ab0..5c9b88487 100644 --- a/crates/hipfire-loader/map.md +++ b/crates/hipfire-loader/map.md @@ -25,14 +25,14 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside |---|---:|---:|---:| | [`src/batch_staging.rs`](src/batch_staging.rs) | 336 | 4 | 0 | | [`src/carriers.rs`](src/carriers.rs) | 2,518 | 11 | 3 | -| [`src/lib.rs`](src/lib.rs) | 4,844 | 95 | 22 | +| [`src/lib.rs`](src/lib.rs) | 5,012 | 99 | 23 | | [`src/spec_build.rs`](src/spec_build.rs) | 233 | 4 | 0 | ### Public API surface - [`src/batch_staging.rs`](src/batch_staging.rs): `BatchStaging`, `qwen_batch_weight_formats_supported`, `qwen_ep_batch_weight_formats_supported`, `stage_continuous_batch` - [`src/carriers.rs`](src/carriers.rs): `Qwen2Carrier`, `Qwen35Carrier`, `LlamaCarrier`, `DotsOcrCarrier`, `Deepseek4Carrier`, `MinimaxCarrier`, `Lfm2MoeCarrier`, `Cohere2MoeCarrier`, `MapleCarrier`, `Gemma4Carrier`, `MuseGlimmerCarrier` -- [`src/lib.rs`](src/lib.rs): `batch_staging`, `carriers`, `spec_build`, `Carrier`, `carrier_for`, `ContinuousBatchRoute`, `continuous_batch_route`, `BenchDecodeRoute`, `bench_decode_route`, `VisionRoute`, `vision_route`, `EpPromptRoute`, +83 more +- [`src/lib.rs`](src/lib.rs): `batch_staging`, `carriers`, `spec_build`, `Carrier`, `carrier_for`, `ContinuousBatchRoute`, `continuous_batch_route`, `BenchDecodeRoute`, `bench_decode_route`, `VisionRoute`, `vision_route`, `EpPromptRoute`, +87 more - [`src/spec_build.rs`](src/spec_build.rs): `Qwen35SlotGuard`, `take`, `model_slot`, `build_speculator` ### Dependencies (from `Cargo.toml`) @@ -48,6 +48,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 4 modules · 7,931 lines · 114 public items · 25 tests · 1 examples +- 4 modules · 8,099 lines · 118 public items · 26 tests · 1 examples diff --git a/crates/hipfire-runtime/map.md b/crates/hipfire-runtime/map.md index 7ba964f5e..509f7cc52 100644 --- a/crates/hipfire-runtime/map.md +++ b/crates/hipfire-runtime/map.md @@ -26,7 +26,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/admission.rs`](src/admission.rs) | 277 | 12 | 8 | | [`src/arch.rs`](src/arch.rs) | 272 | 5 | 0 | | [`src/arch_mapping.rs`](src/arch_mapping.rs) | 117 | 4 | 0 | -| [`src/arch_model.rs`](src/arch_model.rs) | 156 | 1 | 2 | +| [`src/arch_model.rs`](src/arch_model.rs) | 164 | 1 | 2 | | [`src/arch_spec.rs`](src/arch_spec.rs) | 292 | 5 | 0 | | [`src/augmentor.rs`](src/augmentor.rs) | 171 | 5 | 3 | | [`src/bf16_loader.rs`](src/bf16_loader.rs) | 97 | 1 | 4 | @@ -46,17 +46,17 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/ep.rs`](src/ep.rs) | 296 | 2 | 0 | | [`src/eval_common.rs`](src/eval_common.rs) | 231 | 3 | 0 | | [`src/gguf.rs`](src/gguf.rs) | 335 | 20 | 0 | -| [`src/hfq.rs`](src/hfq.rs) | 2,498 | 49 | 11 | +| [`src/hfq.rs`](src/hfq.rs) | 2,530 | 51 | 11 | | [`src/hfq_parallel.rs`](src/hfq_parallel.rs) | 335 | 8 | 2 | | [`src/kv_adaptive.rs`](src/kv_adaptive.rs) | 608 | 24 | 12 | | [`src/kv_backend.rs`](src/kv_backend.rs) | 129 | 1 | 7 | | [`src/kv_mode.rs`](src/kv_mode.rs) | 298 | 10 | 7 | -| [`src/lib.rs`](src/lib.rs) | 78 | 53 | 0 | +| [`src/lib.rs`](src/lib.rs) | 80 | 55 | 0 | | [`src/llama.rs`](src/llama.rs) | 8,738 | 83 | 42 | | [`src/llama_spec.rs`](src/llama_spec.rs) | 617 | 6 | 1 | | [`src/loader_api.rs`](src/loader_api.rs) | 256 | 10 | 4 | | [`src/loop_guard.rs`](src/loop_guard.rs) | 194 | 8 | 4 | -| [`src/model_load.rs`](src/model_load.rs) | 117 | 8 | 1 | +| [`src/model_load.rs`](src/model_load.rs) | 202 | 10 | 3 | | [`src/model_source.rs`](src/model_source.rs) | 92 | 4 | 0 | | [`src/ngram_mod.rs`](src/ngram_mod.rs) | 484 | 11 | 13 | | [`src/paro.rs`](src/paro.rs) | 424 | 9 | 3 | @@ -77,8 +77,10 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | [`src/tool_call.rs`](src/tool_call.rs) | 716 | 7 | 15 | | [`src/tp_shard.rs`](src/tp_shard.rs) | 731 | 25 | 20 | | [`src/triattn.rs`](src/triattn.rs) | 1,355 | 45 | 9 | -| [`src/weight_backend.rs`](src/weight_backend.rs) | 2,038 | 22 | 37 | +| [`src/weight_backend.rs`](src/weight_backend.rs) | 2,149 | 24 | 37 | +| [`src/weight_manifest.rs`](src/weight_manifest.rs) | 1,189 | 35 | 7 | | [`src/weight_pager.rs`](src/weight_pager.rs) | 850 | 31 | 6 | +| [`src/weight_store.rs`](src/weight_store.rs) | 1,314 | 44 | 17 | ### Public API surface @@ -105,17 +107,17 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/ep.rs`](src/ep.rs): `ensure_rank_streams`, `run_layer_program_ep` - [`src/eval_common.rs`](src/eval_common.rs): `verify_ref_sha256`, `verify_slice_md5`, `verify_llama_commit` - [`src/gguf.rs`](src/gguf.rs): `GgmlType`, `from_u32`, `block_size`, `block_bytes`, `tensor_bytes`, `MetaValue`, `as_u32`, `as_f32`, `as_str`, `TensorInfo`, `numel`, `byte_size`, +8 more -- [`src/hfq.rs`](src/hfq.rs): `start_cache_warmup`, `mostly_page_cached`, `mostly_page_cached_memo`, `CacheWarmerGuard`, `HfqTensorInfo`, `HfqTensorManifestEntry`, `HfqSourceIdentity`, `RecommendedSampling`, `HfqFile`, `open`, `open_with_reap_plan`, `attach_overlay`, +37 more +- [`src/hfq.rs`](src/hfq.rs): `start_cache_warmup`, `mostly_page_cached`, `mostly_page_cached_memo`, `CacheWarmerGuard`, `HfqTensorInfo`, `HfqTensorManifestEntry`, `HfqSourceIdentity`, `RecommendedSampling`, `HfqFile`, `open`, `open_with_reap_plan`, `attach_overlay`, +39 more - [`src/hfq_parallel.rs`](src/hfq_parallel.rs): `HFQ_READER_LANES`, `HfqReadJob`, `tensor`, `packed`, `label`, `output_len`, `HfqReadResult`, `read_hfq_jobs_ordered` - [`src/kv_adaptive.rs`](src/kv_adaptive.rs): `KMode`, `bytes_per_head`, `rot_width`, `bits`, `v_bytes_per_head`, `k_buf_bytes_per_layer`, `v_buf_bytes_per_layer`, `cap_min`, `Step`, `Preset`, `KvAdaptive`, `from_preset`, +12 more - [`src/kv_backend.rs`](src/kv_backend.rs): `saddle_core` - [`src/kv_mode.rs`](src/kv_mode.rs): `saddle_core`, `KvModePolicy`, `ResolveResult`, `QWEN35_HFQ_POLICY`, `QWEN35_PARO_POLICY`, `DIR_SAFETENSORS_POLICY`, `LLAMA_HFQ_POLICY`, `HFQ_Q8_ONLY_POLICY`, `QWEN35_PP_POLICY`, `resolve` -- [`src/lib.rs`](src/lib.rs): `admission`, `arch`, `arch_mapping`, `arch_model`, `arch_spec`, `augmentor`, `bf16_loader`, `cache_plan`, `cask`, `config`, `cpu_router`, `ddtree`, +41 more +- [`src/lib.rs`](src/lib.rs): `admission`, `arch`, `arch_mapping`, `arch_model`, `arch_spec`, `augmentor`, `bf16_loader`, `cache_plan`, `cask`, `config`, `cpu_router`, `ddtree`, +43 more - [`src/llama.rs`](src/llama.rs): `ModelArch`, `LlamaConfig`, `from_gguf`, `dequantize_q4_0`, `dequantize_q8_0`, `f16_to_f32`, `f32_to_f16`, `dequantize_q4_k`, `convert_q4k_to_q4f16_g64`, `convert_q4k_to_q4f16_g32`, `dequantize_q6_k`, `ParoRotation`, +71 more - [`src/llama_spec.rs`](src/llama_spec.rs): `verify_block_argmax`, `verify_block_logits`, `verify_block_argmax_capture_gpu`, `verify_block_sampled_capture_gpu`, `verify_tree_logits`, `lm_head_logits_n_rows` - [`src/loader_api.rs`](src/loader_api.rs): `ModelSource`, `from_path`, `arch_id`, `is_dir`, `describe`, `LoadCtx`, `SpecLoadCfg`, `CaskConfig`, `physical_cap`, `physical_cap_with_override` - [`src/loop_guard.rs`](src/loop_guard.rs): `StopReason`, `LoopGuard`, `from_config`, `new`, `off`, `enabled`, `check`, `window_len` -- [`src/model_load.rs`](src/model_load.rs): `Layout`, `single`, `from_gpus`, `device_for_layer`, `output_device`, `LoadedWeights`, `WeightSource`, `load_weights` +- [`src/model_load.rs`](src/model_load.rs): `Layout`, `single`, `from_gpus`, `from_mesh`, `validate`, `device_for_layer`, `output_device`, `LoadedWeights`, `WeightSource`, `load_weights` - [`src/model_source.rs`](src/model_source.rs): `TensorInfo`, `QuantConfig`, `ModelSource`, `open_model` - [`src/ngram_mod.rs`](src/ngram_mod.rs): `HASH_MUL`, `EMPTY`, `NgramModConfig`, `NgramModPool`, `new`, `config`, `occupied`, `clear`, `insert_range`, `draft`, `record_draft_result` - [`src/paro.rs`](src/paro.rs): `repack_awq_to_hfq4g128`, `paro_text_prefix`, `load_paro_weight`, `paro_load_wt`, `paro_load_norm`, `paro_load_f32`, `alias_paro_rotation`, `load_fp16_weight_from_source`, `paro_repack_moe_projection` @@ -136,8 +138,10 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside - [`src/tool_call.rs`](src/tool_call.rs): `ParsedToolCall`, `ToolCallParseResult`, `ToolCallParser`, `HermesJsonParser`, `new`, `Qwen35XmlParser`, `Gemma4NativeParser` - [`src/tp_shard.rs`](src/tp_shard.rs): `ExpertAssign`, `ShardConfig`, `single`, `new`, `new_uneven_experts`, `is_single`, `balanced_range`, `validate`, `q_heads_per_rank`, `kv_heads_per_rank`, `q_head_range`, `kv_head_range`, +13 more - [`src/triattn.rs`](src/triattn.rs): `BandCenter`, `magnitude`, `phase`, `mrl`, `TriAttnCenters`, `new`, `n_bands`, `get`, `set`, `omega`, `save`, `load`, +33 more -- [`src/weight_backend.rs`](src/weight_backend.rs): `hf_name_candidates`, `flat_name_candidates`, `hfq_proj_name`, `hfq_plain_name`, `paro_proj_name`, `paro_plain_name`, `EmbedPlan`, `embed_classify`, `load_embedding`, `embedding_format_dtype`, `load_awq_scale_for`, `f16_bytes_to_f32`, +10 more +- [`src/weight_backend.rs`](src/weight_backend.rs): `hf_name_candidates`, `flat_name_candidates`, `hfq_proj_name`, `hfq_plain_name`, `paro_proj_name`, `paro_plain_name`, `EmbedPlan`, `embed_classify`, `load_embedding`, `embedding_format_dtype`, `load_awq_scale_for`, `f16_bytes_to_f32`, +12 more +- [`src/weight_manifest.rs`](src/weight_manifest.rs): `collective_for_policy`, `PinTarget`, `PlacementHint`, `SourceDType`, `DTypeConstraint`, `any_source`, `source_exact`, `source_from_sources`, `accepts`, `same_source_set`, `FusedQkvLayout`, `ShardPolicy`, +23 more - [`src/weight_pager.rs`](src/weight_pager.rs): `WeightId`, `ExpertRole`, `SharedRole`, `AttnRole`, `NormKind`, `TransferHandle`, `Transport`, `PreadH2DTransport`, `open`, `path`, `PagerConfig`, `WeightPager`, +19 more +- [`src/weight_store.rs`](src/weight_store.rs): `test_support`, `reset`, `arm_fail_after_upload`, `clear_faults`, `resident_allocations`, `resident_releases`, `WeightPlacementKey`, `new`, `WeightProjectionKind`, `WeightProjection`, `WeightHandle`, `WeightOrigin`, +32 more ### Dependencies (from `Cargo.toml`) @@ -152,6 +156,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 56 modules · 50,433 lines · 836 public items · 600 tests · 132 examples +- 58 modules · 53,174 lines · 923 public items · 626 tests · 132 examples diff --git a/crates/hipfire-runtime/src/model_load.rs b/crates/hipfire-runtime/src/model_load.rs index 90b38444e..58f7f1eee 100644 --- a/crates/hipfire-runtime/src/model_load.rs +++ b/crates/hipfire-runtime/src/model_load.rs @@ -7,7 +7,7 @@ use crate::llama::{EmbeddingFormat, WeightTensor}; use hip_bridge::HipResult; -use hipfire_hardware::{DeviceMesh, DimKind, Gpus}; +use hipfire_hardware::{DeviceMesh, DimKind, Gpus, MeshError}; use rdna_compute::{Gpu, GpuTensor}; /// Where each piece of the model lands across a device slice. `single` = the @@ -35,24 +35,24 @@ impl Layout { /// manifest planner owns the full stage grid; this legacy loader view /// selects rank zero for each layer so existing orchestrators continue to /// have one deterministic device index until their typed mesh path lands. - pub fn from_mesh(mesh: &DeviceMesh, n_layers: usize) -> Self { - let mut output_coord = mesh.coord_of(0); + pub fn from_mesh(mesh: &DeviceMesh, n_layers: usize) -> Result { + let mut output_coord = mesh.coord_of(0)?; if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { output_coord[index] = mesh.size_of(DimKind::Pp).saturating_sub(1); } let layer_to_device = (0..n_layers) .map(|layer| { - let mut coord = mesh.coord_of(0); + let mut coord = mesh.coord_of(0)?; if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { coord[index] = mesh.stage_for_layer(layer, n_layers); } mesh.device_of(&coord) }) - .collect(); - Self { - output_device: mesh.device_of(&output_coord), + .collect::, _>>()?; + Ok(Self { + output_device: mesh.device_of(&output_coord)?, layer_to_device, - } + }) } /// Validate the pure layout before any source preparation or GPU upload. @@ -181,7 +181,8 @@ mod tests { fn mesh_layout_selects_stage_rank_zero_without_io() { let mesh = DeviceMesh::rect(&[(DimKind::Pp, 2), (DimKind::Tp, 2)]) .expect("small test mesh construction cannot overflow"); - let layout = Layout::from_mesh(&mesh, 4); + let layout = + Layout::from_mesh(&mesh, 4).expect("valid mesh must produce a deterministic layout"); assert_eq!(layout.output_device(), 2); assert_eq!( (0..4) diff --git a/crates/hipfire-runtime/src/weight_manifest.rs b/crates/hipfire-runtime/src/weight_manifest.rs index 2f2923194..8880a0ac4 100644 --- a/crates/hipfire-runtime/src/weight_manifest.rs +++ b/crates/hipfire-runtime/src/weight_manifest.rs @@ -17,7 +17,7 @@ //! family-local reductions. use crate::tp_shard::ExpertAssign; -use hipfire_hardware::{CollectiveHint, DeviceMesh, DimKind}; +use hipfire_hardware::{CollectiveHint, DeviceMesh, DimKind, MeshError}; use rdna_compute::DType; use std::collections::HashSet; @@ -315,7 +315,11 @@ pub struct ManifestPlan { pub band_xfers: Vec<(usize, CollectiveHint)>, } -fn base_coord_for(entry: &WeightEntry, mesh: &DeviceMesh, n_layers: usize) -> Vec { +fn base_coord_for( + entry: &WeightEntry, + mesh: &DeviceMesh, + n_layers: usize, +) -> Result, MeshError> { let stage = match (entry.placement, &entry.policy, entry.layer) { (PlacementHint::Pin(PinTarget::Embed), _, _) | (PlacementHint::Policy, ShardPolicy::Pin(PinTarget::Embed), _) => 0, @@ -326,18 +330,22 @@ fn base_coord_for(entry: &WeightEntry, mesh: &DeviceMesh, n_layers: usize) -> Ve (PlacementHint::Policy, _, Some(layer)) => mesh.stage_for_layer(layer, n_layers), (PlacementHint::Policy, _, None) => 0, }; - let mut coord = mesh.coord_of(0); + let mut coord = mesh.coord_of(0)?; if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { coord[index] = stage; } - coord + Ok(coord) } /// Compute global placement without touching a source, GPU, or allocator. -pub fn placement_devices(entry: &WeightEntry, mesh: &DeviceMesh, n_layers: usize) -> Vec { - let coord = base_coord_for(entry, mesh, n_layers); +pub fn placement_devices( + entry: &WeightEntry, + mesh: &DeviceMesh, + n_layers: usize, +) -> Result, MeshError> { + let coord = base_coord_for(entry, mesh, n_layers)?; match &entry.policy { - ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } => vec![mesh.device_of(&coord)], + ShardPolicy::Pin(_) | ShardPolicy::Tied { .. } => Ok(vec![mesh.device_of(&coord)?]), ShardPolicy::ExpertSharded { .. } => mesh.group_along(DimKind::Ep, &coord), ShardPolicy::ExpertTensorSharded { .. } => mesh.group_along(DimKind::Tp, &coord), _ => mesh.stage_devices(&coord), @@ -575,23 +583,27 @@ pub fn plan_manifest( .collect(); let weight_placements = weights .iter() - .map(|entry| WeightPlacement { - name: entry.name.clone(), - layer: entry.layer, - devices: placement_devices(entry, mesh, n_layers), + .map(|entry| { + Ok(WeightPlacement { + name: entry.name.clone(), + layer: entry.layer, + devices: placement_devices(entry, mesh, n_layers)?, + }) }) - .collect(); + .collect::, MeshError>>() + .map_err(|error| format!("weight placement failed: {error}"))?; let state_placements = state .iter() .map(|entry| { - let mut coord = mesh.coord_of(0); + let mut coord = mesh.coord_of(0)?; let stage = mesh.stage_for_layer(entry.layer, n_layers); if let Some(index) = mesh.axes().iter().position(|axis| axis.kind == DimKind::Pp) { coord[index] = stage; } - (entry.clone(), mesh.stage_devices(&coord)) + Ok((entry.clone(), mesh.stage_devices(&coord)?)) }) - .collect(); + .collect::, MeshError>>() + .map_err(|error| format!("state placement failed: {error}"))?; let band_xfers = (0..n_layers) .filter_map(|layer| { mesh.band_xfer_after(layer, n_layers) @@ -918,8 +930,8 @@ mod tests { ShardPolicy::Pin(PinTarget::Embed), ); let row = layer_entry("wo", 2, ShardPolicy::RowShard { axis: 1 }); - assert_eq!(placement_devices(&embed, &mesh, 4), vec![0]); - assert_eq!(placement_devices(&row, &mesh, 4), vec![2, 3]); + assert_eq!(placement_devices(&embed, &mesh, 4).unwrap(), vec![0]); + assert_eq!(placement_devices(&row, &mesh, 4).unwrap(), vec![2, 3]); let plan = plan_manifest( &[layer_entry("wo", 0, ShardPolicy::RowShard { axis: 1 }), row], &[], diff --git a/crates/hipfire-runtime/src/weight_store.rs b/crates/hipfire-runtime/src/weight_store.rs index 58bd4512f..08b283f23 100644 --- a/crates/hipfire-runtime/src/weight_store.rs +++ b/crates/hipfire-runtime/src/weight_store.rs @@ -705,7 +705,18 @@ where let origin = WeightOrigin::for_single(mesh, gpu); let mut store = WeightStore::with_origin(origin); for entry in weights { - let devices = placement_devices(entry, mesh, n_layers); + let devices = match placement_devices(entry, mesh, n_layers) { + Ok(devices) => devices, + Err(error) => { + let error = FulfillError { + name: entry.name.clone(), + layer: entry.layer, + device: 0, + reason: format!("device placement failed: {error}"), + }; + return Err(rollback_fulfill_error(store, gpu, error)); + } + }; if devices.as_slice() != [0] { let error = FulfillError { name: entry.name.clone(), @@ -967,7 +978,7 @@ mod tests { // The target guard is pure and can be checked without constructing a // Gpu; the closure would be unreachable on this path. assert!(target_error(&mesh).is_some()); - assert_eq!(placement_devices(&entry, &mesh, 1), vec![0]); + assert_eq!(placement_devices(&entry, &mesh, 1).unwrap(), vec![0]); } #[test] diff --git a/scripts/leanup-thresholds.txt b/scripts/leanup-thresholds.txt index fa8886657..9be510329 100644 --- a/scripts/leanup-thresholds.txt +++ b/scripts/leanup-thresholds.txt @@ -22,11 +22,10 @@ substrate_clean_arch_refs == 0 required_features_daemon == 0 # --- ceilings --- -# Reconciled to the measured f2ea5136 master baseline on 2026-08-28. This -# branch adds no semantic debt; rustfmt makes one existing Qwen35 call visible. -# hipfire-daemon/src/main.rs. Was 43,696 as hipfire-runtime/examples/daemon.rs -# on master; the saddle layering moved it into a crate. -daemon_lines <= 4155 +# Raised for G3 manifest/load lifecycle handling: the daemon now restores an +# intact model after wrong-device unload preflight instead of dropping its only +# retry owner, and reports terminal teardown failures separately. +daemon_lines <= 4194 # Examples compile on every `cargo build --all-targets`. Archived research # probes are gated behind `--features lab`; this is the count still ungated. From 81cf5029bf6d791e93b86eb74e667d54f94558cf Mon Sep 17 00:00:00 2001 From: Kaden Schutt <151092359+Kaden-Schutt@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:24:49 -0700 Subject: [PATCH 13/17] chore(device-mesh): record G3 ratchet trade RATCHET-RAISE: daemon_lines 4155 -> 4194, traded for typed retry restoration and rollback ownership in transactional manifest loading. From 6303b398339da31aee7d31dd2bb0194318620271 Mon Sep 17 00:00:00 2001 From: Kaden Schutt <151092359+Kaden-Schutt@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:54:09 -0700 Subject: [PATCH 14/17] fix(loader): allow text-only VL HFQ trunks --- crates/hipfire-loader/map.md | 4 +-- crates/hipfire-loader/src/carriers.rs | 38 +++++++++++++++++---------- 2 files changed, 26 insertions(+), 16 deletions(-) diff --git a/crates/hipfire-loader/map.md b/crates/hipfire-loader/map.md index 19492c3f8..2a7d9d198 100644 --- a/crates/hipfire-loader/map.md +++ b/crates/hipfire-loader/map.md @@ -24,7 +24,7 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside | File | Lines | Public items | Tests | |---|---:|---:|---:| | [`src/batch_staging.rs`](src/batch_staging.rs) | 336 | 4 | 0 | -| [`src/carriers.rs`](src/carriers.rs) | 2,850 | 11 | 5 | +| [`src/carriers.rs`](src/carriers.rs) | 2,860 | 11 | 6 | | [`src/lib.rs`](src/lib.rs) | 6,773 | 138 | 34 | | [`src/parallel_capability.rs`](src/parallel_capability.rs) | 972 | 9 | 12 | | [`src/spec_build.rs`](src/spec_build.rs) | 236 | 4 | 0 | @@ -50,6 +50,6 @@ _Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside ### Totals -- 5 modules · 11,167 lines · 166 public items · 51 tests · 1 examples +- 5 modules · 11,177 lines · 166 public items · 52 tests · 1 examples diff --git a/crates/hipfire-loader/src/carriers.rs b/crates/hipfire-loader/src/carriers.rs index 5e77a87d5..292dc91ab 100644 --- a/crates/hipfire-loader/src/carriers.rs +++ b/crates/hipfire-loader/src/carriers.rs @@ -172,22 +172,14 @@ fn classify_qwen35(src: &ModelSource) -> Result { _ => unreachable!("arch_id was checked above"), }; - let has_vision_config = config.get("vision_config").is_some(); let has_vision_tensor = source_has_tensor(src, "model.visual.patch_embed.proj.weight"); - let model_type_is_vl = config_model_type(&config) - .map(|model_type| model_type.to_ascii_lowercase().contains("vl")) - .unwrap_or(false); - // Qwen3.5-VL may share arch id 5 or 6 with text checkpoints. A vision - // marker without the actual tower is malformed and must fail closed - // rather than silently turning into a dense text model. - if has_vision_config || has_vision_tensor || model_type_is_vl { - if !has_vision_tensor { - return Err( - "qwen35: vision metadata/model type present but the vision tensor is missing" - .into(), - ); - } + // Vision capability is payload-driven, not metadata-driven. Text-only + // exports legitimately retain the parent VL config/model_type after the + // tower is stripped (including the Qwen3.8 MQV2 XT trunk). Only an actual + // tower tensor opts the load into the vision variant; otherwise keep the + // text backbone and allocate no vision weights. + if has_vision_tensor { return Ok(match backbone { ModelVariant::Qwen35Dense => ModelVariant::Qwen35DenseVl, ModelVariant::Qwen35Moe => ModelVariant::Qwen35MoeVl, @@ -2829,6 +2821,24 @@ mod qwen35_classification_tests { assert_ne!(dense, moe); } + #[test] + fn qwen35_vl_metadata_without_tower_routes_text_backbone() { + let dense = classify_fixture( + 5, + r#"{"config":{"model_type":"qwen3_5_vl","num_experts":0,"vision_config":{}}}"#, + false, + ) + .unwrap(); + let moe = classify_fixture( + 6, + r#"{"config":{"model_type":"qwen3_5_vl_moe","num_experts":8,"vision_config":{}}}"#, + false, + ) + .unwrap(); + assert_eq!(dense, ModelVariant::Qwen35Dense); + assert_eq!(moe, ModelVariant::Qwen35Moe); + } + #[test] fn qwen35_vl_validates_arch_expert_pair_before_vision() { let dense_id_with_experts = classify_fixture( From 69b47d289a9b76e0b76945fb69b41c4ced41b9b5 Mon Sep 17 00:00:00 2001 From: Kaden Schutt <151092359+Kaden-Schutt@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:55:38 -0700 Subject: [PATCH 15/17] fix(daemon): avoid rejecting staged VMM owners --- crates/hipfire-daemon/src/main.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/hipfire-daemon/src/main.rs b/crates/hipfire-daemon/src/main.rs index 190d368fb..d3ab8df91 100644 --- a/crates/hipfire-daemon/src/main.rs +++ b/crates/hipfire-daemon/src/main.rs @@ -678,8 +678,8 @@ impl DaemonLoadState<'_> { } } } else { - hipfire_loader::ensure_vmm_ready_for_load(self.gpu) - .map_err(DaemonLoadOperationError::Internal) + // Prepare already validated VMM before staging candidate allocations. + Ok(()) } } From 8e30750725d1479fbc0a820ae2c21fc4b226a723 Mon Sep 17 00:00:00 2001 From: Kaden Schutt <151092359+Kaden-Schutt@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:58:52 -0700 Subject: [PATCH 16/17] test(hardware): serialize visibility environment --- crates/hipfire-hardware/map.md | 50 ++++++++++++++++++++++++++++++ crates/hipfire-hardware/src/lib.rs | 40 ++++++++++++++++++++++++ 2 files changed, 90 insertions(+) create mode 100644 crates/hipfire-hardware/map.md diff --git a/crates/hipfire-hardware/map.md b/crates/hipfire-hardware/map.md new file mode 100644 index 000000000..8fbfafddf --- /dev/null +++ b/crates/hipfire-hardware/map.md @@ -0,0 +1,50 @@ +# hipfire-hardware — map + +> **Status:** `production` / `research` / `legacy` — pick exactly one +> (vocabulary owned by [`docs/GLOSSARY.md`](../../docs/GLOSSARY.md)). +> **Layer:** see the layering table in +> [`docs/ARCHITECTURE.md`](../../docs/ARCHITECTURE.md) — do not restate it here. + +## Purpose + + + +## Gotchas + + + +## Crate map + + + +_Generated by `scripts/check-crate-maps.py` from the tree — do not edit inside the markers._ + +### Modules + +| File | Lines | Public items | Tests | +|---|---:|---:|---:| +| [`src/lib.rs`](src/lib.rs) | 2,368 | 38 | 12 | +| [`src/mesh.rs`](src/mesh.rs) | 582 | 21 | 8 | + +### Public API surface + +- [`src/lib.rs`](src/lib.rs): `mesh`, `DeviceResolveOpts`, `BoundaryEvent`, `PeerReduceScratchLease`, `peer_reduce_scratch_bytes_per_rank`, `peer_reduce_scratch_total_bytes`, `Gpus`, `init_uniform`, `init_layers`, `init_vram_weighted`, `single`, `init_tp`, +26 more +- [`src/mesh.rs`](src/mesh.rs): `MeshEpoch`, `as_u64`, `DimKind`, `Axis`, `CollectiveHint`, `MeshError`, `DeviceMesh`, `rect`, `single`, `axes`, `epoch`, `n_devices`, +9 more + +### Dependencies (from `Cargo.toml`) + +- path: `hip-bridge`, `hipfire-config`, `rdna-compute` +- external: — +- dev: — +- build: — + +### Reverse dependencies + +- workspace crates with a path dependency on this crate: `hipfire-arch-deepseek4`, `hipfire-arch-llama`, `hipfire-arch-minimax`, `hipfire-arch-qwen35`, `hipfire-generate`, `hipfire-loader`, `hipfire-runtime`, `saddle-lab` + +### Totals + +- 2 modules · 2,950 lines · 59 public items · 22 tests · 0 examples + + diff --git a/crates/hipfire-hardware/src/lib.rs b/crates/hipfire-hardware/src/lib.rs index 89377bccc..2426dc0f5 100644 --- a/crates/hipfire-hardware/src/lib.rs +++ b/crates/hipfire-hardware/src/lib.rs @@ -2067,6 +2067,41 @@ fn preflight_vram_with_opts( mod tests { use super::*; + static VISIBILITY_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + struct VisibilityEnvGuard { + _lock: std::sync::MutexGuard<'static, ()>, + hip: Option, + rocr: Option, + } + + impl VisibilityEnvGuard { + fn acquire() -> Self { + let lock = VISIBILITY_ENV_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Self { + _lock: lock, + hip: std::env::var_os(hipfire_config::HIP_VISIBLE_DEVICES), + rocr: std::env::var_os(hipfire_config::ROCR_VISIBLE_DEVICES), + } + } + } + + impl Drop for VisibilityEnvGuard { + fn drop(&mut self) { + for (name, value) in [ + (hipfire_config::HIP_VISIBLE_DEVICES, self.hip.take()), + (hipfire_config::ROCR_VISIBLE_DEVICES, self.rocr.take()), + ] { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + } + } + #[test] fn uniform_split_basic() { assert_eq!(uniform_split_counts(2, 24), vec![12, 12]); @@ -2240,6 +2275,7 @@ mod tests { #[test] fn resolver_lowers_with_applied_visibility_proof() { + let _env = VisibilityEnvGuard::acquire(); // Simulate successful visibility application: physical 2,3 -> logical 0,1 let vis = hipfire_config::DeviceVisibility { rocr: "2,3".into(), @@ -2263,6 +2299,7 @@ mod tests { #[test] fn resolver_rejects_empty_and_malformed_lists() { + let _env = VisibilityEnvGuard::acquire(); let cases = [ Some("".to_string()), Some(",".to_string()), @@ -2301,6 +2338,9 @@ mod tests { #[test] fn resolver_rejects_stale_visibility_proof() { + let _env = VisibilityEnvGuard::acquire(); + std::env::remove_var(hipfire_config::HIP_VISIBLE_DEVICES); + std::env::remove_var(hipfire_config::ROCR_VISIBLE_DEVICES); // Physical 2,3 but visibility proof is stale (doesn't match env or rocr) let vis = hipfire_config::DeviceVisibility { rocr: "2,3".into(), From f70f343166443e64f1ca44e996b892be95c1f569 Mon Sep 17 00:00:00 2001 From: Kaden Schutt <151092359+Kaden-Schutt@users.noreply.github.com> Date: Tue, 1 Sep 2026 15:59:39 -0700 Subject: [PATCH 17/17] docs(hardware): complete crate map --- crates/hipfire-hardware/map.md | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/crates/hipfire-hardware/map.md b/crates/hipfire-hardware/map.md index 8fbfafddf..516eceb0d 100644 --- a/crates/hipfire-hardware/map.md +++ b/crates/hipfire-hardware/map.md @@ -1,18 +1,20 @@ # hipfire-hardware — map -> **Status:** `production` / `research` / `legacy` — pick exactly one -> (vocabulary owned by [`docs/GLOSSARY.md`](../../docs/GLOSSARY.md)). -> **Layer:** see the layering table in -> [`docs/ARCHITECTURE.md`](../../docs/ARCHITECTURE.md) — do not restate it here. +> **Status:** `production` +> **Layer:** hardware topology and multi-GPU ownership; see +> [`docs/ARCHITECTURE.md`](../../docs/ARCHITECTURE.md). ## Purpose - +Owns typed device meshes, physical-to-logical device resolution, pipeline +layer bands, peer copies, RCCL collectives, and their GPU-lifetime resources. ## Gotchas - +- All HIP work remains on the daemon's owning OS thread. +- Physical selectors lower to logical IDs only with a visibility proof that + matches the current process environment. +- Every `boundary_copy` event must be consumed by `wait_boundary`. ## Crate map