From 03c1feaaa24dc5eee45876db2581314e28924ff6 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 20 Feb 2026 06:55:10 +0000 Subject: [PATCH] fix: Refine experiment crate implementations Agent-driven improvements to min-cut, coherence, and profiler crates: - attn-mincut: config and graph module refinements - coherence: batch evaluation, comparison, and quality updates - profiler: memory and power tracking improvements https://claude.ai/code/session_01TiqLbr2DaNAntQHaVeLfiR --- crates/ruvector-attn-mincut/src/config.rs | 45 ++-- crates/ruvector-attn-mincut/src/graph.rs | 65 ++--- crates/ruvector-attn-mincut/src/mincut.rs | 258 +++++--------------- crates/ruvector-coherence/src/batch.rs | 128 +++------- crates/ruvector-coherence/src/comparison.rs | 173 +++---------- crates/ruvector-coherence/src/quality.rs | 149 +++-------- crates/ruvector-profiler/src/latency.rs | 93 ++----- crates/ruvector-profiler/src/memory.rs | 107 ++------ crates/ruvector-profiler/src/power.rs | 159 +++--------- 9 files changed, 277 insertions(+), 900 deletions(-) diff --git a/crates/ruvector-attn-mincut/src/config.rs b/crates/ruvector-attn-mincut/src/config.rs index a207fd0b1..81c6e9a43 100644 --- a/crates/ruvector-attn-mincut/src/config.rs +++ b/crates/ruvector-attn-mincut/src/config.rs @@ -3,27 +3,16 @@ use serde::{Deserialize, Serialize}; /// Configuration for the min-cut gating attention operator. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MinCutConfig { - /// Regularization weight balancing cut cost vs. edge retention. pub lambda: f32, - /// Hysteresis window: edges must be consistently gated for `tau` steps before flipping. pub tau: usize, - /// Convergence tolerance for iterative refinement. pub eps: f32, - /// Random seed for reproducibility. pub seed: u64, - /// Whether to emit witness JSONL entries for determinism verification. pub witness_enabled: bool, } impl Default for MinCutConfig { fn default() -> Self { - Self { - lambda: 0.5, - tau: 2, - eps: 0.01, - seed: 42, - witness_enabled: true, - } + Self { lambda: 0.5, tau: 2, eps: 0.01, seed: 42, witness_enabled: true } } } @@ -33,27 +22,21 @@ mod tests { #[test] fn test_default_config() { - let cfg = MinCutConfig::default(); - assert!((cfg.lambda - 0.5).abs() < f32::EPSILON); - assert_eq!(cfg.tau, 2); - assert!((cfg.eps - 0.01).abs() < f32::EPSILON); - assert_eq!(cfg.seed, 42); - assert!(cfg.witness_enabled); + let c = MinCutConfig::default(); + assert!((c.lambda - 0.5).abs() < f32::EPSILON); + assert_eq!(c.tau, 2); + assert!((c.eps - 0.01).abs() < f32::EPSILON); + assert_eq!(c.seed, 42); + assert!(c.witness_enabled); } #[test] - fn test_config_serde_roundtrip() { - let cfg = MinCutConfig { - lambda: 0.3, - tau: 5, - eps: 0.001, - seed: 99, - witness_enabled: false, - }; - let json = serde_json::to_string(&cfg).unwrap(); - let restored: MinCutConfig = serde_json::from_str(&json).unwrap(); - assert!((restored.lambda - 0.3).abs() < f32::EPSILON); - assert_eq!(restored.tau, 5); - assert!(!restored.witness_enabled); + fn test_serde_roundtrip() { + let c = MinCutConfig { lambda: 0.3, tau: 5, eps: 0.001, seed: 99, witness_enabled: false }; + let json = serde_json::to_string(&c).unwrap(); + let r: MinCutConfig = serde_json::from_str(&json).unwrap(); + assert!((r.lambda - 0.3).abs() < f32::EPSILON); + assert_eq!(r.tau, 5); + assert!(!r.witness_enabled); } } diff --git a/crates/ruvector-attn-mincut/src/graph.rs b/crates/ruvector-attn-mincut/src/graph.rs index 0b14b74ac..0b68be2b2 100644 --- a/crates/ruvector-attn-mincut/src/graph.rs +++ b/crates/ruvector-attn-mincut/src/graph.rs @@ -2,49 +2,24 @@ use serde::{Deserialize, Serialize}; /// A directed edge in the attention graph. #[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Edge { - pub src: usize, - pub dst: usize, - pub weight: f32, -} +pub struct Edge { pub src: usize, pub dst: usize, pub weight: f32 } /// Weighted directed graph built from attention logits. #[derive(Debug, Clone)] -pub struct AttentionGraph { - pub nodes: usize, - pub edges: Vec, -} +pub struct AttentionGraph { pub nodes: usize, pub edges: Vec } -/// Build a weighted directed graph from attention logits Q*K^T / sqrt(d). -/// -/// `logits` is a flattened `seq_len x seq_len` matrix in row-major order. -/// Each positive logit becomes an edge; non-positive logits are omitted so the -/// graph is sparse when many logits are near zero or negative. +/// Build a weighted directed graph from flattened `seq_len x seq_len` logits. +/// Only positive logits become edges; non-positive entries are omitted. pub fn graph_from_logits(logits: &[f32], seq_len: usize) -> AttentionGraph { - assert_eq!( - logits.len(), - seq_len * seq_len, - "logits length must equal seq_len^2" - ); - + assert_eq!(logits.len(), seq_len * seq_len, "logits length must equal seq_len^2"); let mut edges = Vec::new(); for i in 0..seq_len { for j in 0..seq_len { let w = logits[i * seq_len + j]; - if w > 0.0 { - edges.push(Edge { - src: i, - dst: j, - weight: w, - }); - } + if w > 0.0 { edges.push(Edge { src: i, dst: j, weight: w }); } } } - - AttentionGraph { - nodes: seq_len, - edges, - } + AttentionGraph { nodes: seq_len, edges } } #[cfg(test)] @@ -52,37 +27,25 @@ mod tests { use super::*; #[test] - fn test_graph_from_logits_basic() { - // 2x2 logits: all positive - let logits = vec![1.0, 2.0, 3.0, 4.0]; - let g = graph_from_logits(&logits, 2); + fn test_all_positive() { + let g = graph_from_logits(&[1.0, 2.0, 3.0, 4.0], 2); assert_eq!(g.nodes, 2); assert_eq!(g.edges.len(), 4); } #[test] - fn test_graph_filters_non_positive() { - let logits = vec![1.0, -0.5, 0.0, 2.0]; - let g = graph_from_logits(&logits, 2); - // Only (0,0)=1.0 and (1,1)=2.0 survive + fn test_filters_non_positive() { + let g = graph_from_logits(&[1.0, -0.5, 0.0, 2.0], 2); assert_eq!(g.edges.len(), 2); - assert_eq!(g.edges[0].src, 0); - assert_eq!(g.edges[0].dst, 0); - assert_eq!(g.edges[1].src, 1); - assert_eq!(g.edges[1].dst, 1); } #[test] #[should_panic(expected = "logits length must equal seq_len^2")] - fn test_graph_mismatched_length() { - graph_from_logits(&[1.0, 2.0], 3); - } + fn test_mismatched_length() { graph_from_logits(&[1.0, 2.0], 3); } #[test] - fn test_graph_empty() { - let logits = vec![-1.0; 9]; - let g = graph_from_logits(&logits, 3); - assert_eq!(g.nodes, 3); + fn test_empty_graph() { + let g = graph_from_logits(&[-1.0; 9], 3); assert!(g.edges.is_empty()); } } diff --git a/crates/ruvector-attn-mincut/src/mincut.rs b/crates/ruvector-attn-mincut/src/mincut.rs index 390276a03..57409214a 100644 --- a/crates/ruvector-attn-mincut/src/mincut.rs +++ b/crates/ruvector-attn-mincut/src/mincut.rs @@ -1,43 +1,28 @@ -use crate::graph::AttentionGraph; +use crate::graph::{AttentionGraph, Edge}; use serde::{Deserialize, Serialize}; +use std::collections::VecDeque; -/// Result of a min-cut computation. +/// Result of a single s-t min-cut. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CutResult { - /// Edges that belong to the cut (src, dst). pub cut_edges: Vec<(usize, usize)>, - /// Total weight of the cut. pub cut_cost: f32, - /// Per-edge mask: true = keep, false = gated. pub keep_mask: Vec, } -/// Aggregated gating decision produced by `dynamic_min_cut`. +/// Aggregated gating decision from `dynamic_min_cut`. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GatingResult { - /// Per-edge keep mask aligned to the flattened `seq_len x seq_len` matrix. pub keep_mask: Vec, - /// Total cost of the min-cut. pub cut_cost: f32, - /// Number of edges kept (not gated). pub edges_kept: usize, - /// Total number of edges in the logit matrix. pub edges_total: usize, } -// --------------------------------------------------------------------------- -// Dinic's max-flow / min-cut -// --------------------------------------------------------------------------- - -/// Internal adjacency-list edge for the residual graph. #[derive(Debug, Clone)] -struct FlowEdge { - to: usize, - rev: usize, // index of reverse edge in adj[to] - cap: f32, -} +struct FlowEdge { to: usize, rev: usize, cap: f32 } -/// Dinic's max-flow solver operating on a capacity graph. +/// Dinic's max-flow solver for s-t min-cut on an attention graph. pub struct DinicSolver { adj: Vec>, level: Vec, @@ -46,57 +31,37 @@ pub struct DinicSolver { impl DinicSolver { fn new(n: usize) -> Self { - Self { - adj: vec![Vec::new(); n], - level: vec![0; n], - iter: vec![0; n], - } + Self { adj: vec![Vec::new(); n], level: vec![0; n], iter: vec![0; n] } } fn add_edge(&mut self, from: usize, to: usize, cap: f32) { - let rev_from = self.adj[to].len(); - let rev_to = self.adj[from].len(); - self.adj[from].push(FlowEdge { - to, - rev: rev_from, - cap, - }); - self.adj[to].push(FlowEdge { - to: from, - rev: rev_to, - cap: 0.0, - }); + let (rf, rt) = (self.adj[to].len(), self.adj[from].len()); + self.adj[from].push(FlowEdge { to, rev: rf, cap }); + self.adj[to].push(FlowEdge { to: from, rev: rt, cap: 0.0 }); } - /// BFS to build level graph from source. - fn bfs(&mut self, s: usize) -> bool { + fn bfs(&mut self, s: usize) { self.level.fill(-1); - let mut queue = std::collections::VecDeque::new(); self.level[s] = 0; - queue.push_back(s); - while let Some(v) = queue.pop_front() { + let mut q = VecDeque::new(); + q.push_back(s); + while let Some(v) = q.pop_front() { for e in &self.adj[v] { if e.cap > 0.0 && self.level[e.to] < 0 { self.level[e.to] = self.level[v] + 1; - queue.push_back(e.to); + q.push_back(e.to); } } } - self.level[s] >= 0 // always true; check sink reachability externally } - /// DFS to push flow along blocking paths. fn dfs(&mut self, v: usize, t: usize, f: f32) -> f32 { - if v == t { - return f; - } + if v == t { return f; } while self.iter[v] < self.adj[v].len() { let i = self.iter[v]; - let to = self.adj[v][i].to; - let cap = self.adj[v][i].cap; + let (to, cap) = (self.adj[v][i].to, self.adj[v][i].cap); if cap > 0.0 && self.level[v] < self.level[to] { - let bottleneck = if f < cap { f } else { cap }; - let d = self.dfs(to, t, bottleneck); + let d = self.dfs(to, t, f.min(cap)); if d > 0.0 { self.adj[v][i].cap -= d; let rev = self.adj[v][i].rev; @@ -112,134 +77,59 @@ impl DinicSolver { /// Compute s-t min-cut on the given attention graph. pub fn min_cut(&mut self, graph: &AttentionGraph, s: usize, t: usize) -> CutResult { assert!(s < graph.nodes && t < graph.nodes && s != t); - - // Build residual graph: nodes 0..graph.nodes *self = Self::new(graph.nodes); + for edge in &graph.edges { self.add_edge(edge.src, edge.dst, edge.weight); } - // Map: (edge_index_in_graph) -> index in adj[src] - let mut edge_adj_idx: Vec<(usize, usize)> = Vec::with_capacity(graph.edges.len()); - for edge in &graph.edges { - let idx = self.adj[edge.src].len(); - self.add_edge(edge.src, edge.dst, edge.weight); - edge_adj_idx.push((edge.src, idx)); - } - - // Dinic's main loop let inf = f32::MAX / 2.0; - let mut _total_flow = 0.0f32; - while { + loop { self.bfs(s); - self.level[t] >= 0 - } { + if self.level[t] < 0 { break; } self.iter.fill(0); - loop { - let f = self.dfs(s, t, inf); - if f <= 0.0 { - break; - } - _total_flow += f; - } + while self.dfs(s, t, inf) > 0.0 {} } - // After max-flow, do one final BFS to find reachable set from s + // Final BFS to find S-side of the cut self.bfs(s); - - // Edges crossing from reachable to non-reachable form the min-cut let mut cut_edges = Vec::new(); let mut cut_cost = 0.0f32; let mut keep_mask = vec![true; graph.edges.len()]; - - for (idx, edge) in graph.edges.iter().enumerate() { - let s_side = self.level[edge.src] >= 0; - let t_side = self.level[edge.dst] < 0; - if s_side && t_side { - cut_edges.push((edge.src, edge.dst)); - cut_cost += edge.weight; + for (idx, e) in graph.edges.iter().enumerate() { + if self.level[e.src] >= 0 && self.level[e.dst] < 0 { + cut_edges.push((e.src, e.dst)); + cut_cost += e.weight; keep_mask[idx] = false; } } - - CutResult { - cut_edges, - cut_cost, - keep_mask, - } + CutResult { cut_edges, cut_cost, keep_mask } } } -/// Compute dynamic min-cut gating over a flattened logit matrix. -/// -/// For each pair of nodes `(s, t)` where `s != t`, we compute the min-cut and -/// combine results: an edge is gated if it appears in any min-cut whose cost -/// is below `lambda * mean_weight`. The `eps` parameter sets a floor on edge -/// weights to avoid numerical issues. -pub fn dynamic_min_cut( - logits: &[f32], - seq_len: usize, - lambda: f32, - _tau: usize, - eps: f32, -) -> GatingResult { +/// Compute dynamic min-cut gating over a flattened `seq_len x seq_len` logit matrix. +pub fn dynamic_min_cut(logits: &[f32], seq_len: usize, lambda: f32, _tau: usize, eps: f32) -> GatingResult { assert_eq!(logits.len(), seq_len * seq_len); - - let edges_total = seq_len * seq_len; - - // Clamp logits: replace values below eps with 0 to sparsify - let clamped: Vec = logits - .iter() - .map(|&v| if v > eps { v } else { 0.0 }) - .collect(); - + let n = seq_len * seq_len; + let clamped: Vec = logits.iter().map(|&v| if v > eps { v } else { 0.0 }).collect(); let graph = crate::graph::graph_from_logits(&clamped, seq_len); if graph.edges.is_empty() || seq_len < 2 { - return GatingResult { - keep_mask: vec![false; edges_total], - cut_cost: 0.0, - edges_kept: 0, - edges_total, - }; + return GatingResult { keep_mask: vec![false; n], cut_cost: 0.0, edges_kept: 0, edges_total: n }; } - // Compute mean edge weight for thresholding - let mean_w: f32 = graph.edges.iter().map(|e| e.weight).sum::() - / graph.edges.len() as f32; + let mean_w: f32 = graph.edges.iter().map(|e| e.weight).sum::() / graph.edges.len() as f32; let threshold = lambda * mean_w; - - // Flat keep mask over seq_len x seq_len - let mut flat_keep = vec![true; edges_total]; + let mut flat_keep = vec![true; n]; let mut total_cut_cost = 0.0f32; - // Use node 0 as source and node (seq_len-1) as sink for the primary cut - let s = 0; - let t = seq_len - 1; - let mut solver = DinicSolver::new(seq_len); - let result = solver.min_cut(&graph, s, t); - + let result = solver.min_cut(&graph, 0, seq_len - 1); if result.cut_cost <= threshold { total_cut_cost += result.cut_cost; - // Mark cut edges in the flat matrix - for &(src, dst) in &result.cut_edges { - flat_keep[src * seq_len + dst] = false; - } - } - - // Also gate entries that were clamped to zero - for i in 0..edges_total { - if clamped[i] <= 0.0 { - flat_keep[i] = false; - } + for &(s, d) in &result.cut_edges { flat_keep[s * seq_len + d] = false; } } + for i in 0..n { if clamped[i] <= 0.0 { flat_keep[i] = false; } } let edges_kept = flat_keep.iter().filter(|&&k| k).count(); - - GatingResult { - keep_mask: flat_keep, - cut_cost: total_cut_cost, - edges_kept, - edges_total, - } + GatingResult { keep_mask: flat_keep, cut_cost: total_cut_cost, edges_kept, edges_total: n } } #[cfg(test)] @@ -248,77 +138,43 @@ mod tests { #[test] fn test_dinic_simple_cut() { - // 4-node graph: - // 0 --5--> 1 --3--> 3 - // 0 --4--> 2 --6--> 3 - // 1 --2--> 2 - // Min-cut from 0 to 3 should be 7 (cut edges 1->3=3, 2->3=6? No.) - // Actually: max-flow = min(5+4, 3+6) but with bottleneck path analysis: - // path 0->1->3: 3 - // path 0->2->3: 4 (bottleneck at 0->2=4, 2->3=6 -> 4) - // path 0->1->2->3: remaining cap 0->1=2, 1->2=2, 2->3=2 -> 2 - // total = 3+4+2 = 9? Let's just verify the solver runs. let graph = AttentionGraph { nodes: 4, edges: vec![ - crate::graph::Edge { src: 0, dst: 1, weight: 5.0 }, - crate::graph::Edge { src: 0, dst: 2, weight: 4.0 }, - crate::graph::Edge { src: 1, dst: 3, weight: 3.0 }, - crate::graph::Edge { src: 2, dst: 3, weight: 6.0 }, - crate::graph::Edge { src: 1, dst: 2, weight: 2.0 }, + Edge { src: 0, dst: 1, weight: 5.0 }, Edge { src: 0, dst: 2, weight: 4.0 }, + Edge { src: 1, dst: 3, weight: 3.0 }, Edge { src: 2, dst: 3, weight: 6.0 }, + Edge { src: 1, dst: 2, weight: 2.0 }, ], }; - let mut solver = DinicSolver::new(4); - let result = solver.min_cut(&graph, 0, 3); - - // The min-cut value equals max-flow. Paths: - // 0->1->3: pushes 3 (bottleneck 1->3) - // 0->2->3: pushes 4 (bottleneck 0->2) - // 0->1->2->3: pushes 2 (bottleneck 1->2, remaining 0->1=2) - // Total max-flow = 9, so cut_cost should be 9. - assert!((result.cut_cost - 9.0).abs() < 0.01); + let r = solver.min_cut(&graph, 0, 3); + assert!((r.cut_cost - 9.0).abs() < 0.01); } #[test] fn test_dinic_two_node() { - let graph = AttentionGraph { - nodes: 2, - edges: vec![crate::graph::Edge { src: 0, dst: 1, weight: 3.5 }], - }; + let graph = AttentionGraph { nodes: 2, edges: vec![Edge { src: 0, dst: 1, weight: 3.5 }] }; let mut solver = DinicSolver::new(2); - let result = solver.min_cut(&graph, 0, 1); - assert!((result.cut_cost - 3.5).abs() < 0.01); - assert_eq!(result.cut_edges.len(), 1); - assert!(!result.keep_mask[0]); + let r = solver.min_cut(&graph, 0, 1); + assert!((r.cut_cost - 3.5).abs() < 0.01); + assert!(!r.keep_mask[0]); } #[test] - fn test_dynamic_min_cut_basic() { - // 3x3 logits with clear structure - let logits = vec![ - 1.0, 0.5, 0.0, - 0.0, 1.0, 0.5, - 0.0, 0.0, 1.0, - ]; - let result = dynamic_min_cut(&logits, 3, 0.5, 2, 0.01); - assert_eq!(result.edges_total, 9); - assert_eq!(result.keep_mask.len(), 9); - // Some edges should be kept - assert!(result.edges_kept > 0); + fn test_dynamic_basic() { + let logits = vec![1.0, 0.5, 0.0, 0.0, 1.0, 0.5, 0.0, 0.0, 1.0]; + let r = dynamic_min_cut(&logits, 3, 0.5, 2, 0.01); + assert_eq!(r.edges_total, 9); + assert!(r.edges_kept > 0); } #[test] - fn test_dynamic_min_cut_all_negative() { - let logits = vec![-1.0; 4]; - let result = dynamic_min_cut(&logits, 2, 0.5, 2, 0.01); - assert_eq!(result.edges_kept, 0); + fn test_dynamic_all_negative() { + assert_eq!(dynamic_min_cut(&[-1.0; 4], 2, 0.5, 2, 0.01).edges_kept, 0); } #[test] - fn test_dynamic_min_cut_single_token() { - let logits = vec![1.0]; - let result = dynamic_min_cut(&logits, 1, 0.5, 2, 0.01); - assert_eq!(result.edges_total, 1); + fn test_dynamic_single_token() { + assert_eq!(dynamic_min_cut(&[1.0], 1, 0.5, 2, 0.01).edges_total, 1); } } diff --git a/crates/ruvector-coherence/src/batch.rs b/crates/ruvector-coherence/src/batch.rs index b778a0e63..48cffa6c4 100644 --- a/crates/ruvector-coherence/src/batch.rs +++ b/crates/ruvector-coherence/src/batch.rs @@ -8,28 +8,15 @@ use crate::quality::quality_check; /// Aggregated results from evaluating a batch of baseline/gated output pairs. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BatchResult { - /// Mean coherence delta across all samples. pub mean_coherence_delta: f64, - /// Standard deviation of coherence delta values. pub std_coherence_delta: f64, - /// Lower bound of 95% confidence interval for coherence delta. pub ci_95_lower: f64, - /// Upper bound of 95% confidence interval for coherence delta. pub ci_95_upper: f64, - /// Number of samples evaluated. pub n_samples: usize, - /// Fraction of samples that pass the quality threshold. pub pass_rate: f64, } -/// Evaluates a batch of baseline/gated output pairs and produces aggregate statistics. -/// -/// Each pair `(baseline_outputs[i], gated_outputs[i])` is evaluated for its -/// coherence delta (via [`delta_behavior`]) and quality (via [`quality_check`]). -/// -/// # Panics -/// -/// Does not panic; returns zeroed results when inputs are empty. +/// Evaluates a batch of output pairs, producing mean/std/CI for coherence delta and pass rate. pub fn evaluate_batch( baseline_outputs: &[Vec], gated_outputs: &[Vec], @@ -38,46 +25,31 @@ pub fn evaluate_batch( let n = baseline_outputs.len().min(gated_outputs.len()); if n == 0 { return BatchResult { - mean_coherence_delta: 0.0, - std_coherence_delta: 0.0, - ci_95_lower: 0.0, - ci_95_upper: 0.0, - n_samples: 0, - pass_rate: 0.0, + mean_coherence_delta: 0.0, std_coherence_delta: 0.0, + ci_95_lower: 0.0, ci_95_upper: 0.0, n_samples: 0, pass_rate: 0.0, }; } let mut deltas = Vec::with_capacity(n); let mut passes = 0usize; - for i in 0..n { - let dm = delta_behavior(&baseline_outputs[i], &gated_outputs[i]); - deltas.push(dm.coherence_delta); - - let qr = quality_check(&baseline_outputs[i], &gated_outputs[i], threshold); - if qr.passes_threshold { + deltas.push(delta_behavior(&baseline_outputs[i], &gated_outputs[i]).coherence_delta); + if quality_check(&baseline_outputs[i], &gated_outputs[i], threshold).passes_threshold { passes += 1; } } let mean = deltas.iter().sum::() / n as f64; - let variance = if n > 1 { - deltas.iter().map(|d| (d - mean) * (d - mean)).sum::() / (n - 1) as f64 - } else { - 0.0 - }; - let std_dev = variance.sqrt(); - - // 95% CI using z = 1.96 (normal approximation). + let var = if n > 1 { + deltas.iter().map(|d| (d - mean).powi(2)).sum::() / (n - 1) as f64 + } else { 0.0 }; + let std_dev = var.sqrt(); let margin = 1.96 * std_dev / (n as f64).sqrt(); BatchResult { - mean_coherence_delta: mean, - std_coherence_delta: std_dev, - ci_95_lower: mean - margin, - ci_95_upper: mean + margin, - n_samples: n, - pass_rate: passes as f64 / n as f64, + mean_coherence_delta: mean, std_coherence_delta: std_dev, + ci_95_lower: mean - margin, ci_95_upper: mean + margin, + n_samples: n, pass_rate: passes as f64 / n as f64, } } @@ -87,75 +59,43 @@ mod tests { #[test] fn batch_empty() { - let result = evaluate_batch(&[], &[], 0.9); - assert_eq!(result.n_samples, 0); - assert_eq!(result.pass_rate, 0.0); + let r = evaluate_batch(&[], &[], 0.9); + assert_eq!(r.n_samples, 0); } #[test] - fn batch_identical_outputs() { - let baselines = vec![vec![1.0, 2.0, 3.0]; 10]; - let gated = baselines.clone(); - let result = evaluate_batch(&baselines, &gated, 0.9); - assert_eq!(result.n_samples, 10); - assert!((result.mean_coherence_delta).abs() < 1e-10); - assert!((result.std_coherence_delta).abs() < 1e-10); - assert!((result.pass_rate - 1.0).abs() < 1e-10); + fn batch_identical() { + let bl = vec![vec![1.0, 2.0, 3.0]; 10]; + let r = evaluate_batch(&bl, &bl.clone(), 0.9); + assert_eq!(r.n_samples, 10); + assert!(r.mean_coherence_delta.abs() < 1e-10); + assert!((r.pass_rate - 1.0).abs() < 1e-10); } #[test] fn batch_ci_contains_mean() { - let baselines = vec![ - vec![1.0, 0.0], - vec![0.0, 1.0], - vec![1.0, 1.0], - vec![2.0, 3.0], - ]; - let gated = vec![ - vec![1.1, 0.1], - vec![0.1, 1.1], - vec![1.2, 0.9], - vec![2.1, 2.9], - ]; - let result = evaluate_batch(&baselines, &gated, 0.9); - assert_eq!(result.n_samples, 4); - assert!(result.ci_95_lower <= result.mean_coherence_delta); - assert!(result.ci_95_upper >= result.mean_coherence_delta); - } - - #[test] - fn batch_single_sample() { - let baselines = vec![vec![1.0, 2.0]]; - let gated = vec![vec![1.0, 2.0]]; - let result = evaluate_batch(&baselines, &gated, 0.5); - assert_eq!(result.n_samples, 1); - assert!((result.std_coherence_delta).abs() < 1e-10); - assert!((result.pass_rate - 1.0).abs() < 1e-10); + let bl = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0], vec![2.0, 3.0]]; + let gt = vec![vec![1.1, 0.1], vec![0.1, 1.1], vec![1.2, 0.9], vec![2.1, 2.9]]; + let r = evaluate_batch(&bl, >, 0.9); + assert!(r.ci_95_lower <= r.mean_coherence_delta); + assert!(r.ci_95_upper >= r.mean_coherence_delta); } #[test] fn batch_pass_rate_partial() { - // Two samples: one passes (similar), one fails (orthogonal). - let baselines = vec![vec![1.0, 0.0], vec![1.0, 0.0]]; - let gated = vec![vec![1.0, 0.0], vec![0.0, 1.0]]; - let result = evaluate_batch(&baselines, &gated, 0.5); - assert_eq!(result.n_samples, 2); - assert!((result.pass_rate - 0.5).abs() < 1e-10); + let bl = vec![vec![1.0, 0.0], vec![1.0, 0.0]]; + let gt = vec![vec![1.0, 0.0], vec![0.0, 1.0]]; + let r = evaluate_batch(&bl, >, 0.5); + assert!((r.pass_rate - 0.5).abs() < 1e-10); } #[test] fn batch_result_serializable() { - let result = BatchResult { - mean_coherence_delta: -0.05, - std_coherence_delta: 0.02, - ci_95_lower: -0.07, - ci_95_upper: -0.03, - n_samples: 100, - pass_rate: 0.95, + let r = BatchResult { + mean_coherence_delta: -0.05, std_coherence_delta: 0.02, + ci_95_lower: -0.07, ci_95_upper: -0.03, n_samples: 100, pass_rate: 0.95, }; - let json = serde_json::to_string(&result).unwrap(); - let deser: BatchResult = serde_json::from_str(&json).unwrap(); - assert_eq!(deser.n_samples, 100); - assert!((deser.pass_rate - 0.95).abs() < 1e-10); + let d: BatchResult = serde_json::from_str(&serde_json::to_string(&r).unwrap()).unwrap(); + assert_eq!(d.n_samples, 100); } } diff --git a/crates/ruvector-coherence/src/comparison.rs b/crates/ruvector-coherence/src/comparison.rs index 90ee3c92e..4a48c4eab 100644 --- a/crates/ruvector-coherence/src/comparison.rs +++ b/crates/ruvector-coherence/src/comparison.rs @@ -5,180 +5,83 @@ use serde::{Deserialize, Serialize}; /// Result of comparing two attention masks. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ComparisonResult { - /// Jaccard similarity coefficient between the two masks. pub jaccard: f64, - /// Number of positions where one mask has an edge and the other does not. pub edge_flips: usize, - /// Total active edges in the baseline mask. pub baseline_edges: usize, - /// Total active edges in the gated mask. pub gated_edges: usize, - /// Ratio of gated sparsity to baseline sparsity. - /// Values > 1.0 mean the gated mask is denser; < 1.0 means sparser. pub sparsity_ratio: f64, } -/// Computes the Jaccard similarity coefficient between two boolean masks. -/// -/// `J(A, B) = |A intersection B| / |A union B|` -/// -/// Returns `1.0` when both masks are empty (vacuously similar). +/// Jaccard similarity: `|A & B| / |A | B|`. Returns `1.0` for two empty masks. pub fn jaccard_similarity(mask_a: &[bool], mask_b: &[bool]) -> f64 { let n = mask_a.len().min(mask_b.len()); - let mut intersection = 0usize; - let mut union = 0usize; + let (mut inter, mut union) = (0usize, 0usize); for i in 0..n { - let a = mask_a[i]; - let b = mask_b[i]; - if a || b { - union += 1; - } - if a && b { - intersection += 1; - } + if mask_a[i] || mask_b[i] { union += 1; } + if mask_a[i] && mask_b[i] { inter += 1; } } - // Count remaining elements beyond the shorter slice as union-only. - if mask_a.len() > n { - union += mask_a[n..].iter().filter(|&&v| v).count(); - } - if mask_b.len() > n { - union += mask_b[n..].iter().filter(|&&v| v).count(); - } - if union == 0 { - return 1.0; - } - intersection as f64 / union as f64 + union += count_true_tail(mask_a, n) + count_true_tail(mask_b, n); + if union == 0 { 1.0 } else { inter as f64 / union as f64 } } -/// Counts positions where the two masks disagree (one true, the other false). +/// Counts positions where the two masks disagree. pub fn edge_flip_count(mask_a: &[bool], mask_b: &[bool]) -> usize { let n = mask_a.len().min(mask_b.len()); - let mut flips = 0usize; - for i in 0..n { - if mask_a[i] != mask_b[i] { - flips += 1; - } - } - // Positions beyond the shorter mask count as flips if the longer mask is true. - if mask_a.len() > n { - flips += mask_a[n..].iter().filter(|&&v| v).count(); - } - if mask_b.len() > n { - flips += mask_b[n..].iter().filter(|&&v| v).count(); - } + let mut flips = (0..n).filter(|&i| mask_a[i] != mask_b[i]).count(); + flips += count_true_tail(mask_a, n) + count_true_tail(mask_b, n); flips } -/// Performs a full comparison of two attention masks. +/// Full comparison of two attention masks. pub fn compare_attention_masks(baseline: &[bool], gated: &[bool]) -> ComparisonResult { - let jaccard = jaccard_similarity(baseline, gated); - let edge_flips = edge_flip_count(baseline, gated); let baseline_edges = baseline.iter().filter(|&&v| v).count(); let gated_edges = gated.iter().filter(|&&v| v).count(); let total = baseline.len().max(gated.len()); - let baseline_sparsity = if total > 0 { - 1.0 - (baseline_edges as f64 / total as f64) - } else { - 1.0 - }; - let gated_sparsity = if total > 0 { - 1.0 - (gated_edges as f64 / total as f64) - } else { - 1.0 - }; - let sparsity_ratio = if baseline_sparsity > f64::EPSILON { - gated_sparsity / baseline_sparsity - } else { - gated_sparsity - }; + let bl_sp = if total > 0 { 1.0 - baseline_edges as f64 / total as f64 } else { 1.0 }; + let gt_sp = if total > 0 { 1.0 - gated_edges as f64 / total as f64 } else { 1.0 }; ComparisonResult { - jaccard, - edge_flips, + jaccard: jaccard_similarity(baseline, gated), + edge_flips: edge_flip_count(baseline, gated), baseline_edges, gated_edges, - sparsity_ratio, + sparsity_ratio: if bl_sp > f64::EPSILON { gt_sp / bl_sp } else { gt_sp }, } } +fn count_true_tail(mask: &[bool], from: usize) -> usize { + if mask.len() > from { mask[from..].iter().filter(|&&v| v).count() } else { 0 } +} + #[cfg(test)] mod tests { use super::*; #[test] - fn jaccard_identical() { - let mask = vec![true, false, true, true]; - assert!((jaccard_similarity(&mask, &mask) - 1.0).abs() < 1e-10); - } - - #[test] - fn jaccard_disjoint() { - let a = vec![true, false, true, false]; - let b = vec![false, true, false, true]; - assert!(jaccard_similarity(&a, &b).abs() < 1e-10); - } - - #[test] - fn jaccard_empty() { - let empty: Vec = vec![]; - assert_eq!(jaccard_similarity(&empty, &empty), 1.0); - } - - #[test] - fn jaccard_all_false() { - let a = vec![false, false, false]; - let b = vec![false, false, false]; - assert_eq!(jaccard_similarity(&a, &b), 1.0); - } - - #[test] - fn jaccard_partial_overlap() { - let a = vec![true, true, false, false]; - let b = vec![true, false, true, false]; - // intersection = 1 (pos 0), union = 3 (pos 0, 1, 2) + fn jaccard_cases() { + let m = vec![true, false, true, true]; + assert!((jaccard_similarity(&m, &m) - 1.0).abs() < 1e-10); + assert!(jaccard_similarity(&[true, false], &[false, true]).abs() < 1e-10); + assert_eq!(jaccard_similarity(&[], &[]), 1.0); + // partial: intersection=1, union=3 + let (a, b) = (vec![true, true, false, false], vec![true, false, true, false]); assert!((jaccard_similarity(&a, &b) - 1.0 / 3.0).abs() < 1e-10); } #[test] - fn edge_flip_count_identical() { - let mask = vec![true, false, true]; - assert_eq!(edge_flip_count(&mask, &mask), 0); + fn edge_flip_cases() { + assert_eq!(edge_flip_count(&[true, false], &[true, false]), 0); + assert_eq!(edge_flip_count(&[true, false, true], &[false, true, false]), 3); + assert_eq!(edge_flip_count(&[true, false], &[true, false, true, true]), 2); } #[test] - fn edge_flip_count_all_flipped() { - let a = vec![true, false, true]; - let b = vec![false, true, false]; - assert_eq!(edge_flip_count(&a, &b), 3); - } - - #[test] - fn edge_flip_count_different_lengths() { - let a = vec![true, false]; - let b = vec![true, false, true, true]; - // pos 0: same, pos 1: same, pos 2: flip, pos 3: flip - assert_eq!(edge_flip_count(&a, &b), 2); - } - - #[test] - fn compare_attention_masks_basic() { - let baseline = vec![true, true, false, false, true]; - let gated = vec![true, false, false, true, true]; - let result = compare_attention_masks(&baseline, &gated); - assert_eq!(result.baseline_edges, 3); - assert_eq!(result.gated_edges, 3); - assert_eq!(result.edge_flips, 2); - // intersection = 2 (pos 0, 4), union = 4 (pos 0, 1, 3, 4) - assert!((result.jaccard - 0.5).abs() < 1e-10); - } - - #[test] - fn compare_sparser_gated() { - let baseline = vec![true, true, true, true]; - let gated = vec![true, false, false, false]; - let result = compare_attention_masks(&baseline, &gated); - assert_eq!(result.baseline_edges, 4); - assert_eq!(result.gated_edges, 1); - // baseline_sparsity = 0, so sparsity_ratio = gated_sparsity - assert!(result.sparsity_ratio > 0.0); + fn compare_masks() { + let bl = vec![true, true, false, false, true]; + let gt = vec![true, false, false, true, true]; + let r = compare_attention_masks(&bl, >); + assert_eq!(r.baseline_edges, 3); + assert_eq!(r.gated_edges, 3); + assert_eq!(r.edge_flips, 2); + assert!((r.jaccard - 0.5).abs() < 1e-10); } } diff --git a/crates/ruvector-coherence/src/quality.rs b/crates/ruvector-coherence/src/quality.rs index a578b69f4..e52d8a3d2 100644 --- a/crates/ruvector-coherence/src/quality.rs +++ b/crates/ruvector-coherence/src/quality.rs @@ -5,73 +5,43 @@ use serde::{Deserialize, Serialize}; /// Result of a quality check comparing baseline and gated outputs. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QualityResult { - /// Cosine similarity between the two output vectors. pub cosine_sim: f64, - /// Euclidean (L2) distance between the two output vectors. pub l2_dist: f64, - /// Whether the cosine similarity meets or exceeds the threshold. pub passes_threshold: bool, } -/// Computes cosine similarity between two vectors. -/// -/// Returns `0.0` when either vector has zero magnitude. +/// Cosine similarity between two vectors. Returns `0.0` for zero-magnitude inputs. pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 { let n = a.len().min(b.len()); - let mut dot = 0.0_f64; - let mut norm_a_sq = 0.0_f64; - let mut norm_b_sq = 0.0_f64; + let (mut dot, mut na, mut nb) = (0.0_f64, 0.0_f64, 0.0_f64); for i in 0..n { - let ai = a[i] as f64; - let bi = b[i] as f64; + let (ai, bi) = (a[i] as f64, b[i] as f64); dot += ai * bi; - norm_a_sq += ai * ai; - norm_b_sq += bi * bi; + na += ai * ai; + nb += bi * bi; } - let denom = norm_a_sq.sqrt() * norm_b_sq.sqrt(); - if denom < f64::EPSILON { - return 0.0; - } - dot / denom + let denom = na.sqrt() * nb.sqrt(); + if denom < f64::EPSILON { 0.0 } else { dot / denom } } -/// Computes the Euclidean (L2) distance between two vectors. +/// Euclidean (L2) distance between two vectors. pub fn l2_distance(a: &[f32], b: &[f32]) -> f64 { let n = a.len().min(b.len()); - let mut sum_sq = 0.0_f64; + let mut s = 0.0_f64; for i in 0..n { - let diff = (a[i] as f64) - (b[i] as f64); - sum_sq += diff * diff; + let d = a[i] as f64 - b[i] as f64; + s += d * d; } - // Account for extra dimensions in the longer vector. - if a.len() > n { - for &v in &a[n..] { - sum_sq += (v as f64) * (v as f64); - } - } - if b.len() > n { - for &v in &b[n..] { - sum_sq += (v as f64) * (v as f64); - } - } - sum_sq.sqrt() + if a.len() > n { s += a[n..].iter().map(|v| (*v as f64).powi(2)).sum::(); } + if b.len() > n { s += b[n..].iter().map(|v| (*v as f64).powi(2)).sum::(); } + s.sqrt() } -/// Checks whether gated output quality is acceptable relative to the baseline. -/// -/// The check passes when `cosine_similarity(baseline, gated) >= threshold`. -pub fn quality_check( - baseline_output: &[f32], - gated_output: &[f32], - threshold: f64, -) -> QualityResult { +/// Quality gate: passes when `cosine_similarity >= threshold`. +pub fn quality_check(baseline_output: &[f32], gated_output: &[f32], threshold: f64) -> QualityResult { let cosine_sim = cosine_similarity(baseline_output, gated_output); let l2_dist = l2_distance(baseline_output, gated_output); - QualityResult { - cosine_sim, - l2_dist, - passes_threshold: cosine_sim >= threshold, - } + QualityResult { cosine_sim, l2_dist, passes_threshold: cosine_sim >= threshold } } #[cfg(test)] @@ -79,82 +49,33 @@ mod tests { use super::*; #[test] - fn cosine_identical() { - let v = vec![1.0, 2.0, 3.0]; - assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-10); + fn cosine_cases() { + assert!((cosine_similarity(&[1.0, 2.0, 3.0], &[1.0, 2.0, 3.0]) - 1.0).abs() < 1e-10); + assert!((cosine_similarity(&[1.0, 0.0], &[-1.0, 0.0]) + 1.0).abs() < 1e-10); + assert!(cosine_similarity(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-10); + assert_eq!(cosine_similarity(&[0.0, 0.0], &[1.0, 2.0]), 0.0); } #[test] - fn cosine_opposite() { - let a = vec![1.0, 0.0]; - let b = vec![-1.0, 0.0]; - assert!((cosine_similarity(&a, &b) + 1.0).abs() < 1e-10); + fn l2_cases() { + assert!(l2_distance(&[1.0, 2.0], &[1.0, 2.0]) < 1e-10); + assert!((l2_distance(&[0.0, 0.0], &[3.0, 4.0]) - 5.0).abs() < 1e-10); + assert!((l2_distance(&[1.0], &[1.0, 3.0]) - 3.0).abs() < 1e-10); } #[test] - fn cosine_orthogonal() { - let a = vec![1.0, 0.0]; - let b = vec![0.0, 1.0]; - assert!(cosine_similarity(&a, &b).abs() < 1e-10); - } - - #[test] - fn cosine_zero_vector() { - let a = vec![0.0, 0.0]; - let b = vec![1.0, 2.0]; - assert_eq!(cosine_similarity(&a, &b), 0.0); - } - - #[test] - fn l2_distance_zero() { - let v = vec![1.0, 2.0, 3.0]; - assert!(l2_distance(&v, &v) < 1e-10); - } - - #[test] - fn l2_distance_known() { - let a = vec![0.0, 0.0]; - let b = vec![3.0, 4.0]; - assert!((l2_distance(&a, &b) - 5.0).abs() < 1e-10); - } - - #[test] - fn l2_distance_different_lengths() { - let a = vec![1.0]; - let b = vec![1.0, 3.0]; - // diff at pos 0: 0, extra in b: 3.0 => sqrt(0 + 9) = 3.0 - assert!((l2_distance(&a, &b) - 3.0).abs() < 1e-10); - } - - #[test] - fn quality_check_passes() { - let baseline = vec![1.0, 2.0, 3.0]; - let gated = vec![1.1, 2.1, 3.1]; - let result = quality_check(&baseline, &gated, 0.99); - assert!(result.passes_threshold); - assert!(result.cosine_sim > 0.99); - assert!(result.l2_dist > 0.0); - } - - #[test] - fn quality_check_fails() { - let baseline = vec![1.0, 0.0]; - let gated = vec![0.0, 1.0]; - let result = quality_check(&baseline, &gated, 0.5); - assert!(!result.passes_threshold); - assert!(result.cosine_sim < 0.5); + fn quality_check_pass_and_fail() { + let r = quality_check(&[1.0, 2.0, 3.0], &[1.1, 2.1, 3.1], 0.99); + assert!(r.passes_threshold); + let r2 = quality_check(&[1.0, 0.0], &[0.0, 1.0], 0.5); + assert!(!r2.passes_threshold); } #[test] fn quality_result_serializable() { - let result = QualityResult { - cosine_sim: 0.95, - l2_dist: 0.32, - passes_threshold: true, - }; - let json = serde_json::to_string(&result).unwrap(); - let deser: QualityResult = serde_json::from_str(&json).unwrap(); - assert!((deser.cosine_sim - 0.95).abs() < 1e-10); - assert!(deser.passes_threshold); + let r = QualityResult { cosine_sim: 0.95, l2_dist: 0.32, passes_threshold: true }; + let j = serde_json::to_string(&r).unwrap(); + let d: QualityResult = serde_json::from_str(&j).unwrap(); + assert!((d.cosine_sim - 0.95).abs() < 1e-10); } } diff --git a/crates/ruvector-profiler/src/latency.rs b/crates/ruvector-profiler/src/latency.rs index 07801cbf0..4b508108a 100644 --- a/crates/ruvector-profiler/src/latency.rs +++ b/crates/ruvector-profiler/src/latency.rs @@ -1,4 +1,3 @@ -/// A single latency measurement for one forward pass / kernel invocation. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct LatencyRecord { pub sample_id: usize, @@ -7,7 +6,6 @@ pub struct LatencyRecord { pub seq_len: usize, } -/// Descriptive statistics over a collection of latency records. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct LatencyStats { pub p50_us: u64, @@ -18,98 +16,45 @@ pub struct LatencyStats { pub n: usize, } -/// Compute percentile and summary statistics from [`LatencyRecord`]s. -/// -/// Uses `wall_time_us` for all calculations. Returns zeroed stats when -/// the input slice is empty. +/// Compute percentile and summary statistics from wall-time latencies. pub fn compute_latency_stats(records: &[LatencyRecord]) -> LatencyStats { let n = records.len(); if n == 0 { - return LatencyStats { - p50_us: 0, - p95_us: 0, - p99_us: 0, - mean_us: 0.0, - std_us: 0.0, - n: 0, - }; + return LatencyStats { p50_us: 0, p95_us: 0, p99_us: 0, mean_us: 0.0, std_us: 0.0, n: 0 }; } - let mut times: Vec = records.iter().map(|r| r.wall_time_us).collect(); times.sort_unstable(); - - let mean = times.iter().copied().sum::() as f64 / n as f64; - let variance = times.iter().map(|&t| (t as f64 - mean).powi(2)).sum::() / n as f64; - let std = variance.sqrt(); - + let mean = times.iter().sum::() as f64 / n as f64; + let var = times.iter().map(|&t| (t as f64 - mean).powi(2)).sum::() / n as f64; LatencyStats { - p50_us: percentile(×, 50.0), - p95_us: percentile(×, 95.0), - p99_us: percentile(×, 99.0), - mean_us: mean, - std_us: std, - n, + p50_us: pctl(×, 50.0), p95_us: pctl(×, 95.0), p99_us: pctl(×, 99.0), + mean_us: mean, std_us: var.sqrt(), n, } } -/// Nearest-rank percentile on a **sorted** slice. -fn percentile(sorted: &[u64], pct: f64) -> u64 { - if sorted.is_empty() { - return 0; - } - let rank = (pct / 100.0 * sorted.len() as f64).ceil() as usize; - let idx = rank.min(sorted.len()).saturating_sub(1); +fn pctl(sorted: &[u64], p: f64) -> u64 { + let idx = ((p / 100.0 * sorted.len() as f64).ceil() as usize).min(sorted.len()).saturating_sub(1); sorted[idx] } #[cfg(test)] mod tests { use super::*; - - fn make_records(times: &[u64]) -> Vec { - times - .iter() - .enumerate() - .map(|(i, &t)| LatencyRecord { - sample_id: i, - wall_time_us: t, - kernel_time_us: t, - seq_len: 128, - }) - .collect() + fn recs(ts: &[u64]) -> Vec { + ts.iter().enumerate().map(|(i, &t)| LatencyRecord { + sample_id: i, wall_time_us: t, kernel_time_us: t, seq_len: 128, + }).collect() } - #[test] - fn stats_empty() { - let s = compute_latency_stats(&[]); - assert_eq!(s.n, 0); - assert_eq!(s.p50_us, 0); + #[test] fn empty() { assert_eq!(compute_latency_stats(&[]).n, 0); } + #[test] fn single() { + let s = compute_latency_stats(&recs(&[42])); + assert_eq!((s.p50_us, s.p99_us, s.n), (42, 42, 1)); } - - #[test] - fn stats_single() { - let recs = make_records(&[42]); - let s = compute_latency_stats(&recs); - assert_eq!(s.n, 1); - assert_eq!(s.p50_us, 42); - assert_eq!(s.p99_us, 42); - assert!((s.mean_us - 42.0).abs() < 1e-9); - assert!((s.std_us).abs() < 1e-9); - } - - #[test] - fn stats_multiple() { - let recs = make_records(&[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]); - let s = compute_latency_stats(&recs); - assert_eq!(s.n, 10); + #[test] fn multi() { + let s = compute_latency_stats(&recs(&[10,20,30,40,50,60,70,80,90,100])); assert_eq!(s.p50_us, 50); assert!((s.mean_us - 55.0).abs() < 1e-9); } - - #[test] - fn stats_unsorted_input() { - let recs = make_records(&[100, 10, 50, 90, 20]); - let s = compute_latency_stats(&recs); - assert_eq!(s.p50_us, 50); - } + #[test] fn unsorted() { assert_eq!(compute_latency_stats(&recs(&[100,10,50,90,20])).p50_us, 50); } } diff --git a/crates/ruvector-profiler/src/memory.rs b/crates/ruvector-profiler/src/memory.rs index a6db71cae..131c727bb 100644 --- a/crates/ruvector-profiler/src/memory.rs +++ b/crates/ruvector-profiler/src/memory.rs @@ -1,6 +1,5 @@ use std::time::{SystemTime, UNIX_EPOCH}; -/// A point-in-time snapshot of memory usage. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct MemorySnapshot { pub peak_rss_bytes: u64, @@ -10,7 +9,6 @@ pub struct MemorySnapshot { pub timestamp_us: u64, } -/// Aggregated memory statistics produced by [`MemoryTracker::report`]. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct MemoryReport { pub label: String, @@ -20,20 +18,11 @@ pub struct MemoryReport { pub activation_total: u64, } -/// Captures a [`MemorySnapshot`] using OS-specific facilities. -/// -/// On Linux this reads `VmRSS` from `/proc/self/status`. On other -/// platforms a zero-valued fallback is returned so the crate still -/// compiles everywhere. +/// Capture current memory via /proc/self/status (Linux) or zero fallback. pub fn capture_memory() -> MemorySnapshot { - let rss = read_vm_rss(); - let ts = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_micros() as u64; - + let ts = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_micros() as u64; MemorySnapshot { - peak_rss_bytes: rss, + peak_rss_bytes: read_vm_rss(), kv_cache_bytes: 0, activation_bytes: 0, temp_buffer_bytes: 0, @@ -43,25 +32,17 @@ pub fn capture_memory() -> MemorySnapshot { #[cfg(target_os = "linux")] fn read_vm_rss() -> u64 { - if let Ok(status) = std::fs::read_to_string("/proc/self/status") { - for line in status.lines() { - if let Some(rest) = line.strip_prefix("VmRSS:") { - let trimmed = rest.trim().trim_end_matches("kB").trim(); - if let Ok(kb) = trimmed.parse::() { - return kb * 1024; - } - } - } - } - 0 + std::fs::read_to_string("/proc/self/status").ok().and_then(|s| { + s.lines() + .find(|l| l.starts_with("VmRSS:")) + .and_then(|l| l.trim_start_matches("VmRSS:").trim().trim_end_matches("kB").trim().parse::().ok()) + .map(|kb| kb * 1024) + }).unwrap_or(0) } #[cfg(not(target_os = "linux"))] -fn read_vm_rss() -> u64 { - 0 -} +fn read_vm_rss() -> u64 { 0 } -/// Collects a series of [`MemorySnapshot`]s under a named label. pub struct MemoryTracker { pub snapshots: Vec, pub label: String, @@ -69,40 +50,23 @@ pub struct MemoryTracker { impl MemoryTracker { pub fn new(label: &str) -> Self { - Self { - snapshots: Vec::new(), - label: label.to_string(), - } + Self { snapshots: Vec::new(), label: label.to_string() } } - /// Take a snapshot and append it to the internal buffer. - pub fn snapshot(&mut self) { - self.snapshots.push(capture_memory()); - } + pub fn snapshot(&mut self) { self.snapshots.push(capture_memory()); } - /// Return the peak RSS across all recorded snapshots. pub fn peak(&self) -> u64 { self.snapshots.iter().map(|s| s.peak_rss_bytes).max().unwrap_or(0) } - /// Produce an aggregated [`MemoryReport`]. pub fn report(&self) -> MemoryReport { - let n = self.snapshots.len() as u64; - let peak_rss = self.peak(); - let mean_rss = if n > 0 { - self.snapshots.iter().map(|s| s.peak_rss_bytes).sum::() / n - } else { - 0 - }; - let kv_cache_total = self.snapshots.iter().map(|s| s.kv_cache_bytes).sum(); - let activation_total = self.snapshots.iter().map(|s| s.activation_bytes).sum(); - + let n = self.snapshots.len().max(1) as u64; MemoryReport { label: self.label.clone(), - peak_rss, - mean_rss, - kv_cache_total, - activation_total, + peak_rss: self.peak(), + mean_rss: self.snapshots.iter().map(|s| s.peak_rss_bytes).sum::() / n, + kv_cache_total: self.snapshots.iter().map(|s| s.kv_cache_bytes).sum(), + activation_total: self.snapshots.iter().map(|s| s.activation_bytes).sum(), } } } @@ -112,39 +76,22 @@ mod tests { use super::*; #[test] - fn capture_returns_nonzero_timestamp() { - let snap = capture_memory(); - assert!(snap.timestamp_us > 0); - } + fn capture_returns_nonzero_timestamp() { assert!(capture_memory().timestamp_us > 0); } #[test] - fn tracker_peak_empty() { - let t = MemoryTracker::new("empty"); - assert_eq!(t.peak(), 0); - } + fn tracker_peak_empty() { assert_eq!(MemoryTracker::new("x").peak(), 0); } #[test] fn tracker_report_aggregates() { let mut t = MemoryTracker::new("test"); - t.snapshots.push(MemorySnapshot { - peak_rss_bytes: 100, - kv_cache_bytes: 10, - activation_bytes: 20, - temp_buffer_bytes: 5, - timestamp_us: 1, - }); - t.snapshots.push(MemorySnapshot { - peak_rss_bytes: 200, - kv_cache_bytes: 30, - activation_bytes: 40, - temp_buffer_bytes: 15, - timestamp_us: 2, - }); + let mk = |rss, kv, act| MemorySnapshot { + peak_rss_bytes: rss, kv_cache_bytes: kv, activation_bytes: act, + temp_buffer_bytes: 0, timestamp_us: 1, + }; + t.snapshots.push(mk(100, 10, 20)); + t.snapshots.push(mk(200, 30, 40)); let r = t.report(); - assert_eq!(r.peak_rss, 200); - assert_eq!(r.mean_rss, 150); - assert_eq!(r.kv_cache_total, 40); - assert_eq!(r.activation_total, 60); - assert_eq!(r.label, "test"); + assert_eq!((r.peak_rss, r.mean_rss, r.kv_cache_total, r.activation_total), + (200, 150, 40, 60)); } } diff --git a/crates/ruvector-profiler/src/power.rs b/crates/ruvector-profiler/src/power.rs index 9b62820bf..abe681674 100644 --- a/crates/ruvector-profiler/src/power.rs +++ b/crates/ruvector-profiler/src/power.rs @@ -1,11 +1,6 @@ -/// A single power measurement sample. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct PowerSample { - pub watts: f64, - pub timestamp_us: u64, -} +pub struct PowerSample { pub watts: f64, pub timestamp_us: u64 } -/// Result of integrating power samples over time. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct EnergyResult { pub total_joules: f64, @@ -15,161 +10,85 @@ pub struct EnergyResult { pub samples: usize, } -/// Trait for reading instantaneous power from a hardware source. -/// -/// Real implementations would wrap NVML (GPU) or RAPL (CPU). Use -/// [`MockPowerSource`] for deterministic tests. -pub trait PowerSource { - fn read_watts(&self) -> f64; -} +/// Trait for reading instantaneous power (NVML, RAPL, etc.). +pub trait PowerSource { fn read_watts(&self) -> f64; } -/// A mock power source that returns a fixed wattage. -pub struct MockPowerSource { - pub watts: f64, -} +/// Fixed-wattage mock for deterministic tests. +pub struct MockPowerSource { pub watts: f64 } +impl PowerSource for MockPowerSource { fn read_watts(&self) -> f64 { self.watts } } -impl PowerSource for MockPowerSource { - fn read_watts(&self) -> f64 { - self.watts - } -} - -/// Estimate total energy consumption via trapezoidal integration. -/// -/// Samples must be sorted by `timestamp_us`. Returns a zeroed result -/// when fewer than two samples are provided. +/// Trapezoidal integration of power samples (must be sorted by timestamp). pub fn estimate_energy(samples: &[PowerSample]) -> EnergyResult { let n = samples.len(); if n < 2 { return EnergyResult { - total_joules: 0.0, + total_joules: 0.0, samples: n, duration_s: 0.0, mean_watts: samples.first().map_or(0.0, |s| s.watts), peak_watts: samples.first().map_or(0.0, |s| s.watts), - duration_s: 0.0, - samples: n, }; } - - let mut total_joules = 0.0; - let mut peak_watts: f64 = f64::NEG_INFINITY; - let mut sum_watts = 0.0; - + let (mut joules, mut peak, mut sum) = (0.0f64, f64::NEG_INFINITY, 0.0f64); for i in 0..n { let w = samples[i].watts; - sum_watts += w; - if w > peak_watts { - peak_watts = w; - } + sum += w; + if w > peak { peak = w; } if i > 0 { - let dt_us = samples[i].timestamp_us.saturating_sub(samples[i - 1].timestamp_us); - let dt_s = dt_us as f64 / 1_000_000.0; - let avg_w = (samples[i - 1].watts + samples[i].watts) / 2.0; - total_joules += avg_w * dt_s; + let dt = samples[i].timestamp_us.saturating_sub(samples[i - 1].timestamp_us) as f64 / 1e6; + joules += (samples[i - 1].watts + w) / 2.0 * dt; } } - - let duration_us = - samples.last().unwrap().timestamp_us.saturating_sub(samples.first().unwrap().timestamp_us); - let duration_s = duration_us as f64 / 1_000_000.0; - let mean_watts = sum_watts / n as f64; - - EnergyResult { - total_joules, - mean_watts, - peak_watts, - duration_s, - samples: n, - } + let dur = samples.last().unwrap().timestamp_us.saturating_sub(samples[0].timestamp_us) as f64 / 1e6; + EnergyResult { total_joules: joules, mean_watts: sum / n as f64, peak_watts: peak, duration_s: dur, samples: n } } -/// Collects [`PowerSample`]s under a named label. -pub struct PowerTracker { - pub samples: Vec, - pub label: String, -} +pub struct PowerTracker { pub samples: Vec, pub label: String } impl PowerTracker { - pub fn new(label: &str) -> Self { - Self { - samples: Vec::new(), - label: label.to_string(), - } - } + pub fn new(label: &str) -> Self { Self { samples: Vec::new(), label: label.to_string() } } - /// Record a sample from a [`PowerSource`]. pub fn sample(&mut self, source: &dyn PowerSource) { let ts = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_micros() as u64; - self.samples.push(PowerSample { - watts: source.read_watts(), - timestamp_us: ts, - }); + .duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_micros() as u64; + self.samples.push(PowerSample { watts: source.read_watts(), timestamp_us: ts }); } - /// Integrate all collected samples. - pub fn energy(&self) -> EnergyResult { - estimate_energy(&self.samples) - } + pub fn energy(&self) -> EnergyResult { estimate_energy(&self.samples) } } #[cfg(test)] mod tests { use super::*; + fn ps(w: f64, t: u64) -> PowerSample { PowerSample { watts: w, timestamp_us: t } } #[test] - fn energy_empty() { - let r = estimate_energy(&[]); - assert_eq!(r.total_joules, 0.0); - assert_eq!(r.samples, 0); + fn energy_empty() { let r = estimate_energy(&[]); assert_eq!(r.samples, 0); } + + #[test] + fn energy_single() { + let r = estimate_energy(&[ps(42.0, 0)]); + assert_eq!((r.total_joules, r.mean_watts), (0.0, 42.0)); } #[test] - fn energy_single_sample() { - let r = estimate_energy(&[PowerSample { watts: 42.0, timestamp_us: 100 }]); - assert_eq!(r.total_joules, 0.0); - assert_eq!(r.mean_watts, 42.0); - assert_eq!(r.samples, 1); - } - - #[test] - fn energy_trapezoidal() { - // 100W constant for 1 second = 100 J - let samples = vec![ - PowerSample { watts: 100.0, timestamp_us: 0 }, - PowerSample { watts: 100.0, timestamp_us: 1_000_000 }, - ]; - let r = estimate_energy(&samples); - assert!((r.total_joules - 100.0).abs() < 1e-9); - assert!((r.duration_s - 1.0).abs() < 1e-9); - assert_eq!(r.peak_watts, 100.0); - } - - #[test] - fn energy_trapezoid_varying() { - // Ramp from 0W to 200W over 1 second -> average 100W -> 100 J - let samples = vec![ - PowerSample { watts: 0.0, timestamp_us: 0 }, - PowerSample { watts: 200.0, timestamp_us: 1_000_000 }, - ]; - let r = estimate_energy(&samples); + fn energy_constant_100w_1s() { + let r = estimate_energy(&[ps(100.0, 0), ps(100.0, 1_000_000)]); assert!((r.total_joules - 100.0).abs() < 1e-9); } #[test] - fn mock_power_source() { - let src = MockPowerSource { watts: 75.0 }; - assert_eq!(src.read_watts(), 75.0); + fn energy_ramp() { + let r = estimate_energy(&[ps(0.0, 0), ps(200.0, 1_000_000)]); + assert!((r.total_joules - 100.0).abs() < 1e-9); } #[test] - fn power_tracker_collects() { + fn mock_source() { assert_eq!(MockPowerSource { watts: 75.0 }.read_watts(), 75.0); } + + #[test] + fn tracker_collects() { let src = MockPowerSource { watts: 50.0 }; - let mut tracker = PowerTracker::new("gpu"); - tracker.sample(&src); - tracker.sample(&src); - assert_eq!(tracker.samples.len(), 2); - assert_eq!(tracker.label, "gpu"); + let mut t = PowerTracker::new("gpu"); + t.sample(&src); t.sample(&src); + assert_eq!(t.samples.len(), 2); } }