avarok_core/kimi_k3/
greedy.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Token-greedy helper for C1: prefill the prompt, then decode one token at a time.
4
5use super::cache::HybridCache;
6use super::cpu_forward::{forward_token, logits};
7use super::cpu_weights::{Ablation, K3CpuModel};
8use super::ops::argmax;
9
10/// Greedy decode. Returns `prompt || new_tokens` (length `prompt.len() + max_new`).
11pub fn greedy_decode(
12    model: &K3CpuModel,
13    prompt: &[u32],
14    max_new: usize,
15    ablation: Ablation,
16) -> Vec<u32> {
17    assert!(!prompt.is_empty(), "C1 greedy needs a non-empty prompt");
18    let mut cache = HybridCache::from_graph(&model.graph, &model.kda);
19    let mut tokens = prompt.to_vec();
20    let mut h = Vec::new();
21    for (pos, &tok) in prompt.iter().enumerate() {
22        h = forward_token(model, tok, pos, &mut cache, ablation);
23    }
24    for _ in 0..max_new {
25        let next = argmax(&logits(model, &h));
26        tokens.push(next);
27        let pos = tokens.len() - 1;
28        h = forward_token(model, next, pos, &mut cache, ablation);
29    }
30    tokens
31}
32
33#[cfg(test)]
34mod tests {
35    use super::*;
36
37    #[test]
38    fn greedy_emits_requested_new_tokens() {
39        let model = K3CpuModel::synthetic_tiny();
40        let out = greedy_decode(&model, &[1, 2], 4, Ablation::default());
41        assert_eq!(out.len(), 6);
42        assert_eq!(&out[..2], &[1, 2]);
43        assert!(out[2..].iter().all(|&t| (t as usize) < model.vocab));
44    }
45
46    #[test]
47    fn greedy_is_deterministic() {
48        let model = K3CpuModel::synthetic_tiny();
49        let a = greedy_decode(&model, &[3], 8, Ablation::default());
50        let b = greedy_decode(&model, &[3], 8, Ablation::default());
51        assert_eq!(a, b);
52    }
53}
avarok_core/kimi_k3/greedy.rs

avarok_core/kimi_k3/
greedy.rs
