forked from lance-format/lance-graph
-
Notifications
You must be signed in to change notification settings - Fork 0
D1.3 decode-kernel + residual composition (Phase 1 scaffold complete, 104/104 tests) #235
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,339 @@ | ||
| //! **LAB-ONLY.** D1.3 — residual PQ via decode-kernel composition. | ||
| //! | ||
| //! Scope correction (per `cognitive-shader-architecture.md`): this module | ||
| //! sits on the **hydration / calibration path**, not the cascade inference | ||
| //! path. The inference cascade uses `p64_bridge::CognitiveShader::cascade` | ||
| //! at Layer 2 (line 582 of that doc); decode kernels here are for offline | ||
| //! codec-candidate calibration. A codec that passes the token-agreement | ||
| //! cert gate (D2.x) then runs at **weight hydration time** (GGUF → palette | ||
| //! + Fingerprint<256> + holographic residual), never per-inference. | ||
| //! | ||
| //! This module defines: | ||
| //! | ||
| //! - [`DecodeKernel`] — the codec decode/encode trait, signature-keyed | ||
| //! into `CodecKernelCache<H>` at the `H` slot where `H: DecodeKernel`. | ||
| //! - [`StubDecodeKernel`] — deterministic reference for tests (byte-level | ||
| //! round-trip, no quantization loss, matches Rule F "serialise once at | ||
| //! edge" — the decode/encode IS the edge). | ||
| //! - [`ResidualComposer`] — composes two decoders with subtract/add: | ||
| //! `decode(enc) = base.decode(enc[..k]) + residual.decode(enc[k..])` | ||
| //! `encode(v) = [base.encode(v); residual.encode(v - base.decode(base.encode(v)))]` | ||
| //! Depth `d > 1` recurses: the residual field itself is a `ResidualComposer`. | ||
|
|
||
| use std::collections::hash_map::DefaultHasher; | ||
| use std::hash::{Hash, Hasher}; | ||
|
|
||
| /// Error from encode / decode — size mismatch or downstream kernel failure. | ||
| #[derive(Debug, Clone, PartialEq, Eq)] | ||
| pub enum DecodeError { | ||
| /// Input slice size doesn't match the kernel's expected `bytes_per_row` | ||
| /// (for decode) or `dim` (for encode). | ||
| SizeMismatch { expected: usize, actual: usize }, | ||
| /// A composed-kernel stage failed with a downstream reason. | ||
| Stage { stage: &'static str, detail: String }, | ||
| } | ||
|
|
||
| impl std::fmt::Display for DecodeError { | ||
| fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||
| match self { | ||
| Self::SizeMismatch { expected, actual } => { | ||
| write!(f, "decode size mismatch: expected {expected}, got {actual}") | ||
| } | ||
| Self::Stage { stage, detail } => write!(f, "stage {stage} failed: {detail}"), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl std::error::Error for DecodeError {} | ||
|
|
||
| /// A codec decode/encode kernel. | ||
| /// | ||
| /// Object-safe so `ResidualComposer` can hold `Box<dyn DecodeKernel>` for | ||
| /// each stage; keyed by `signature()` into `CodecKernelCache<H>` at the | ||
| /// `H = Box<dyn DecodeKernel>` slot. | ||
| pub trait DecodeKernel: Send + Sync + std::fmt::Debug { | ||
| /// Decode `bytes` (exactly `bytes_per_row` long) into `dim` `f32` values. | ||
| fn decode(&self, bytes: &[u8]) -> Result<Vec<f32>, DecodeError>; | ||
| /// Encode `vec` (exactly `dim` `f32` values) into `bytes_per_row` bytes. | ||
| fn encode(&self, vec: &[f32]) -> Result<Vec<u8>, DecodeError>; | ||
| /// Bytes per encoded row — composers rely on this to split payloads. | ||
| fn bytes_per_row(&self) -> u32; | ||
| /// Dimension of the decoded `f32` vector. | ||
| fn dim(&self) -> u32; | ||
| /// Stable hash for JIT cache keying. | ||
| fn signature(&self) -> u64; | ||
| /// Backend tier ("avx512" | "amx" | "stub") — never "scalar" on SoA. | ||
| fn backend(&self) -> &'static str; | ||
| } | ||
|
|
||
| // ─── Stub decoder (byte-exact round-trip for testing) ──────────────────── | ||
|
|
||
| /// Deterministic byte-exact decoder. `dim` `f32` values ⇔ `dim * 4` bytes | ||
| /// via native-endian reinterpret. No quantization, no compression — the | ||
| /// round-trip is exact. | ||
| /// | ||
| /// This IS NOT a real codec. It exists so the `ResidualComposer` and | ||
| /// `CodecKernelCache` composition tests can verify the plumbing without | ||
| /// depending on a trained palette. | ||
| #[derive(Debug, Clone, Copy)] | ||
| pub struct StubDecodeKernel { | ||
| pub dim: u32, | ||
| pub tag: u64, | ||
| } | ||
|
|
||
| impl StubDecodeKernel { | ||
| pub const fn new(dim: u32, tag: u64) -> Self { Self { dim, tag } } | ||
| } | ||
|
|
||
| impl DecodeKernel for StubDecodeKernel { | ||
| fn decode(&self, bytes: &[u8]) -> Result<Vec<f32>, DecodeError> { | ||
| let expected = self.bytes_per_row() as usize; | ||
| if bytes.len() != expected { | ||
| return Err(DecodeError::SizeMismatch { expected, actual: bytes.len() }); | ||
| } | ||
| let mut out = Vec::with_capacity(self.dim as usize); | ||
| for chunk in bytes.chunks_exact(4) { | ||
| out.push(f32::from_ne_bytes(chunk.try_into().unwrap())); | ||
| } | ||
| Ok(out) | ||
| } | ||
|
|
||
| fn encode(&self, vec: &[f32]) -> Result<Vec<u8>, DecodeError> { | ||
| let expected = self.dim as usize; | ||
| if vec.len() != expected { | ||
| return Err(DecodeError::SizeMismatch { expected, actual: vec.len() }); | ||
| } | ||
| let mut out = Vec::with_capacity(expected * 4); | ||
| for &v in vec { | ||
| out.extend_from_slice(&v.to_ne_bytes()); | ||
| } | ||
| Ok(out) | ||
| } | ||
|
|
||
| fn bytes_per_row(&self) -> u32 { self.dim * 4 } | ||
| fn dim(&self) -> u32 { self.dim } | ||
|
|
||
| fn signature(&self) -> u64 { | ||
| let mut h = DefaultHasher::new(); | ||
| "stub_decode".hash(&mut h); | ||
| self.dim.hash(&mut h); | ||
| self.tag.hash(&mut h); | ||
| h.finish() | ||
| } | ||
|
|
||
| fn backend(&self) -> &'static str { "stub" } | ||
| } | ||
|
|
||
| // ─── Residual composer ─────────────────────────────────────────────────── | ||
|
|
||
| /// Composes a `base` decoder with a `residual` decoder: the encoded payload | ||
| /// is the concatenation of `base.encode(v)` + `residual.encode(v - | ||
| /// base.decode(base.encode(v)))`; decode reverses by decoding each stage's | ||
| /// own byte range and summing. | ||
| /// | ||
| /// **Contract:** `base.dim() == residual.dim()`. Enforced by `new()`. | ||
| /// | ||
| /// Deeper residual chains (depth > 1) nest `ResidualComposer`s — the | ||
| /// `residual` slot itself is a `Box<dyn DecodeKernel>` which can be | ||
| /// another `ResidualComposer`. | ||
| #[derive(Debug)] | ||
| pub struct ResidualComposer { | ||
| base: Box<dyn DecodeKernel>, | ||
| residual: Box<dyn DecodeKernel>, | ||
| } | ||
|
|
||
| impl ResidualComposer { | ||
| /// Build a two-stage residual composer. | ||
| /// | ||
| /// Returns `Err(SizeMismatch)` when `base.dim() != residual.dim()`. | ||
| pub fn new( | ||
| base: Box<dyn DecodeKernel>, | ||
| residual: Box<dyn DecodeKernel>, | ||
| ) -> Result<Self, DecodeError> { | ||
| if base.dim() != residual.dim() { | ||
| return Err(DecodeError::SizeMismatch { | ||
| expected: base.dim() as usize, | ||
| actual: residual.dim() as usize, | ||
| }); | ||
| } | ||
| Ok(Self { base, residual }) | ||
| } | ||
| } | ||
|
|
||
| impl DecodeKernel for ResidualComposer { | ||
| fn decode(&self, bytes: &[u8]) -> Result<Vec<f32>, DecodeError> { | ||
| let base_b = self.base.bytes_per_row() as usize; | ||
| let expected = self.bytes_per_row() as usize; | ||
| if bytes.len() != expected { | ||
| return Err(DecodeError::SizeMismatch { expected, actual: bytes.len() }); | ||
| } | ||
| let base_v = self | ||
| .base | ||
| .decode(&bytes[..base_b]) | ||
| .map_err(|e| DecodeError::Stage { stage: "base::decode", detail: e.to_string() })?; | ||
| let residual_v = self | ||
| .residual | ||
| .decode(&bytes[base_b..]) | ||
| .map_err(|e| DecodeError::Stage { stage: "residual::decode", detail: e.to_string() })?; | ||
| let mut out = base_v; | ||
| for (dst, &r) in out.iter_mut().zip(&residual_v) { | ||
| *dst += r; | ||
| } | ||
| Ok(out) | ||
| } | ||
|
|
||
| fn encode(&self, vec: &[f32]) -> Result<Vec<u8>, DecodeError> { | ||
| let expected = self.dim() as usize; | ||
| if vec.len() != expected { | ||
| return Err(DecodeError::SizeMismatch { expected, actual: vec.len() }); | ||
| } | ||
| // First-pass encode + its self-reconstruction. | ||
| let base_bytes = self | ||
| .base | ||
| .encode(vec) | ||
| .map_err(|e| DecodeError::Stage { stage: "base::encode", detail: e.to_string() })?; | ||
| let base_reconstructed = self | ||
| .base | ||
| .decode(&base_bytes) | ||
| .map_err(|e| DecodeError::Stage { stage: "base::decode", detail: e.to_string() })?; | ||
| // Residual = original − base.decode(base.encode(original)). | ||
| let residual_vec: Vec<f32> = | ||
| vec.iter().zip(&base_reconstructed).map(|(a, b)| a - b).collect(); | ||
| let residual_bytes = self | ||
| .residual | ||
| .encode(&residual_vec) | ||
| .map_err(|e| DecodeError::Stage { stage: "residual::encode", detail: e.to_string() })?; | ||
| // Concat. | ||
| let mut out = Vec::with_capacity(base_bytes.len() + residual_bytes.len()); | ||
| out.extend_from_slice(&base_bytes); | ||
| out.extend_from_slice(&residual_bytes); | ||
| Ok(out) | ||
| } | ||
|
|
||
| fn bytes_per_row(&self) -> u32 { | ||
| self.base.bytes_per_row() + self.residual.bytes_per_row() | ||
| } | ||
|
|
||
| fn dim(&self) -> u32 { self.base.dim() } | ||
|
|
||
| fn signature(&self) -> u64 { | ||
| let mut h = DefaultHasher::new(); | ||
| "residual_compose".hash(&mut h); | ||
| self.base.signature().hash(&mut h); | ||
| self.residual.signature().hash(&mut h); | ||
| h.finish() | ||
| } | ||
|
|
||
| fn backend(&self) -> &'static str { | ||
| // Composer backend = the less-optimized stage's backend (weakest link | ||
| // drives actual latency). For the stub case this is "stub"; for real | ||
| // D1.1b kernels, the OPQ matmul stage dominates. | ||
| if self.base.backend() == "stub" || self.residual.backend() == "stub" { | ||
| "stub" | ||
| } else { | ||
| self.base.backend() | ||
| } | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
|
|
||
| #[test] | ||
| fn stub_round_trip_is_exact() { | ||
| let k = StubDecodeKernel::new(8, 1); | ||
| let v = vec![1.0_f32, -2.5, 3.25, 0.0, -7.875, 100.0, -0.125, 42.0]; | ||
| let enc = k.encode(&v).unwrap(); | ||
| assert_eq!(enc.len(), 32); | ||
| let dec = k.decode(&enc).unwrap(); | ||
| assert_eq!(dec, v); | ||
| } | ||
|
|
||
| #[test] | ||
| fn stub_rejects_wrong_input_size() { | ||
| let k = StubDecodeKernel::new(4, 0); | ||
| let err = k.encode(&[1.0, 2.0, 3.0]).unwrap_err(); | ||
| assert!(matches!(err, DecodeError::SizeMismatch { expected: 4, actual: 3 })); | ||
| let err = k.decode(&[0u8; 10]).unwrap_err(); | ||
| assert!(matches!(err, DecodeError::SizeMismatch { expected: 16, actual: 10 })); | ||
| } | ||
|
|
||
| #[test] | ||
| fn residual_compose_round_trip_is_exact_when_both_stubs() { | ||
| let base = Box::new(StubDecodeKernel::new(4, 1)); | ||
| let residual = Box::new(StubDecodeKernel::new(4, 2)); | ||
| let comp = ResidualComposer::new(base, residual).unwrap(); | ||
|
|
||
| let v = vec![1.5_f32, -2.0, 3.125, -0.5]; | ||
| let enc = comp.encode(&v).unwrap(); | ||
| assert_eq!(enc.len(), 32, "4 dim × 4 bytes × 2 stages = 32"); | ||
| let dec = comp.decode(&enc).unwrap(); | ||
| // Both stages byte-exact ⇒ residual is all zeros ⇒ decoded = base_reconstructed = v. | ||
| assert_eq!(dec, v); | ||
| } | ||
|
|
||
| #[test] | ||
| fn residual_compose_mismatched_dims_rejected() { | ||
| let base = Box::new(StubDecodeKernel::new(4, 0)); | ||
| let residual = Box::new(StubDecodeKernel::new(8, 0)); | ||
| let err = ResidualComposer::new(base, residual).unwrap_err(); | ||
| assert!(matches!(err, DecodeError::SizeMismatch { expected: 4, actual: 8 })); | ||
| } | ||
|
|
||
| #[test] | ||
| fn residual_compose_bytes_per_row_sums_stages() { | ||
| let base = Box::new(StubDecodeKernel::new(6, 1)); // 24 bytes | ||
| let residual = Box::new(StubDecodeKernel::new(6, 2)); // 24 bytes | ||
| let comp = ResidualComposer::new(base, residual).unwrap(); | ||
| assert_eq!(comp.bytes_per_row(), 48); | ||
| assert_eq!(comp.dim(), 6); | ||
| } | ||
|
|
||
| #[test] | ||
| fn residual_compose_nested_depth_two_round_trip() { | ||
| // depth=2: ResidualComposer whose `residual` slot is itself a ResidualComposer. | ||
| let inner_base = Box::new(StubDecodeKernel::new(4, 1)); | ||
| let inner_residual = Box::new(StubDecodeKernel::new(4, 2)); | ||
| let inner = Box::new(ResidualComposer::new(inner_base, inner_residual).unwrap()); | ||
| let outer_base = Box::new(StubDecodeKernel::new(4, 3)); | ||
| let outer = ResidualComposer::new(outer_base, inner).unwrap(); | ||
|
|
||
| assert_eq!(outer.bytes_per_row(), 48, "4 dim × 4 bytes × 3 stages"); | ||
| let v = vec![1.0_f32, 2.0, 3.0, 4.0]; | ||
| let enc = outer.encode(&v).unwrap(); | ||
| let dec = outer.decode(&enc).unwrap(); | ||
| assert_eq!(dec, v); | ||
| } | ||
|
|
||
| #[test] | ||
| fn signatures_distinguish_composer_from_stages() { | ||
| let base = Box::new(StubDecodeKernel::new(4, 1)); | ||
| let residual = Box::new(StubDecodeKernel::new(4, 2)); | ||
| let s_base = base.signature(); | ||
| let s_res = residual.signature(); | ||
| let comp = ResidualComposer::new(base, residual).unwrap(); | ||
| let s_comp = comp.signature(); | ||
| assert_ne!(s_comp, s_base); | ||
| assert_ne!(s_comp, s_res); | ||
| } | ||
|
|
||
| #[test] | ||
| fn signature_depends_on_stage_order() { | ||
| let a = Box::new(StubDecodeKernel::new(4, 1)); | ||
| let b = Box::new(StubDecodeKernel::new(4, 2)); | ||
| let ab = ResidualComposer::new(a, b).unwrap(); | ||
| let a2 = Box::new(StubDecodeKernel::new(4, 1)); | ||
| let b2 = Box::new(StubDecodeKernel::new(4, 2)); | ||
| let ba = ResidualComposer::new(b2, a2).unwrap(); | ||
| assert_ne!(ab.signature(), ba.signature(), "base/residual order is part of identity"); | ||
| } | ||
|
|
||
| #[test] | ||
| fn composer_backend_reports_stub_when_any_stage_is_stub() { | ||
| let base = Box::new(StubDecodeKernel::new(4, 1)); | ||
| let residual = Box::new(StubDecodeKernel::new(4, 2)); | ||
| let comp = ResidualComposer::new(base, residual).unwrap(); | ||
| assert_eq!(comp.backend(), "stub"); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ResidualComposer::decodeadds stage outputs withzip, which truncates to the shorter vector. If either stage returns fewer thandim()elements, this returnsOk(...)with silently corrupted output instead of surfacing an error, even though the trait contract says decode should produce full-dimension vectors. Add explicit length checks for both decoded stage vectors before the accumulation loop so malformed/buggy stage implementations fail fast.Useful? React with 👍 / 👎.