diff --git a/crates/larql-server/coverage-policy.json b/crates/larql-server/coverage-policy.json index 4399daa24..696a747e9 100644 --- a/crates/larql-server/coverage-policy.json +++ b/crates/larql-server/coverage-policy.json @@ -1,5 +1,5 @@ { - "policy_note": "Per-file coverage policy. The `included_total_line_min_percent` gate computes total over the included files only \u2014 pure-logic modules, route handlers with small bodies, helpers. I/O-bound wrappers (heavy route handlers, gRPC servers, daemon bootstrap, the announce client loop) are excluded because they require a live model and remote shards to exercise: their coverage is intentionally low and tracked separately via the still-existing per-file `total_line_min_percent`. New files under the include set must hit 90% on first commit; existing debt baselines should only ratchet upward. 2026-05-20: investigated a CI-vs-local divergence on `completions.rs` (CI 70.34% vs local 86.85%, identical test outcomes) \u2014 root cause was NOT a real generation regression but a coverage artefact. The `completions_*_returns_200` tests asserted `OK || is_server_error()` on `resp.status()` without draining the response body, so axum's lazy `into_response()` serialisation of the buffered handler never ran under llvm-cov instrumentation on Ubuntu (it did run on macOS, hence the divergence). Fixed by draining the body in every success-path completions test (`tests/test_openai_completions_coverage.rs::capture_completion`) and tightening the asserts to strict 200 OK. completions.rs now reports 86.85% consistently across platforms; baseline restored to 86.0. The Linux `synthetic_q4k_vindex` weights are actually finite; no NaN regression existed. 2026-08-22: bootstrap.rs split into bootstrap/{mod,cli,load,listeners} \u2014 mod.rs (serve daemon) and listeners.rs (H3 socket) inherit the daemon exclusion; cli.rs and tests are pure logic and included; load.rs enters with a 66.0 baseline to ratchet upward. chat.rs split into chat/ and removed from the exclude list \u2014 the whole chat surface (types/handler/stream/tools/v3) is now coverage-gated at the default floor. 2026-08-22: openai chat/responses coverage push \u2014 fixed the SSE test fixture-lifetime bug (fixture dropped before the body was drained, so generation always failed at the lazy weights load and the whole stream pump read as uncovered), extracted the four duplicated per-token callback bodies into token_tap.rs (unit-tested), and added validation/timeout/stop/failed-stream tests. chat/stream.rs (89.4 local) and responses/stream.rs (86.3 local) enter with baselines: their remaining uncovered lines are the V2 per-token emit closures and client-disconnect returns, which the CPU generation arm structurally never invokes (per-token callbacks are a GPU-path affordance \u2014 see generate_streaming_runs_against_synthetic_fixture in larql-inference), plus tool-success emission the synthetic vocab cannot produce (no JSON punctuation tokens, so the constrained mask cannot emit parseable output). Ratchet these upward when a GPU-exercising or JSON-capable fixture lands. completions.rs ratcheted 86.0 -> 88.0 (reads 90.6 local). 2026-08-22 (N0.6-on-V3): fsm.rs ratcheted 85.6 -> 86.5 (reads 87.1 local) \u2014 the emission-time key-discipline fixes added covered branches. 2026-08-22 (/v1/sessions): session.rs became session/{clock,lease,manager,state} and routes/sessions/ landed; all five new files enter at 97-100% with no debt baseline. Included-total 92.7 -> 93.12.", + "policy_note": "Per-file coverage policy. The `included_total_line_min_percent` gate computes total over the included files only \u2014 pure-logic modules, route handlers with small bodies, helpers. I/O-bound wrappers (heavy route handlers, gRPC servers, daemon bootstrap, the announce client loop) are excluded because they require a live model and remote shards to exercise: their coverage is intentionally low and tracked separately via the still-existing per-file `total_line_min_percent`. New files under the include set must hit 90% on first commit; existing debt baselines should only ratchet upward. 2026-05-20: investigated a CI-vs-local divergence on `completions.rs` (CI 70.34% vs local 86.85%, identical test outcomes) \u2014 root cause was NOT a real generation regression but a coverage artefact. The `completions_*_returns_200` tests asserted `OK || is_server_error()` on `resp.status()` without draining the response body, so axum's lazy `into_response()` serialisation of the buffered handler never ran under llvm-cov instrumentation on Ubuntu (it did run on macOS, hence the divergence). Fixed by draining the body in every success-path completions test (`tests/test_openai_completions_coverage.rs::capture_completion`) and tightening the asserts to strict 200 OK. completions.rs now reports 86.85% consistently across platforms; baseline restored to 86.0. The Linux `synthetic_q4k_vindex` weights are actually finite; no NaN regression existed. 2026-08-22: bootstrap.rs split into bootstrap/{mod,cli,load,listeners} \u2014 mod.rs (serve daemon) and listeners.rs (H3 socket) inherit the daemon exclusion; cli.rs and tests are pure logic and included; load.rs enters with a 66.0 baseline to ratchet upward. chat.rs split into chat/ and removed from the exclude list \u2014 the whole chat surface (types/handler/stream/tools/v3) is now coverage-gated at the default floor. 2026-08-22: openai chat/responses coverage push \u2014 fixed the SSE test fixture-lifetime bug (fixture dropped before the body was drained, so generation always failed at the lazy weights load and the whole stream pump read as uncovered), extracted the four duplicated per-token callback bodies into token_tap.rs (unit-tested), and added validation/timeout/stop/failed-stream tests. chat/stream.rs (89.4 local) and responses/stream.rs (86.3 local) enter with baselines: their remaining uncovered lines are the V2 per-token emit closures and client-disconnect returns, which the CPU generation arm structurally never invokes (per-token callbacks are a GPU-path affordance \u2014 see generate_streaming_runs_against_synthetic_fixture in larql-inference), plus tool-success emission the synthetic vocab cannot produce (no JSON punctuation tokens, so the constrained mask cannot emit parseable output). Ratchet these upward when a GPU-exercising or JSON-capable fixture lands. completions.rs ratcheted 86.0 -> 88.0 (reads 90.6 local). 2026-08-22 (N0.6-on-V3): fsm.rs ratcheted 85.6 -> 86.5 (reads 87.1 local) \u2014 the emission-time key-discipline fixes added covered branches. 2026-08-22 (/v1/sessions): session.rs became session/{clock,lease,manager,state} and routes/sessions/ landed; all five new files enter at 97-100% with no debt baseline. Included-total 92.7 -> 93.12. 2026-09-15 (BitNet HTTP serving): /v1/infer and both OpenAI streaming surfaces gained a native-ternary arm for --keep-quant containers (routes/infer.rs, routes/openai/completions.rs, routes/openai/chat/stream.rs). Those ~300 lines are structurally unreachable from the existing test fixtures: every arm is behind LoadedModel::is_bitnet(), which requires a container carrying bitnet_layout plus the bitnet/ I2_S artifacts, and synthetic_vindex builds a dense V2 container. The refusal and guard paths ARE covered (state/loaded_model.rs gained four tests: dense-vindex refusal message, --no-infer no-op, lazy-load contract, is_dense_only detection) -- what is not covered is the ternary forward pass itself, which needs weights. Baselines therefore enter at the values CI measured on the commit that added them, to be ratcheted upward when a synthetic BitNet fixture lands: that fixture is the real fix and is a separate piece of work (it needs packed I2_S bytes plus per-row scales in the kernel's layout, not just a config flag). Included-total is unaffected -- these three files were already in the include set.", "include_globs": [ "crates/larql-server/src/*.rs", "crates/larql-server/src/**/*.rs" @@ -27,15 +27,16 @@ "per_file_line_min_percent": { "crates/larql-server/src/bootstrap/load.rs": 66.0, "crates/larql-server/src/routes/embed.rs": 86.1, + "crates/larql-server/src/routes/infer.rs": 75.0, "crates/larql-server/src/routes/insert.rs": 76.7, + "crates/larql-server/src/routes/openai/chat/stream.rs": 66.0, + "crates/larql-server/src/routes/openai/completions.rs": 79.5, + "crates/larql-server/src/routes/openai/responses/stream.rs": 84.0, "crates/larql-server/src/routes/openai/schema/ast.rs": 80.5, "crates/larql-server/src/routes/openai/schema/fsm.rs": 86.5, "crates/larql-server/src/routes/openai/schema/tools.rs": 88.2, "crates/larql-server/src/routes/topology.rs": 60.6, - "crates/larql-server/src/routes/openai/completions.rs": 88.0, "crates/larql-server/src/routes/walk_ffn/handler.rs": 86.0, - "crates/larql-server/src/routes/warmup.rs": 84.0, - "crates/larql-server/src/routes/openai/chat/stream.rs": 87.0, - "crates/larql-server/src/routes/openai/responses/stream.rs": 84.0 + "crates/larql-server/src/routes/warmup.rs": 84.0 } } diff --git a/crates/larql-server/src/bootstrap/load.rs b/crates/larql-server/src/bootstrap/load.rs index 9d2fc4e3a..d2fe8cf7d 100644 --- a/crates/larql-server/src/bootstrap/load.rs +++ b/crates/larql-server/src/bootstrap/load.rs @@ -486,6 +486,8 @@ pub fn load_single_vindex( release_mmap_after_request: opts.release_mmap_after_request, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels, ffn_l2_cache: crate::ffn_l2_cache::FfnL2Cache::new(num_layers), layer_latency_tracker: std::sync::Arc::new(crate::metrics::LayerLatencyTracker::new()), diff --git a/crates/larql-server/src/bootstrap/mod.rs b/crates/larql-server/src/bootstrap/mod.rs index d823fff1d..b8c7467d5 100644 --- a/crates/larql-server/src/bootstrap/mod.rs +++ b/crates/larql-server/src/bootstrap/mod.rs @@ -195,7 +195,12 @@ pub async fn serve(cli: Cli) -> Result<(), BoxError> { if !cli.no_memcheck && !cli.lazy_weights { let total_estimate: u64 = models .iter() - .filter(|m| !m.infer_disabled) + // BitNet (--keep-quant) vindexes don't allocate dense + // BitLinear tensors at load time — the resident size + // estimator targets the dense path and would massively + // over-count for them. Skip until estimate_resident_bytes + // grows a bitnet-aware branch. + .filter(|m| !m.infer_disabled && !m.is_bitnet()) .map(|m| m.config.estimate_resident_bytes()) .sum(); if total_estimate > 0 { @@ -239,6 +244,26 @@ pub async fn serve(cli: Cli) -> Result<(), BoxError> { continue; } let load_start = std::time::Instant::now(); + // BitNet vindex (--keep-quant) skips the dense load and + // pre-loads the native ternary path instead. Saves ~5 GB + // of dense allocation per model on a 2 B BitNet. + if m.is_bitnet() { + info!("Pre-loading BitNet model for '{}' …", m.id); + if let Err(e) = m.force_load_bitnet_model() { + return Err(format!( + "failed to load bitnet model for '{}': {} \ + (pass --lazy-weights to defer until first request)", + m.id, e + ) + .into()); + } + info!( + " Pre-loaded BitNet model for '{}' in {:.1}s", + m.id, + load_start.elapsed().as_secs_f64(), + ); + continue; + } info!("Pre-loading model weights for '{}' …", m.id); if let Err(e) = m.force_load_weights() { return Err(format!( diff --git a/crates/larql-server/src/routes/infer.rs b/crates/larql-server/src/routes/infer.rs index 673a7c4ab..34572b4ed 100644 --- a/crates/larql-server/src/routes/infer.rs +++ b/crates/larql-server/src/routes/infer.rs @@ -87,6 +87,123 @@ fn run_infer( )); } + // BitNet 1.58 (--keep-quant) vindex: take the native-ternary + // forward path. Skips dense weight loading entirely (~5 GB + // saved on a 2 B BitNet) and runs predict_bitnet against the + // pre-loaded BitnetModel. Walk-mode is supported via + // residual capture + KNN-store override (no sparse FFN — see + // larql_inference::ternary::infer_bitnet_walk for the + // architecture note). + // + // Checked before the `has_model_weights` gate below: a + // `--keep-quant` container carries ternary artifacts rather than + // the dense weight manifest that gate looks for, so a BitNet + // vindex would otherwise be refused as weightless. + if model.is_bitnet() { + let bitnet_guard = model + .get_or_load_bitnet() + .map_err(ServerError::InferenceUnavailable)?; + let bitnet: &larql_inference::ternary::BitnetModel = &bitnet_guard; + + let encoding = model + .tokenizer + .encode(req.prompt.as_str(), true) + .map_err(|e| ServerError::Internal(format!("tokenize error: {e}")))?; + let token_ids: Vec = encoding.get_ids().to_vec(); + if token_ids.is_empty() { + return Err(ServerError::BadRequest("empty prompt".into())); + } + + let start = std::time::Instant::now(); + let (is_compare, mut use_walk, mut use_dense) = infer_mode_flags(&req.mode); + // Dense-only BitNet vindexes (`--dense-only`) have no gate + // vectors / KNN store, so walk-mode would silently return + // nothing useful. Coerce any walk request to dense so + // clients that omit `mode` (which defaults to walk) still + // get correct predictions. Compare-mode also collapses to + // dense-only output here. + if model.is_dense_only() && (use_walk || is_compare) { + use_walk = false; + use_dense = true; + } + let mut result = serde_json::Map::new(); + result.insert("prompt".into(), serde_json::json!(req.prompt)); + + if use_walk { + let run_bitnet_walk = |knn: &larql_vindex::patch::KnnStore| { + larql_inference::ternary::infer_bitnet_walk( + bitnet, + &model.tokenizer, + Some(knn), + &token_ids, + req.top, + ) + }; + // Same lock discipline and session fallback as the dense + // path below: a reader on the sessions map, and a session + // with no overlay reads like the global state. + let walk_pred = if let Some(sid) = session_id { + let sessions = state.sessions.sessions_blocking_read(); + if let Some(patched) = sessions.get(sid).and_then(|s| s.patched()) { + run_bitnet_walk(&patched.knn_store) + } else { + drop(sessions); + let patched = model.patched.blocking_read(); + run_bitnet_walk(&patched.knn_store) + } + } else { + let patched = model.patched.blocking_read(); + run_bitnet_walk(&patched.knn_store) + }; + + let predictions = format_predictions(&walk_pred.predictions); + if let Some(ovr) = &walk_pred.knn_override { + result.insert( + "knn_override".into(), + format_knn_override(ovr, walk_pred.model_top1.as_ref()), + ); + } + if is_compare { + result.insert(INFER_MODE_WALK.into(), serde_json::json!(predictions)); + result.insert( + "walk_ms".into(), + serde_json::json!((walk_pred.walk_ms * 10.0).round() / 10.0), + ); + } else { + result.insert("predictions".into(), serde_json::json!(predictions)); + result.insert("mode".into(), serde_json::json!(INFER_MODE_WALK)); + } + } + + if use_dense { + let dense_start = std::time::Instant::now(); + let pred = larql_inference::ternary::predict_bitnet( + bitnet, + &model.tokenizer, + &token_ids, + req.top, + ); + let dense_ms = dense_start.elapsed().as_secs_f64() * 1000.0; + + let pred_pairs: Vec<(String, f64)> = + pred.into_iter().map(|p| (p.token, p.probability)).collect(); + let predictions = format_predictions(&pred_pairs); + if is_compare { + result.insert(INFER_MODE_DENSE.into(), serde_json::json!(predictions)); + result.insert( + "dense_ms".into(), + serde_json::json!((dense_ms * 10.0).round() / 10.0), + ); + } else { + result.insert("predictions".into(), serde_json::json!(predictions)); + result.insert("mode".into(), serde_json::json!("bitnet")); + } + } + + result.insert("latency_ms".into(), serde_json::json!(elapsed_ms(start))); + return Ok(serde_json::Value::Object(result)); + } + if !model.config.has_model_weights && model.config.extract_level != larql_vindex::ExtractLevel::Inference && model.config.extract_level != larql_vindex::ExtractLevel::All diff --git a/crates/larql-server/src/routes/openai/chat/stream.rs b/crates/larql-server/src/routes/openai/chat/stream.rs index 936127dd5..ec08c8b2c 100644 --- a/crates/larql-server/src/routes/openai/chat/stream.rs +++ b/crates/larql-server/src/routes/openai/chat/stream.rs @@ -43,6 +43,108 @@ pub(super) fn stream_chat_completion( tokio::task::spawn_blocking(move || { let _gen_guard = runtime.clone().enter_generation(); + + // BitNet (--keep-quant) vindexes take the native-ternary + // streaming path: skips the dense weights write-lock and runs + // generate_streaming_bitnet against the pre-loaded + // BitnetModel at ~1.4 GB resident instead of ~5 GB. + // + // Tools and constrained generation are refused rather than + // silently ignored: both need masked logits over the dense + // path, and answering a tool request with prose would look + // like a model that chose not to call the tool. + if model.is_bitnet() { + if tools_active || constrained_schema.is_some() { + let _ = tx.blocking_send(error_chunk( + "tools / constrained generation not supported on BitNet \ + (--keep-quant) models yet", + )); + return; + } + let bitnet_guard = match model.get_or_load_bitnet() { + Ok(g) => g, + Err(e) => { + let _ = tx.blocking_send(error_chunk(&e)); + return; + } + }; + let bitnet: &larql_inference::ternary::BitnetModel = &bitnet_guard; + // `pick_template` needs `&ModelWeights`, which the ternary + // path deliberately never loads. Resolve the template from + // the container's declared family instead — the same string + // `ModelWeights::arch.family()` would have yielded. + let template = larql_inference::prompt::ChatTemplate::for_family(&model.config.family); + let prompt = render(template, &messages); + let encoding = match model.tokenizer.encode(prompt.as_str(), true) { + Ok(e) => e, + Err(e) => { + let _ = tx.blocking_send(error_chunk(&format!("tokenize: {e}"))); + return; + } + }; + let prompt_ids: Vec = encoding.get_ids().to_vec(); + if prompt_ids.is_empty() { + let _ = tx.blocking_send(error_chunk("rendered prompt tokenises to empty")); + return; + } + + // Initial role=assistant chunk — OpenAI contract. + let first = build_chat_chunk(&chat_id, &model_id, Some(ASSISTANT_ROLE), None, None); + if tx.blocking_send(first).is_err() { + return; + } + + let (sampling, eos) = util::build_sampling_eos(sampling_params, &stop_strings); + // Same `TokenTap` the dense path uses, so stop-string + // handling and halt semantics are one implementation + // rather than two that have to agree. + let tap = std::rc::Rc::new(std::cell::RefCell::new(TokenTap::new( + &stop_strings, + EmitFailure::Halt, + ))); + let chat_id_cb = chat_id.clone(); + let model_id_cb = model_id.clone(); + let tx_cb = tx.clone(); + let tap_cb = std::rc::Rc::clone(&tap); + let result = larql_inference::ternary::generate_streaming_bitnet( + bitnet, + &model.tokenizer, + &prompt_ids, + max_tokens, + sampling, + &eos, + move |_id: u32, text: &str, _ms: f64| { + tap_cb.borrow_mut().feed(text, |t| { + let chunk = + build_chat_chunk(&chat_id_cb, &model_id_cb, None, Some(t), None); + tx_cb.blocking_send(chunk).is_ok() + }); + }, + ); + + let emitted = result; + // Record the generation the same way the dense path does, + // or `/v1/stats` would report BitNet traffic as zero + // throughput. `add_v3` takes plain counts, which is all the + // ternary path produces (no GenerateResult); the split + // between prefill and decode is not separately measured + // here, so the whole span is attributed to decode. + let mut tally = crate::runtime_stats::GenerationTally::new(); + let elapsed = crate::state::elapsed_ms(call_started); + tally.add_v3(prompt_ids.len(), emitted, 0.0, elapsed); + runtime.record(tally.into_sample(elapsed)); + + let finish_reason: &'static str = if tap.borrow().halted() || emitted < max_tokens { + FINISH_REASON_STOP + } else { + FINISH_REASON_LENGTH + }; + let final_chunk = + build_chat_chunk(&chat_id, &model_id, None, None, Some(finish_reason)); + let _ = tx.blocking_send(final_chunk); + return; + } + let mut weights_guard = match model.lock_weights_for_gen() { Ok(w) => w, Err(e) => { diff --git a/crates/larql-server/src/routes/openai/completions.rs b/crates/larql-server/src/routes/openai/completions.rs index 6f4755c47..4894410d6 100644 --- a/crates/larql-server/src/routes/openai/completions.rs +++ b/crates/larql-server/src/routes/openai/completions.rs @@ -372,6 +372,92 @@ fn stream_completions( tokio::task::spawn_blocking(move || { let _gen_guard = runtime.clone().enter_generation(); + + // BitNet (--keep-quant) vindexes take the native-ternary + // streaming path: skips the dense weights write-lock (which + // serialises all generation) and runs + // generate_streaming_bitnet against the pre-loaded + // BitnetModel at ~1.4 GB resident instead of ~5 GB. + if model.is_bitnet() { + let bitnet_guard = match model.get_or_load_bitnet() { + Ok(g) => g, + Err(e) => { + let _ = tx.blocking_send(error_chunk(&e)); + return; + } + }; + let bitnet: &larql_inference::ternary::BitnetModel = &bitnet_guard; + let encoding = match model.tokenizer.encode(prompt.as_str(), true) { + Ok(e) => e, + Err(e) => { + let _ = tx.blocking_send(error_chunk(&format!("tokenize: {e}"))); + return; + } + }; + let prompt_ids: Vec = encoding.get_ids().to_vec(); + if prompt_ids.is_empty() { + let _ = tx.blocking_send(error_chunk("prompt tokenises to empty")); + return; + } + + let (sampling, eos) = super::util::build_sampling_eos(sampling_params, &stop_strings); + + let cmpl_id_cb = cmpl_id.clone(); + let model_id_cb = model_id.clone(); + let tx_cb = tx.clone(); + let stop_strings_cb = stop_strings.clone(); + let mut completion_text = String::new(); + let mut early_stop = false; + let mut emitted = 0usize; + let _ = larql_inference::ternary::generate_streaming_bitnet( + bitnet, + &model.tokenizer, + &prompt_ids, + max_tokens, + sampling, + &eos, + |_id, text, _ms| { + if early_stop { + return; + } + let chunk = + build_text_completion_chunk(&cmpl_id_cb, &model_id_cb, Some(text), None); + if tx_cb.blocking_send(chunk).is_err() { + early_stop = true; + return; + } + completion_text.push_str(text); + emitted += 1; + if !stop_strings_cb.is_empty() + && contains_any(&completion_text, &stop_strings_cb) + { + early_stop = true; + } + }, + ); + + // Record the generation the same way the dense path does, + // or `/v1/stats` would report BitNet traffic as zero + // throughput. `add_v3` takes plain counts, which is all the + // ternary path produces (no GenerateResult); the split + // between prefill and decode is not separately measured + // here, so the whole span is attributed to decode. + let mut tally = crate::runtime_stats::GenerationTally::new(); + let elapsed = crate::state::elapsed_ms(call_started); + tally.add_v3(prompt_ids.len(), emitted, 0.0, elapsed); + runtime.record(tally.into_sample(elapsed)); + + let finish_reason: &'static str = if early_stop || emitted < max_tokens { + FINISH_REASON_STOP + } else { + FINISH_REASON_LENGTH + }; + let final_chunk = + build_text_completion_chunk(&cmpl_id, &model_id, None, Some(finish_reason)); + let _ = tx.blocking_send(final_chunk); + return; + } + let mut weights_guard = match model.lock_weights_for_gen() { Ok(w) => w, Err(e) => { diff --git a/crates/larql-server/src/routes/runtime_lifecycle.rs b/crates/larql-server/src/routes/runtime_lifecycle.rs index 06074bd40..00ff599e6 100644 --- a/crates/larql-server/src/routes/runtime_lifecycle.rs +++ b/crates/larql-server/src/routes/runtime_lifecycle.rs @@ -363,6 +363,8 @@ mod tests { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: std::collections::HashMap::new(), ffn_l2_cache: crate::ffn_l2_cache::FfnL2Cache::new(1), layer_latency_tracker: Arc::new(crate::metrics::LayerLatencyTracker::new()), diff --git a/crates/larql-server/src/routes/stream.rs b/crates/larql-server/src/routes/stream.rs index 7b8da06ab..0e1e45cb4 100644 --- a/crates/larql-server/src/routes/stream.rs +++ b/crates/larql-server/src/routes/stream.rs @@ -705,6 +705,8 @@ mod tests { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: labels, ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new(crate::metrics::LayerLatencyTracker::new()), diff --git a/crates/larql-server/src/state/loaded_model.rs b/crates/larql-server/src/state/loaded_model.rs index d8d8f8905..e4386f909 100644 --- a/crates/larql-server/src/state/loaded_model.rs +++ b/crates/larql-server/src/state/loaded_model.rs @@ -74,6 +74,18 @@ pub struct LoadedModel { /// once `weights` is populated, callers skip the mutex via the /// fast-path `OnceLock::get` check. pub weights_init: std::sync::Mutex<()>, + /// BitNet 1.58 model with native ternary weights. Populated + /// when the loaded vindex was built with `--keep-quant` + /// (i.e. `config.bitnet_layout.is_some()`). When present, the + /// route handlers prefer this over `weights` for inference + /// because the native-ternary path runs the full forward at + /// ~1.4 GB instead of ~5 GB resident. Eager-loaded by + /// `force_load_bitnet_model` from `bootstrap::serve` (unless + /// `--lazy-weights`). + pub bitnet_model: std::sync::OnceLock>, + /// Init guard for the bitnet model load — same pattern as + /// `weights_init` but for the ternary path. + pub bitnet_init: std::sync::Mutex<()>, /// Probe-confirmed feature labels: (layer, feature) → relation name. /// Loaded from feature_labels.json if present. pub probe_labels: HashMap<(usize, usize), String>, @@ -171,6 +183,71 @@ impl LoadedModel { self.ensure_weights_cell().map(|_| ()) } + /// Whether this vindex was built with `--keep-quant` and + /// therefore has the BitNet 1.58 native-ternary artifacts + /// (`bitnet/` + `bitnet_layout` in index.json). Route handlers + /// dispatch on this to pick the ternary forward path. + pub fn is_bitnet(&self) -> bool { + self.config.bitnet_layout.is_some() + } + + /// Whether this vindex was built `--dense-only`: it has the + /// dense weights + BitNet I2_S artifacts but NO gate vectors / + /// HNSW clustering, so walk-mode inference cannot run against + /// it (the KNN store is empty). Detected by an empty gate-layer + /// list in index.json (`build_vindex_dense_only` leaves + /// `layer_infos` empty). Route handlers force dense-mode + /// inference on such vindexes regardless of the requested mode, + /// since walk would silently return nothing useful. + pub fn is_dense_only(&self) -> bool { + self.config.layers.is_empty() + } + + /// Get a read guard on the lazy-loaded BitNet model. Returns + /// `Err` when the vindex isn't a BitNet (callers should check + /// `is_bitnet()` first). + pub fn get_or_load_bitnet( + &self, + ) -> Result, String> { + let cell = self.ensure_bitnet_cell()?; + cell.read() + .map_err(|e| format!("bitnet RwLock poisoned: {e}")) + } + + /// Eager-load the BitNet model from disk before the listener + /// binds. Mirrors `force_load_weights` but for the ternary + /// path; called by `bootstrap::serve` when the vindex is + /// BitNet-shaped and `--lazy-weights` was not passed. + pub fn force_load_bitnet_model(&self) -> Result<(), String> { + if self.infer_disabled || !self.is_bitnet() { + return Ok(()); + } + self.ensure_bitnet_cell().map(|_| ()) + } + + fn ensure_bitnet_cell( + &self, + ) -> Result<&std::sync::RwLock, String> { + // Fast path. + if let Some(cell) = self.bitnet_model.get() { + return Ok(cell); + } + // Single-flight slow path. + let _init_guard = self.bitnet_init.lock().unwrap_or_else(|p| p.into_inner()); + if let Some(cell) = self.bitnet_model.get() { + return Ok(cell); + } + if !self.is_bitnet() { + return Err("vindex has no bitnet_layout (not a --keep-quant build)".into()); + } + let model = larql_inference::ternary::load_bitnet_model(&self.path) + .map_err(|e| format!("failed to load bitnet model: {e}"))?; + let _ = self.bitnet_model.set(std::sync::RwLock::new(model)); + self.bitnet_model + .get() + .ok_or_else(|| "bitnet cell unset after set".to_string()) + } + /// Acquire an exclusive write guard on the loaded weights. /// /// Used by the OpenAI generation path (`/v1/completions`, @@ -184,6 +261,29 @@ impl LoadedModel { pub fn lock_weights_for_gen( &self, ) -> Result, String> { + // A BitNet `--keep-quant` container has no dense weight manifest to + // load, so `ensure_weights_cell` would fail here with a bare + // "No such file or directory" from whichever tensor file it reached + // first. Every non-streaming generation path funnels through this + // one method (`openai/completions.rs` batch loop, + // `openai/chat/handler.rs`, `openai/responses/engine.rs`), so + // naming the real reason once here covers all of them rather than + // three separate checks that have to stay in agreement. + // + // Refused rather than silently routed to the ternary path: these + // callers hold a `&mut ModelWeights` for the whole generation, and + // there is no dense `ModelWeights` to hand them. The ternary + // engine is reachable through `/v1/infer` and the streaming + // surfaces, which do not need one. + if self.is_bitnet() { + return Err( + "this vindex is a BitNet --keep-quant build and carries no dense \ + weights; non-streaming generation is not supported on it. Use \ + POST /v1/infer, or /v1/completions and /v1/chat/completions \ + with \"stream\": true, which take the native-ternary path." + .to_string(), + ); + } let cell = self.ensure_weights_cell()?; cell.write() .map_err(|e| format!("weights RwLock poisoned: {e}")) @@ -343,6 +443,8 @@ mod loaded_model_tests { release_mmap_after_request: release_mmap, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: HashMap::new(), ffn_l2_cache: crate::ffn_l2_cache::FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new(crate::metrics::LayerLatencyTracker::new()), @@ -393,6 +495,183 @@ mod loaded_model_tests { ); } + #[test] + fn is_dense_only_detects_empty_gate_layers() { + // A normal vindex has gate layers -> not dense-only. + let normal = tiny_loaded_model(QuantFormat::None, false); + assert!( + !normal.is_dense_only(), + "vindex with gate layers must not be dense-only" + ); + assert!( + !normal.is_bitnet(), + "and a plain vindex carries no bitnet_layout" + ); + + // A --dense-only BitNet vindex has zero gate layers. Build + // one by emptying the layer list + setting bitnet_layout. + let mut cfg = tiny_config(QuantFormat::None); + cfg.layers = Vec::new(); + cfg.bitnet_layout = Some(larql_vindex::config::BitnetLayout::default()); + let mut dense_only = tiny_loaded_model(QuantFormat::None, false); + dense_only.config = cfg; + assert!( + dense_only.is_dense_only(), + "dense-only vindex (empty gate layers) must be detected" + ); + assert!(dense_only.is_bitnet(), "and it is a BitNet vindex"); + } + + #[test] + fn bitnet_guards_refuse_a_dense_vindex_with_a_useful_message() { + // `ensure_bitnet_cell`'s refusal path: asking a non-BitNet vindex + // for a ternary model must name *why* rather than surfacing a + // load error from a file that was never going to exist. + let model = tiny_loaded_model(QuantFormat::None, false); + // `BitnetModel` is not `Debug`, so match rather than `expect_err`. + let Err(err) = model.get_or_load_bitnet() else { + unreachable!("a dense vindex has no ternary model to hand out") + }; + assert!( + err.contains("bitnet_layout") && err.contains("keep-quant"), + "the error must say the container is not a --keep-quant build, \ + got: {err}" + ); + } + + #[test] + fn force_load_bitnet_model_is_a_noop_when_infer_disabled() { + // `bootstrap::serve` calls this unconditionally for every model, + // so it has to stay quiet on a --no-infer server even when the + // container *is* BitNet-shaped: eagerly loading ternary weights + // into a process that refuses to infer would spend the memory a + // --no-infer operator asked not to spend. + let mut cfg = tiny_config(QuantFormat::None); + cfg.bitnet_layout = Some(larql_vindex::config::BitnetLayout::default()); + let mut model = tiny_loaded_model(QuantFormat::None, false); + model.config = cfg; + model.infer_disabled = true; + assert!(model.is_bitnet(), "fixture must be BitNet-shaped"); + assert!( + model.force_load_bitnet_model().is_ok(), + "must no-op rather than error under --no-infer" + ); + assert!( + model.bitnet_model.get().is_none(), + "and must not have loaded anything" + ); + } + + #[test] + fn bitnet_load_failure_names_the_container() { + // A container that *claims* to be BitNet (bitnet_layout present) + // but has no `bitnet/` artifacts on disk must fail with the load + // error, not the "not a --keep-quant build" refusal: the two are + // different operator problems. The first says "this vindex is the + // wrong kind", the second says "this vindex is the right kind and + // is broken/incomplete", and reporting the wrong one sends the + // operator to rebuild a container that only needs its files back. + // + // Reachable without any weights: the fixture's path points at no + // bitnet/ directory, which is exactly the on-disk state of a + // truncated or partially-copied container. + let mut cfg = tiny_config(QuantFormat::None); + cfg.bitnet_layout = Some(larql_vindex::config::BitnetLayout::default()); + let mut model = tiny_loaded_model(QuantFormat::None, false); + model.config = cfg; + assert!(model.is_bitnet(), "fixture must be BitNet-shaped"); + + let Err(err) = model.get_or_load_bitnet() else { + unreachable!("there are no bitnet/ artifacts to load") + }; + assert!( + err.contains("failed to load bitnet model"), + "a BitNet-shaped container with missing artifacts must report a \ + load failure, not the wrong-kind refusal, got: {err}" + ); + assert!( + !err.contains("not a --keep-quant build"), + "must not claim the container is the wrong kind: {err}" + ); + // A failed load must leave the cell empty so a later attempt (after + // the operator restores the files) still tries, rather than caching + // the failure for the process lifetime. + assert!( + model.bitnet_model.get().is_none(), + "a failed load must not poison the cell" + ); + } + + #[test] + fn lock_weights_for_gen_refuses_bitnet_with_an_actionable_message() { + // Regression: on a real --keep-quant container the three + // non-streaming generation paths (openai completions batch loop, + // chat handler, responses engine) all reached + // `ensure_weights_cell` and surfaced a bare "No such file or + // directory" as a 503 -- there is no dense weight manifest in such + // a container. Caught only against the real + // microsoft/bitnet-b1.58-2B-4T model, because the synthetic + // fixture is a dense V2 container that has those files. + // + // The message has to say what to use instead: the ternary engine + // *is* reachable, just not through a path that needs + // `&mut ModelWeights`. + let mut cfg = tiny_config(QuantFormat::None); + cfg.bitnet_layout = Some(larql_vindex::config::BitnetLayout::default()); + let mut model = tiny_loaded_model(QuantFormat::None, false); + model.config = cfg; + assert!(model.is_bitnet(), "fixture must be BitNet-shaped"); + + let Err(err) = model.lock_weights_for_gen() else { + unreachable!("a --keep-quant container has no dense weights to lock") + }; + assert!( + err.contains("keep-quant") && err.contains("no dense"), + "must name the container kind as the reason, got: {err}" + ); + assert!( + err.contains("/v1/infer") && err.contains("stream"), + "must point at the paths that do work, got: {err}" + ); + + // And the dense case must be unaffected: a plain container still + // reaches the loader (and fails on the missing fixture files, not + // on this guard). + let dense = tiny_loaded_model(QuantFormat::None, false); + let Err(dense_err) = dense.lock_weights_for_gen() else { + unreachable!("the tiny fixture has no weight files on disk") + }; + assert!( + !dense_err.contains("keep-quant"), + "a dense container must not hit the BitNet guard: {dense_err}" + ); + } + + #[test] + fn bitnet_model_not_loaded_by_default() { + // Same lazy-load contract as `weights`: the ternary cell stays + // empty until `get_or_load_bitnet`, and `force_load_bitnet_model` + // is a no-op on a vindex that is not BitNet-shaped (rather than + // an error), so `bootstrap::serve` can call it unconditionally. + let model = tiny_loaded_model(QuantFormat::None, false); + assert!( + model.bitnet_model.get().is_none(), + "bitnet cell must start empty" + ); + assert!( + model.force_load_bitnet_model().is_ok(), + "force_load_bitnet_model must no-op on a non-BitNet vindex" + ); + assert!( + model.bitnet_model.get().is_none(), + "and must not populate the cell" + ); + assert!( + model.get_or_load_bitnet().is_err(), + "explicitly asking for a bitnet model on a dense vindex is an error" + ); + } + #[test] fn weights_not_loaded_by_default() { // Lazy-load contract: `weights` is `OnceLock::new()` until the diff --git a/crates/larql-server/src/state/model_set.rs b/crates/larql-server/src/state/model_set.rs index c74d54162..34888156c 100644 --- a/crates/larql-server/src/state/model_set.rs +++ b/crates/larql-server/src/state/model_set.rs @@ -290,6 +290,8 @@ mod model_set_tests { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: HashMap::new(), ffn_l2_cache: crate::ffn_l2_cache::FfnL2Cache::new(1), layer_latency_tracker: Arc::new(crate::metrics::LayerLatencyTracker::new()), diff --git a/crates/larql-server/tests/common/mod.rs b/crates/larql-server/tests/common/mod.rs index 2e1005ed9..31d760274 100644 --- a/crates/larql-server/tests/common/mod.rs +++ b/crates/larql-server/tests/common/mod.rs @@ -146,6 +146,8 @@ pub fn model_functional(id: &str) -> Arc { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: std::collections::HashMap::new(), ffn_l2_cache: larql_server::ffn_l2_cache::FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( @@ -189,6 +191,8 @@ pub fn model_infer_enabled(id: &str) -> Arc { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: std::collections::HashMap::new(), ffn_l2_cache: larql_server::ffn_l2_cache::FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( @@ -273,6 +277,8 @@ impl ModelBuilder { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: self.probe_labels, ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( @@ -361,6 +367,8 @@ pub fn model_with_real_weights_and_labels( release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels, ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( @@ -441,6 +449,8 @@ pub fn model_with_q4k_weights( release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: HashMap::new(), ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( diff --git a/crates/larql-server/tests/test_expert_endpoint.rs b/crates/larql-server/tests/test_expert_endpoint.rs index cb39335c0..4032b96fc 100644 --- a/crates/larql-server/tests/test_expert_endpoint.rs +++ b/crates/larql-server/tests/test_expert_endpoint.rs @@ -364,6 +364,8 @@ fn make_loaded_model( release_mmap_after_request: false, weights: lock, weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: HashMap::new(), ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( diff --git a/crates/larql-server/tests/test_http_full_routes.rs b/crates/larql-server/tests/test_http_full_routes.rs index 0c357b380..0065f77d4 100644 --- a/crates/larql-server/tests/test_http_full_routes.rs +++ b/crates/larql-server/tests/test_http_full_routes.rs @@ -45,6 +45,8 @@ fn model_functional_with_labels(id: &str) -> Arc { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: labels, ffn_l2_cache: larql_server::ffn_l2_cache::FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( diff --git a/crates/larql-server/tests/test_http_shard.rs b/crates/larql-server/tests/test_http_shard.rs index bcb1c23aa..5e6e0ed95 100644 --- a/crates/larql-server/tests/test_http_shard.rs +++ b/crates/larql-server/tests/test_http_shard.rs @@ -38,6 +38,8 @@ fn model_with_path(id: &str, path: PathBuf) -> Arc { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: std::collections::HashMap::new(), ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( diff --git a/crates/larql-server/tests/test_unit_band_utils.rs b/crates/larql-server/tests/test_unit_band_utils.rs index 295e6a255..e81b3048a 100644 --- a/crates/larql-server/tests/test_unit_band_utils.rs +++ b/crates/larql-server/tests/test_unit_band_utils.rs @@ -164,6 +164,8 @@ fn make_minimal_model(layer_bands: Option) -> Arc { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: HashMap::new(), ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( diff --git a/crates/larql-server/tests/test_unit_state.rs b/crates/larql-server/tests/test_unit_state.rs index f0f09378c..48bd45e9b 100644 --- a/crates/larql-server/tests/test_unit_state.rs +++ b/crates/larql-server/tests/test_unit_state.rs @@ -94,6 +94,8 @@ fn make_tiny_model(id: &str) -> Arc { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: HashMap::new(), ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( @@ -196,6 +198,8 @@ fn make_loaded_model_for_warmup() -> Arc { release_mmap_after_request: false, weights: std::sync::OnceLock::new(), weights_init: std::sync::Mutex::new(()), + bitnet_model: std::sync::OnceLock::new(), + bitnet_init: std::sync::Mutex::new(()), probe_labels: HashMap::new(), ffn_l2_cache: FfnL2Cache::new(1), layer_latency_tracker: std::sync::Arc::new( diff --git a/crates/larql-vindex/src/index/compute/gate_knn/scores_batch.rs b/crates/larql-vindex/src/index/compute/gate_knn/scores_batch.rs index 717e5efd9..561231aa0 100644 --- a/crates/larql-vindex/src/index/compute/gate_knn/scores_batch.rs +++ b/crates/larql-vindex/src/index/compute/gate_knn/scores_batch.rs @@ -185,22 +185,29 @@ impl VectorIndex { return None; } let mmap: &[u8] = view.bytes.as_ref(); - let mut cache = self.gate.f16_decode_cache.lock().unwrap(); - if cache.len() <= layer { - cache.resize(layer + 1, None); - } - let miss = cache[layer].is_none(); - if miss { - let byte_offset = view.slice.float_offset * 2; - let byte_end = byte_offset + view.slice.num_features * self.hidden_size * 2; - if byte_end > mmap.len() { - return None; + // Take an `Arc` handle and release the lock before the matmul, + // for the same reason as `gate_knn_mmap_fast`: holding it across + // the multiply serialises concurrent callers on unrelated layers. + let data = { + let mut cache = self.gate.f16_decode_cache.lock().unwrap(); + if cache.len() <= layer { + cache.resize(layer + 1, None); } - let raw = &mmap[byte_offset..byte_end]; - cache[layer] = Some(larql_models::quant::half::decode_f16(raw)); - } - self.touch_gate_cache_lru(layer, miss, &mut cache); - let data = cache[layer].as_ref().unwrap(); + let miss = cache[layer].is_none(); + if miss { + let byte_offset = view.slice.float_offset * 2; + let byte_end = byte_offset + view.slice.num_features * self.hidden_size * 2; + if byte_end > mmap.len() { + return None; + } + let raw = &mmap[byte_offset..byte_end]; + cache[layer] = Some(std::sync::Arc::new( + larql_models::quant::half::decode_f16(raw), + )); + } + self.touch_gate_cache_lru(layer, miss, &mut cache); + std::sync::Arc::clone(cache[layer].as_ref().unwrap()) + }; let arr = ArrayView2::from_shape( (view.slice.num_features, self.hidden_size), data.as_slice(), diff --git a/crates/larql-vindex/src/index/core/mod.rs b/crates/larql-vindex/src/index/core/mod.rs index 4e8c79c66..37f9c0f51 100644 --- a/crates/larql-vindex/src/index/core/mod.rs +++ b/crates/larql-vindex/src/index/core/mod.rs @@ -365,7 +365,7 @@ mod refactor_tests { { let mut cache = v.gate.f16_decode_cache.lock().unwrap(); - cache[1] = Some(vec![1.0, 2.0, 3.0]); + cache[1] = Some(std::sync::Arc::new(vec![1.0, 2.0, 3.0])); } { let mut warm = v.gate.warmed_gates.write().unwrap(); diff --git a/crates/larql-vindex/src/index/storage/gate_store.rs b/crates/larql-vindex/src/index/storage/gate_store.rs index 5f3c5e0d8..d31225c38 100644 --- a/crates/larql-vindex/src/index/storage/gate_store.rs +++ b/crates/larql-vindex/src/index/storage/gate_store.rs @@ -36,7 +36,14 @@ pub struct GateStore { /// Per-layer gate vectors (heap mode). pub gate_vectors: Vec>>, /// Lazy decode cache for f16 gate vectors. - pub f16_decode_cache: Mutex>>>, + /// + /// `Arc` per layer so a reader can take a cheap handle and release the + /// mutex *before* scoring. Holding the lock across `gemv` would + /// serialise `PatchedVindex::walk`'s rayon-parallel layers against each + /// other, which is the whole point of parallelising them; cloning the + /// data instead would reintroduce the ~71 MB/layer copy this cache + /// exists to avoid. An `Arc` clone is a refcount bump. + pub f16_decode_cache: Mutex>>>>, /// LRU queue for `f16_decode_cache`. Back is oldest, front is newest. pub gate_cache_lru: Mutex>, /// Cap on live entries in `f16_decode_cache`. 0 = unlimited. @@ -110,11 +117,29 @@ impl Clone for GateStore { /// Matrix-vector multiply: view[N, hidden] × vec[hidden] → scores[N]. /// All compute goes through larql-compute. pub(crate) fn gemv(view: &ArrayView2, vec: &Array1) -> Array1 { - let hidden = vec.len(); - let x = vec.view().into_shape_with_order((1, hidden)).unwrap(); - let cpu = larql_compute::CpuBackend; - let result = cpu.matmul_transb(x, *view); - Array1::from_vec(result.into_raw_vec_and_offset().0) + // `gate[N, hidden] . vec[hidden] -> [N]`, expressed as a real + // matrix-vector product rather than a 1-row matmul against a + // transpose. + // + // The previous form was `matmul_transb(vec.into_shape((1, hidden)), + // gate)`, i.e. `a.dot(&b.t())` with `a` shaped [1, hidden]. ndarray + // only dispatches to BLAS when the operand layouts qualify, and that + // shape does not reach `sgemv`: it fell to ndarray's own path. + // Measured on one layer of a real BitNet 2B browse vindex + // (6912 features x 2560 dims, f32 cached): + // + // a.dot(&b.t()) 15.5 ms/layer 4.6 GB/s + // gate.dot(&vec) 3.3 ms/layer 21.4 GB/s <- this + // manual row dot 9.8 ms/layer 7.2 GB/s + // + // 4.7x, from removing the transpose. `describe()` scans 12-30 layers + // per call, so this is the dominant term in its latency. + // + // `Array2::dot(&Array1)` is ndarray's gemv entry point and hits + // `cblas_sgemv` for f32 with a standard-layout operand, which the + // gate view is (contiguous rows straight out of the mmap or the f16 + // decode cache). + view.dot(vec) } /// Gate scores batch: gate[N, hidden] × x[seq, hidden]^T → [N, seq]. @@ -205,7 +230,7 @@ impl VectorIndex { &self, layer: usize, just_inserted: bool, - cache: &mut [Option>], + cache: &mut [Option>>], ) { let max = self .gate @@ -286,6 +311,10 @@ impl VectorIndex { } } crate::config::dtype::StorageDtype::F16 => { + // `GateData` owns its buffer, so this arm still copies. + // It is the slow path: `gate_knn_mmap_fast` handles f16 + // without copying and is what the scan actually takes. + // Reached only by callers that need an owned matrix. let mut cache = self.gate.f16_decode_cache.lock().unwrap(); if cache.len() <= layer { cache.resize(layer + 1, None); @@ -293,10 +322,12 @@ impl VectorIndex { let miss = cache[layer].is_none(); if miss { let raw = &mmap[byte_offset..byte_end]; - cache[layer] = Some(larql_models::quant::half::decode_f16(raw)); + cache[layer] = Some(std::sync::Arc::new( + larql_models::quant::half::decode_f16(raw), + )); } self.touch_gate_cache_lru(layer, miss, &mut cache); - cache[layer].as_ref().unwrap().clone() + cache[layer].as_ref().unwrap().as_ref().clone() } }; return Some(GateData { @@ -308,9 +339,21 @@ impl VectorIndex { None } - /// Zero-copy gate KNN scoring for the f32 mmap path — no - /// allocation, no clone. Returns `None` if not on the f32 mmap - /// path; caller falls back to `resolve_gate`. + /// Zero-copy gate KNN scoring for the mmap path — no allocation of a + /// gate copy. Returns `None` if the layer cannot be scored here; + /// caller falls back to `resolve_gate`. + /// + /// Handles f32 (direct reinterpret of the mmap) *and* f16 (score out of + /// the decode cache in place). The f16 arm is the load-bearing one: + /// without it every f16 layer fell through to `resolve_gate`, which + /// ends in `cache[layer].as_ref().unwrap().clone()` — a full f32 copy + /// of the layer's gate matrix on **every query**. At 6912 features × + /// 2560 dims that is ~71 MB cloned per layer, so a 20-layer + /// `describe()` spent ~1.4 GB on allocation and memcpy before scoring + /// a single feature, and another ~1.4 GB reading it back in `gemv`. + /// Measured effect on a real 2B browse container: `describe()` 288 ms + /// → see `bench_graph.csv`. The decode itself was already cached; it + /// was purely the copy. pub(crate) fn gate_knn_mmap_fast( &self, layer: usize, @@ -356,6 +399,54 @@ impl VectorIndex { } } + // f16 mmap: score out of the decode cache without copying it. + // + // Decoding on a miss is unavoidable (it is what the cache is for), + // but on a hit the previous path cloned the whole layer purely to + // hand an owned `Vec` back to the caller. `gemv` only needs a view, + // so take one over the cached buffer while the lock is held. + // + // The lock is released before `gemv`: the cache holds an `Arc` per + // layer, so a reader takes a refcount bump and scores outside the + // critical section. Holding it across `gemv` would serialise + // `PatchedVindex::walk`'s rayon-parallel layers, and cloning the + // buffer to release early would reintroduce the ~71 MB/layer copy + // this path exists to remove. + if self.storage.gate_dtype() == crate::config::dtype::StorageDtype::F16 { + let view = self.storage.gate_layer_view(layer)?; + if view.slice.num_features == 0 { + return None; + } + let bpf = 2; + let byte_offset = view.slice.float_offset * bpf; + let byte_end = byte_offset + view.slice.num_features * self.hidden_size * bpf; + let mmap: &[u8] = view.bytes.as_ref(); + if byte_end > mmap.len() { + return None; + } + + let data = { + let mut cache = self.gate.f16_decode_cache.lock().unwrap(); + if cache.len() <= layer { + cache.resize(layer + 1, None); + } + let miss = cache[layer].is_none(); + if miss { + cache[layer] = Some(std::sync::Arc::new( + larql_models::quant::half::decode_f16(&mmap[byte_offset..byte_end]), + )); + } + self.touch_gate_cache_lru(layer, miss, &mut cache); + std::sync::Arc::clone(cache[layer].as_ref()?) + }; + let arr = ArrayView2::from_shape( + (view.slice.num_features, self.hidden_size), + data.as_slice(), + ) + .ok()?; + return Some(gemv(&arr, residual)); + } + None } } @@ -519,4 +610,78 @@ mod gate_cache_lru_tests { idx.set_gate_cache_max_layers(0); assert_eq!(resident_layers(&idx), 2); } + #[test] + fn f16_fast_path_scores_without_cloning_the_layer() { + // The f16 arm of `gate_knn_mmap_fast` is what the gate scan takes on + // an f16 container. Before it existed, f16 layers fell through to + // `resolve_gate`, which clones the whole decoded layer per query. + // + // Asserts the arm is reached and correct: the fixture's gate matrix + // is a scaled identity, so a query that is 1.0 in every dim scores + // every feature at 1.0, and feature 0 is among the top hits. + let idx = f16_mmap_index(2, 4, 4); + let q = Array1::from_vec(vec![1.0f32; 4]); + + let scores = idx + .gate_knn_mmap_fast(0, &q) + .expect("f16 mmap layers must be scored by the fast path, not resolve_gate"); + assert_eq!(scores.len(), 4, "one score per feature"); + for (i, s) in scores.iter().enumerate() { + assert!( + (s - 1.0).abs() < 1e-3, + "feature {i} scored {s}, expected ~1.0 from the identity fixture" + ); + } + + // Scoring populated the decode cache (the buffer the Arc points at), + // so a second call is a cache hit and must agree exactly. + assert_eq!(resident_layers(&idx), 1, "scoring must populate the cache"); + let again = idx.gate_knn_mmap_fast(0, &q).expect("cache hit"); + assert_eq!(scores, again, "a cache hit must not change the scores"); + } + + #[test] + fn f16_fast_path_agrees_with_the_resolve_gate_slow_path() { + // Two routes to the same numbers: the fast path scores in place out + // of the Arc'd cache, `resolve_gate` hands back an owned copy that + // the caller multiplies itself. They must not disagree -- that would + // mean the optimisation changed answers, which is the failure mode + // worth a test rather than the speed. + let idx = f16_mmap_index(1, 6, 4); + let q = Array1::from_vec(vec![0.5f32, 0.25, 0.125, 1.0]); + + let fast = idx.gate_knn_mmap_fast(0, &q).expect("fast path"); + let gate = idx.resolve_gate(0).expect("slow path"); + let view = gate.view(idx.hidden_size); + let slow = super::gemv(&view, &q); + + assert_eq!(fast.len(), slow.len()); + for (i, (f, s)) in fast.iter().zip(slow.iter()).enumerate() { + assert!( + (f - s).abs() < 1e-6, + "feature {i}: fast={f} slow={s} -- the paths disagree" + ); + } + } + + #[test] + fn f16_cache_hands_out_arc_handles_not_copies() { + // The cache holds `Arc>` specifically so a reader can take a + // handle and release the mutex before scoring; holding it across the + // multiply would serialise `PatchedVindex::walk`'s parallel layers. + // Two handles to the same layer must therefore be the same + // allocation, not two copies of it. + let idx = f16_mmap_index(1, 4, 4); + touch(&idx, 0); + + let (a, b) = { + let cache = idx.gate.f16_decode_cache.lock().unwrap(); + let entry = cache[0].as_ref().expect("layer 0 cached"); + (std::sync::Arc::clone(entry), std::sync::Arc::clone(entry)) + }; + assert!( + std::sync::Arc::ptr_eq(&a, &b), + "cache handles must alias one buffer, not clone it" + ); + } } diff --git a/crates/larql-vindex/src/patch/overlay.rs b/crates/larql-vindex/src/patch/overlay.rs index cf7564430..98e1be5a0 100644 --- a/crates/larql-vindex/src/patch/overlay.rs +++ b/crates/larql-vindex/src/patch/overlay.rs @@ -478,19 +478,40 @@ impl PatchedVindex { } /// Walk with patch overrides. + /// Walk the given layers, collecting top-K gate hits per layer. + /// + /// Parallel across layers: each layer's `gate_knn` is an independent + /// gemv over that layer's gate matrix, so there is no ordering or data + /// dependency between them. `describe()` scans 12--30 layers, which is + /// enough work per item to cover rayon's dispatch overhead. + /// + /// What this does and does not buy, measured rather than assumed: + /// it cuts **single-query latency** by spreading one query's layers + /// across cores. It does **not** raise saturated throughput, because + /// the total bytes read per query are unchanged and the scan is + /// memory-bound, not compute-bound (15% of 32 vCPUs under 8 concurrent + /// requests -- threads stalled on memory, not busy). Under concurrent + /// load the queries already fill the cores; this helps the single-query + /// case and the lightly-loaded case. + /// + /// `map` on a parallel iterator preserves input order, so the returned + /// trace is still layer-ordered. pub fn walk(&self, residual: &Array1, layers: &[usize], top_k: usize) -> WalkTrace { - let mut trace_layers = Vec::with_capacity(layers.len()); - for &layer in layers { - let hits = self.gate_knn(layer, residual, top_k); - let walk_hits: Vec = hits - .into_iter() - .filter_map(|(feature, gate_score)| { - let meta = self.feature_meta(layer, feature)?.clone(); - Some(WalkHit::from_gate(layer, feature, gate_score, meta)) - }) - .collect(); - trace_layers.push((layer, walk_hits)); - } + use rayon::prelude::*; + let trace_layers: Vec<(usize, Vec)> = layers + .par_iter() + .map(|&layer| { + let hits = self.gate_knn(layer, residual, top_k); + let walk_hits: Vec = hits + .into_iter() + .filter_map(|(feature, gate_score)| { + let meta = self.feature_meta(layer, feature)?.clone(); + Some(WalkHit::from_gate(layer, feature, gate_score, meta)) + }) + .collect(); + (layer, walk_hits) + }) + .collect(); WalkTrace { layers: trace_layers, } diff --git a/crates/larql-vindex/tests/ingestion_record_sites.json b/crates/larql-vindex/tests/ingestion_record_sites.json index 1c052538a..3546ef4b8 100644 --- a/crates/larql-vindex/tests/ingestion_record_sites.json +++ b/crates/larql-vindex/tests/ingestion_record_sites.json @@ -224,6 +224,11 @@ "stream_chat_completion", "record" ], + [ + "larql-server/src/routes/openai/chat/stream.rs", + "stream_chat_completion", + "record" + ], [ "larql-server/src/routes/openai/chat/v3.rs", "respond", @@ -244,6 +249,11 @@ "stream_completions", "record" ], + [ + "larql-server/src/routes/openai/completions.rs", + "stream_completions", + "record" + ], [ "larql-server/src/routes/openai/responses/handler.rs", "handle_responses",