qualia_core_db/wgsl_forge/calibration/
mod.rs

1//! W10 — Forge calibration pipeline (the training-related upgrade).
2//!
3//! This is an **upgrade of the existing forge**, not a new forge: the forge already PRODUCES +
4//! CERTIFIES artifacts (it certifies GPU kernels against a CPU oracle and transcodes GGUF→p64).
5//! Calibration/adaptation artifacts — AWQ activation scales, int8-KV scales, sparse KV
6//! dictionaries — are the same produce-and-certify pattern applied to a new artifact class, so they
7//! live here as a `calibration` concern beside the kernel/transcode entry points. The engine still
8//! only RUNS certified artifacts.
9//!
10//! Pipeline (5 stages): **corpus → capture → learn → certify → package.**
11//! - **corpus** ([`corpus`]) — assemble/expand the calibration text. *Local Ollama is a legitimate
12//!   resource HERE* (offline domain-diverse synthesis), strictly forge-side — it never enters the
13//!   inference runtime (CLAUDE.md §1 holds).
14//! - **capture** — run OUR engine over the corpus with instrumentation on (reuses `llm_awq`'s
15//!   activation hooks; KV capture arrives with W5a). This CANNOT come from Ollama — the artifacts
16//!   compress our engine's own tensors (GQA layout, RoPE convention, layer shapes are engine-specific).
17//! - **learn** — fit the artifact (AWQ scale fold reuses the existing AWQ pipeline; int8-KV scales =
18//!   W5a; dictionary/Top-K SAE = W5b).
19//! - **certify** ([`certify`]) — the ΔPPL ≤ gate via the existing [`perplexity_eval_blocking`] oracle.
20//! - **package** ([`package`]) — certified artifact + provenance (corpus hash, engine version, gate
21//!   numbers) as a CBOR-framed sidecar, so the engine can refuse uncertified artifacts.
22//!
23//! Native-only: the pipeline drives the real inference stack (GGUF, GPU, PPL oracle).
24
25#![cfg(not(target_arch = "wasm32"))]
26
27pub mod corpus;
28pub mod kv_dictionary;
29pub mod package;
30
31/// Re-export the runtime KV-dictionary install/reconstruct — it now lives in core
32/// ([`crate::kv_dict_runtime`]) so the engine can run a certified artifact without the forge feature.
33pub use crate::kv_dict_runtime;
34
35pub use corpus::CorpusSpec;
36pub use kv_dictionary::{learn_dictionary, KvDictionary, SparseCode};
37pub use package::Provenance;
38
39use std::path::PathBuf;
40
41/// Which calibration artifact to produce.
42#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
43pub enum ArtifactKind {
44    /// AWQ per-input-channel activation scales for the FFN projections (implemented — reuses the
45    /// existing AWQ capture+fold+sweep).
46    AwqScales,
47    /// int8 K/V-cache scales (per head-slot). Gated on W5a (the int8 KV cache).
48    KvInt8Scales,
49    /// Sparse KV dictionary / Top-K SAE (Lexico-style). Gated on W5b (needs a custom corpus + OMP/k-SVD).
50    KvDictionary,
51}
52
53impl ArtifactKind {
54    /// Stable machine label (diagnostics / provenance display).
55    #[allow(dead_code)]
56    pub fn label(self) -> &'static str {
57        match self {
58            ArtifactKind::AwqScales => "awq_scales",
59            ArtifactKind::KvInt8Scales => "kv_int8_scales",
60            ArtifactKind::KvDictionary => "kv_dictionary",
61        }
62    }
63}
64
65/// The ΔPPL acceptance gate for a lossy artifact (fraction, e.g. 0.05 = 5%).
66#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
67pub struct GateSpec {
68    pub max_delta_ppl: f64,
69}
70
71impl Default for GateSpec {
72    fn default() -> Self {
73        // The project-wide compression gate (5%), same constant the AWQ/ternary work uses.
74        Self {
75            max_delta_ppl: crate::llm_eval::MAX_DELTA_PPL,
76        }
77    }
78}
79
80/// A calibration job: produce `artifact` for `model_path`, calibrated on `corpus`, gated by `gate`.
81#[derive(Debug, Clone)]
82pub struct CalibrationJob {
83    pub model_path: PathBuf,
84    pub artifact: ArtifactKind,
85    pub corpus: CorpusSpec,
86    pub gate: GateSpec,
87    /// Token budget per PPL pass (0 = the oracle's default full corpus).
88    pub max_tok: usize,
89}
90
91/// The outcome of a calibration run.
92#[derive(Debug, Clone)]
93pub struct CalibrationReport {
94    pub artifact: ArtifactKind,
95    /// Content hash of the assembled calibration corpus (provenance).
96    pub corpus_hash: u64,
97    pub corpus_docs: usize,
98    /// Reference (uncompressed) perplexity.
99    pub ref_ppl: f64,
100    /// Candidate (artifact-applied) perplexity.
101    pub cand_ppl: f64,
102    /// (cand - ref) / ref.
103    pub delta_ppl: f64,
104    /// Whether `delta_ppl <= gate.max_delta_ppl`.
105    pub passed: bool,
106    /// The packaged artifact bytes (artifact + CBOR provenance frame) — only when `passed`.
107    pub packaged: Option<Vec<u8>>,
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub enum CalibrationError {
112    CorpusEmpty,
113    OllamaUnavailable(String),
114    CaptureFailed(String),
115    CertifyFailed(String),
116    PackageFailed(String),
117    /// The artifact kind exists in the taxonomy but its learner is not built yet (honest gate, not a
118    /// silent stub) — the workstream that lands it is named.
119    NotYetImplemented(&'static str),
120}
121
122impl std::fmt::Display for CalibrationError {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        match self {
125            CalibrationError::CorpusEmpty => write!(f, "calibration corpus is empty"),
126            CalibrationError::OllamaUnavailable(e) => {
127                write!(f, "ollama corpus synthesis unavailable: {e}")
128            }
129            CalibrationError::CaptureFailed(e) => write!(f, "activation capture failed: {e}"),
130            CalibrationError::CertifyFailed(e) => write!(f, "certification (PPL) failed: {e}"),
131            CalibrationError::PackageFailed(e) => write!(f, "artifact packaging failed: {e}"),
132            CalibrationError::NotYetImplemented(w) => {
133                write!(f, "artifact learner not yet implemented ({w})")
134            }
135        }
136    }
137}
138
139/// Run the calibration pipeline for one job. The forge's third produce-and-certify entry point,
140/// beside kernel certification and GGUF→p64 transcode.
141pub fn run_calibration(job: &CalibrationJob) -> Result<CalibrationReport, CalibrationError> {
142    // Stage 1 — corpus. Assembled + hashed for provenance regardless of artifact kind.
143    let docs = corpus::assemble(&job.corpus)?;
144    if docs.is_empty() {
145        return Err(CalibrationError::CorpusEmpty);
146    }
147    let corpus_hash = corpus::content_hash(&docs);
148    let corpus_docs = docs.len();
149
150    // Stages 2–4 — capture + learn + certify, per artifact kind.
151    match job.artifact {
152        ArtifactKind::AwqScales => run_awq(job, corpus_hash, corpus_docs),
153        // int8-KV scale calibration needs the W5a int8 KV cache (write-side amax → per-head-slot f16
154        // scale). The seam lands with W5a; visible-not-stubbed until then.
155        ArtifactKind::KvInt8Scales => Err(CalibrationError::NotYetImplemented("W5a int8 KV cache")),
156        // Sparse KV dictionary (Lexico / Top-K SAE): W5b. Phase 3 = the go/no-go — capture the
157        // engine's real KV vectors, learn a per-layer dictionary, and compare its reconstruction to
158        // int8's at footprint. Packaging a certified runtime artifact is Phase 4 and only justified if
159        // this gate says GO.
160        ArtifactKind::KvDictionary => run_kv_dictionary(job, corpus_hash, corpus_docs),
161    }
162}
163
164/// Per-(layer, stream) rate-distortion comparison. int8 (the W5a incumbent, 8 bits/elem) is a strong,
165/// accurate baseline; a sparse dictionary trades accuracy for a much smaller footprint. So the decision
166/// isn't "does the dictionary beat int8's accuracy" (it won't — int8 has ~4× the bits) but "at the
167/// dictionary's OWN low bit rate, does the learned basis beat NAIVE uniform quantization?" — i.e. is the
168/// learned codebook worth more than just quantizing more coarsely.
169#[derive(Debug, Clone)]
170pub struct KvLayerVerdict {
171    pub layer: usize,
172    /// "K" or "V".
173    pub stream: &'static str,
174    pub n_vectors: usize,
175    /// int8 incumbent (8-bit) reconstruction error and footprint — context, not the gate.
176    pub recon_int8: f64,
177    pub int8_bits: f64,
178    /// k-sparse dictionary reconstruction error and asymptotic code footprint (bits/vec, dictionary
179    /// amortized away as at deployment scale).
180    pub recon_dict: f64,
181    pub dict_code_bits: f64,
182    /// Naive uniform quantization at the dictionary's matched bit rate — the head-to-head baseline.
183    pub matched_bits: u32,
184    pub recon_uniform_matched: f64,
185    /// GO here: the learned dictionary reconstructs at least as well as uniform quantization AT THE
186    /// SAME bit rate (the learned basis earns its keep).
187    pub go: bool,
188}
189
190/// The W5b Phase-3 decision: does a learned sparse dictionary beat int8 on the engine's real KV
191/// vectors? Per-layer detail plus an overall verdict.
192#[derive(Debug, Clone)]
193pub struct KvDictReport {
194    pub head_dim: usize,
195    pub n_atoms: usize,
196    pub sparsity: usize,
197    pub layers: Vec<KvLayerVerdict>,
198    /// GO iff a majority of analyzed (layer, stream) pairs have the dictionary dominating int8.
199    pub overall_go: bool,
200}
201
202/// W5b Phase 3 — the sparse-KV-dictionary go/no-go on **real engine KV vectors**.
203///
204/// Enables the [`crate::kv_capture`] hook, runs a calibration forward (the native attention path routes
205/// through the CPU SDPA that the hook taps), then per sampled layer learns a dictionary over the
206/// captured K and V vectors and compares its reconstruction error + footprint to per-vector int8. This
207/// is a **measurement**, not a packaged artifact — the artifact (runtime dictionary decode + ΔPPL
208/// certify) is Phase 4, built only if this returns GO.
209#[cfg(not(target_arch = "wasm32"))]
210#[allow(clippy::too_many_arguments)]
211pub fn kv_dictionary_go_no_go(
212    model_path: &str,
213    n_layer_hint: usize,
214    max_per_layer: usize,
215    n_atoms: usize,
216    sparsity: usize,
217    iters: usize,
218    max_tok: usize,
219    layer_stride: usize,
220) -> Result<KvDictReport, CalibrationError> {
221    // The KV hook taps `cpu_attention_pass`, which the fast GPU decode path bypasses. To route the
222    // whole attention block through the CPU reference SDPA (so K/V pass through the hook AND the hidden
223    // state stays correct for deeper layers), force three flags for the capture and restore them after:
224    //   * cpu_attention ON  — route K/V/Q projections through `cpu_attention_pass`;
225    //   * preproject   OFF  — skip the fused GPU K/V pre-projection that writes straight to VRAM;
226    //   * o_fuse       OFF  — skip the fused GPU Q+O tail so attention output comes from the CPU SDPA.
227    // This is the engine's certified correctness reference, so the captured K/V are the real vectors,
228    // just computed on the reference path rather than the fast one.
229    let prev_cpu_attn = crate::llm_bench::cpu_attention_enabled();
230    let prev_preproject = crate::llm_bench::attention_preproject_enabled();
231    let prev_o_fuse = crate::llm_bench::attention_o_fuse_enabled();
232    crate::llm_bench::set_cpu_attention(true);
233    crate::llm_bench::set_attention_preproject(false);
234    crate::llm_bench::set_attention_o_fuse(false);
235    crate::kv_capture::enable(n_layer_hint, max_per_layer);
236    let ev = crate::llm_bench::perplexity_eval_blocking(model_path, max_tok);
237    let cap = crate::kv_capture::snapshot();
238    crate::kv_capture::disable();
239    crate::kv_capture::clear();
240    crate::llm_bench::set_cpu_attention(prev_cpu_attn);
241    crate::llm_bench::set_attention_preproject(prev_preproject);
242    crate::llm_bench::set_attention_o_fuse(prev_o_fuse);
243    ev.map_err(CalibrationError::CaptureFailed)?;
244    let cap = cap.ok_or_else(|| {
245        CalibrationError::CaptureFailed(
246            "no KV captured — the calibration forward never hit the CPU attention path".into(),
247        )
248    })?;
249
250    Ok(analyze_kv_capture(
251        &cap,
252        n_atoms,
253        sparsity,
254        iters,
255        layer_stride,
256    ))
257}
258
259/// W5b Phase 3 — the same go/no-go via the **GPU-readback** capture route ([`crate::llm_bench::
260/// capture_kv_gpu_readback`]): reads the real fast-decode-path K/V straight from VRAM instead of forcing
261/// the CPU reference SDPA. Running this alongside [`kv_dictionary_go_no_go`] cross-checks the measured
262/// geometry — if the two routes agree, the verdict is trustworthy.
263#[cfg(not(target_arch = "wasm32"))]
264#[allow(clippy::too_many_arguments)]
265pub fn kv_dictionary_go_no_go_gpu(
266    model_path: &str,
267    max_per_layer: usize,
268    n_atoms: usize,
269    sparsity: usize,
270    iters: usize,
271    max_tok: usize,
272    layer_stride: usize,
273) -> Result<KvDictReport, CalibrationError> {
274    let cap = crate::llm_bench::capture_kv_gpu_readback(model_path, max_tok, max_per_layer)
275        .map_err(CalibrationError::CaptureFailed)?;
276    Ok(analyze_kv_capture(
277        &cap,
278        n_atoms,
279        sparsity,
280        iters,
281        layer_stride,
282    ))
283}
284
285/// Shared analysis: per sampled layer + stream, learn a dictionary over the captured vectors and score
286/// it against int8 (context) and matched-rate uniform quantization (the gate). Pure CPU; identical for
287/// both capture routes so their verdicts are directly comparable.
288#[cfg(not(target_arch = "wasm32"))]
289fn analyze_kv_capture(
290    cap: &crate::kv_capture::KvCapture,
291    n_atoms: usize,
292    sparsity: usize,
293    iters: usize,
294    layer_stride: usize,
295) -> KvDictReport {
296    use kv_dictionary::{
297        dict_code_bits_per_vector, int8_bits_per_vector, int8_reconstruction_error,
298        learn_dictionary, uniform_reconstruction_error,
299    };
300
301    let head_dim = cap.head_dim;
302    let int8_bits = int8_bits_per_vector(head_dim);
303    // Asymptotic code rate of the dictionary (dictionary amortized away, as at deployment scale) and
304    // the integer bit-rate of uniform quantization that matches it — the head-to-head baseline.
305    let dict_code_bits = dict_code_bits_per_vector(n_atoms, sparsity, 16);
306    let matched_bits = ((dict_code_bits / head_dim.max(1) as f64).round() as u32).clamp(2, 8);
307    let mut layers = Vec::new();
308    for li in (0..cap.k.len()).step_by(layer_stride.max(1)) {
309        for (stream, vecs) in [("K", &cap.k[li]), ("V", &cap.v[li])] {
310            // Need enough vectors for the dictionary to be meaningful (≥ atom count).
311            if vecs.len() < n_atoms.max(16) {
312                continue;
313            }
314            let dict = learn_dictionary(vecs, head_dim, n_atoms, sparsity, iters);
315            let recon_dict = dict.reconstruction_error(vecs, sparsity);
316            let recon_int8 = int8_reconstruction_error(vecs);
317            let recon_uniform_matched = uniform_reconstruction_error(vecs, matched_bits);
318            // GO: the learned basis beats naive uniform quantization at the same bit rate.
319            let go = recon_dict <= recon_uniform_matched;
320            layers.push(KvLayerVerdict {
321                layer: li,
322                stream,
323                n_vectors: vecs.len(),
324                recon_int8,
325                int8_bits,
326                recon_dict,
327                dict_code_bits,
328                matched_bits,
329                recon_uniform_matched,
330                go,
331            });
332        }
333    }
334    let go_count = layers.iter().filter(|l| l.go).count();
335    let overall_go = !layers.is_empty() && go_count * 2 >= layers.len();
336    KvDictReport {
337        head_dim,
338        n_atoms,
339        sparsity,
340        layers,
341        overall_go,
342    }
343}
344
345/// `run_calibration` arm for [`ArtifactKind::KvDictionary`]: the full W5b Phase-4 certify — learn the
346/// per-layer dictionaries, gate them by real ΔPPL, and package if they pass. Defaults: 256 atoms,
347/// 4-sparse, 20 learn iterations.
348#[cfg(not(target_arch = "wasm32"))]
349fn run_kv_dictionary(
350    job: &CalibrationJob,
351    corpus_hash: u64,
352    corpus_docs: usize,
353) -> Result<CalibrationReport, CalibrationError> {
354    let model = job.model_path.to_string_lossy().to_string();
355    certify_kv_dictionary(
356        &model,
357        256,
358        4,
359        20,
360        job.max_tok,
361        job.gate,
362        corpus_hash,
363        corpus_docs,
364    )
365}
366
367/// RAII: engage the CPU-reference attention path + f32 KV (so the reconstruct hook is on the PPL path)
368/// and restore the prior flags on drop — so any early return during certification can't leave the
369/// engine's global toggles flipped.
370#[cfg(not(target_arch = "wasm32"))]
371struct CpuRefPathGuard {
372    cpu: bool,
373    pre: bool,
374    of: bool,
375    int8: bool,
376}
377
378#[cfg(not(target_arch = "wasm32"))]
379impl CpuRefPathGuard {
380    fn engage() -> Self {
381        let g = Self {
382            cpu: crate::llm_bench::cpu_attention_enabled(),
383            pre: crate::llm_bench::attention_preproject_enabled(),
384            of: crate::llm_bench::attention_o_fuse_enabled(),
385            int8: crate::llm_bench::kv_int8_enabled(),
386        };
387        crate::llm_bench::set_cpu_attention(true);
388        crate::llm_bench::set_attention_preproject(false);
389        crate::llm_bench::set_attention_o_fuse(false);
390        crate::llm_bench::set_kv_int8(false);
391        g
392    }
393}
394
395#[cfg(not(target_arch = "wasm32"))]
396impl Drop for CpuRefPathGuard {
397    fn drop(&mut self) {
398        crate::llm_bench::set_cpu_attention(self.cpu);
399        crate::llm_bench::set_attention_preproject(self.pre);
400        crate::llm_bench::set_attention_o_fuse(self.of);
401        crate::llm_bench::set_kv_int8(self.int8);
402    }
403}
404
405/// Learn one config's dictionaries over an already-captured KV set and measure the candidate PPL, then
406/// gate + package. Assumes the CPU-reference/f32 path is already engaged and `ref_ppl` was measured
407/// under the same conditions — the delta is purely the dictionary. Config-specific work only, so a
408/// sweep pays the capture + reference cost once and this per config.
409#[cfg(not(target_arch = "wasm32"))]
410#[allow(clippy::too_many_arguments)]
411fn certify_one_config(
412    cap: &crate::kv_capture::KvCapture,
413    ref_ppl: f64,
414    model_path: &str,
415    n_atoms: usize,
416    sparsity: usize,
417    iters: usize,
418    max_tok: usize,
419    gate: GateSpec,
420    corpus_hash: u64,
421    corpus_docs: usize,
422) -> Result<CalibrationReport, CalibrationError> {
423    use kv_dictionary::{learn_dictionary, KvDictionary};
424
425    let head_dim = cap.head_dim;
426    let learn = |vecs: &Vec<Vec<f32>>| -> Option<KvDictionary> {
427        (vecs.len() >= n_atoms.max(16))
428            .then(|| learn_dictionary(vecs, head_dim, n_atoms, sparsity, iters))
429    };
430    let k_dicts: Vec<Option<KvDictionary>> = cap.k.iter().map(learn).collect();
431    let v_dicts: Vec<Option<KvDictionary>> = cap.v.iter().map(learn).collect();
432    if k_dicts.iter().all(|d| d.is_none()) && v_dicts.iter().all(|d| d.is_none()) {
433        return Err(CalibrationError::CaptureFailed(
434            "no layer had enough KV to learn a dictionary".into(),
435        ));
436    }
437
438    kv_dict_runtime::enable(k_dicts.clone(), v_dicts.clone(), sparsity);
439    let cand_res = crate::llm_bench::perplexity_eval_blocking(model_path, max_tok);
440    kv_dict_runtime::disable();
441    kv_dict_runtime::clear();
442    let (cand_ppl, _) = cand_res.map_err(CalibrationError::CertifyFailed)?;
443
444    let delta_ppl = crate::llm_eval::delta_ppl(ref_ppl, cand_ppl);
445    let passed = delta_ppl <= gate.max_delta_ppl;
446
447    let packaged = if passed {
448        // The engine's core loader (`kv_dict_runtime::load_certified`) deserializes exactly this struct,
449        // so the forge and engine share one on-disk dictionary format.
450        let art = crate::kv_dict_runtime::KvDictArtifact {
451            sparsity,
452            head_dim,
453            k: k_dicts,
454            v: v_dicts,
455        };
456        let mut cbor = Vec::new();
457        ciborium::into_writer(&art, &mut cbor)
458            .map_err(|e| CalibrationError::PackageFailed(format!("cbor encode: {e}")))?;
459        let prov = Provenance::new(
460            ArtifactKind::KvDictionary,
461            corpus_hash,
462            corpus_docs,
463            ref_ppl,
464            cand_ppl,
465            delta_ppl,
466            true,
467        );
468        Some(package::frame_artifact(&cbor, &prov))
469    } else {
470        None
471    };
472
473    Ok(CalibrationReport {
474        artifact: ArtifactKind::KvDictionary,
475        corpus_hash,
476        corpus_docs,
477        ref_ppl,
478        cand_ppl,
479        delta_ppl,
480        passed,
481        packaged,
482    })
483}
484
485/// W5b Phase 4 — certify one or more KV-dictionary configs by **ΔPPL**, sharing the expensive capture +
486/// reference-PPL across all of them.
487///
488/// Captures the engine's real KV once (fast GPU-readback route), measures the reference perplexity once
489/// on the CPU-reference/f32 attention path (dictionary OFF), then for each `(n_atoms, sparsity)` config
490/// learns the per-layer dictionaries, measures the candidate PPL (dictionary ON, same path), gates at
491/// `ΔPPL ≤ gate`, and packages the CBOR dictionaries + provenance if it passes (fail-closed). Because
492/// the capture and reference are config-independent, an N-config sweep costs one capture + one reference
493/// + N candidate passes, not N full certifications. Realizing the memory saving in the GPU cache layout
494/// is Phase 4b.
495#[cfg(not(target_arch = "wasm32"))]
496#[allow(clippy::too_many_arguments)]
497pub fn sweep_kv_dictionary(
498    model_path: &str,
499    configs: &[(usize, usize)],
500    iters: usize,
501    max_tok: usize,
502    gate: GateSpec,
503    corpus_hash: u64,
504    corpus_docs: usize,
505) -> Result<Vec<(usize, usize, CalibrationReport)>, CalibrationError> {
506    // Capture once, on the FAST path (before engaging the CPU-reference path below).
507    let cap = crate::llm_bench::capture_kv_gpu_readback(model_path, max_tok, 2048)
508        .map_err(CalibrationError::CaptureFailed)?;
509
510    // Engage the CPU-reference/f32 path for ALL the PPL runs (restored on drop / any early return).
511    let _guard = CpuRefPathGuard::engage();
512
513    // Reference perplexity once (dictionary OFF), shared across configs.
514    kv_dict_runtime::disable();
515    let (ref_ppl, _) = crate::llm_bench::perplexity_eval_blocking(model_path, max_tok)
516        .map_err(CalibrationError::CertifyFailed)?;
517
518    let mut out = Vec::with_capacity(configs.len());
519    for &(n_atoms, sparsity) in configs {
520        let report = certify_one_config(
521            &cap,
522            ref_ppl,
523            model_path,
524            n_atoms,
525            sparsity,
526            iters,
527            max_tok,
528            gate,
529            corpus_hash,
530            corpus_docs,
531        )?;
532        out.push((n_atoms, sparsity, report));
533    }
534    Ok(out)
535}
536
537/// Single-config convenience wrapper over [`sweep_kv_dictionary`].
538#[cfg(not(target_arch = "wasm32"))]
539#[allow(clippy::too_many_arguments)]
540pub fn certify_kv_dictionary(
541    model_path: &str,
542    n_atoms: usize,
543    sparsity: usize,
544    iters: usize,
545    max_tok: usize,
546    gate: GateSpec,
547    corpus_hash: u64,
548    corpus_docs: usize,
549) -> Result<CalibrationReport, CalibrationError> {
550    let mut results = sweep_kv_dictionary(
551        model_path,
552        &[(n_atoms, sparsity)],
553        iters,
554        max_tok,
555        gate,
556        corpus_hash,
557        corpus_docs,
558    )?;
559    results
560        .pop()
561        .map(|(_, _, r)| r)
562        .ok_or_else(|| CalibrationError::CertifyFailed("no config produced a report".into()))
563}
564
565/// AWQ scales: the real capture→fold→sweep pipeline (reuses [`awq_sweep_blocking`]), certified
566/// against the Q8 reference and packaged as an AWQ-folded Q4_0 FFN p64 with a provenance frame.
567#[cfg(not(target_arch = "wasm32"))]
568fn run_awq(
569    job: &CalibrationJob,
570    corpus_hash: u64,
571    corpus_docs: usize,
572) -> Result<CalibrationReport, CalibrationError> {
573    use crate::p64_weight::FfnQuant;
574    let model = job.model_path.to_string_lossy().to_string();
575    // Coarse α-sweep (0 / 0.5 / 1) over the Q4_0 FFN — AWQ's design regime. `awq_sweep_blocking`
576    // runs the capture + fold + PPL certify internally and returns (ref_ppl, [(alpha, ppl, uniq)]).
577    let alphas = [0.0f32, 0.5, 1.0];
578    let (ref_ppl, results) =
579        crate::llm_bench::awq_sweep_blocking(&model, &alphas, job.max_tok, FfnQuant::Q4_0)
580            .map_err(CalibrationError::CertifyFailed)?;
581    let best = results
582        .iter()
583        .filter(|(_, p, _)| p.is_finite() && *p > 1.0)
584        .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
585        .copied()
586        .ok_or_else(|| {
587            CalibrationError::CertifyFailed("AWQ sweep produced no finite PPL".into())
588        })?;
589    let (best_alpha, cand_ppl, _uniq) = best;
590    let delta_ppl = crate::llm_eval::delta_ppl(ref_ppl, cand_ppl);
591    let passed = delta_ppl <= job.gate.max_delta_ppl;
592
593    let packaged = if passed {
594        // Package: recompute the AWQ scales at the winning α + fold into a Q4_0 FFN p64, then frame
595        // it with the CBOR provenance. Capture the scales via a fresh AWQ pass over the reference.
596        let bytes = std::fs::read(&job.model_path)
597            .map_err(|e| CalibrationError::PackageFailed(format!("read gguf: {e}")))?;
598        let scales = capture_awq_scales(&model, job.max_tok)?;
599        let p64 = crate::p64_weight::compile_gguf_to_q42_ffn_quant_awq(
600            &bytes,
601            14,
602            Some(&scales),
603            best_alpha,
604            FfnQuant::Q4_0,
605        )
606        .map_err(|e| CalibrationError::PackageFailed(format!("awq compile: {e}")))?;
607        let prov = Provenance::new(
608            ArtifactKind::AwqScales,
609            corpus_hash,
610            corpus_docs,
611            ref_ppl,
612            cand_ppl,
613            delta_ppl,
614            true,
615        );
616        Some(package::frame_artifact(&p64, &prov))
617    } else {
618        None
619    };
620
621    Ok(CalibrationReport {
622        artifact: ArtifactKind::AwqScales,
623        corpus_hash,
624        corpus_docs,
625        ref_ppl,
626        cand_ppl,
627        delta_ppl,
628        passed,
629        packaged,
630    })
631}
632
633/// Capture per-FFN-layer per-input-channel AWQ activation scales by running a calibration forward
634/// over the reference model with the `llm_awq` hooks enabled. Native (drives the GPU forward).
635#[cfg(not(target_arch = "wasm32"))]
636fn capture_awq_scales(model: &str, max_tok: usize) -> Result<Vec<Vec<f32>>, CalibrationError> {
637    // SmolLM2-360M shape; the hook self-sizes on the first record. Layers/channels are read from the
638    // snapshot, so an over-estimate here is harmless (the AWQ module clamps).
639    crate::llm_awq::enable(64, 4096).map_err(CalibrationError::CaptureFailed)?;
640    let cap = crate::llm_bench::perplexity_eval_blocking(model, max_tok);
641    let scales = crate::llm_awq::snapshot();
642    crate::llm_awq::disable();
643    cap.map_err(CalibrationError::CaptureFailed)?;
644    if scales.is_empty() || scales.iter().all(|l| l.iter().all(|&v| v == 0.0)) {
645        return Err(CalibrationError::CaptureFailed(
646            "AWQ hooks captured no activations".into(),
647        ));
648    }
649    Ok(scales)
650}
651
652#[cfg(test)]
653mod tests {
654    use super::*;
655
656    #[test]
657    fn unimplemented_kinds_are_visible_not_stubbed() {
658        // A tiny in-memory corpus so we reach the artifact dispatch.
659        let job = |kind| CalibrationJob {
660            model_path: PathBuf::from("/nonexistent.gguf"),
661            artifact: kind,
662            corpus: CorpusSpec::Files(vec![]),
663            gate: GateSpec::default(),
664            max_tok: 0,
665        };
666        // Empty Files corpus → CorpusEmpty (the corpus stage gates before artifact dispatch).
667        assert_eq!(
668            run_calibration(&job(ArtifactKind::KvInt8Scales)).unwrap_err(),
669            CalibrationError::CorpusEmpty
670        );
671    }
672
673    #[test]
674    fn gate_default_is_the_project_ppl_gate() {
675        assert_eq!(
676            GateSpec::default().max_delta_ppl,
677            crate::llm_eval::MAX_DELTA_PPL
678        );
679    }
680
681    #[test]
682    fn artifact_kind_labels_stable() {
683        assert_eq!(ArtifactKind::AwqScales.label(), "awq_scales");
684        assert_eq!(ArtifactKind::KvInt8Scales.label(), "kv_int8_scales");
685        assert_eq!(ArtifactKind::KvDictionary.label(), "kv_dictionary");
686    }
687}
qualia_core_db/wgsl_forge/calibration/mod.rs

qualia_core_db/wgsl_forge/calibration/
mod.rs
