poly_kv/
replay.rs

1use crate::codec::create_codec;
2use crate::error::{PolyKvError, Result};
3use crate::policy::CODEC_TURBO_8BIT;
4use crate::{AgentShell, SharedKVPool};
5use serde::{Deserialize, Serialize};
6
7/// Schema for model-shaped compressed-attention replay receipts.
8pub const MODEL_REPLAY_RECEIPT_SCHEMA: &str = "poly_kv_model_replay_receipt_v1";
9
10/// One replay query plus a synthetic label used for PPL-proxy comparison.
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub struct ModelReplayQuery {
13    /// Query vector for one layer/head.
14    pub query: Vec<f32>,
15    /// Label token in the deterministic projection vocabulary.
16    pub label_token: usize,
17}
18
19/// Configuration for model-shaped replay.
20#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
21pub struct ModelReplayConfig {
22    /// Layer index to replay.
23    pub layer: usize,
24    /// KV head index to replay.
25    pub head: usize,
26    /// Candidate budgets to sweep. First passing value is selected.
27    pub candidate_ks: Vec<usize>,
28    /// Synthetic projection vocabulary size.
29    pub vocab_size: usize,
30    /// Seed for deterministic output projection.
31    pub projection_seed: u64,
32    /// Minimum mean output cosine for a candidate budget.
33    pub min_output_cosine: f64,
34    /// Maximum mean attention-output MSE for a candidate budget.
35    pub max_output_mse: f64,
36    /// Maximum mean KL(exact || compressed) for projected logits.
37    pub max_kl_divergence: f64,
38    /// Maximum allowed PPL-proxy delta.
39    pub max_ppl_delta: f64,
40    /// Minimum top-1 logit agreement rate.
41    pub min_top1_agreement: f64,
42}
43
44impl Default for ModelReplayConfig {
45    fn default() -> Self {
46        Self {
47            layer: 0,
48            head: 0,
49            candidate_ks: Vec::new(),
50            vocab_size: 64,
51            projection_seed: 0,
52            min_output_cosine: 0.75,
53            max_output_mse: 0.25,
54            max_kl_divergence: 0.50,
55            max_ppl_delta: 1.0,
56            min_top1_agreement: 0.25,
57        }
58    }
59}
60
61/// Metrics for one candidate-k replay pass.
62#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
63pub struct CandidateReplayMetrics {
64    pub candidate_k: usize,
65    pub output_cosine_mean: f64,
66    pub output_mse_mean: f64,
67    pub kl_divergence_mean: f64,
68    pub top1_agreement: f64,
69    pub ppl_proxy_exact: f64,
70    pub ppl_proxy_compressed: f64,
71    pub ppl_proxy_delta: f64,
72    pub decoded_values_total: u64,
73    pub full_decode_value_count: u64,
74    pub decode_reduction: f64,
75    pub passed: bool,
76    pub blockers: Vec<String>,
77}
78
79/// Aggregate replay metrics for the selected candidate budget.
80#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
81pub struct ModelReplayMetrics {
82    pub query_count: usize,
83    pub exact_attention_outputs: u64,
84    pub logit_vectors_compared: u64,
85    pub output_cosine_mean: f64,
86    pub output_mse_mean: f64,
87    pub kl_divergence_mean: f64,
88    pub top1_agreement: f64,
89    pub ppl_proxy_exact: f64,
90    pub ppl_proxy_compressed: f64,
91    pub ppl_proxy_delta: f64,
92    pub decoded_values_total: u64,
93    pub full_decode_value_count: u64,
94    pub decode_reduction: f64,
95}
96
97/// Model-shaped replay receipt.
98#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
99pub struct ModelReplayReceipt {
100    pub schema_version: String,
101    pub claim_boundary: String,
102    pub config: ModelReplayConfig,
103    pub selected_candidate_k: usize,
104    pub candidate_results: Vec<CandidateReplayMetrics>,
105    pub metrics: ModelReplayMetrics,
106    pub passed: bool,
107    pub blockers: Vec<String>,
108}
109
110/// Schema for captured-tensor model replay receipts.
111pub const CAPTURED_MODEL_REPLAY_RECEIPT_SCHEMA: &str = "poly_kv_captured_model_replay_v1";
112
113/// One captured query/key/value/logit sample.
114#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
115pub struct CapturedReplayQuery {
116    pub query: Vec<f32>,
117    pub keys: Vec<Vec<f32>>,
118    pub values: Vec<Vec<f32>>,
119    pub exact_attention_output: Vec<f32>,
120    pub exact_logits: Vec<f32>,
121    pub label_token: usize,
122}
123
124/// Captured tensor replay fixture.
125#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
126pub struct CapturedReplayFixture {
127    pub schema_version: String,
128    pub model_id: String,
129    pub head_dim: usize,
130    pub shared_tokens: usize,
131    pub seed: u64,
132    pub output_projection: Vec<Vec<f32>>,
133    pub queries: Vec<CapturedReplayQuery>,
134}
135
136/// Captured replay gate config.
137#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
138pub struct CapturedReplayConfig {
139    pub candidate_ks: Vec<usize>,
140    pub min_output_cosine: f64,
141    pub max_output_mse: f64,
142    pub max_kl_divergence: f64,
143    pub max_ppl_delta: f64,
144    pub min_top1_agreement: f64,
145}
146
147/// Captured tensor replay receipt.
148#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
149pub struct CapturedReplayReceipt {
150    pub schema_version: String,
151    pub model_id: String,
152    pub claim_boundary: String,
153    pub config: CapturedReplayConfig,
154    pub selected_candidate_k: usize,
155    pub candidate_results: Vec<CandidateReplayMetrics>,
156    pub metrics: ModelReplayMetrics,
157    pub passed: bool,
158    pub blockers: Vec<String>,
159}
160
161#[derive(Clone)]
162struct ExactCandidate {
163    key: Vec<f32>,
164    value: Vec<f32>,
165}
166
167/// Run a deterministic model-shaped replay over the compressed pool+shell path.
168///
169/// This compares compressed candidate selection against exact full-decode
170/// attention and then projects both outputs through a deterministic synthetic
171/// output head to produce logit/KL/PPL-proxy metrics. It is not real model PPL.
172pub fn run_model_replay(
173    pool: &SharedKVPool,
174    shell: &AgentShell,
175    queries: &[ModelReplayQuery],
176    config: ModelReplayConfig,
177) -> Result<ModelReplayReceipt> {
178    if config.candidate_ks.is_empty() {
179        return Err(PolyKvError::InvalidPolicy(
180            "candidate_ks must not be empty".to_string(),
181        ));
182    }
183    if queries.is_empty() {
184        return Err(PolyKvError::InvalidPolicy(
185            "queries must not be empty".to_string(),
186        ));
187    }
188    if config.vocab_size == 0 {
189        return Err(PolyKvError::InvalidPolicy(
190            "vocab_size must be greater than zero".to_string(),
191        ));
192    }
193
194    let head_dim = pool.manifest.shape.head_dim;
195    for query in queries {
196        if query.query.len() != head_dim {
197            return Err(PolyKvError::DimensionMismatch {
198                expected: head_dim,
199                got: query.query.len(),
200            });
201        }
202        if query.label_token >= config.vocab_size {
203            return Err(PolyKvError::InvalidPolicy(format!(
204                "label_token {} >= vocab_size {}",
205                query.label_token, config.vocab_size
206            )));
207        }
208    }
209
210    let exact_candidates = exact_candidates(pool, shell, config.layer, config.head)?;
211    let full_decode_value_count = (exact_candidates.len() * queries.len()) as u64;
212    let projection = output_projection(config.vocab_size, head_dim, config.projection_seed);
213    let mut candidate_results = Vec::with_capacity(config.candidate_ks.len());
214
215    for &candidate_k in &config.candidate_ks {
216        candidate_results.push(eval_candidate_k(
217            pool,
218            shell,
219            queries,
220            &config,
221            &exact_candidates,
222            &projection,
223            candidate_k,
224            full_decode_value_count,
225        )?);
226    }
227
228    let selected_idx = candidate_results
229        .iter()
230        .position(|r| r.passed)
231        .unwrap_or(candidate_results.len() - 1);
232    let selected = candidate_results[selected_idx].clone();
233    let passed = selected.passed;
234    let blockers = if passed {
235        Vec::new()
236    } else {
237        selected.blockers.clone()
238    };
239
240    Ok(ModelReplayReceipt {
241        schema_version: MODEL_REPLAY_RECEIPT_SCHEMA.to_string(),
242        claim_boundary: "deterministic model-shaped replay over synthetic projection; not real model PPL, not production KV-cache preservation, and not provider/framework KV-cache byte-reduction evidence".to_string(),
243        config,
244        selected_candidate_k: selected.candidate_k,
245        metrics: ModelReplayMetrics {
246            query_count: queries.len(),
247            exact_attention_outputs: queries.len() as u64,
248            logit_vectors_compared: queries.len() as u64,
249            output_cosine_mean: selected.output_cosine_mean,
250            output_mse_mean: selected.output_mse_mean,
251            kl_divergence_mean: selected.kl_divergence_mean,
252            top1_agreement: selected.top1_agreement,
253            ppl_proxy_exact: selected.ppl_proxy_exact,
254            ppl_proxy_compressed: selected.ppl_proxy_compressed,
255            ppl_proxy_delta: selected.ppl_proxy_delta,
256            decoded_values_total: selected.decoded_values_total,
257            full_decode_value_count,
258            decode_reduction: selected.decode_reduction,
259        },
260        candidate_results,
261        passed,
262        blockers,
263    })
264}
265
266/// Run captured-tensor replay through the existing compressed pool+shell path.
267pub fn run_captured_model_replay(
268    fixture: &CapturedReplayFixture,
269    config: CapturedReplayConfig,
270) -> Result<CapturedReplayReceipt> {
271    validate_captured_fixture(fixture, &config)?;
272    let full_decode_value_count = (fixture
273        .queries
274        .iter()
275        .map(|query| query.values.len())
276        .sum::<usize>()
277        * config.candidate_ks.len().max(1)
278        / config.candidate_ks.len().max(1)) as u64;
279    let mut candidate_results = Vec::with_capacity(config.candidate_ks.len());
280    for &candidate_k in &config.candidate_ks {
281        candidate_results.push(eval_captured_candidate_k(fixture, &config, candidate_k)?);
282    }
283    let selected_idx = candidate_results
284        .iter()
285        .position(|result| result.passed)
286        .unwrap_or(candidate_results.len() - 1);
287    let selected = candidate_results[selected_idx].clone();
288    let passed = selected.passed;
289    let blockers = if passed {
290        Vec::new()
291    } else {
292        selected.blockers.clone()
293    };
294    Ok(CapturedReplayReceipt {
295        schema_version: CAPTURED_MODEL_REPLAY_RECEIPT_SCHEMA.to_string(),
296        model_id: fixture.model_id.clone(),
297        claim_boundary: "captured tensor replay against fixture Q/K/V/logits; not pretrained LLM PPL, not production KV-cache preservation, and not provider/framework KV-cache byte-reduction evidence".to_string(),
298        config,
299        selected_candidate_k: selected.candidate_k,
300        metrics: ModelReplayMetrics {
301            query_count: fixture.queries.len(),
302            exact_attention_outputs: fixture.queries.len() as u64,
303            logit_vectors_compared: fixture.queries.len() as u64,
304            output_cosine_mean: selected.output_cosine_mean,
305            output_mse_mean: selected.output_mse_mean,
306            kl_divergence_mean: selected.kl_divergence_mean,
307            top1_agreement: selected.top1_agreement,
308            ppl_proxy_exact: selected.ppl_proxy_exact,
309            ppl_proxy_compressed: selected.ppl_proxy_compressed,
310            ppl_proxy_delta: selected.ppl_proxy_delta,
311            decoded_values_total: selected.decoded_values_total,
312            full_decode_value_count,
313            decode_reduction: selected.decode_reduction,
314        },
315        candidate_results,
316        passed,
317        blockers,
318    })
319}
320
321#[allow(clippy::too_many_arguments)]
322fn eval_candidate_k(
323    pool: &SharedKVPool,
324    shell: &AgentShell,
325    queries: &[ModelReplayQuery],
326    config: &ModelReplayConfig,
327    exact_candidates: &[ExactCandidate],
328    projection: &[Vec<f32>],
329    candidate_k: usize,
330    full_decode_value_count: u64,
331) -> Result<CandidateReplayMetrics> {
332    let mut cosines = Vec::with_capacity(queries.len());
333    let mut mses = Vec::with_capacity(queries.len());
334    let mut kls = Vec::with_capacity(queries.len());
335    let mut exact_nlls = Vec::with_capacity(queries.len());
336    let mut compressed_nlls = Vec::with_capacity(queries.len());
337    let mut top1_matches = 0usize;
338    let mut decoded_values_total = 0u64;
339
340    for query in queries {
341        let exact = exact_attention_output(&query.query, exact_candidates);
342        let exact_logits = project_logits(&exact, projection);
343        let exact_probs = softmax(&exact_logits);
344        let exact_top1 = argmax(&exact_logits);
345        let exact_nll = nll(&exact_probs, query.label_token);
346
347        let compressed = shell.attention_topk_compressed(
348            pool,
349            config.layer,
350            config.head,
351            &query.query,
352            candidate_k,
353        )?;
354        decoded_values_total += compressed.receipt.decoded_value_vectors;
355        let compressed_output = compressed_attention_output(&compressed.hits);
356        let compressed_logits = project_logits(&compressed_output, projection);
357        let compressed_probs = softmax(&compressed_logits);
358        let compressed_top1 = argmax(&compressed_logits);
359        let compressed_nll = nll(&compressed_probs, query.label_token);
360
361        cosines.push(cosine(&exact, &compressed_output));
362        mses.push(mse(&exact, &compressed_output));
363        kls.push(kl_divergence(&exact_probs, &compressed_probs));
364        exact_nlls.push(exact_nll);
365        compressed_nlls.push(compressed_nll);
366        if exact_top1 == compressed_top1 {
367            top1_matches += 1;
368        }
369    }
370
371    let output_cosine_mean = mean(&cosines);
372    let output_mse_mean = mean(&mses);
373    let kl_divergence_mean = mean(&kls);
374    let exact_nll = mean(&exact_nlls);
375    let compressed_nll = mean(&compressed_nlls);
376    let ppl_proxy_exact = exact_nll.exp();
377    let ppl_proxy_compressed = compressed_nll.exp();
378    let ppl_proxy_delta = ppl_proxy_compressed - ppl_proxy_exact;
379    let top1_agreement = top1_matches as f64 / queries.len() as f64;
380    let decode_reduction = full_decode_value_count as f64 / decoded_values_total.max(1) as f64;
381
382    let mut blockers = Vec::new();
383    if output_cosine_mean < config.min_output_cosine {
384        blockers.push(format!(
385            "output_cosine_mean {output_cosine_mean:.4} < {:.4}",
386            config.min_output_cosine
387        ));
388    }
389    if output_mse_mean > config.max_output_mse {
390        blockers.push(format!(
391            "output_mse_mean {output_mse_mean:.4} > {:.4}",
392            config.max_output_mse
393        ));
394    }
395    if kl_divergence_mean > config.max_kl_divergence {
396        blockers.push(format!(
397            "kl_divergence_mean {kl_divergence_mean:.4} > {:.4}",
398            config.max_kl_divergence
399        ));
400    }
401    let ppl_proxy_delta_abs = ppl_proxy_delta.abs();
402    if ppl_proxy_delta_abs > config.max_ppl_delta {
403        blockers.push(format!(
404            "abs(ppl_proxy_delta) {ppl_proxy_delta_abs:.4} > {:.4}",
405            config.max_ppl_delta
406        ));
407    }
408    if top1_agreement < config.min_top1_agreement {
409        blockers.push(format!(
410            "top1_agreement {top1_agreement:.4} < {:.4}",
411            config.min_top1_agreement
412        ));
413    }
414
415    Ok(CandidateReplayMetrics {
416        candidate_k,
417        output_cosine_mean,
418        output_mse_mean,
419        kl_divergence_mean,
420        top1_agreement,
421        ppl_proxy_exact,
422        ppl_proxy_compressed,
423        ppl_proxy_delta,
424        decoded_values_total,
425        full_decode_value_count,
426        decode_reduction,
427        passed: blockers.is_empty(),
428        blockers,
429    })
430}
431
432fn validate_captured_fixture(
433    fixture: &CapturedReplayFixture,
434    config: &CapturedReplayConfig,
435) -> Result<()> {
436    if config.candidate_ks.is_empty() {
437        return Err(PolyKvError::InvalidPolicy(
438            "candidate_ks must not be empty".to_string(),
439        ));
440    }
441    if fixture.queries.is_empty() {
442        return Err(PolyKvError::InvalidPolicy(
443            "captured fixture must contain at least one query".to_string(),
444        ));
445    }
446    if fixture.head_dim == 0 {
447        return Err(PolyKvError::InvalidPolicy(
448            "head_dim must be greater than zero".to_string(),
449        ));
450    }
451    if fixture.output_projection.is_empty() {
452        return Err(PolyKvError::InvalidPolicy(
453            "output_projection must not be empty".to_string(),
454        ));
455    }
456    for row in &fixture.output_projection {
457        if row.len() != fixture.head_dim {
458            return Err(PolyKvError::DimensionMismatch {
459                expected: fixture.head_dim,
460                got: row.len(),
461            });
462        }
463    }
464    for query in &fixture.queries {
465        if query.query.len() != fixture.head_dim
466            || query.exact_attention_output.len() != fixture.head_dim
467        {
468            return Err(PolyKvError::DimensionMismatch {
469                expected: fixture.head_dim,
470                got: query.query.len(),
471            });
472        }
473        if query.keys.len() != query.values.len() || query.keys.is_empty() {
474            return Err(PolyKvError::InvalidPolicy(
475                "captured keys/values must be non-empty and same length".to_string(),
476            ));
477        }
478        if fixture.shared_tokens == 0 || fixture.shared_tokens >= query.keys.len() {
479            return Err(PolyKvError::InvalidPolicy(
480                "shared_tokens must split captured rows into non-empty pool and shell tiers"
481                    .to_string(),
482            ));
483        }
484        if query.exact_logits.len() != fixture.output_projection.len() {
485            return Err(PolyKvError::DimensionMismatch {
486                expected: fixture.output_projection.len(),
487                got: query.exact_logits.len(),
488            });
489        }
490        if query.label_token >= query.exact_logits.len() {
491            return Err(PolyKvError::InvalidPolicy(format!(
492                "label_token {} >= logits len {}",
493                query.label_token,
494                query.exact_logits.len()
495            )));
496        }
497        for row in query.keys.iter().chain(query.values.iter()) {
498            if row.len() != fixture.head_dim {
499                return Err(PolyKvError::DimensionMismatch {
500                    expected: fixture.head_dim,
501                    got: row.len(),
502                });
503            }
504        }
505    }
506    Ok(())
507}
508
509fn eval_captured_candidate_k(
510    fixture: &CapturedReplayFixture,
511    config: &CapturedReplayConfig,
512    candidate_k: usize,
513) -> Result<CandidateReplayMetrics> {
514    let mut cosines = Vec::with_capacity(fixture.queries.len());
515    let mut mses = Vec::with_capacity(fixture.queries.len());
516    let mut kls = Vec::with_capacity(fixture.queries.len());
517    let mut exact_nlls = Vec::with_capacity(fixture.queries.len());
518    let mut compressed_nlls = Vec::with_capacity(fixture.queries.len());
519    let mut top1_matches = 0usize;
520    let mut decoded_values_total = 0u64;
521    let mut full_decode_value_count = 0u64;
522
523    for (query_idx, query) in fixture.queries.iter().enumerate() {
524        let (pool, shell) = build_captured_pool_shell(fixture, query, query_idx)?;
525        let compressed = shell.attention_topk_compressed(&pool, 0, 0, &query.query, candidate_k)?;
526        decoded_values_total += compressed.receipt.decoded_value_vectors;
527        full_decode_value_count += query.values.len() as u64;
528
529        let compressed_output = compressed_attention_output(&compressed.hits);
530        let compressed_logits = project_logits(&compressed_output, &fixture.output_projection);
531        let exact_probs = softmax(&query.exact_logits);
532        let compressed_probs = softmax(&compressed_logits);
533        let exact_nll = nll(&exact_probs, query.label_token);
534        let compressed_nll = nll(&compressed_probs, query.label_token);
535        cosines.push(cosine(&query.exact_attention_output, &compressed_output));
536        mses.push(mse(&query.exact_attention_output, &compressed_output));
537        kls.push(kl_divergence(&exact_probs, &compressed_probs));
538        exact_nlls.push(exact_nll);
539        compressed_nlls.push(compressed_nll);
540        if argmax(&query.exact_logits) == argmax(&compressed_logits) {
541            top1_matches += 1;
542        }
543    }
544
545    let output_cosine_mean = mean(&cosines);
546    let output_mse_mean = mean(&mses);
547    let kl_divergence_mean = mean(&kls);
548    let ppl_proxy_exact = mean(&exact_nlls).exp();
549    let ppl_proxy_compressed = mean(&compressed_nlls).exp();
550    let ppl_proxy_delta = ppl_proxy_compressed - ppl_proxy_exact;
551    let top1_agreement = top1_matches as f64 / fixture.queries.len() as f64;
552    let decode_reduction = full_decode_value_count as f64 / decoded_values_total.max(1) as f64;
553
554    let mut blockers = Vec::new();
555    if output_cosine_mean < config.min_output_cosine {
556        blockers.push(format!(
557            "output_cosine_mean {output_cosine_mean:.4} < {:.4}",
558            config.min_output_cosine
559        ));
560    }
561    if output_mse_mean > config.max_output_mse {
562        blockers.push(format!(
563            "output_mse_mean {output_mse_mean:.4} > {:.4}",
564            config.max_output_mse
565        ));
566    }
567    if kl_divergence_mean > config.max_kl_divergence {
568        blockers.push(format!(
569            "kl_divergence_mean {kl_divergence_mean:.4} > {:.4}",
570            config.max_kl_divergence
571        ));
572    }
573    let ppl_proxy_delta_abs = ppl_proxy_delta.abs();
574    if ppl_proxy_delta_abs > config.max_ppl_delta {
575        blockers.push(format!(
576            "abs(ppl_proxy_delta) {ppl_proxy_delta_abs:.4} > {:.4}",
577            config.max_ppl_delta
578        ));
579    }
580    if top1_agreement < config.min_top1_agreement {
581        blockers.push(format!(
582            "top1_agreement {top1_agreement:.4} < {:.4}",
583            config.min_top1_agreement
584        ));
585    }
586
587    Ok(CandidateReplayMetrics {
588        candidate_k,
589        output_cosine_mean,
590        output_mse_mean,
591        kl_divergence_mean,
592        top1_agreement,
593        ppl_proxy_exact,
594        ppl_proxy_compressed,
595        ppl_proxy_delta,
596        decoded_values_total,
597        full_decode_value_count,
598        decode_reduction,
599        passed: blockers.is_empty(),
600        blockers,
601    })
602}
603
604fn build_captured_pool_shell(
605    fixture: &CapturedReplayFixture,
606    query: &CapturedReplayQuery,
607    query_idx: usize,
608) -> Result<(SharedKVPool, AgentShell)> {
609    let shape = crate::KvTensorShape {
610        attention_type: crate::AttentionType::MHA,
611        num_layers: 1,
612        num_heads: 1,
613        num_kv_heads: 1,
614        head_dim: fixture.head_dim,
615        hidden_size: fixture.head_dim,
616    };
617    let rows: Vec<Vec<f32>> = query
618        .keys
619        .iter()
620        .zip(&query.values)
621        .map(|(key, value)| {
622            let mut row = Vec::with_capacity(fixture.head_dim * 2);
623            row.extend_from_slice(key);
624            row.extend_from_slice(value);
625            row
626        })
627        .collect();
628    let shared: Vec<(String, Vec<f32>)> = rows
629        .iter()
630        .take(fixture.shared_tokens)
631        .enumerate()
632        .map(|(idx, row)| (format!("q{query_idx}_shared_{idx}"), row.clone()))
633        .collect();
634    let hot: Vec<(String, Vec<f32>)> = rows
635        .iter()
636        .skip(fixture.shared_tokens)
637        .enumerate()
638        .map(|(idx, row)| (format!("q{query_idx}_hot_{idx}"), row.clone()))
639        .collect();
640    let (pool, _) = SharedKVPool::build(&shared, &shape, fixture.seed + query_idx as u64)?;
641    let (shell, _) = pool.materialize_shell(
642        &format!("captured_{}_{query_idx}", fixture.model_id),
643        &hot,
644        fixture.seed + 10_000 + query_idx as u64,
645    )?;
646    Ok((pool, shell))
647}
648
649fn exact_candidates(
650    pool: &SharedKVPool,
651    shell: &AgentShell,
652    layer_idx: usize,
653    head_idx: usize,
654) -> Result<Vec<ExactCandidate>> {
655    let layer = pool.decompress_layer(layer_idx)?;
656    if head_idx >= layer.num_heads {
657        return Err(PolyKvError::Internal(format!(
658            "head_idx {head_idx} out of range (have {})",
659            layer.num_heads
660        )));
661    }
662    let head_dim = layer.head_dim;
663    let mut out = Vec::with_capacity(layer.num_tokens);
664    let pool_keys = &layer.keys[head_idx];
665    let pool_values = &layer.values[head_idx];
666    for token_idx in 0..layer.num_tokens {
667        let start = token_idx * head_dim;
668        out.push(ExactCandidate {
669            key: pool_keys[start..start + head_dim].to_vec(),
670            value: pool_values[start..start + head_dim].to_vec(),
671        });
672    }
673
674    if let Some(shell_layer) = shell
675        .unique_layers
676        .iter()
677        .find(|l| l.layer_index == layer_idx as u32)
678    {
679        let num_heads = pool.manifest.shape.num_kv_heads as usize;
680        let shell_tokens = shell_layer.key_blocks.len() / num_heads;
681        let turbo_codec = create_codec(
682            CODEC_TURBO_8BIT,
683            head_dim,
684            None,
685            Some(&pool.policy.turbo_config),
686        )?;
687        for token_idx in 0..shell_tokens {
688            let block_idx = token_idx * num_heads + head_idx;
689            out.push(ExactCandidate {
690                key: turbo_codec.decode(
691                    &shell_layer.key_blocks[block_idx].encoded_payload,
692                    shell.build_seed,
693                )?,
694                value: turbo_codec.decode(
695                    &shell_layer.value_blocks[block_idx].encoded_payload,
696                    shell.build_seed,
697                )?,
698            });
699        }
700    }
701    Ok(out)
702}
703
704fn exact_attention_output(query: &[f32], candidates: &[ExactCandidate]) -> Vec<f32> {
705    let scores: Vec<f32> = candidates.iter().map(|c| dot(query, &c.key)).collect();
706    let weights = softmax_f32(&scores);
707    let dim = candidates.first().map(|c| c.value.len()).unwrap_or(0);
708    let mut out = vec![0.0f32; dim];
709    for (weight, cand) in weights.iter().zip(candidates) {
710        for (dst, value) in out.iter_mut().zip(&cand.value) {
711            *dst += *weight * *value;
712        }
713    }
714    out
715}
716
717fn compressed_attention_output(hits: &[crate::CompressedShellAttentionHit]) -> Vec<f32> {
718    if hits.is_empty() {
719        return Vec::new();
720    }
721    let scores: Vec<f32> = hits.iter().map(|h| h.score).collect();
722    let weights = softmax_f32(&scores);
723    let dim = hits[0].value.len();
724    let mut out = vec![0.0f32; dim];
725    for (weight, hit) in weights.iter().zip(hits) {
726        for (dst, value) in out.iter_mut().zip(&hit.value) {
727            *dst += *weight * *value;
728        }
729    }
730    out
731}
732
733fn output_projection(vocab_size: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
734    (0..vocab_size)
735        .map(|token| {
736            (0..dim)
737                .map(|i| {
738                    let x = (token as f32 + 1.0) * 0.013
739                        + (i as f32 + 1.0) * 0.017
740                        + seed as f32 * 0.0001;
741                    x.sin() * 0.25 + x.cos() * 0.10
742                })
743                .collect()
744        })
745        .collect()
746}
747
748fn project_logits(output: &[f32], projection: &[Vec<f32>]) -> Vec<f32> {
749    projection.iter().map(|row| dot(output, row)).collect()
750}
751
752fn dot(a: &[f32], b: &[f32]) -> f32 {
753    a.iter().zip(b).map(|(x, y)| x * y).sum()
754}
755
756fn softmax_f32(scores: &[f32]) -> Vec<f32> {
757    let max = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
758    let mut exps: Vec<f32> = scores.iter().map(|v| (*v - max).exp()).collect();
759    let sum: f32 = exps.iter().sum();
760    if sum <= f32::EPSILON || !sum.is_finite() {
761        return vec![1.0 / scores.len().max(1) as f32; scores.len()];
762    }
763    for v in &mut exps {
764        *v /= sum;
765    }
766    exps
767}
768
769fn softmax(scores: &[f32]) -> Vec<f64> {
770    softmax_f32(scores).into_iter().map(f64::from).collect()
771}
772
773fn nll(probs: &[f64], label: usize) -> f64 {
774    -probs[label].max(1e-12).ln()
775}
776
777fn kl_divergence(p: &[f64], q: &[f64]) -> f64 {
778    p.iter()
779        .zip(q)
780        .map(|(pi, qi)| {
781            if *pi <= 0.0 {
782                0.0
783            } else {
784                pi * (pi / qi.max(1e-12)).ln()
785            }
786        })
787        .sum()
788}
789
790fn argmax(values: &[f32]) -> usize {
791    values
792        .iter()
793        .enumerate()
794        .max_by(|a, b| a.1.total_cmp(b.1))
795        .map(|(idx, _)| idx)
796        .unwrap_or(0)
797}
798
799fn cosine(a: &[f32], b: &[f32]) -> f64 {
800    let dot: f64 = a
801        .iter()
802        .zip(b)
803        .map(|(x, y)| f64::from(*x) * f64::from(*y))
804        .sum();
805    let na: f64 = a
806        .iter()
807        .map(|x| f64::from(*x) * f64::from(*x))
808        .sum::<f64>()
809        .sqrt();
810    let nb: f64 = b
811        .iter()
812        .map(|x| f64::from(*x) * f64::from(*x))
813        .sum::<f64>()
814        .sqrt();
815    if na <= f64::EPSILON || nb <= f64::EPSILON {
816        0.0
817    } else {
818        dot / (na * nb)
819    }
820}
821
822fn mse(a: &[f32], b: &[f32]) -> f64 {
823    if a.is_empty() {
824        return 0.0;
825    }
826    a.iter()
827        .zip(b)
828        .map(|(x, y)| {
829            let d = f64::from(*x) - f64::from(*y);
830            d * d
831        })
832        .sum::<f64>()
833        / a.len() as f64
834}
835
836fn mean(values: &[f64]) -> f64 {
837    if values.is_empty() {
838        0.0
839    } else {
840        values.iter().sum::<f64>() / values.len() as f64
841    }
842}
poly_kv/replay.rs

poly_kv/
replay.rs
