diff --git a/crates/larql-cli/src/commands/primary/run_cmd.rs b/crates/larql-cli/src/commands/primary/run_cmd.rs index 100fda41c..129e447db 100644 --- a/crates/larql-cli/src/commands/primary/run_cmd.rs +++ b/crates/larql-cli/src/commands/primary/run_cmd.rs @@ -137,6 +137,19 @@ pub struct RunArgs { #[arg(long, value_name = "URL")] pub ffn: Option, + /// Serve the routed expert banks from a VINDEX3 container, keeping the + /// rest of the model (tokenizer, config, embeddings, attention, norms, + /// routers, dense/shared FFN, LM head) from the VINDEX2 `MODEL` argument. + /// + /// Exactly one operand source is replaced — spec §4 classes 4 and 5 — + /// so a comparison against the same prompt without this flag is a + /// statement about the routed bytes and nothing else. + /// + /// Never falls back: if the container cannot serve every routed layer the + /// model needs, the run is refused before the prompt is encoded. + #[arg(long, value_name = "DIR")] + pub routed_from: Option, + /// HTTP timeout in seconds for --ffn. #[arg(long, default_value = "60")] pub ffn_timeout_secs: u64, @@ -308,6 +321,14 @@ pub fn run(args: RunArgs) -> Result<(), Box> { return run_bitnet(&vindex_path, &args); } + if let Some(ref routed_dir) = args.routed_from { + let prompt = args + .prompt + .as_deref() + .ok_or("--routed-from requires a prompt argument (chat mode not yet supported)")?; + return run_with_routed_container(&vindex_path, routed_dir, prompt, args.max_tokens); + } + if let Some(ref ffn_url) = args.ffn { let prompt = args.prompt.as_deref().ok_or( "--ffn requires a prompt argument (chat mode not yet supported with --ffn-dispatch batch)", @@ -758,6 +779,8 @@ fn run_with_moe_shards( ); } let started = std::time::Instant::now(); + // Fatal by policy: a shard failure aborts the run rather than + // finishing the sentence from a model missing an expert layer. let toks = generate_kquant_cpu_remote( &mut weights, &tokenizer, @@ -765,7 +788,8 @@ fn run_with_moe_shards( max_tokens, &index, &remote, - ); + ) + .map_err(|e| format!("remote MoE dispatch failed, generation aborted: {e}"))?; let total_ms = started.elapsed().as_secs_f64() * 1000.0; let strings: Vec = toks.into_iter().map(|(s, _)| s).collect(); let n = strings.len(); @@ -887,6 +911,82 @@ fn run_with_moe_shards( Ok(()) } +/// `--routed-from DIR` — routed expert banks served from a VINDEX3 container. +/// +/// The composition, precisely: +/// +/// ```text +/// VINDEX2 model tokenizer, config, embeddings, attention, norms, +/// routers, dense/shared FFN, LM head +/// VINDEX3 dir routed gate/up and routed down banks (spec §4 classes 4-5) +/// ``` +/// +/// Everything but the routed banks is read exactly as an ordinary run reads +/// it, so the same prompt without the flag is a controlled comparison: the +/// only variable is where the expert bytes came from. +/// +/// This is a *composed* run, not a VINDEX3 model. A container holding only +/// routed banks has no tokenizer and no spine; `larql run ` is +/// still correctly refused. Container completeness is a separate rung. +fn run_with_routed_container( + vindex_path: &std::path::Path, + routed_dir: &str, + prompt: &str, + max_tokens: usize, +) -> Result<(), Box> { + let routed_path = std::path::Path::new(routed_dir); + + let mut cb = larql_vindex::SilentLoadCallbacks; + let mut weights = larql_vindex::load_model_weights_kquant(vindex_path, &mut cb) + .map_err(|e| format!("failed to load spine weights: {e}"))?; + let tokenizer = larql_vindex::load_vindex_tokenizer(vindex_path) + .map_err(|e| format!("failed to load tokenizer: {e}"))?; + let mut index = larql_vindex::VectorIndex::load_vindex(vindex_path, &mut cb) + .map_err(|e| format!("failed to load vindex: {e}"))?; + index + .load_attn_kquant(vindex_path) + .map_err(|e| format!("failed to load attn Q4K: {e}"))?; + index + .load_interleaved_kquant(vindex_path) + .map_err(|e| format!("failed to load interleaved Q4K: {e}"))?; + let _ = index.load_lm_head_kquant(vindex_path); + + // Compose *before* the prompt is encoded. Every shape, count and region is + // checked here, so a mismatch is reported against two named artifacts + // rather than surfacing as a wrong number seventeen layers into a forward + // pass that has already printed part of an answer. + let routed = larql_inference::ffn::ContainerRoutedBackend::open(routed_path, &weights, true) + .map_err(|e| format!("--routed-from refused: {e}"))?; + eprintln!("{}", routed.describe(vindex_path)); + + let wrapped_prompt = + larql_inference::chat::render_user_prompt(vindex_path, weights.arch.family(), prompt)?; + let prompt_ids = larql_inference::encode_prompt(&tokenizer, &*weights.arch, &wrapped_prompt) + .map_err(|e| format!("failed to tokenise prompt: {e}"))?; + + let started = std::time::Instant::now(); + let toks = larql_inference::vindex::generate_kquant_cpu_routed( + &mut weights, + &tokenizer, + &prompt_ids, + max_tokens, + &index, + &routed, + ) + .map_err(|e| format!("routed container dispatch failed, generation aborted: {e}"))?; + let total_ms = started.elapsed().as_secs_f64() * 1000.0; + + let text: String = toks.iter().map(|(t, _)| t.as_str()).collect(); + println!("{text}"); + let n = toks.len(); + eprintln!( + "\n {n} token(s) in {:.0} ms ({:.0} ms/token)", + total_ms, + if n == 0 { 0.0 } else { total_ms / n as f64 } + ); + Ok(()) +} + /// `--ffn URL` dispatch path for dense models. /// /// Metal runs attention on the local GPU. Every layer's FFN is a round trip diff --git a/crates/larql-cli/src/main.rs b/crates/larql-cli/src/main.rs index ba9320b37..e257c9ea0 100644 --- a/crates/larql-cli/src/main.rs +++ b/crates/larql-cli/src/main.rs @@ -317,6 +317,7 @@ struct ChatArgs { /// Route FFN to a remote larql-server. #[arg(long, value_name = "URL")] ffn: Option, + routed_from: Option, /// HTTP timeout in seconds for --ffn. #[arg(long, default_value = "60")] @@ -338,6 +339,7 @@ impl From for run_cmd::RunArgs { context_window: 0, engine: None, ffn: c.ffn, + routed_from: c.routed_from, ffn_timeout_secs: c.ffn_timeout_secs, metal: false, verbose: c.verbose, diff --git a/crates/larql-inference/src/ffn/local_moe.rs b/crates/larql-inference/src/ffn/local_moe.rs index 47518397d..9e713fe38 100644 --- a/crates/larql-inference/src/ffn/local_moe.rs +++ b/crates/larql-inference/src/ffn/local_moe.rs @@ -80,17 +80,23 @@ impl<'a> FfnBackend for LocalMoeFfn<'a> { ) -> Result>, larql_execution::BoxRefusal> { // Local dispatch over resident weights: there is no operand this could // fail to reach, so it never refuses. - Ok(Some(moe_ffn_block_cpu_with_index( - self.weights, - h_post_attn, - layer, - &WeightFfn { - weights: self.weights, - }, - None, - None, - self.index, - ))) + Ok(Some( + moe_ffn_block_cpu_with_index( + self.weights, + h_post_attn, + layer, + &WeightFfn { + weights: self.weights, + }, + None, + None, + self.index, + ) + // No route is bound, so the refusal branch is unreachable rather than + // ignored — stated so a future caller that *does* bind one here has to + // decide what a failure means instead of inheriting silence. + .expect("no MoE route is bound, so no refusal is reachable"), + )) } } diff --git a/crates/larql-inference/src/ffn/mod.rs b/crates/larql-inference/src/ffn/mod.rs index ec1abe20d..add0ebb30 100644 --- a/crates/larql-inference/src/ffn/mod.rs +++ b/crates/larql-inference/src/ffn/mod.rs @@ -17,6 +17,7 @@ pub mod graph_backend; pub mod local_moe; pub mod moe_backend; pub mod moe_bound; +pub mod moe_container; pub mod moe_remote; pub mod remote; pub mod sparse; @@ -34,8 +35,11 @@ pub use larql_compute::ffn::{ // ── Re-exports ── pub use local_moe::LocalMoeFfn; -pub use moe_backend::{InProcessMoeBackend, MoeBackendError, MoeExpertBackend}; +pub use moe_backend::{ + InProcessMoeBackend, MoeBackendError, MoeExpertBackend, MoeFailurePolicy, MoeRoute, +}; pub use moe_bound::BoundMoeBackend; +pub use moe_container::{CompositionError, ContainerRoutedBackend}; pub use moe_remote::{ MoeFfn, MoeRouterWeights, RecordedRefusal, RefusalPolicy, RemoteMoeBackend, RemoteMoeError, RemoteMoeFfn, ShardConfig, diff --git a/crates/larql-inference/src/ffn/moe_backend.rs b/crates/larql-inference/src/ffn/moe_backend.rs index 1ed555331..6775f8e34 100644 --- a/crates/larql-inference/src/ffn/moe_backend.rs +++ b/crates/larql-inference/src/ffn/moe_backend.rs @@ -54,6 +54,71 @@ pub enum MoeBackendError { Remote(#[from] RemoteMoeError), #[error("bound expert execution failed: {0}")] Bound(#[from] larql_vindex::runtime::ExecutionError), + /// A routed operand could not be sourced from its VINDEX3 container. + /// + /// Distinct from [`Self::Bound`]: that is the executor rejecting operands + /// it was given, this is not having them. Collapsing the two would report + /// a missing bank as a kernel failure and send the reader to the wrong + /// half of the system. Reaching this at generation time is itself a bug — + /// composition is validated when the container is opened — so the message + /// names the layer and expert rather than assuming a retry. + #[error("routed container operand unavailable: {0}")] + Container(String), +} + +/// What an operation does when a MoE route refuses. +/// +/// This is a property of the **operation**, not of the backend. The same +/// backend is legitimately used both to generate (where a failed layer must +/// abort) and to probe (where a refusal is the measurement). A backend cannot +/// distinguish those callers, so asking it would bake operation policy into an +/// operand provider — and would imply a false taxonomy in which some backends' +/// failures are tolerable. They are not: a remote shard failing mid-generation +/// and contributing zero is exactly as invalid as a missing container region. +/// +/// The default for anything that produces a continuation, a score, a parity +/// number or a benchmark is [`Self::Fatal`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MoeFailurePolicy { + /// Abort the operation. No token, score or measurement is produced from a + /// forward pass in which any expert contribution was not computed. + Fatal, + /// Record the refusal and continue with an incomplete result. + /// + /// Analysis-only. A caller selecting this is declaring that it *intends* + /// to inspect partial execution and will report the incompleteness; it may + /// not present the outcome as an ordinary model continuation. + RecordRefusal, +} + +/// A route plus the policy the calling operation applies to its refusals. +/// +/// Paired rather than passed separately so a caller cannot supply a backend +/// and forget to state what a failure means — the omission that let a failed +/// layer contribute silent zeros for as long as this branch has existed. +#[derive(Clone, Copy)] +pub struct MoeRoute<'a> { + pub backend: &'a dyn MoeExpertBackend, + pub policy: MoeFailurePolicy, +} + +impl<'a> MoeRoute<'a> { + /// A route whose failures abort the operation. The correct choice for + /// generation, scoring, parity and benchmarking. + pub fn fatal(backend: &'a dyn MoeExpertBackend) -> Self { + Self { + backend, + policy: MoeFailurePolicy::Fatal, + } + } + + /// A route whose failures are recorded and tolerated. Analysis only. + pub fn recording(backend: &'a dyn MoeExpertBackend) -> Self { + Self { + backend, + policy: MoeFailurePolicy::RecordRefusal, + } + } } /// A route that computes one hybrid-MoE layer's expert contribution. diff --git a/crates/larql-inference/src/ffn/moe_container.rs b/crates/larql-inference/src/ffn/moe_container.rs new file mode 100644 index 000000000..122adc1d1 --- /dev/null +++ b/crates/larql-inference/src/ffn/moe_container.rs @@ -0,0 +1,459 @@ +//! The routed-bank override: expert bytes sourced from a VINDEX3 container. +//! +//! # What this closes +//! +//! Every VINDEX3 execution result before this one bound its operands out of a +//! **VINDEX2** file. `moe_bound` says so plainly, and the container ladder +//! (c8/c9) proved only that the same bytes could be *written* to a VINDEX3 +//! container — not that anything would ever *read* one to compute with. That +//! left a gap nothing on the execution path could cross: `Vindex3Container` +//! appeared nowhere in this crate. +//! +//! ```text +//! before VINDEX2 index ──> MoeLayerWeights ──> bound executor +//! now VINDEX3 container ──> expert regions ─┘ +//! (everything else still VINDEX2) +//! ``` +//! +//! # Exactly one operand source is replaced +//! +//! The spine — tokenizer, config, embeddings, attention, norms, routers, +//! shared/dense FFN, LM head — is read from the VINDEX2 model, unchanged and +//! byte-for-byte. Only classes 4 and 5 (routed gate/up and routed down, spec +//! §4) come from the container. That is what makes a comparison against a +//! plain VINDEX2 run a statement about *the routed bytes* and nothing else: if +//! this route substituted anything further, a divergence would have somewhere +//! else to hide. +//! +//! This is also the K3 deployment topology in miniature — a small resident +//! spine plus huge paged routed banks — but it is **not** artifact identity. A +//! composed run is two directories, and a VINDEX3 model that still needs a +//! VINDEX2 directory to run is not yet a model. Container completeness (the +//! `control/`, `dense/`, `shared/` and sidecar classes of spec §5) is a +//! separate job and this type does not pretend to stand in for it. +//! +//! # Public API only, deliberately +//! +//! Regions are reached through `Vindex3Container -> segment -> region_bytes` +//! and never by parsing LYRW descriptors here. The binary layout is mid-change +//! (24 B → 28 B bank descriptors with an explicit `group_width`); a backend +//! that read the descriptor itself would have to be rewritten alongside it, +//! and would have been a second implementation of the reader in the meantime. +//! +//! # No fallback, ever +//! +//! If a layer, an expert or a role is missing, this refuses. It does not quietly +//! serve the VINDEX2 bytes that are sitting right there in the same process — +//! doing so would make "the model ran from VINDEX3" unfalsifiable, which is the +//! only claim the route exists to support. + +use std::path::Path; + +use larql_models::ModelWeights; +use ndarray::Array2; + +use larql_compute::pipeline_layer::build_moe_weights; +use larql_compute::MoeLayerWeights; +use larql_vindex::format::lyrw2::region_role::RegionRole; +use larql_vindex::format::vindex3::import::routed_storage_key; +use larql_vindex::format::vindex3::Vindex3Container; + +use super::moe_backend::{MoeBackendError, MoeExpertBackend}; +use super::moe_bound::BoundMoeBackend; + +/// Route name for diagnostics. Never branched on. +const ROUTE_NAME: &str = "moe-vindex3-container"; + +/// Bank ordinal of the routed bank within a segment. c8/c9 write one bank per +/// segment file, and spec §6 (draft-3) makes that a MUST. +const ROUTED_BANK: u16 = 0; + +/// Expert bytes served from a VINDEX3 container, everything else from VINDEX2. +pub struct ContainerRoutedBackend { + container: Vindex3Container, + inner: BoundMoeBackend, +} + +/// Why a container could not be composed with a loaded model. +/// +/// Every variant names both sides. A composition refusal that said only +/// "incompatible" would send the reader to the wrong one of two artifacts. +#[derive(Debug)] +pub enum CompositionError { + Open(String), + /// The container describes a different model than the spine does. + Identity { + spine: String, + container: String, + }, + /// A routed layer the model needs is not in the container. + MissingLayer { + layer: usize, + }, + /// The container's segment for a layer cannot be resolved or read. + UnreadableLayer { + layer: usize, + why: String, + }, + /// A shape the container declares disagrees with the model's. + Shape { + layer: usize, + what: &'static str, + spine: u64, + container: u64, + }, +} + +impl std::fmt::Display for CompositionError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Open(why) => write!(f, "cannot open the routed container: {why}"), + Self::Identity { spine, container } => write!( + f, + "the container describes `{container}` but the model is `{spine}` — \ + composing them would serve one model's experts inside another" + ), + Self::MissingLayer { layer } => write!( + f, + "layer {layer} routes to experts but the container has no bank for \ + it; `--routed-from` never falls back to the VINDEX2 bytes, because \ + a run that silently did could not support the claim it exists for" + ), + Self::UnreadableLayer { layer, why } => { + write!(f, "layer {layer}'s bank is unreadable: {why}") + } + Self::Shape { + layer, + what, + spine, + container, + } => write!( + f, + "layer {layer}: the model expects {what} {spine}, the container \ + declares {container}" + ), + } + } +} + +impl std::error::Error for CompositionError {} + +impl ContainerRoutedBackend { + /// Open `root` and check it can serve every routed layer `weights` needs. + /// + /// All validation happens here — before a prompt is encoded, before a + /// token is generated. A composition that failed on layer 17 of the first + /// forward pass would have already printed part of an answer produced by a + /// model half of which was not the one asked for. + pub fn open( + root: &Path, + weights: &ModelWeights, + production: bool, + ) -> Result { + let container = + Vindex3Container::open(root).map_err(|e| CompositionError::Open(e.to_string()))?; + + let backend = Self { + container, + inner: if production { + BoundMoeBackend::production() + } else { + BoundMoeBackend::reference() + }, + }; + backend.check_composable(weights)?; + Ok(backend) + } + + /// Refuse any container that cannot serve this model's routed layers. + /// + /// Coverage is checked against the **model**, not the container: a + /// container with thirty banks still cannot serve a model with thirty-one + /// routed layers, and asking the container what it has would let the + /// missing one pass unnoticed. + fn check_composable(&self, weights: &ModelWeights) -> Result<(), CompositionError> { + let arch = &*weights.arch; + let spine_model = weights.arch.family().to_string(); + let container_model = self.container.index().family.clone(); + if spine_model != container_model { + return Err(CompositionError::Identity { + spine: spine_model, + container: container_model, + }); + } + + for layer in 0..weights.num_layers { + let Some(moe) = build_moe_weights(weights, arch, layer) else { + continue; // A dense layer routes nowhere; the container owes it nothing. + }; + let declared = self + .container + .layer(layer as u32) + .ok_or(CompositionError::MissingLayer { layer })?; + + let experts = declared.routed_bank.experts as usize; + if experts != moe.num_experts { + return Err(CompositionError::Shape { + layer, + what: "experts", + spine: moe.num_experts as u64, + container: experts as u64, + }); + } + if let Some(dims) = declared.routed_bank.expert_dims.as_ref() { + if dims.input as usize != weights.hidden_size { + return Err(CompositionError::Shape { + layer, + what: "hidden size", + spine: weights.hidden_size as u64, + container: dims.input as u64, + }); + } + if dims.intermediate as usize != moe.intermediate_size { + return Err(CompositionError::Shape { + layer, + what: "semantic intermediate width", + spine: moe.intermediate_size as u64, + container: dims.intermediate as u64, + }); + } + } + + // Resolve and read the bank now rather than at first token: a key + // that resolves to a missing file is a composition failure, and + // discovering it mid-generation would be one too late. + let expected = expert_region_sizes(&moe); + self.check_layer_regions(layer, experts, expected)?; + } + Ok(()) + } + + /// Every expert's two regions must be present and the size the model expects. + fn check_layer_regions( + &self, + layer: usize, + experts: usize, + expected: (usize, usize), + ) -> Result<(), CompositionError> { + let reader = self + .container + .segment(&routed_storage_key(layer as u32)) + .map_err(|e| CompositionError::UnreadableLayer { + layer, + why: e.to_string(), + })?; + let (want_gate_up, want_down) = expected; + + for expert in 0..experts as u32 { + for (role, want) in [ + (RegionRole::GateUpFused, want_gate_up), + (RegionRole::Down, want_down), + ] { + let got = reader + .region_bytes(ROUTED_BANK, expert, role) + .map_err(|e| CompositionError::UnreadableLayer { + layer, + why: format!("expert {expert} {role:?}: {e}"), + })? + .ok_or_else(|| CompositionError::UnreadableLayer { + layer, + why: format!("expert {expert} has no {role:?} region"), + })?; + if got.len() != want { + return Err(CompositionError::Shape { + layer, + what: "routed region bytes", + spine: want as u64, + container: got.len() as u64, + }); + } + } + } + Ok(()) + } + + /// One line describing what this run actually composed. + pub fn describe(&self, spine: &Path) -> String { + format!( + "composed run: VINDEX2 spine {} + VINDEX3 routed banks {}", + spine.display(), + self.container.root().display() + ) + } +} + +/// The byte length each of a layer's two region kinds must have. +/// +/// Taken from the incumbent's own slices rather than recomputed from shapes: +/// the model already holds the authoritative lengths, and deriving them here +/// would be a second implementation of the layout to disagree with. +fn expert_region_sizes(moe: &MoeLayerWeights<'_>) -> (usize, usize) { + ( + moe.experts_gate_up.first().map_or(0, |s| s.len()), + moe.experts_down.first().map_or(0, |s| s.len()), + ) +} + +impl MoeExpertBackend for ContainerRoutedBackend { + fn forward_moe_seq( + &self, + weights: &ModelWeights, + layer: usize, + h: &Array2, + norm_offset: f32, + eps: f32, + ) -> Result, MoeBackendError> { + let arch = &*weights.arch; + let Some(mut moe) = build_moe_weights(weights, arch, layer) else { + // No experts to route into. Zeros, matching the in-process path — + // erroring here would change the model rather than the route. + return Ok(Array2::zeros((h.nrows(), h.ncols()))); + }; + + // Replace *only* the routed operands. Every other field of `moe` — + // router projection and scales, the norms, the policy flags — stays as + // the VINDEX2 model produced it. + let reader = self + .container + .segment(&routed_storage_key(layer as u32)) + .map_err(|e| MoeBackendError::Container(format!("layer {layer} bank: {e}")))?; + + let mut gate_up = Vec::with_capacity(moe.num_experts); + let mut down = Vec::with_capacity(moe.num_experts); + for expert in 0..moe.num_experts as u32 { + for (role, sink) in [ + (RegionRole::GateUpFused, &mut gate_up), + (RegionRole::Down, &mut down), + ] { + let bytes = reader + .region_bytes(ROUTED_BANK, expert, role) + .map_err(|e| { + MoeBackendError::Container(format!( + "layer {layer} expert {expert} {role:?}: {e}" + )) + })? + .ok_or_else(|| { + MoeBackendError::Container(format!( + "layer {layer} expert {expert} has no {role:?} region" + )) + })?; + sink.push(bytes); + } + } + moe.experts_gate_up = gate_up; + moe.experts_down = down; + + Ok(self.inner.run_layer(layer, &moe, h, norm_offset, eps)?) + } + + fn name(&self) -> &'static str { + ROUTE_NAME + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::make_test_gemma4_moe_weights; + use larql_vindex::format::lyrw2::region_format::RegionFormat; + use larql_vindex::format::vindex3::{ContainerBuilder, MoeLayerSource}; + use tempfile::TempDir; + + /// Build a container from the fixture model's own expert bytes. + /// + /// `skip` omits that layer, so a test can produce a container that is + /// well-formed but cannot serve the model it is composed with — the + /// distinction the coverage check exists to make. + fn container_for(weights: &larql_models::ModelWeights, skip: Option) -> TempDir { + let dir = TempDir::new().unwrap(); + let mut builder = ContainerBuilder::create(dir.path()).unwrap(); + let arch = &*weights.arch; + let mut wrote = false; + for layer in 0..weights.num_layers { + if Some(layer) == skip { + continue; + } + let Some(moe) = build_moe_weights(weights, arch, layer) else { + continue; + }; + let source = MoeLayerSource { + layer: layer as u32, + experts_gate_up: moe.experts_gate_up.clone(), + experts_down: moe.experts_down.clone(), + format: RegionFormat::Q4K, + hidden_size: weights.hidden_size as u32, + gate_up_stored_intermediate: moe.intermediate_size as u32, + down_stored_intermediate: moe.inter_padded() as u32, + semantic_intermediate: moe.intermediate_size as u32, + top_k: moe.top_k as u32, + }; + if builder.add_moe_layer(&source).is_ok() { + wrote = true; + } + } + assert!(wrote, "fixture produced no routed layers to import"); + builder + .finish( + weights.arch.family(), + weights.arch.family(), + weights.hidden_size, + weights.num_layers, + ) + .unwrap(); + dir + } + + #[test] + fn a_container_missing_a_routed_layer_is_refused_before_generation() { + // Coverage is checked against the model, so an omitted layer must be + // caught at open — not at the token where that layer first runs, by + // which point part of an answer has already been produced. + let weights = make_test_gemma4_moe_weights(); + let full = container_for(&weights, None); + ContainerRoutedBackend::open(full.path(), &weights, true) + .expect("a complete container must compose"); + + let missing = container_for(&weights, Some(0)); + let err = ContainerRoutedBackend::open(missing.path(), &weights, true) + .err() + .expect("a container missing a routed layer must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("no bank for it") || msg.contains("layer 0"), + "the refusal must name the missing layer, got: {msg}" + ); + } + + #[test] + fn a_container_describing_another_model_is_refused() { + let weights = make_test_gemma4_moe_weights(); + let dir = container_for(&weights, None); + + // Rewrite the family in place: same banks, different identity. + let index_path = dir.path().join("index.json"); + let text = std::fs::read_to_string(&index_path).unwrap(); + let mut json: serde_json::Value = serde_json::from_str(&text).unwrap(); + json["family"] = serde_json::Value::String("not-this-model".into()); + std::fs::write(&index_path, serde_json::to_string_pretty(&json).unwrap()).unwrap(); + + let err = ContainerRoutedBackend::open(dir.path(), &weights, true) + .err() + .expect("composing a container for another model must be refused"); + assert!( + err.to_string().contains("not-this-model"), + "the refusal must name both sides, got: {err}" + ); + } + + #[test] + fn a_composed_route_reports_both_artifacts() { + // Provenance is not decoration: a composed run is two directories and + // a result recorded against only one of them is unreproducible. + let weights = make_test_gemma4_moe_weights(); + let dir = container_for(&weights, None); + let backend = ContainerRoutedBackend::open(dir.path(), &weights, true).unwrap(); + let line = backend.describe(std::path::Path::new("/models/spine.vindex")); + assert!(line.contains("/models/spine.vindex"), "{line}"); + assert!(line.contains(&dir.path().display().to_string()), "{line}"); + } +} diff --git a/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs b/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs index c7e4de6a8..fb93e2f76 100644 --- a/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs +++ b/crates/larql-inference/src/ffn/moe_remote/ffn_adapter.rs @@ -182,6 +182,12 @@ impl crate::ffn::MoeExpertBackend for RefusalRecorder<'_> { crate::ffn::MoeBackendError::Bound(inner) => inner.refusal(), // A remote dispatch failure is the operand not being here. crate::ffn::MoeBackendError::Remote(_) => larql_execution::RefusalKind::Residency, + // So is a routed region missing from its container — same + // diagnosis, different distance: the bytes this layer needs + // are not reachable, rather than reachable and rejected. + crate::ffn::MoeBackendError::Container(_) => { + larql_execution::RefusalKind::Residency + } }; // First only: the earliest refusal is the diagnosis, and later ones // are usually the same cause repeating per layer. @@ -289,8 +295,14 @@ impl FfnBackend for MoeFfn<'_> { weights: self.weights, }, None, - Some(&recorder), - ); + // Recording, because this adapter has its own `RefusalPolicy` one + // level up: it needs the refusal *captured* so `Strict` can turn it + // into an error and `BestEffort` can degrade having said so. + // Making the block itself fatal here would bypass that decision. + Some(crate::ffn::MoeRoute::recording(&recorder)), + ) + // Unreachable: `recording` never propagates out of the block. + .expect("RecordRefusal never returns Err"); // `moe_ffn_block_cpu` has already logged the refusal and left the // expert contribution at zero, so `out` is the dense half wearing the // shape of an answer. Under `Strict` it must not escape — which is the diff --git a/crates/larql-inference/src/vindex/kquant_forward/generation.rs b/crates/larql-inference/src/vindex/kquant_forward/generation.rs index e75b0e0cf..368d0ab4c 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/generation.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/generation.rs @@ -4,7 +4,7 @@ use tokenizers::Tokenizer; use crate::forward::PredictResult; -use super::hidden::predict_kquant_hidden; +use super::hidden::{predict_kquant_hidden, predict_kquant_hidden_checked}; /// End-to-end predict on a Q4_K/Q6_K vindex. pub fn predict_kquant( @@ -65,6 +65,12 @@ pub fn generate_kquant_cpu( /// Like [`generate_kquant_cpu`] but dispatches MoE expert matmuls to remote shard /// servers via [`crate::ffn::RemoteMoeBackend`]. +/// +/// Generation is [`crate::ffn::MoeFailurePolicy::Fatal`]: a shard that fails +/// mid-decode aborts the request. It does not continue with that layer's +/// expert contribution zeroed — a network failure cannot produce a valid +/// continuation, and one that looks valid is worse than none. Tokens already +/// emitted stay emitted; no further token is produced. pub fn generate_kquant_cpu_remote( weights: &mut ModelWeights, tokenizer: &Tokenizer, @@ -72,11 +78,40 @@ pub fn generate_kquant_cpu_remote( max_tokens: usize, index: &VectorIndex, moe_remote: &crate::ffn::RemoteMoeBackend, -) -> Vec<(String, u32)> { +) -> Result, crate::ffn::MoeBackendError> { + generate_kquant_cpu_routed( + weights, tokenizer, prompt_ids, max_tokens, index, moe_remote, + ) +} + +/// Greedy generation with the expert half served by any bound route. +/// +/// The route is the only variable: embeddings, attention, norms, routers, the +/// dense slab and the LM head all come from the loaded model exactly as the +/// default path reads them. That is what lets a run through a VINDEX3 +/// container be compared against a plain one and have the difference mean +/// *the routed bytes*. +/// +/// Always [`crate::ffn::MoeFailurePolicy::Fatal`] — this produces a +/// continuation, and a continuation computed with a layer's experts missing is +/// not a continuation of the model that was asked for. +pub fn generate_kquant_cpu_routed( + weights: &mut ModelWeights, + tokenizer: &Tokenizer, + prompt_ids: &[u32], + max_tokens: usize, + index: &VectorIndex, + backend: &dyn crate::ffn::MoeExpertBackend, +) -> Result, crate::ffn::MoeBackendError> { let mut ids = prompt_ids.to_vec(); let mut out: Vec<(String, u32)> = Vec::with_capacity(max_tokens); for _ in 0..max_tokens { - let h = predict_kquant_hidden(weights, &ids, index, Some(moe_remote)); + let h = predict_kquant_hidden_checked( + weights, + &ids, + index, + Some(crate::ffn::MoeRoute::fatal(backend)), + )?; let last = h.nrows().saturating_sub(1); let h_last = h.slice(ndarray::s![last..last + 1, ..]).to_owned(); let logits = crate::forward::hidden_to_raw_logits(weights, &h_last); @@ -96,7 +131,7 @@ pub fn generate_kquant_cpu_remote( break; } } - out + Ok(out) } /// KV-cached autoregressive generation: one prefill over the prompt, then @@ -484,16 +519,34 @@ mod tests { /// the remote returns Err (see hidden.rs's `Some(remote)` branch); /// the generation loop still picks tokens off the dense path. #[test] - fn generate_kquant_cpu_remote_runs_against_disconnected_backend() { + fn generate_kquant_cpu_remote_aborts_against_a_disconnected_backend() { + // Fault injection at the operation boundary. This test previously + // asserted the opposite — that generation "still picks tokens off the + // dense path" with the backend disconnected — which is precisely the + // defect: every expert contribution was zero and the loop emitted + // fluent tokens from a model that had lost its entire MoE half, with + // one stderr line as the only evidence. + // + // A disconnected shard cannot produce a valid continuation, so the + // only correct outcome is no continuation. use crate::ffn::RemoteMoeBackend; use crate::test_utils::{make_test_gemma4_moe_weights, make_test_q4k_vindex}; let mut weights = make_test_gemma4_moe_weights(); let index = make_test_q4k_vindex(&weights); let tokenizer = make_test_tokenizer(weights.vocab_size); let remote = RemoteMoeBackend::new_disconnected(); + let out = generate_kquant_cpu_remote(&mut weights, &tokenizer, &[0u32, 1], 2, &index, &remote); - assert!(out.len() <= 2); + + let err = out.expect_err("a disconnected shard must abort generation, not degrade it"); + assert!( + matches!( + err, + crate::ffn::MoeBackendError::Remote(_) | crate::ffn::MoeBackendError::Container(_) + ), + "the refusal must name the unreachable operand, got: {err}" + ); } /// Parity gate for the KV-cached decode loop: on the same fixture, diff --git a/crates/larql-inference/src/vindex/kquant_forward/hidden.rs b/crates/larql-inference/src/vindex/kquant_forward/hidden.rs index be25b0223..f21201c6e 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/hidden.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/hidden.rs @@ -14,12 +14,38 @@ use super::tensors::{insert_q4k_layer_tensors, remove_layer_tensors}; /// Compute the final hidden state for `token_ids` against a Q4_K/Q6_K /// vindex, dequantising attn + FFN one layer at a time. Returns the /// `[seq_len, hidden]` array; caller owns the lm_head step. +/// Analysis-mode forward: a refusing MoE route is recorded, not fatal. +/// +/// Kept at this signature because the overwhelming majority of its callers +/// pass `None`, where no refusal is reachable at all. Any caller that *does* +/// pass a backend and produces a continuation, a score, a parity number or a +/// benchmark must use [`predict_kquant_hidden_checked`] with +/// [`crate::ffn::MoeFailurePolicy::Fatal`] instead — under `RecordRefusal` a +/// failed layer contributes zeros and the result is analysis-only. pub fn predict_kquant_hidden( weights: &ModelWeights, token_ids: &[u32], index: &VectorIndex, moe: Option<&dyn crate::ffn::MoeExpertBackend>, ) -> Array2 { + let route = moe.map(crate::ffn::MoeRoute::recording); + predict_kquant_hidden_checked(weights, token_ids, index, route) + // `RecordRefusal` never propagates an error out of the block, so this + // is unreachable rather than swallowed. Stated as an expectation so a + // future policy change cannot turn it back into a silent zero. + .expect("RecordRefusal never returns Err") +} + +/// Forward pass whose MoE refusals are governed by the caller's policy. +/// +/// This is the form every operation that produces a result about the model +/// should use: it returns the refusal rather than folding it into the tensor. +pub fn predict_kquant_hidden_checked( + weights: &ModelWeights, + token_ids: &[u32], + index: &VectorIndex, + moe: Option>, +) -> Result, crate::ffn::MoeBackendError> { let num_layers = weights.num_layers; let mut scratch = larql_models::DequantScratch::new(); let mut h = embed_tokens_pub(weights, token_ids); @@ -57,7 +83,7 @@ pub fn predict_kquant_hidden( ple_inputs.get(layer), shared_kv, moe, - ) { + )? { h = h_new; if let Some(kv) = kv_out { kv_cache.insert(layer, kv); @@ -90,7 +116,7 @@ pub fn predict_kquant_hidden( } } - h + Ok(h) } /// Build `MoeRouterWeights` for a single layer from the model's vector store. @@ -127,6 +153,14 @@ pub fn build_moe_router_weights<'a>( }) } +/// What one MoE layer produced: the new residual, plus any K/V worth caching. +/// +/// The outer `Option` is an **absence** — attention had nothing to run for this +/// layer — while the `Result` around it is a **failure**, a refused expert +/// route. Keeping them distinct is the point: collapsing "nothing to do" into +/// "something went wrong" is how a missing operand becomes a zeroed layer. +type MoeLayerOutcome = Option<(Array2, Option)>; + /// CPU forward for one hybrid-MoE layer (Gemma 4 26B A4B). fn run_moe_layer_cpu( weights: larql_models::WeightsView, @@ -135,15 +169,24 @@ fn run_moe_layer_cpu( ffn: &dyn crate::ffn::FfnBackend, ple_input: Option<&Array2>, shared_kv: Option<&SharedKV>, - moe: Option<&dyn crate::ffn::MoeExpertBackend>, -) -> Option<(Array2, Option)> { + moe: Option>, +) -> Result { + // Attention returning `None` means this layer has nothing to run — an + // absence, not a failure, so it stays `Ok(None)` and is distinct from a + // refused expert route. let (h_post_attn, kv_out) = if let Some(shared) = shared_kv { - let (h_pa, _, _) = - crate::attention::run_attention_block_shared(weights, h, layer, false, Some(shared))?; + let Some((h_pa, _, _)) = + crate::attention::run_attention_block_shared(weights, h, layer, false, Some(shared)) + else { + return Ok(None); + }; (h_pa, None) } else { - let (h_pa, _, _, k_rope, v_final) = - crate::attention::run_attention_block_with_kv_out(weights, h, layer, false, None)?; + let Some((h_pa, _, _, k_rope, v_final)) = + crate::attention::run_attention_block_with_kv_out(weights, h, layer, false, None) + else { + return Ok(None); + }; (h_pa, Some((k_rope, v_final))) }; @@ -154,8 +197,8 @@ fn run_moe_layer_cpu( ffn, ple_input, moe, - ); - Some((h_out, kv_out)) + )?; + Ok(Some((h_out, kv_out))) } /// CPU MoE FFN block for one hybrid-MoE layer, given the **post-attention** @@ -179,8 +222,8 @@ pub fn moe_ffn_block_cpu( layer: usize, ffn: &dyn crate::ffn::FfnBackend, ple_input: Option<&Array2>, - moe: Option<&dyn crate::ffn::MoeExpertBackend>, -) -> Array2 { + moe: Option>, +) -> Result, crate::ffn::MoeBackendError> { moe_ffn_block_cpu_with_index(weights, h_post_attn, layer, ffn, ple_input, moe, None) } @@ -206,9 +249,9 @@ pub fn moe_ffn_block_cpu_with_index( layer: usize, ffn: &dyn crate::ffn::FfnBackend, ple_input: Option<&Array2>, - moe: Option<&dyn crate::ffn::MoeExpertBackend>, + moe: Option>, index: Option<&larql_vindex::VectorIndex>, -) -> Array2 { +) -> Result, crate::ffn::MoeBackendError> { let arch = &*weights.arch; let norm_offset = arch.norm_weight_offset(); let eps = arch.norm_eps(); @@ -255,19 +298,34 @@ pub fn moe_ffn_block_cpu_with_index( let seq_len = h_post_attn.nrows(); let mut h2 = Array2::::zeros((seq_len, hidden)); - if let Some(backend) = moe { + if let Some(route) = moe { // Every non-default route goes through one call. Which route it is — // remote shards, a VINDEX3 bound plan — is the backend's business, and // the block loop's only job is to place the contribution. let _t_expert = std::time::Instant::now(); - let out = backend.forward_moe_seq(weights, layer, h_post_attn, norm_offset, eps); + let out = route + .backend + .forward_moe_seq(weights, layer, h_post_attn, norm_offset, eps); crate::decode_stages::record_expert(_t_expert.elapsed().as_nanos()); match out { Ok(out) => h2 = out, - Err(e) => eprintln!( - "[moe_ffn_block_cpu] {} dispatch error L{layer}: {e}", - backend.name() - ), + // The failure is the caller's to interpret. This branch used to + // print and continue unconditionally, which left `h2` at zeros — + // so a broken layer produced a fluent continuation from a model + // with that layer's entire expert contribution removed, and the + // only trace was a line on stderr. + Err(e) => match route.policy { + crate::ffn::MoeFailurePolicy::Fatal => return Err(e), + crate::ffn::MoeFailurePolicy::RecordRefusal => { + // Analysis-only: `h2` stays zero and the caller has + // declared it will report the incompleteness. + eprintln!( + "[moe_ffn_block_cpu] {} refused L{layer} (analysis mode, \ + expert contribution omitted): {e}", + route.backend.name() + ); + } + }, } } else { // Local experts count toward the expert stage too (`LARQL_DECODE_STAGES`) @@ -296,7 +354,7 @@ pub fn moe_ffn_block_cpu_with_index( let mut h_ple = crate::forward::ple::apply_per_layer_embedding(weights, &out, layer, ple_input); crate::forward::layer::apply_layer_scalar(weights, &mut h_ple, layer); - return h_ple; + return Ok(h_ple); } else { // Pure MoE with no expert weights would otherwise fall through to // `h_post_ffn_dense`, which here *is* `h_post_attn` — an identity @@ -364,7 +422,7 @@ pub fn moe_ffn_block_cpu_with_index( } } - h_out + Ok(h_out) } #[cfg(test)] diff --git a/crates/larql-inference/src/vindex/kquant_forward/mod.rs b/crates/larql-inference/src/vindex/kquant_forward/mod.rs index 7f4b9929e..d978c6245 100644 --- a/crates/larql-inference/src/vindex/kquant_forward/mod.rs +++ b/crates/larql-inference/src/vindex/kquant_forward/mod.rs @@ -31,11 +31,11 @@ pub use generation::{ generate_kquant_cpu, generate_kquant_cpu_cached, generate_kquant_cpu_constrained, generate_kquant_cpu_constrained_cached, generate_kquant_cpu_constrained_cached_streaming, generate_kquant_cpu_constrained_streaming, generate_kquant_cpu_constrained_streaming_sampled, - generate_kquant_cpu_remote, is_end_of_turn, predict_kquant, + generate_kquant_cpu_remote, generate_kquant_cpu_routed, is_end_of_turn, predict_kquant, }; pub use hidden::{ build_moe_router_weights, moe_ffn_block_cpu, moe_ffn_block_cpu_with_index, - predict_kquant_hidden, + predict_kquant_hidden, predict_kquant_hidden_checked, }; pub use hooks::predict_kquant_hidden_hooked; pub use interventions::{ diff --git a/crates/larql-inference/src/vindex/mod.rs b/crates/larql-inference/src/vindex/mod.rs index f40e02624..f19b865ad 100644 --- a/crates/larql-inference/src/vindex/mod.rs +++ b/crates/larql-inference/src/vindex/mod.rs @@ -19,11 +19,12 @@ pub use kquant_forward::{ generate_kquant_cpu_cached, generate_kquant_cpu_constrained, generate_kquant_cpu_constrained_cached, generate_kquant_cpu_constrained_cached_streaming, generate_kquant_cpu_constrained_streaming, generate_kquant_cpu_constrained_streaming_sampled, - generate_kquant_cpu_remote, insert_q4k_layer_tensors, insert_q4k_layer_tensors_resident, - is_end_of_turn, kquant_ffn_forward_layer, kquant_ffn_forward_layer_q8k, moe_ffn_block_cpu, - moe_ffn_block_cpu_with_index, predict_kquant, predict_kquant_decode_step, - predict_kquant_decode_step_direct, predict_kquant_decode_step_direct_with_state, - predict_kquant_hidden, predict_kquant_hidden_hooked, predict_kquant_hidden_with_ffn, + generate_kquant_cpu_remote, generate_kquant_cpu_routed, insert_q4k_layer_tensors, + insert_q4k_layer_tensors_resident, is_end_of_turn, kquant_ffn_forward_layer, + kquant_ffn_forward_layer_q8k, moe_ffn_block_cpu, moe_ffn_block_cpu_with_index, predict_kquant, + predict_kquant_decode_step, predict_kquant_decode_step_direct, + predict_kquant_decode_step_direct_with_state, predict_kquant_hidden, + predict_kquant_hidden_checked, predict_kquant_hidden_hooked, predict_kquant_hidden_with_ffn, predict_kquant_hidden_with_mapped_head_residual_delta, predict_kquant_hidden_with_mapped_pre_o_head, predict_kquant_hidden_with_original_head_residual_delta, diff --git a/crates/larql-vindex/examples/vindex3_composed_parity.rs b/crates/larql-vindex/examples/vindex3_composed_parity.rs new file mode 100644 index 000000000..eb4641f3b --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_composed_parity.rs @@ -0,0 +1,182 @@ +//! Teacher-forced parity: VINDEX2 experts vs the same experts from a VINDEX3 +//! container. +//! +//! # Why a greedy continuation is not the gate +//! +//! Running both ways and seeing "Paris." twice proves almost nothing. A short +//! greedy match survives large divergence: the argmax is stable long after the +//! distribution has moved, so agreement on a handful of tokens licenses a claim +//! about *those tokens* and not about the forward pass. This compares the +//! objects that actually differ. +//! +//! # The control is the same loop, not the same answer +//! +//! Both arms call `predict_kquant_hidden_checked` with a bound `MoeRoute`, over +//! identical token ids, on one loaded model. The only difference between them +//! is where the routed expert bytes are read from: +//! +//! ```text +//! control InProcessMoeBackend experts from the VINDEX2 mapped index +//! test ContainerRoutedBackend experts from the VINDEX3 container +//! ``` +//! +//! Comparing against a plain `larql run` instead would change the decode loop +//! as well as the byte source, and a difference could then be either. +//! +//! # What the expected result is, and why +//! +//! c9 established that every routed region in the container is byte-identical +//! to its VINDEX2 source (7680/7680 on `gemma4-26b-a4b`). The same bytes +//! through the same kernels must therefore give **bit-identical** hidden +//! states. Not "close" — identical. A max absolute difference of 1e-7 would be +//! a finding, not a pass: it would mean the two paths are not doing the same +//! arithmetic, and the container would be introducing something. +//! +//! ```text +//! cargo run --release --example vindex3_composed_parity -- [prompt] +//! ``` + +use larql_inference::ffn::{ContainerRoutedBackend, InProcessMoeBackend, MoeRoute}; +use larql_inference::vindex::predict_kquant_hidden_checked; + +const DEFAULT_PROMPT: &str = "The capital of France is"; + +fn main() -> Result<(), String> { + let mut args = std::env::args().skip(1); + let v2 = args + .next() + .ok_or("usage: vindex3_composed_parity [prompt]")?; + let v3 = args + .next() + .ok_or("usage: vindex3_composed_parity [prompt]")?; + let prompt = args.next().unwrap_or_else(|| DEFAULT_PROMPT.to_string()); + + let v2_path = std::path::Path::new(&v2); + let v3_path = std::path::Path::new(&v3); + + println!("teacher-forced parity — VINDEX2 experts vs VINDEX3 container experts"); + println!(" spine {v2}"); + println!(" container {v3}"); + println!(" prompt {prompt:?}\n"); + + let mut cb = larql_vindex::SilentLoadCallbacks; + let weights = larql_vindex::load_model_weights_kquant(v2_path, &mut cb) + .map_err(|e| format!("load weights: {e}"))?; + let tokenizer = + larql_vindex::load_vindex_tokenizer(v2_path).map_err(|e| format!("load tokenizer: {e}"))?; + let mut index = larql_vindex::VectorIndex::load_vindex(v2_path, &mut cb) + .map_err(|e| format!("load vindex: {e}"))?; + index + .load_attn_kquant(v2_path) + .map_err(|e| format!("load attn Q4K: {e}"))?; + index + .load_interleaved_kquant(v2_path) + .map_err(|e| format!("load interleaved Q4K: {e}"))?; + let _ = index.load_lm_head_kquant(v2_path); + + // Wrap exactly as `larql run` does. Scoring the raw string instead would + // teacher-force a token sequence the deployed path never sees, and a parity + // result on inputs nobody serves is a weaker claim than it appears. + let wrapped = + larql_inference::chat::render_user_prompt(v2_path, weights.arch.family(), &prompt) + .map_err(|e| format!("render prompt: {e}"))?; + let ids = larql_inference::encode_prompt(&tokenizer, &*weights.arch, &wrapped) + .map_err(|e| format!("tokenise: {e}"))?; + println!(" {} token(s) teacher-forced\n", ids.len()); + + let routed = ContainerRoutedBackend::open(v3_path, &weights, true) + .map_err(|e| format!("compose: {e}"))?; + + // Control first: if the in-process route itself cannot run, a difference + // afterwards would have two candidate causes rather than one. + let control = predict_kquant_hidden_checked( + &weights, + &ids, + &index, + Some(MoeRoute::fatal(&InProcessMoeBackend)), + ) + .map_err(|e| format!("control run: {e}"))?; + + let test = + predict_kquant_hidden_checked(&weights, &ids, &index, Some(MoeRoute::fatal(&routed))) + .map_err(|e| format!("container run: {e}"))?; + + if control.shape() != test.shape() { + return Err(format!( + "shape differs: control {:?}, container {:?}", + control.shape(), + test.shape() + )); + } + + let mut identical = 0usize; + let mut max_abs = 0.0f32; + for (a, b) in control.iter().zip(test.iter()) { + if a.to_bits() == b.to_bits() { + identical += 1; + } + let d = (a - b).abs(); + if d > max_abs { + max_abs = d; + } + } + let total = control.len(); + + println!("hidden state"); + println!(" elements {total}"); + println!(" bit-identical {identical} / {total}"); + println!(" max |difference| {max_abs:e}"); + + // The decision the model would actually make. Only the **last** position + // predicts the next token, so the argmax has to be taken over that row — + // over the whole flattened `[seq_len, vocab]` it is the largest logit + // anywhere in the prompt, which is not a decision the model ever makes. + let last = |h: &ndarray::Array2| -> ndarray::Array2 { + let r = h.nrows().saturating_sub(1); + h.slice(ndarray::s![r..r + 1, ..]).to_owned() + }; + let logits_control = larql_inference::forward::hidden_to_raw_logits(&weights, &last(&control)); + let logits_test = larql_inference::forward::hidden_to_raw_logits(&weights, &last(&test)); + let argmax = |v: &[f32]| -> usize { + v.iter() + .enumerate() + .filter(|(_, x)| x.is_finite()) + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(i, _)| i) + .unwrap_or(0) + }; + let top_control = argmax(&logits_control); + let top_test = argmax(&logits_test); + let logit_max_abs = logits_control + .iter() + .zip(logits_test.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0f32, f32::max); + + println!("\nfinal logits"); + println!(" max |difference| {logit_max_abs:e}"); + println!( + " top-1 control {top_control} ({:?}), container {top_test} ({:?})", + tokenizer + .decode(&[top_control as u32], true) + .unwrap_or_default(), + tokenizer + .decode(&[top_test as u32], true) + .unwrap_or_default() + ); + + if identical != total || max_abs != 0.0 || logit_max_abs != 0.0 { + return Err(format!( + "NOT bit-identical — {} of {total} hidden elements differ, max |Δ| {max_abs:e}. \ + The container's regions are byte-identical to source (c9), so the same bytes \ + through the same kernels must give the same numbers; a difference here is the \ + binding, not the storage.", + total - identical + )); + } + + println!("\nPASS — bit-identical. The routed experts were read from the VINDEX3"); + println!("container and the forward pass is indistinguishable from the VINDEX2 one."); + println!("Composed run only: the container still has no spine and no tokenizer."); + Ok(()) +} diff --git a/crates/larql-vindex/examples/vindex3_composed_perf.rs b/crates/larql-vindex/examples/vindex3_composed_perf.rs new file mode 100644 index 000000000..f7925396a --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_composed_perf.rs @@ -0,0 +1,227 @@ +//! Where does a composed-run token actually go? +//! +//! The first composed run measured 9.2 s/token, which is a number nobody +//! should quote until it is decomposed. Two candidate stories: +//! +//! ```text +//! A the loop full-recompute (PLE, no KV cache) is simply this slow, +//! and VINDEX2 pays the same +//! B the container 7680 region_bytes lookups per token cost real time +//! ``` +//! +//! They are separable, because both routes run the *same* loop and the +//! decode-stage instrumentation already splits attention / dense / expert / +//! lm_head. Everything outside the expert stage is identical code on identical +//! bytes, so: +//! +//! ```text +//! container tax = expert_ms(container) - expert_ms(in-process) +//! ``` +//! +//! Anything else that differs would be a bug, not a cost. +//! +//! # Reading the result honestly +//! +//! Run it on a cool machine on AC. This project has been bitten by 1.5–3× +//! phantom regressions measured on a hot box, and a composed-vs-control +//! comparison is *less* exposed to that than an absolute number is — both arms +//! see the same thermal state — but the absolute ms/token is not. +//! +//! Arms alternate rather than running one after the other for the same reason: +//! if the machine drifts during the run, drift lands on both arms instead of +//! whichever went second. +//! +//! ```text +//! LARQL_DECODE_STAGES=1 cargo run --release --example vindex3_composed_perf \ +//! -- [tokens] [rounds] +//! ``` + +use larql_inference::ffn::{ContainerRoutedBackend, InProcessMoeBackend, MoeExpertBackend}; +use larql_inference::vindex::generate_kquant_cpu_routed; + +const DEFAULT_TOKENS: usize = 4; +const DEFAULT_ROUNDS: usize = 2; +const PROMPT: &str = "The capital of France is"; + +struct Sample { + wall_ms: f64, + attn_ms: f64, + dense_ms: f64, + expert_ms: f64, + lmhead_ms: f64, + tokens: usize, +} + +fn main() -> Result<(), String> { + let mut args = std::env::args().skip(1); + let v2 = args + .next() + .ok_or("usage: vindex3_composed_perf [tokens] [rounds]")?; + let v3 = args + .next() + .ok_or("usage: vindex3_composed_perf [tokens] [rounds]")?; + let tokens: usize = args + .next() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_TOKENS); + let rounds: usize = args + .next() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_ROUNDS); + + if !larql_inference::decode_stages::is_enabled() { + eprintln!("note: set LARQL_DECODE_STAGES=1 for the per-stage split"); + } + + let v2_path = std::path::Path::new(&v2); + let v3_path = std::path::Path::new(&v3); + + let mut cb = larql_vindex::SilentLoadCallbacks; + let mut weights = larql_vindex::load_model_weights_kquant(v2_path, &mut cb) + .map_err(|e| format!("load weights: {e}"))?; + let tokenizer = + larql_vindex::load_vindex_tokenizer(v2_path).map_err(|e| format!("load tokenizer: {e}"))?; + let mut index = larql_vindex::VectorIndex::load_vindex(v2_path, &mut cb) + .map_err(|e| format!("load vindex: {e}"))?; + index + .load_attn_kquant(v2_path) + .map_err(|e| format!("load attn Q4K: {e}"))?; + index + .load_interleaved_kquant(v2_path) + .map_err(|e| format!("load interleaved Q4K: {e}"))?; + let _ = index.load_lm_head_kquant(v2_path); + + let wrapped = larql_inference::chat::render_user_prompt(v2_path, weights.arch.family(), PROMPT) + .map_err(|e| format!("render prompt: {e}"))?; + let ids = larql_inference::encode_prompt(&tokenizer, &*weights.arch, &wrapped) + .map_err(|e| format!("tokenise: {e}"))?; + + let routed = ContainerRoutedBackend::open(v3_path, &weights, true) + .map_err(|e| format!("compose: {e}"))?; + let in_process = InProcessMoeBackend; + + println!("composed-run cost decomposition"); + println!(" prompt {} token(s)", ids.len()); + println!(" generate {tokens} token(s) x {rounds} round(s), arms alternating\n"); + + let mut control: Vec = Vec::new(); + let mut container: Vec = Vec::new(); + + // Round 0 is a discarded warm-up. The first pass faults in ~12 GB of + // container and ~15 GB of model through the page cache, and its cost + // decays for several rounds afterwards — measuring into that decay + // attributes warm-up to whichever arm happened to go first. + for round in 0..=rounds { + // ABBA: swap which arm leads on alternate rounds. With a fixed order + // the second arm is always the warmer one, which is enough on its own + // to invent a difference in the direction of whoever runs second. + let control_first = round % 2 == 0; + for slot in 0..2 { + let is_control = (slot == 0) == control_first; + let backend: &dyn MoeExpertBackend = if is_control { &in_process } else { &routed }; + let label = if is_control { "vindex2" } else { "vindex3" }; + + larql_inference::decode_stages::reset(); + let started = std::time::Instant::now(); + let out = + generate_kquant_cpu_routed(&mut weights, &tokenizer, &ids, tokens, &index, backend) + .map_err(|e| format!("{label} arm: {e}"))?; + let wall_ms = started.elapsed().as_secs_f64() * 1000.0; + let (attn_ms, dense_ms, expert_ms, lmhead_ms) = + larql_inference::decode_stages::snapshot_ms(); + + let s = Sample { + wall_ms, + attn_ms, + dense_ms, + expert_ms, + lmhead_ms, + tokens: out.len().max(1), + }; + let tag = if round == 0 { + " (warm-up, discarded)" + } else { + "" + }; + println!( + " round {round} {label:8} {:8.0} ms ({:7.0} ms/token){tag}", + s.wall_ms, + s.wall_ms / s.tokens as f64 + ); + if round == 0 { + continue; + } + if is_control { + control.push(s); + } else { + container.push(s); + } + } + } + + fn best(v: &[Sample]) -> &Sample { + // Fastest round: the minimum is the least contaminated by whatever + // else the machine was doing, and the comparison is like-for-like. + v.iter() + .min_by(|a, b| { + (a.wall_ms / a.tokens as f64) + .partial_cmp(&(b.wall_ms / b.tokens as f64)) + .unwrap_or(std::cmp::Ordering::Equal) + }) + .expect("at least one round") + } + let c = best(&control); + let k = best(&container); + + let per = |s: &Sample, v: f64| v / s.tokens as f64; + println!("\n best round, ms/token"); + println!( + " {:<10} {:>10} {:>10} {:>10} {:>10} {:>10}", + "arm", "wall", "attn", "dense", "expert", "lm_head" + ); + for (label, s) in [("vindex2", c), ("vindex3", k)] { + println!( + " {:<10} {:>10.0} {:>10.0} {:>10.0} {:>10.0} {:>10.0}", + label, + per(s, s.wall_ms), + per(s, s.attn_ms), + per(s, s.dense_ms), + per(s, s.expert_ms), + per(s, s.lmhead_ms) + ); + } + + let tax_ms = per(k, k.expert_ms) - per(c, c.expert_ms); + let wall_delta = per(k, k.wall_ms) - per(c, c.wall_ms); + println!("\n container tax (expert stage) {tax_ms:+.0} ms/token"); + println!(" wall delta {wall_delta:+.0} ms/token"); + if per(c, c.wall_ms) > 0.0 { + println!( + " container / vindex2 wall {:.3}x", + per(k, k.wall_ms) / per(c, c.wall_ms) + ); + } + + // What the instrumented stages do *not* explain. On the full-recompute + // path most of a token is per-layer Q4_K dequantisation + // (`insert_q4k_layer_tensors`), which no stage counter covers — so a + // decomposition that only reported the four stages would silently account + // for a minority of the time and invite the reader to divide the rest + // among them. + for (label, s) in [("vindex2", c), ("vindex3", k)] { + let accounted = + per(s, s.attn_ms) + per(s, s.dense_ms) + per(s, s.expert_ms) + per(s, s.lmhead_ms); + let wall = per(s, s.wall_ms); + println!( + " {label:<10} unaccounted {:>8.0} ms/token ({:.0}% of wall — not in any stage counter)", + wall - accounted, + 100.0 * (wall - accounted) / wall.max(1.0) + ); + } + + println!( + "\n Everything outside the expert stage is identical code on identical\n \ + bytes; a large difference there would be a bug, not a cost." + ); + Ok(()) +} diff --git a/crates/larql-vindex/examples/vindex3_import_gemma_all_layers.rs b/crates/larql-vindex/examples/vindex3_import_gemma_all_layers.rs new file mode 100644 index 000000000..c769842ca --- /dev/null +++ b/crates/larql-vindex/examples/vindex3_import_gemma_all_layers.rs @@ -0,0 +1,191 @@ +//! Container ladder **c9** — every MoE layer, so the model *is* a VINDEX3 +//! container rather than one layer of one being. +//! +//! c8 proved a real Gemma layer survives the round trip byte for byte. It left +//! two things open, and they are different in kind: +//! +//! ```text +//! quantitative one layer -> thirty. A loop. +//! qualitative thirty layers do not fit in RAM the way one did. +//! ``` +//! +//! The second is why this is not `vindex3_import_gemma_layer` with a `for` +//! around it. `import_one_layer` returns a `ContainerSpec` that *owns* its +//! segment bytes; at ~421 MB per layer, describing thirty would hold ~12 GB of +//! already-mmapped weights a second time before writing any of them. +//! [`ContainerBuilder`] streams instead — each layer goes straight to its final +//! path and is forgotten — so peak memory is one expert's slice regardless of +//! model size. +//! +//! # What passing here does and does not license +//! +//! It licenses: *this model's routed layers exist as a VINDEX3 container, and +//! every region in it is byte-identical to the VINDEX2 source it came from.* +//! +//! It does **not** license "Gemma runs from VINDEX3". Execution still binds +//! over VINDEX2 bytes (`vindex3_gemma_layer_parity`); what c8 and c9 establish +//! is that those are the same bytes, which is what lets that result carry over. +//! The remaining rungs are the production expert kernel, full-layer residual +//! parity, and greedy token parity through `larql run` — and `extract` still +//! writes VINDEX2 (spec §12.1: new extractions default to VINDEX3 only once +//! the ABI freezes *and* the E0 preservation matrix passes). +//! +//! # Usage +//! +//! ```text +//! cargo run --release --example vindex3_import_gemma_all_layers -- [out-dir] +//! ``` + +use larql_vindex::format::lyrw2::region_role::RegionRole; +use larql_vindex::format::vindex3::import::{region_format_for, routed_storage_key}; +use larql_vindex::format::vindex3::{ContainerBuilder, MoeLayerSource, Vindex3Container}; + +use larql_compute::pipeline_layer::build_moe_weights; + +/// Bytes per GiB, for reporting only. +const GIB: f64 = (1024 * 1024 * 1024) as f64; + +fn main() -> Result<(), String> { + let mut args = std::env::args().skip(1); + let vindex = args + .next() + .ok_or("usage: vindex3_import_gemma_all_layers [out-dir]")?; + let out = args.next().unwrap_or_else(|| "./vindex3-gemma".to_string()); + let out_dir = std::path::Path::new(&out); + + println!("c9 — import every MoE layer into one VINDEX3 container"); + println!(" vindex {vindex}"); + println!(" out {out}\n"); + + let mut callbacks = larql_vindex::SilentLoadCallbacks; + let weights = + larql_vindex::load_model_weights_kquant(std::path::Path::new(&vindex), &mut callbacks) + .map_err(|e| format!("load weights: {e}"))?; + let arch = &*weights.arch; + let num_layers = weights.num_layers; + let hidden = weights.hidden_size as u32; + + // ── Write ──────────────────────────────────────────────────────────── + // + // Non-MoE layers are skipped, not refused: a hybrid model's dense spine is + // a legitimate part of the architecture, and this rung is about the routed + // banks. `num_layers` below still records the model's real depth so a + // reader can tell a routed-only container from a shallower model. + let mut builder = ContainerBuilder::create(out_dir).map_err(|e| format!("create: {e}"))?; + let mut imported: Vec = Vec::new(); + let mut skipped: Vec = Vec::new(); + + for layer in 0..num_layers as u32 { + let Some(moe) = build_moe_weights(&weights, arch, layer as usize) else { + skipped.push(layer); + continue; + }; + let source = MoeLayerSource { + layer, + experts_gate_up: moe.experts_gate_up.clone(), + experts_down: moe.experts_down.clone(), + format: region_format_for(moe.expert_data_format) + .map_err(|e| format!("layer {layer}: {e}"))?, + hidden_size: hidden, + // gate_up is never padded; only down is. See `MoeLayerSource`. + gate_up_stored_intermediate: moe.intermediate_size as u32, + down_stored_intermediate: moe.inter_padded() as u32, + semantic_intermediate: moe.intermediate_size as u32, + top_k: moe.top_k as u32, + }; + let size = builder + .add_moe_layer(&source) + .map_err(|e| format!("layer {layer}: {e}"))?; + imported.push(layer); + println!( + " layer {layer:>3} {} experts, top-{}, {:.2} GiB (running {:.2} GiB)", + moe.num_experts, + moe.top_k, + size as f64 / GIB, + builder.bytes_written() as f64 / GIB + ); + } + + if imported.is_empty() { + return Err("no MoE layers found — this model has no routed banks to import".into()); + } + let total = builder.bytes_written(); + builder + .finish( + weights.arch.family(), + weights.arch.family(), + hidden as usize, + num_layers, + ) + .map_err(|e| format!("finish: {e}"))?; + + println!( + "\n wrote {} routed layer(s), {:.2} GiB; skipped {} non-MoE layer(s)", + imported.len(), + total as f64 / GIB, + skipped.len() + ); + + // ── Verify ─────────────────────────────────────────────────────────── + // + // Reopen from disk rather than trusting the builder: the claim is about + // what a reader finds, not about what the writer intended. + let container = Vindex3Container::open(out_dir).map_err(|e| format!("reopen: {e}"))?; + let defects = container.verify(); + if !defects.is_empty() { + for d in &defects { + eprintln!(" DEFECT {d}"); + } + return Err(format!("{} structural defect(s)", defects.len())); + } + println!(" verify no structural defects"); + + // ── The gate ───────────────────────────────────────────────────────── + // + // Every region of every imported layer, byte for byte, against the source + // it came from. Re-deriving each layer's source slices here rather than + // holding them from the write loop keeps the comparison honest: it reads + // the incumbent afresh, exactly as a consumer of the VINDEX2 file would. + let mut checked = 0usize; + for &layer in &imported { + let moe = build_moe_weights(&weights, arch, layer as usize) + .ok_or_else(|| format!("layer {layer} stopped being MoE between passes"))?; + let reader = container + .segment(&routed_storage_key(layer)) + .map_err(|e| format!("layer {layer} segment: {e}"))?; + + for expert in 0..moe.num_experts as u32 { + for (role, expected) in [ + ( + RegionRole::GateUpFused, + moe.experts_gate_up[expert as usize], + ), + (RegionRole::Down, moe.experts_down[expert as usize]), + ] { + let got = reader + .region_bytes(0, expert, role) + .map_err(|e| format!("layer {layer} expert {expert} {role:?}: {e}"))? + .ok_or_else(|| { + format!("layer {layer} expert {expert} {role:?} missing from container") + })?; + if got != expected { + return Err(format!( + "layer {layer} expert {expert} {role:?}: {} bytes on disk differ \ + from the source's {}", + got.len(), + expected.len() + )); + } + checked += 1; + } + } + } + + println!(" bytes {checked} regions identical to the VINDEX2 source"); + println!( + "\nc9 PASS — {} routed layers are one VINDEX3 container, byte-identical to source.", + imported.len() + ); + println!("Execution still binds VINDEX2 bytes; `extract` still writes VINDEX2."); + Ok(()) +} diff --git a/crates/larql-vindex/examples/vindex3_import_gemma_layer.rs b/crates/larql-vindex/examples/vindex3_import_gemma_layer.rs index b14f187e2..f24b5c460 100644 --- a/crates/larql-vindex/examples/vindex3_import_gemma_layer.rs +++ b/crates/larql-vindex/examples/vindex3_import_gemma_layer.rs @@ -88,7 +88,9 @@ fn main() -> Result<(), String> { experts_down: moe.experts_down.clone(), format: region_format_for(moe.expert_data_format)?, hidden_size: hidden, - stored_intermediate: stored, + // gate_up is never padded; only down is. See `MoeLayerSource`. + gate_up_stored_intermediate: semantic, + down_stored_intermediate: stored, semantic_intermediate: semantic, top_k: moe.top_k as u32, }; @@ -162,20 +164,11 @@ fn main() -> Result<(), String> { /// Map the source's quantisation to the region format that describes it. /// -/// Refuses rather than guesses: a format this container cannot name is one a -/// reader could not interpret, and silently labelling it something else is the -/// transcode this whole module exists to avoid. +/// Thin wrapper over the importer's own mapping so both drivers agree by +/// construction — two copies of this table could drift, and a drift here is a +/// mislabelled region rather than a visible error. fn region_format_for( q: larql_compute::QuantFormat, ) -> Result { - use larql_compute::QuantFormat; - use larql_vindex::format::lyrw2::region_format::RegionFormat; - match q { - QuantFormat::Q4_K => Ok(RegionFormat::Q4K), - QuantFormat::F32 => Ok(RegionFormat::F32), - other => Err(format!( - "expert format {other:?} has no VINDEX3 region format yet — import \ - would have to transcode, which the container forbids" - )), - } + larql_vindex::format::vindex3::import::region_format_for(q).map_err(|e| e.to_string()) } diff --git a/crates/larql-vindex/src/format/load.rs b/crates/larql-vindex/src/format/load.rs index 6c7ac02dd..4e0270312 100644 --- a/crates/larql-vindex/src/format/load.rs +++ b/crates/larql-vindex/src/format/load.rs @@ -44,6 +44,13 @@ impl VectorIndex { callbacks: &mut dyn IndexLoadCallbacks, layer_range: Option<(usize, usize)>, ) -> Result { + // Same gate, same reason as `load_vindex_config`: this builds the v1 + // vector index over a v1 layout, and a VINDEX3 directory reaching it + // would be read against weights that live somewhere else entirely. + // `larql serve` enters here via `bootstrap`, so an ungated parse would + // put the mis-detection on the serving path rather than the CLI's. + detect_generation(dir)?.require(ContainerGeneration::V2)?; + // Read config let config_path = dir.join(INDEX_JSON); let config_text = std::fs::read_to_string(&config_path)?; @@ -384,7 +391,18 @@ fn synthesize_gate_from_q4k( } /// Load embeddings from a .vindex directory. +/// +/// Gated on the generation for the reason [`load_vindex_config`] states, and +/// with more urgency: this is the **first** thing the walk/run path touches +/// (`walk_cmd` reads embeddings before it reads the config), so an ungated +/// parse here is the one that decides whether a VINDEX3 directory is refused +/// by name or wanders into the v1 layout. It refused only by luck — VINDEX3's +/// `index.json` happens to omit `intermediate_size`, so serde rejected it with +/// a field-level message that names nothing about generations. Any +/// `#[serde(default)]` added for schema-1 compatibility would have opened it +/// silently, which is exactly the failure this check exists to prevent. pub fn load_vindex_embeddings(dir: &Path) -> Result<(Array2, f32), VindexError> { + detect_generation(dir)?.require(ContainerGeneration::V2)?; let config_text = std::fs::read_to_string(dir.join(INDEX_JSON))?; let config: VindexConfig = serde_json::from_str(&config_text).map_err(|e| VindexError::Parse(e.to_string()))?; @@ -522,6 +540,70 @@ mod tests { assert_eq!(cfg.family, "llama"); } + /// A VINDEX3 `index.json` carrying **every** field `VindexConfig` needs. + /// + /// The real one omits `intermediate_size`, which is the only reason the + /// ungated v1 entry points refused it — an accident of field overlap, not + /// a decision. This fixture removes that accident so the test measures the + /// generation gate itself. Without the gate, these parse cleanly and the + /// v1 loader proceeds against a layout whose weights are not there. + fn write_v3_index_json_that_would_parse_as_v1(dir: &Path) { + let json = serde_json::json!({ + "version": 3, + "model": "test/v3", + "family": "gemma4", + "num_layers": 2, + "hidden_size": 8, + "intermediate_size": 4, + "vocab_size": 16, + "embed_scale": 1.0, + "layers": [], + "down_top_k": 5, + "has_model_weights": false, + "extract_level": "browse", + "dtype": "f32", + "quant": "none", + "moe_manifest": "moe_manifest.json", + "segments": { "routed/layer_000": 1 } + }); + std::fs::write(dir.join("index.json"), json.to_string()).unwrap(); + } + + #[test] + fn every_v1_entry_point_refuses_a_v3_container_by_generation() { + // The regression this guards: each of these reads `index.json` for + // itself, and `walk_cmd` reaches `load_vindex_embeddings` *before* the + // gated `load_vindex_config`. If any one of them loses its gate, a + // VINDEX3 directory is served against v1 offsets rather than refused. + let dir = TempDir::new().unwrap(); + write_v3_index_json_that_would_parse_as_v1(dir.path()); + + for (entry, err) in [ + ("load_vindex_config", load_vindex_config(dir.path()).err()), + ( + "load_vindex_embeddings", + load_vindex_embeddings(dir.path()).err(), + ), + ( + "load_vindex_with_range", + VectorIndex::load_vindex_with_range( + dir.path(), + &mut crate::SilentLoadCallbacks, + None, + ) + .err(), + ), + ] { + let err = err.unwrap_or_else(|| panic!("{entry} accepted a VINDEX3 container")); + let msg = err.to_string(); + assert!( + msg.contains("VINDEX3") || msg.contains("VINDEX2 loader"), + "{entry} refused for the wrong reason — the message must name \ + the generation, not a missing field. Got: {msg}" + ); + } + } + #[test] fn load_vindex_config_missing_file_errors() { let dir = TempDir::new().unwrap(); diff --git a/crates/larql-vindex/src/format/vindex3/build.rs b/crates/larql-vindex/src/format/vindex3/build.rs new file mode 100644 index 000000000..641907d83 --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/build.rs @@ -0,0 +1,168 @@ +//! Stream a many-layer VINDEX3 container to disk (container ladder c9). +//! +//! # Why c9 is not just c8 in a loop +//! +//! [`super::write::write_container`] takes a [`ContainerSpec`] whose segments +//! own their bytes. That is the right shape for one fixture-sized bank and the +//! wrong shape for a model: layer 0 of `gemma4-26b-a4b` is ~421 MB, so +//! describing thirty of them the c8 way would hold ~12 GB of already-mapped +//! weights a second time, in RAM, purely to hand them to a writer that puts +//! them straight back on disk. +//! +//! This builder inverts that. Each layer is written **directly at its final +//! path** as it arrives, and only the manifest entries — a few hundred bytes +//! per layer — accumulate: +//! +//! ```text +//! c8 describe all layers -> hold all bytes -> write everything +//! c9 describe one layer -> write it -> forget it -> next +//! ``` +//! +//! Peak memory is therefore one expert's slice, whatever the model's size. +//! +//! # The write order is the crash contract +//! +//! Segments, then the manifest, then `index.json` **last** — the same ordering +//! [`super::write::write_container`] states, and for the same reason. +//! `index.json` is the discriminator every reader dispatches on, so a crash +//! part-way through a thirty-layer import must leave a directory that is *not +//! yet* a VINDEX3 container, rather than one that announces itself as VINDEX3 +//! and is missing two thirds of its banks. That distinction matters more here +//! than at c8, because a multi-gigabyte import is long enough to actually be +//! interrupted. +//! +//! # Still verbatim +//! +//! Nothing here transcodes, requantises or repacks. The builder places bytes +//! and declares them, exactly as §9.1 requires of the one-layer path — a +//! streaming assembler that also converted formats would hide the silent +//! conversion one level below where c8 already forbids it. + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +use super::import::{manifest_layer, routed_storage_key, write_segment_file, MoeLayerSource}; +use super::index::Vindex3Index; +use super::write::{segment_path, MOE_MANIFEST_JSON}; +use crate::format::filenames::INDEX_JSON; +use crate::format::moe_manifest::layer::MoeLayer; +use crate::format::moe_manifest::MoeManifest; +use crate::VindexError; + +/// Every segment key a container declares appears exactly once. A repeat means +/// the caller would have silently overwritten a bank it had already written. +const SEGMENT_DECLARED_ONCE: u32 = 1; + +/// Assembles a VINDEX3 container one layer at a time. +/// +/// Construct with [`ContainerBuilder::create`], feed layers with +/// [`ContainerBuilder::add_moe_layer`], and close with +/// [`ContainerBuilder::finish`]. Dropping the builder without finishing leaves +/// the segments on disk but writes no `index.json`, so the directory is not a +/// container and no reader will treat it as one. +pub struct ContainerBuilder { + root: PathBuf, + layers: Vec, + declared: BTreeMap, + bytes_written: u64, +} + +impl ContainerBuilder { + /// Prepare `root` to receive segments. Creates the directory if absent. + pub fn create(root: &Path) -> Result { + std::fs::create_dir_all(root).map_err(VindexError::Io)?; + Ok(Self { + root: root.to_path_buf(), + layers: Vec::new(), + declared: BTreeMap::new(), + bytes_written: 0, + }) + } + + /// Write one routed MoE layer's bank and record how to interpret it. + /// + /// Returns the segment's size on disk. Refuses a layer whose storage key + /// this container already declares: importing it would overwrite the + /// earlier bank while leaving both manifest entries in place, which + /// produces a container that verifies clean and executes the wrong + /// weights — the one failure mode worth spending a check on. + pub fn add_moe_layer(&mut self, source: &MoeLayerSource<'_>) -> Result { + let storage = routed_storage_key(source.layer); + if self.declared.contains_key(&storage) { + return Err(VindexError::Parse(format!( + "layer {} is already in this container as `{storage}` — importing \ + it twice would overwrite the first bank and leave two manifest \ + entries pointing at the survivor", + source.layer + ))); + } + + let path = segment_path(&self.root, &storage); + write_segment_file(source, &path)?; + let size = std::fs::metadata(&path).map_err(VindexError::Io)?.len(); + + self.layers.push(manifest_layer(source, &storage)); + self.declared.insert(storage, SEGMENT_DECLARED_ONCE); + self.bytes_written += size; + Ok(size) + } + + /// How many layers have been written so far. + pub fn layers_added(&self) -> usize { + self.layers.len() + } + + /// Total segment bytes written so far. + pub fn bytes_written(&self) -> u64 { + self.bytes_written + } + + /// Write the manifest and then `index.json`, making `root` a container. + /// + /// `num_layers` is the **model's** layer count, not the number imported: + /// a container holding only the MoE layers of a hybrid model still has to + /// describe the model it came from, or a reader cannot tell a partial + /// import from a shallower architecture. + pub fn finish( + self, + model: impl Into, + family: impl Into, + hidden_size: usize, + num_layers: usize, + ) -> Result<(), VindexError> { + if self.layers.is_empty() { + return Err(VindexError::Parse( + "a VINDEX3 container needs at least one segment; an index \ + declaring none cannot be bound and would fail at execution \ + instead of here" + .into(), + )); + } + + // Validate before publishing: a manifest this builder produced and + // cannot itself parse would push the failure into the container it + // just spent minutes writing, discovered only by whoever opens it. + let manifest = MoeManifest::new(self.layers); + let manifest_json = serde_json::to_string_pretty(&manifest) + .map_err(|e| VindexError::Parse(format!("serialise moe_manifest: {e}")))?; + MoeManifest::parse(&manifest_json)?; + std::fs::write(self.root.join(MOE_MANIFEST_JSON), &manifest_json) + .map_err(VindexError::Io)?; + + let index = Vindex3Index::new( + model, + family, + hidden_size, + num_layers, + MOE_MANIFEST_JSON, + self.declared, + ); + let index_json = serde_json::to_string_pretty(&index) + .map_err(|e| VindexError::Parse(format!("serialise index.json: {e}")))?; + std::fs::write(self.root.join(INDEX_JSON), index_json).map_err(VindexError::Io) + } +} + +#[cfg(test)] +#[path = "build_tests.rs"] +mod tests; diff --git a/crates/larql-vindex/src/format/vindex3/build_tests.rs b/crates/larql-vindex/src/format/vindex3/build_tests.rs new file mode 100644 index 000000000..353cfdcd1 --- /dev/null +++ b/crates/larql-vindex/src/format/vindex3/build_tests.rs @@ -0,0 +1,246 @@ +//! Tests for the c9 streaming builder. +//! +//! c8's gate was "one layer's regions survive the round trip unchanged". c9 +//! adds two things that only exist once there is more than one layer: every +//! layer must round-trip (not just the first or last), and the container must +//! not *announce itself* as a container until all of them are down. + +use super::*; +use crate::format::generation::{detect_generation, ContainerGeneration}; +use crate::format::lyrw2::region_format::RegionFormat; +use crate::format::lyrw2::region_role::RegionRole; +use crate::format::vindex3::read::Vindex3Container; +use tempfile::tempdir; + +const HIDDEN: u32 = 8; +const STORED_INTER: u32 = 12; +const SEMANTIC_INTER: u32 = 10; +const EXPERTS: usize = 3; +const TOP_K: u32 = 2; +const LAYERS: u32 = 4; +const MODEL_LAYERS: usize = 6; + +/// Distinct per (layer, expert) so a bank read from the wrong layer — the +/// failure a loop introduces and a one-layer test cannot see — cannot pass. +/// +/// Sized at the **semantic** width, not the stored one: gate_up is never +/// padded, and a fixture that pads both regions equally cannot tell a correct +/// importer from one that describes gate_up with down's width. +fn gate_up_bytes(layer: u32, expert: usize) -> Vec { + let n = (SEMANTIC_INTER * 2 * HIDDEN) as usize * 4; + (0..n) + .map(|i| (layer as usize * 101 + expert * 31 + i) as u8) + .collect() +} + +fn down_bytes(layer: u32, expert: usize) -> Vec { + let n = (HIDDEN * STORED_INTER) as usize * 4; + (0..n) + .map(|i| (layer as usize * 53 + expert * 17 + i + 7) as u8) + .collect() +} + +struct Owned { + gate_up: Vec>, + down: Vec>, +} + +fn owned(layer: u32) -> Owned { + Owned { + gate_up: (0..EXPERTS).map(|e| gate_up_bytes(layer, e)).collect(), + down: (0..EXPERTS).map(|e| down_bytes(layer, e)).collect(), + } +} + +fn source(layer: u32, o: &Owned) -> MoeLayerSource<'_> { + MoeLayerSource { + layer, + experts_gate_up: o.gate_up.iter().map(|v| v.as_slice()).collect(), + experts_down: o.down.iter().map(|v| v.as_slice()).collect(), + format: RegionFormat::F32, + hidden_size: HIDDEN, + gate_up_stored_intermediate: SEMANTIC_INTER, + down_stored_intermediate: STORED_INTER, + semantic_intermediate: SEMANTIC_INTER, + top_k: TOP_K, + } +} + +/// Build a complete `LAYERS`-layer container at `root`. +fn build_all(root: &Path) -> Vec { + let mut builder = ContainerBuilder::create(root).unwrap(); + let mut kept = Vec::new(); + for layer in 0..LAYERS { + let o = owned(layer); + builder.add_moe_layer(&source(layer, &o)).unwrap(); + kept.push(o); + } + builder + .finish("gemma-fixture", "gemma", HIDDEN as usize, MODEL_LAYERS) + .unwrap(); + kept +} + +// ── The verbatim guarantee, now across layers ──────────────────────────── + +#[test] +fn every_layer_round_trips_byte_for_byte() { + let dir = tempdir().unwrap(); + let kept = build_all(dir.path()); + + let container = Vindex3Container::open(dir.path()).unwrap(); + assert!( + container.verify().is_empty(), + "structural defects: {:?}", + container.verify() + ); + + let mut checked = 0usize; + for layer in 0..LAYERS { + let reader = container.segment(&routed_storage_key(layer)).unwrap(); + for expert in 0..EXPERTS { + for (role, expected) in [ + ( + RegionRole::GateUpFused, + kept[layer as usize].gate_up[expert].as_slice(), + ), + ( + RegionRole::Down, + kept[layer as usize].down[expert].as_slice(), + ), + ] { + let got = reader + .region_bytes(0, expert as u32, role) + .unwrap() + .unwrap_or_else(|| panic!("layer {layer} expert {expert} {role:?} missing")); + assert_eq!( + got, expected, + "layer {layer} expert {expert} {role:?} differs from its source" + ); + checked += 1; + } + } + } + assert_eq!(checked, LAYERS as usize * EXPERTS * 2); +} + +#[test] +fn the_container_reports_v3_and_declares_one_segment_per_layer() { + let dir = tempdir().unwrap(); + build_all(dir.path()); + + assert_eq!( + detect_generation(dir.path()).unwrap(), + ContainerGeneration::V3 + ); + + let container = Vindex3Container::open(dir.path()).unwrap(); + for layer in 0..LAYERS { + assert!( + container.segment(&routed_storage_key(layer)).is_ok(), + "layer {layer} is not resolvable from the written container" + ); + } +} + +// ── The crash contract ─────────────────────────────────────────────────── + +#[test] +fn a_directory_is_not_a_container_until_finish_writes_the_index() { + // The property that makes a long import safe to interrupt: segments land + // first, and until `index.json` exists no reader will dispatch here. + let dir = tempdir().unwrap(); + let mut builder = ContainerBuilder::create(dir.path()).unwrap(); + let o = owned(0); + builder.add_moe_layer(&source(0, &o)).unwrap(); + + assert!( + segment_path(dir.path(), &routed_storage_key(0)).exists(), + "the segment should be on disk before the index is" + ); + assert!( + !dir.path().join(INDEX_JSON).exists(), + "index.json must not exist mid-import — a reader would treat a partial \ + directory as a complete container" + ); + assert!(detect_generation(dir.path()).is_err()); + + builder + .finish("gemma-fixture", "gemma", HIDDEN as usize, MODEL_LAYERS) + .unwrap(); + assert!(dir.path().join(INDEX_JSON).exists()); +} + +// ── Refusals ───────────────────────────────────────────────────────────── + +#[test] +fn importing_the_same_layer_twice_is_refused() { + let dir = tempdir().unwrap(); + let mut builder = ContainerBuilder::create(dir.path()).unwrap(); + let o = owned(0); + builder.add_moe_layer(&source(0, &o)).unwrap(); + + let err = builder.add_moe_layer(&source(0, &o)).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("already in this container"), + "refusal should name the collision, got: {msg}" + ); + assert_eq!(builder.layers_added(), 1); +} + +#[test] +fn finishing_with_no_layers_is_refused() { + let dir = tempdir().unwrap(); + let builder = ContainerBuilder::create(dir.path()).unwrap(); + let err = builder + .finish("gemma-fixture", "gemma", HIDDEN as usize, MODEL_LAYERS) + .unwrap_err(); + assert!(err.to_string().contains("at least one segment")); + assert!( + !dir.path().join(INDEX_JSON).exists(), + "a refused finish must not leave an index behind" + ); +} + +#[test] +fn a_layer_the_importer_refuses_leaves_the_builder_usable() { + // `check()` runs before any byte is written, so a bad layer must not + // consume its storage key or half-write its bank. + let dir = tempdir().unwrap(); + let mut builder = ContainerBuilder::create(dir.path()).unwrap(); + let o = owned(0); + + let mut bad = source(0, &o); + bad.top_k = EXPERTS as u32 + 1; + assert!(builder.add_moe_layer(&bad).is_err()); + assert_eq!(builder.layers_added(), 0); + + builder.add_moe_layer(&source(0, &o)).unwrap(); + assert_eq!(builder.layers_added(), 1); +} + +// ── Accounting ─────────────────────────────────────────────────────────── + +#[test] +fn bytes_written_tracks_the_segments_on_disk() { + let dir = tempdir().unwrap(); + build_all(dir.path()); + + let on_disk: u64 = (0..LAYERS) + .map(|l| { + std::fs::metadata(segment_path(dir.path(), &routed_storage_key(l))) + .unwrap() + .len() + }) + .sum(); + + let mut builder = ContainerBuilder::create(tempdir().unwrap().path()).unwrap(); + let mut total = 0u64; + for layer in 0..LAYERS { + let o = owned(layer); + total += builder.add_moe_layer(&source(layer, &o)).unwrap(); + } + assert_eq!(total, on_disk); + assert_eq!(builder.bytes_written(), on_disk); +} diff --git a/crates/larql-vindex/src/format/vindex3/import.rs b/crates/larql-vindex/src/format/vindex3/import.rs index 3376adead..43b17189a 100644 --- a/crates/larql-vindex/src/format/vindex3/import.rs +++ b/crates/larql-vindex/src/format/vindex3/import.rs @@ -73,10 +73,21 @@ pub struct MoeLayerSource<'a> { pub format: RegionFormat, /// Residual width entering the layer. pub hidden_size: u32, - /// Per-expert intermediate width, as **stored** — Gemma pads this, and the - /// padded extent is what the bytes actually contain. The semantic width is - /// a view over it, recorded by the manifest's `expert_dims`, not here. - pub stored_intermediate: u32, + /// Intermediate width as stored in the **gate/up** region. + /// + /// Separate from [`Self::down_stored_intermediate`] because the two are + /// genuinely different on Gemma, and describing them with one number + /// over-declares the gate/up region. `gate_up` is `[2 × inter, hidden]` + /// and is *never* padded — `hidden` is already a 256-multiple, so Q4_K + /// quantises it cleanly — while `down` is `[hidden, inter]` with `inter` + /// padded up to the next super-block (704 → 768 on `gemma4-26b-a4b`). + /// See `larql_compute::cpu::ops::moe::forward`, which states both. + pub gate_up_stored_intermediate: u32, + /// Intermediate width as stored in the **down** region, padding included. + /// + /// This is the extent a kernel contracts over, so it is also the bank + /// descriptor's `intermediate_dim`. + pub down_stored_intermediate: u32, /// Semantic intermediate width the operation means. pub semantic_intermediate: u32, /// Experts activated per token. @@ -108,12 +119,17 @@ impl MoeLayerSource<'_> { self.num_experts() )); } - if self.semantic_intermediate > self.stored_intermediate { - return refuse(format!( - "semantic intermediate {} exceeds the stored {} — a view can \ - narrow the stored extent, never widen it", - self.semantic_intermediate, self.stored_intermediate - )); + for (which, stored) in [ + ("gate_up", self.gate_up_stored_intermediate), + ("down", self.down_stored_intermediate), + ] { + if self.semantic_intermediate > stored { + return refuse(format!( + "semantic intermediate {} exceeds the {which} region's stored \ + {stored} — a view can narrow the stored extent, never widen it", + self.semantic_intermediate + )); + } } // Ragged expert slices mean the source's own layout is inconsistent; // importing them would bake that into a container that then fails at @@ -138,15 +154,24 @@ impl MoeLayerSource<'_> { } } -/// Write `source`'s regions into one LYRW v2 segment, verbatim. +/// Write `source`'s regions into one LYRW v2 segment file at `dest`, verbatim. /// -/// `staging` is where the writer builds the file; the bytes are read back and -/// returned, so the caller decides where the container finally lives. -pub fn segment_bytes_for_layer( +/// The bytes are streamed region by region and never held whole, so peak memory +/// is one expert's slice regardless of how large the layer is. That matters at +/// c9: a 26B layer is ~421 MB, and materialising thirty of them to describe a +/// model would cost more RAM than loading the model did. +/// +/// Parent directories are created — a segment key is a *path* (`routed/…`), +/// composed rather than globbed (§12.1), so its directory is part of the key's +/// meaning rather than something the caller should have to anticipate. +pub fn write_segment_file( source: &MoeLayerSource<'_>, - staging: &std::path::Path, -) -> Result, VindexError> { + dest: &std::path::Path, +) -> Result<(), VindexError> { source.check()?; + if let Some(parent) = dest.parent() { + std::fs::create_dir_all(parent).map_err(VindexError::Io)?; + } let experts = source.num_experts(); let bank = BankDescriptor { @@ -154,21 +179,26 @@ pub fn segment_bytes_for_layer( kind: BankKind::Routed, num_entries: experts, input_dim: source.hidden_size, - intermediate_dim: source.stored_intermediate, + intermediate_dim: source.down_stored_intermediate, output_dim: source.hidden_size, region_schema_count: REGIONS_PER_EXPERT, browse: BrowseMode::None, }; - // Shapes are the **stored** ones. `gate_up` is one fused range of - // 2 x intermediate rows; `down` contracts the intermediate back to hidden. + // Shapes are the **stored** ones, and the two regions do not share a + // stored width. `gate_up` is one fused range of 2 x its own intermediate + // rows; `down` contracts its own (padded) intermediate back to hidden. + // Using one number for both over-declares gate_up by the padding, which + // nothing downstream would catch: regions are copied verbatim, so the + // declared shape never enters the byte length, and `verify` checks + // structure rather than shape-against-length. let schemas = vec![ RegionSchema::unpaired( SCHEMA_GATE_UP, RegionRole::GateUpFused, source.format, Packing::RowMajor, - source.stored_intermediate * 2, + source.gate_up_stored_intermediate * 2, source.hidden_size, ), RegionSchema::unpaired( @@ -177,12 +207,12 @@ pub fn segment_bytes_for_layer( source.format, Packing::RowMajor, source.hidden_size, - source.stored_intermediate, + source.down_stored_intermediate, ), ]; let plan = Lyrw2Plan::single_segment(source.layer, bank, schemas); - let mut writer = Lyrw2Writer::create(staging, plan) + let mut writer = Lyrw2Writer::create(dest, plan) .map_err(|e| VindexError::Parse(format!("create LYRW v2 writer: {e}")))?; // Entry order is (expert, then schema), matching the plan's declaration. @@ -196,8 +226,20 @@ pub fn segment_bytes_for_layer( } writer .finish() - .map_err(|e| VindexError::Parse(format!("finish LYRW v2 segment: {e}")))?; + .map_err(|e| VindexError::Parse(format!("finish LYRW v2 segment: {e}"))) +} +/// Write `source`'s regions into one LYRW v2 segment, verbatim, and return them. +/// +/// `staging` is where the writer builds the file; the bytes are read back and +/// returned, so the caller decides where the container finally lives. This is +/// the one-layer (c8) shape — it costs a full copy of the layer in RAM, which +/// is why [`write_segment_file`] exists for the all-layer path. +pub fn segment_bytes_for_layer( + source: &MoeLayerSource<'_>, + staging: &std::path::Path, +) -> Result, VindexError> { + write_segment_file(source, staging)?; std::fs::read(staging).map_err(VindexError::Io) } @@ -275,6 +317,24 @@ pub fn routed_storage_key(layer: u32) -> String { format!("routed/layer_{layer:03}") } +/// Map a source's expert quantisation to the region format that describes it. +/// +/// Refuses rather than guesses. A format this container cannot name is one a +/// reader could not interpret, and labelling it as something else would be the +/// silent transcode this module exists to prevent — the failure would surface +/// as wrong numbers at execution, not as an error at import. +pub fn region_format_for(q: larql_compute::QuantFormat) -> Result { + use larql_compute::QuantFormat; + match q { + QuantFormat::Q4_K => Ok(RegionFormat::Q4K), + QuantFormat::F32 => Ok(RegionFormat::F32), + other => Err(VindexError::Parse(format!( + "expert format {other:?} has no VINDEX3 region format yet — import \ + would have to transcode, which the container forbids" + ))), + } +} + #[cfg(test)] #[path = "import_tests.rs"] mod tests; diff --git a/crates/larql-vindex/src/format/vindex3/import_tests.rs b/crates/larql-vindex/src/format/vindex3/import_tests.rs index d9520a6c0..db430a5e6 100644 --- a/crates/larql-vindex/src/format/vindex3/import_tests.rs +++ b/crates/larql-vindex/src/format/vindex3/import_tests.rs @@ -19,8 +19,12 @@ const EXPERTS: usize = 3; const TOP_K: u32 = 2; /// Distinct, position-dependent bytes so a wrong offset cannot pass. +/// +/// Sized at the **semantic** width: gate_up is never padded (only down is), so +/// a fixture padding both equally could not distinguish a correct importer +/// from one that describes gate_up using down's stored width. fn gate_up_bytes(expert: usize) -> Vec { - let n = (STORED_INTER * 2 * HIDDEN) as usize * 4; + let n = (SEMANTIC_INTER * 2 * HIDDEN) as usize * 4; (0..n).map(|i| (expert * 31 + i) as u8).collect() } @@ -48,7 +52,8 @@ fn source(o: &Owned) -> MoeLayerSource<'_> { experts_down: o.down.iter().map(|v| v.as_slice()).collect(), format: RegionFormat::F32, hidden_size: HIDDEN, - stored_intermediate: STORED_INTER, + gate_up_stored_intermediate: SEMANTIC_INTER, + down_stored_intermediate: STORED_INTER, semantic_intermediate: SEMANTIC_INTER, top_k: TOP_K, } @@ -238,5 +243,136 @@ fn a_semantic_width_wider_than_the_stored_one_is_refused() { let staging = tempdir().unwrap(); let err = segment_bytes_for_layer(&src, &staging.path().join("s.lyrw")) .expect_err("a widening view must not import"); - assert!(format!("{err}").contains("exceeds the stored"), "{err}"); + let msg = format!("{err}"); + // The refusal names which region it is too wide for, because the two + // regions carry different stored widths and "the stored one" is ambiguous. + assert!(msg.contains("exceeds the"), "{msg}"); + assert!( + msg.contains("gate_up") || msg.contains("down"), + "the refusal must name the region whose stored width was exceeded: {msg}" + ); +} + +#[test] +fn each_region_is_declared_at_its_own_stored_width() { + // The defect this pins: `gate_up` and `down` do not share a stored width. + // gate_up is `[2 x inter, hidden]` and never padded; down is + // `[hidden, inter]` with inter padded to the next super-block. Describing + // both with down's width over-declares gate_up by the padding — and + // nothing downstream catches it, because regions are copied verbatim (so + // the declared shape never enters the byte length) and `verify` checks + // structure rather than shape-against-length. + let o = owned(); + let src = source(&o); + assert_ne!( + src.gate_up_stored_intermediate, src.down_stored_intermediate, + "the fixture must exercise the asymmetry or this test proves nothing" + ); + + let staging = tempdir().unwrap(); + let path = staging.path().join("s.lyrw"); + let bytes = segment_bytes_for_layer(&src, &path).unwrap(); + let reader = Lyrw2Reader::parse(&bytes).unwrap(); + + let schemas = reader.schemas_for(0).expect("bank 0 has schemas"); + let gate_up = schemas + .iter() + .find(|s| s.role == RegionRole::GateUpFused) + .expect("a gate_up schema"); + let down = schemas + .iter() + .find(|s| s.role == RegionRole::Down) + .expect("a down schema"); + + assert_eq!( + gate_up.rows, + SEMANTIC_INTER * 2, + "gate_up must be declared at its own (unpadded) width" + ); + assert_eq!(gate_up.cols, HIDDEN); + assert_eq!(down.rows, HIDDEN); + assert_eq!( + down.cols, STORED_INTER, + "down must be declared at its own (padded) width" + ); + + // And the declaration must match what is actually there. + let region = reader + .region_bytes(0, 0, RegionRole::GateUpFused) + .unwrap() + .unwrap(); + assert_eq!( + region.len(), + (gate_up.rows * gate_up.cols) as usize * 4, + "declared gate_up shape does not account for the region's bytes" + ); +} + +// ── Format naming ──────────────────────────────────────────────────────── + +#[test] +fn a_representable_expert_format_maps_to_its_region_format() { + use larql_compute::QuantFormat; + assert_eq!( + region_format_for(QuantFormat::Q4_K).unwrap(), + RegionFormat::Q4K + ); + assert_eq!( + region_format_for(QuantFormat::F32).unwrap(), + RegionFormat::F32 + ); +} + +#[test] +fn an_unrepresentable_expert_format_is_refused_not_relabelled() { + // The whole importer exists to avoid silent conversion, so a format the + // container cannot name must stop the import rather than be written under + // a neighbouring tag. Mislabelling here would surface as wrong numbers at + // execution, with nothing pointing back to the import. + use larql_compute::QuantFormat; + for unnameable in [QuantFormat::Q6_K, QuantFormat::BF16, QuantFormat::Q8_0] { + let err = region_format_for(unnameable) + .expect_err("a format with no VINDEX3 region format must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("no VINDEX3 region format") && msg.contains("transcode"), + "the refusal must say why it cannot proceed, got: {msg}" + ); + } +} + +// ── Per-region width refusals ──────────────────────────────────────────── + +#[test] +fn a_semantic_width_wider_than_the_down_region_is_refused_naming_down() { + // The gate_up arm is checked first, so a source that only violates the + // down width is the one that proves the second arm is reachable — and + // that the message names the region rather than saying "the stored one". + let o = owned(); + let mut src = source(&o); + src.gate_up_stored_intermediate = STORED_INTER * 2; + src.down_stored_intermediate = SEMANTIC_INTER - 1; + let staging = tempdir().unwrap(); + let err = segment_bytes_for_layer(&src, &staging.path().join("s.lyrw")) + .expect_err("a view wider than the down region must not import"); + let msg = format!("{err}"); + assert!(msg.contains("down"), "must name the down region: {msg}"); +} + +// ── Segment placement ──────────────────────────────────────────────────── + +#[test] +fn writing_a_segment_creates_the_directories_its_key_implies() { + // A storage key is a path (`routed/layer_000`), composed rather than + // globbed, so its directory is part of the key's meaning — the c9 builder + // relies on that instead of pre-creating a tree it would have to keep in + // step with the key format. + let o = owned(); + let src = source(&o); + let root = tempdir().unwrap(); + let nested = root.path().join("routed").join("deeper").join("s.lyrw"); + assert!(!nested.parent().unwrap().exists()); + + write_segment_file(&src, &nested).unwrap(); + assert!(nested.exists(), "segment file was not placed at its key"); } diff --git a/crates/larql-vindex/src/format/vindex3/mod.rs b/crates/larql-vindex/src/format/vindex3/mod.rs index 0b097d844..32ff9bfb4 100644 --- a/crates/larql-vindex/src/format/vindex3/mod.rs +++ b/crates/larql-vindex/src/format/vindex3/mod.rs @@ -36,6 +36,7 @@ //! the silent conversion §9.1 forbids, buried one layer below where anyone //! would look for it. Regions are placed exactly as their producer wrote them. +pub mod build; pub mod import; pub mod index; pub mod profile; @@ -47,7 +48,8 @@ pub mod variants; pub mod verify; pub mod write; -pub use import::{import_one_layer, MoeLayerSource}; +pub use build::ContainerBuilder; +pub use import::{import_one_layer, write_segment_file, MoeLayerSource}; pub use index::{Vindex3Index, PROFILE_EXACT}; pub use profile::{Profile, ProfileSelectionError, ResolvedProfile}; pub use read::Vindex3Container;