aprender-serve 0.71.3

//! M32c.2.2.2.1.2 + M32d — autoregressive loop for Qwen3-MoE with KV cache.
//!
//! Composes `OwnedQuantizedModel::forward_single_qwen3_moe_with_cache`
//! (M32d) into a per-token decode loop. This is the sibling of
//! `run_gguf_generate` for `qwen3_moe` arch.
//!
//! ## Design
//! Per `qwen3-moe-serve-dispatch-v1` v1.2.0 + M32d playbook:
//!   1. Read MoE config (num_experts, k, intermediate) from GGUF metadata.
//!   2. Build per-layer `Qwen3MoeQuantizedLayer` descriptors via
//!      `load_qwen3_moe_layer` once at start.
//!   3. Allocate `OwnedQuantizedKVCache` sized to `prompt_len + max_tokens`.
//!   4. Prefill: per prompt token, call
//!      `forward_single_qwen3_moe_with_cache`. Cache builds incrementally.
//!      The final iteration's logits are the seed for decode.
//!   5. Decode: per output token, greedy-argmax + call
//!      `forward_single_qwen3_moe_with_cache` for the next-token logits.
//!      Stop on `stop_tokens` or `max_tokens` exhausted.
//!
//! ## Performance
//! Post-M32d: 5-15 tok/s sustained on Qwen3-Coder-30B-A3B (vs ~0.5 tok/s
//! pre-M32d full-prefill-per-token). Each output token amortizes to one
//! per-layer attention (cached K/V read) + one per-layer MoE FFN
//! dispatch — no re-prefill.
//!
//! ## What's NOT in scope
//! - Top-p / top-k / temperature sampling (greedy-only for V1_001 +
//!   V1_004 discharge; sampling is M32 follow-up)
//! - Streaming SSE (cache exposes natural emit-per-token point; one-line
//!   addition once needed — separate contract `qwen3-moe-streaming-sse-v1`)
//! - GPU MoE (separate `qwen3-moe-forward-gpu-v1` track)
//! - Cache rollback / beam search (cache.rollback_to exists; not wired)

use crate::error::{RealizarError, Result};
use crate::gguf::moe_session::{Qwen3MoeForward, Qwen3MoeSession};
use crate::gguf::{MappedGGUFModel, OwnedQuantizedModel, QuantizedGenerateConfig};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};

/// Sample the next token from logits per `QuantizedGenerateConfig`.
///
/// Discharges `qwen3-moe-sampling-v1.yaml`:
/// - greedy fallback when `temperature == 0` OR `top_k == 1` (V1_001 + V1_004)
/// - seeded RNG → deterministic across runs with same seed (V1_002)
/// - seed differences produce different outputs (V1_003)
///
/// Mirrors the deleted dense `sample_advanced` (#4266) but uses a seeded
/// `StdRng` instead of `rand::rng()` for reproducibility.
pub(crate) fn sample_from_logits(
    logits: &[f32],
    config: &QuantizedGenerateConfig,
    rng: &mut StdRng,
    recent_tokens: &[u32],
) -> Result<u32> {
    if logits.is_empty() {
        return Err(RealizarError::InvalidShape {
            reason: "sample_from_logits: empty logits vector".to_string(),
        });
    }

    let penalized = penalize_recent(logits, config, recent_tokens);

    // Greedy fallback: temperature == 0 OR top_k == 1 (after repetition penalty)
    if config.temperature == 0.0 || config.top_k == 1 {
        return Ok(crate::sampling::argmax(&penalized));
    }

    // Temperature scaling, then the top-k and top-p (nucleus) filters
    let scaled: Vec<f32> = penalized.iter().map(|&x| x / config.temperature).collect();
    let mut indexed = top_k_sorted(&scaled, config.top_k);
    if config.top_p > 0.0 && config.top_p < 1.0 {
        truncate_top_p(&mut indexed, config.top_p);
    }

    Ok(multinomial_draw(&indexed, rng))
}

/// Step 1: Repetition penalty (qwen3-moe-repetition-penalty-v1).
/// Applied BEFORE temperature scaling. The dense helper, so the MoE and
/// dense paths cannot drift: once per distinct recent token, as Candle's
/// apply_repeat_penalty and llama.cpp do (PMAT-383/384).
/// No-op when repeat_penalty == 1.0 OR repeat_last_n == 0.
fn penalize_recent(
    logits: &[f32],
    config: &QuantizedGenerateConfig,
    recent_tokens: &[u32],
) -> Vec<f32> {
    let mut p: Vec<f32> = logits.to_vec();
    if config.repeat_penalty != 1.0 && config.repeat_last_n > 0 && !recent_tokens.is_empty() {
        crate::gguf::OwnedQuantizedModel::apply_repeat_penalty(
            &mut p,
            recent_tokens,
            config.repeat_penalty,
            config.repeat_last_n,
        );
    }
    p
}

/// Top-k filter: `(index, value)` pairs sorted by value, descending, cut to
/// `top_k` when `0 < top_k < len`.
fn top_k_sorted(scaled: &[f32], top_k: usize) -> Vec<(usize, f32)> {
    let mut indexed: Vec<(usize, f32)> = scaled.iter().copied().enumerate().collect();
    indexed.sort_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
    if top_k > 0 && top_k < indexed.len() {
        indexed.truncate(top_k);
    }
    indexed
}

/// Top-p (nucleus): keep the smallest prefix whose cumulative softmax >= `top_p`.
/// `indexed` is sorted descending, so its first value is the max.
fn truncate_top_p(indexed: &mut Vec<(usize, f32)>, top_p: f32) {
    let max_val = indexed.first().map_or(0.0, |(_, v)| *v);
    let exp_vals: Vec<f32> = indexed.iter().map(|(_, v)| (v - max_val).exp()).collect();
    let total: f32 = exp_vals.iter().sum();
    if total <= 0.0 {
        return;
    }
    let mut cumulative = 0.0;
    let cutoff = exp_vals
        .iter()
        .position(|&ev| {
            cumulative += ev / total;
            cumulative >= top_p
        })
        .map_or(indexed.len(), |i| i + 1);
    indexed.truncate(cutoff);
}

/// Softmax over the filtered set + one multinomial draw from `rng`.
fn multinomial_draw(indexed: &[(usize, f32)], rng: &mut StdRng) -> u32 {
    let max_val = indexed.first().map_or(0.0, |(_, v)| *v);
    let exp_sum: f32 = indexed.iter().map(|(_, v)| (v - max_val).exp()).sum();
    if exp_sum <= 0.0 {
        // Degenerate softmax: fall back to argmax of filtered set
        return indexed.first().map_or(0, |(i, _)| *i as u32);
    }

    let r: f32 = rng.random();
    let mut cumulative = 0.0;
    for (idx, v) in indexed {
        cumulative += (v - max_val).exp() / exp_sum;
        if cumulative >= r {
            return *idx as u32;
        }
    }
    indexed.last().map_or(0, |(i, _)| *i as u32)
}

/// Run autoregressive token generation for a Qwen3-MoE GGUF model.
///
/// PMAT-4269 (M1): a thin entry into the one engine — a
/// [`Qwen3MoeSession`](crate::gguf::moe_session::Qwen3MoeSession) owns the
/// prefill, token choice, stop tokens and context budget; this function only
/// opens it and runs one turn.
///
/// # Arguments
/// * `mapped` — the mmapped GGUF (the per-layer expert tensors borrow from it
///   for the length of the call).
/// * `model` — the standard `OwnedQuantizedModel` constructed via
///   `OwnedQuantizedModel::from_mapped`.
/// * `input_tokens` — the prompt token IDs.
/// * `gen_config` — generation config (max_tokens, sampling, stop tokens).
///
/// # Returns
/// Full token sequence including prompt: `[prompt..., generated...]`.
///
/// # Errors
/// - An empty prompt, or one the declared context cannot hold.
/// - Architecture isn't qwen3_moe (caller should dispatch correctly).
/// - MoE config metadata missing (`expert_count`, `expert_used_count`,
///   `expert_feed_forward_length`).
/// - Per-layer MoE descriptor load failure, or a forward pass error.
pub fn run_qwen3_moe_generate(
    mapped: &MappedGGUFModel,
    model: &OwnedQuantizedModel,
    input_tokens: &[u32],
    gen_config: &QuantizedGenerateConfig,
) -> Result<Vec<u32>> {
    let mut session = Qwen3MoeSession::new(Qwen3MoeForward::cpu(mapped, model)?);
    Ok(session
        .generate(input_tokens, gen_config, &mut |_| true)?
        .tokens)
}

/// Streaming variant of `run_qwen3_moe_generate` — discharges
/// `qwen3-moe-streaming-sse-v1.yaml` per-token emit requirement.
///
/// `on_token(next_token)` fires for every chosen token, including the one that
/// triggered a stop, before the stop is honoured; returning `false` ends the
/// turn (e.g. client disconnect). The same session turn as the non-streaming
/// variant, so the two emit the same tokens.
pub fn run_qwen3_moe_generate_streaming(
    mapped: &MappedGGUFModel,
    model: &OwnedQuantizedModel,
    input_tokens: &[u32],
    gen_config: &QuantizedGenerateConfig,
    mut on_token: impl FnMut(u32) -> bool,
) -> Result<()> {
    let mut session = Qwen3MoeSession::new(Qwen3MoeForward::cpu(mapped, model)?);
    session.generate(input_tokens, gen_config, &mut on_token)?;
    Ok(())
}

#[cfg(test)]
mod sample_from_logits_tests {
    //! Unit tests for the `qwen3-moe-sampling-v1.yaml` falsifiers
    //! against `sample_from_logits` directly. Run without a real
    //! Qwen3-MoE GGUF (uses synthetic logits arrays). Complements the
    //! env-gated integration tests in
    //! `crates/aprender-serve/tests/qwen3_moe_sampling_v1.rs` which
    //! validate the same invariants on top of a real model.
    use super::*;

    fn mk_config(temperature: f32, top_k: usize, top_p: f32, seed: u64) -> QuantizedGenerateConfig {
        QuantizedGenerateConfig {
            max_tokens: 1,
            temperature,
            top_k,
            top_p,
            seed,
            stop_tokens: Vec::new(),
            ..QuantizedGenerateConfig::default()
        }
    }

    /// V1_001: greedy fallback (temperature == 0) returns argmax deterministically.
    #[test]
    fn v1_001_temperature_zero_is_argmax_deterministic() {
        let logits = vec![1.0, 5.0, 2.0, 4.0, 3.0]; // argmax = index 1
        let cfg = mk_config(0.0, 50, 1.0, 42);
        for _ in 0..5 {
            let mut rng = StdRng::seed_from_u64(cfg.seed);
            let token = sample_from_logits(&logits, &cfg, &mut rng, &[]).unwrap();
            assert_eq!(token, 1, "V1_001: temperature=0 must return argmax");
        }
    }

    /// V1_001: top_k == 1 ALSO triggers greedy fallback (independent path).
    #[test]
    fn v1_001_top_k_one_is_argmax_deterministic() {
        let logits = vec![3.0, 1.0, 7.0, 2.0, 5.0]; // argmax = index 2
        let cfg = mk_config(5.0 /* high temp ignored */, 1, 1.0, 42);
        for _ in 0..5 {
            let mut rng = StdRng::seed_from_u64(cfg.seed);
            let token = sample_from_logits(&logits, &cfg, &mut rng, &[]).unwrap();
            assert_eq!(token, 2, "V1_001: top_k=1 must return argmax");
        }
    }

    /// V1_002: temperature > 0 with fixed seed returns the same token across runs.
    #[test]
    fn v1_002_seeded_rng_is_reproducible() {
        let logits = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let cfg = mk_config(0.7, 50, 0.95, 42);

        let mut tokens = Vec::new();
        for _ in 0..5 {
            let mut rng = StdRng::seed_from_u64(cfg.seed);
            tokens.push(sample_from_logits(&logits, &cfg, &mut rng, &[]).unwrap());
        }
        let first = tokens[0];
        for (i, &t) in tokens.iter().enumerate() {
            assert_eq!(
                t, first,
                "V1_002: seed=42 must produce same token; iter {i} got {t}, expected {first}"
            );
        }
    }

    /// V1_003: different seeds produce different tokens on average. Single-token
    /// inevitably has collisions, so probe across 32 seeds and assert at least
    /// 3 distinct tokens (very loose bound; collisions on 1-of-8 logits with
    /// reasonable temp are rare).
    #[test]
    fn v1_003_different_seeds_diverge() {
        let logits = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let cfg_template = mk_config(1.5 /* spread mass */, 0, 1.0, 0);

        let mut tokens = std::collections::HashSet::new();
        for seed in 0..32u64 {
            let mut cfg = cfg_template.clone();
            cfg.seed = seed;
            let mut rng = StdRng::seed_from_u64(cfg.seed);
            tokens.insert(sample_from_logits(&logits, &cfg, &mut rng, &[]).unwrap());
        }
        assert!(
            tokens.len() >= 3,
            "V1_003: 32 seeds must produce ≥ 3 distinct tokens (got {})",
            tokens.len()
        );
    }

    /// V1_004: top_k=1 with HIGH temperature == greedy (regardless of RNG state).
    #[test]
    fn v1_004_top_k_one_equals_pure_greedy() {
        let logits = vec![0.1, 0.2, 0.3, 0.4, 0.5, 99.0, 0.6, 0.7]; // argmax = 5
        let high_temp_top_k_one = mk_config(50.0, 1, 1.0, 12345);
        let pure_greedy = mk_config(0.0, 1, 1.0, 999_999);

        let mut rng_a = StdRng::seed_from_u64(high_temp_top_k_one.seed);
        let mut rng_b = StdRng::seed_from_u64(pure_greedy.seed);
        let a = sample_from_logits(&logits, &high_temp_top_k_one, &mut rng_a, &[]).unwrap();
        let b = sample_from_logits(&logits, &pure_greedy, &mut rng_b, &[]).unwrap();
        assert_eq!(
            a, b,
            "V1_004: top_k=1 == pure greedy regardless of temperature"
        );
        assert_eq!(a, 5, "V1_004: argmax of logits is at index 5");
    }

    /// Edge case: empty logits returns InvalidShape error (no panic).
    #[test]
    fn empty_logits_returns_error() {
        let cfg = mk_config(0.7, 50, 0.95, 42);
        let mut rng = StdRng::seed_from_u64(cfg.seed);
        let result = sample_from_logits(&[], &cfg, &mut rng, &[]);
        assert!(result.is_err(), "empty logits must error, not panic");
    }

    /// Edge case: top_p=1.0 has no effect (just regular sampling).
    #[test]
    fn top_p_one_is_no_op() {
        let logits = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let cfg_with_top_p = mk_config(0.7, 0 /* no top_k cap */, 1.0, 42);
        let cfg_no_top_p = mk_config(0.7, 0, 0.0 /* sentinel: not active */, 42);

        let mut rng_a = StdRng::seed_from_u64(cfg_with_top_p.seed);
        let mut rng_b = StdRng::seed_from_u64(cfg_no_top_p.seed);
        let a = sample_from_logits(&logits, &cfg_with_top_p, &mut rng_a, &[]).unwrap();
        let b = sample_from_logits(&logits, &cfg_no_top_p, &mut rng_b, &[]).unwrap();
        assert_eq!(a, b, "top_p=1.0 must equal top_p=0.0 (both no-op)");
    }

    // ========================================================================
    // qwen3-moe-repetition-penalty-v1 falsifier tests
    // ========================================================================

    fn mk_config_with_penalty(
        temperature: f32,
        top_k: usize,
        repeat_penalty: f32,
        repeat_last_n: usize,
        seed: u64,
    ) -> QuantizedGenerateConfig {
        QuantizedGenerateConfig {
            max_tokens: 1,
            temperature,
            top_k,
            top_p: 1.0,
            repeat_penalty,
            repeat_last_n,
            seed,
            stop_tokens: Vec::new(),
            ..QuantizedGenerateConfig::default()
        }
    }

    /// V1_001 (repetition penalty): repeat_penalty == 1.0 is a no-op even with
    /// non-empty recent_tokens.
    #[test]
    fn rep_penalty_v1_001_no_op_at_one() {
        let logits = vec![3.0, 5.0, 2.0, 4.0]; // argmax = 1
        let recent = vec![1, 1, 1]; // many repetitions of token 1
        let cfg = mk_config_with_penalty(0.0, 1, 1.0 /* no-op */, 100, 42);

        let mut rng = StdRng::seed_from_u64(cfg.seed);
        let token = sample_from_logits(&logits, &cfg, &mut rng, &recent).unwrap();
        // Without penalty, argmax stays at index 1 even though token 1 is in recent.
        assert_eq!(
            token, 1,
            "V1_001: repeat_penalty=1.0 must be a no-op (argmax stays at 1)"
        );
    }

    /// V1_001 (repetition penalty): repeat_last_n == 0 is a no-op.
    #[test]
    fn rep_penalty_v1_001_no_op_when_repeat_last_n_zero() {
        let logits = vec![3.0, 5.0, 2.0, 4.0];
        let recent = vec![1, 1, 1];
        let cfg = mk_config_with_penalty(
            0.0, 1, 2.0, /* would penalize */
            0,   /* no-op */
            42,
        );

        let mut rng = StdRng::seed_from_u64(cfg.seed);
        let token = sample_from_logits(&logits, &cfg, &mut rng, &recent).unwrap();
        assert_eq!(
            token, 1,
            "V1_001: repeat_last_n=0 must be a no-op (argmax stays at 1)"
        );
    }

    /// V1_002: repeat_penalty > 1.0 down-weights repeated tokens (positive logit branch).
    #[test]
    fn rep_penalty_v1_002_down_weights_repeated() {
        // logits[1] = 5.0 (would be argmax); penalty 2 → 2.5. New argmax = 3 (4.0).
        let logits = vec![3.0, 5.0, 2.0, 4.0];
        let recent = vec![1, 1]; // token 1 repeated twice: still ONE application
        let cfg = mk_config_with_penalty(0.0, 1, 2.0, 100, 42);

        let mut rng = StdRng::seed_from_u64(cfg.seed);
        let token = sample_from_logits(&logits, &cfg, &mut rng, &recent).unwrap();
        assert_eq!(
            token, 3,
            "V1_002: repeat_penalty must shift argmax away from repeated token 1"
        );
    }

    /// V1_002 (contract text): recent = [2, 2, 2], penalty 2 halves logit 2 once, 10 → 5, not
    /// 10 / 2³ = 1.25. A token repeated k times is penalised once, as llama.cpp and Candle do.
    #[test]
    fn rep_penalty_v1_002_repeats_penalise_once() {
        let logits = vec![1.0, 10.0, 10.0, 4.0];
        let cfg = mk_config_with_penalty(0.0, 1, 2.0, 100, 42);
        // Once: [1, 10, 5, 4] → argmax 1. Compounded ×8: logit 2 → 0.039, argmax still 1, so
        // probe the value directly through the shared helper as well.
        let mut rng = StdRng::seed_from_u64(cfg.seed);
        let token = sample_from_logits(&logits, &cfg, &mut rng, &[2; 8]).unwrap();
        assert_eq!(token, 1);
        let mut p = logits.clone();
        crate::gguf::OwnedQuantizedModel::apply_repeat_penalty(&mut p, &[2, 2, 2], 2.0, 100);
        assert_eq!(
            p,
            vec![1.0, 10.0, 5.0, 4.0],
            "a repeated token is penalised once"
        );
        // Where it decides the token: 6.0 / 1.5 = 4.0 > 3.0 once; 6.0 / 1.5⁴ = 1.19 < 3.0 compounded.
        let close = vec![3.0, 6.0, 0.0, 0.0];
        let cfg = mk_config_with_penalty(0.0, 1, 1.5, 100, 42);
        let mut rng = StdRng::seed_from_u64(cfg.seed);
        let token = sample_from_logits(&close, &cfg, &mut rng, &[1, 1, 1, 1]).unwrap();
        assert_eq!(
            token, 1,
            "four repeats of token 1 must not compound to 1.5⁴"
        );
    }

    /// V1_002: negative logits get MULTIPLIED by penalty (Candle's convention).
    #[test]
    fn rep_penalty_v1_002_negative_logit_branch() {
        // Mix: logit[2] is negative. Penalty multiplies it (more negative).
        let logits = vec![3.0, 1.0, -2.0, 4.0]; // argmax = 3
        let recent = vec![2]; // token 2 has negative logit
        let cfg = mk_config_with_penalty(0.0, 1, 2.0, 100, 42);

        let mut rng = StdRng::seed_from_u64(cfg.seed);
        let token = sample_from_logits(&logits, &cfg, &mut rng, &recent).unwrap();
        // After penalty: [3.0, 1.0, -4.0, 4.0] → argmax still = 3, but logit[2]
        // is now more strongly suppressed. Confirms the branch ran without
        // accidentally amplifying.
        assert_eq!(token, 3, "V1_002 negative branch: argmax stays at 3");
    }

    /// V1_003: repeat_last_n bounds the penalty window correctly.
    #[test]
    fn rep_penalty_v1_003_window_bounds() {
        // recent = [1, 3, 3]: the last 2 hold only token 3; the last 3 hold tokens 1 and 3.
        // penalty 1.5, logit[1] = 10.0 vs logit[2] = 9.0.
        let logits = vec![1.0, 10.0, 9.0, 3.0];
        let recent = vec![1, 3, 3];

        let cfg_n2 = mk_config_with_penalty(0.0, 1, 1.5, 2, 42);
        let mut rng = StdRng::seed_from_u64(42);
        let token_n2 = sample_from_logits(&logits, &cfg_n2, &mut rng, &recent).unwrap();
        // Token 1 is outside the window: argmax stays 1 (10.0).
        assert_eq!(token_n2, 1, "V1_003 n=2: token 1 is outside the window");

        let cfg_n3 = mk_config_with_penalty(0.0, 1, 1.5, 3, 42);
        let mut rng = StdRng::seed_from_u64(42);
        let token_n3 = sample_from_logits(&logits, &cfg_n3, &mut rng, &recent).unwrap();
        // 10.0 / 1.5 = 6.67 < 9.0: argmax = 2.
        assert_eq!(token_n3, 2, "V1_003 n=3: token 1 is penalised");

        // n=0 means "no penalty" (per backwards-compat invariant).
        let logits_close = logits;
        let cfg_n0 = mk_config_with_penalty(0.0, 1, 1.5, 0, 42);
        let mut rng = StdRng::seed_from_u64(42);
        let token_n0 = sample_from_logits(&logits_close, &cfg_n0, &mut rng, &recent).unwrap();
        // No penalty: argmax = 1 (10.0).
        assert_eq!(token_n0, 1, "V1_003 n=0: no-op, argmax = 1");
    }
}