Speculative decoding

/// Speculative decoding generator
///
/// Uses a draft model to speculate multiple tokens ahead, then verifies them
/// with the target model in a single forward pass.
pub struct SpeculativeGenerate<'a, M, D, C>
where
    M: SpeculativeModel,
    D: SpeculativeModel,
    C: KeyValueCache + Default,
{
    /// Target (large) model
    target_model: &'a mut M,
    /// Draft (small) model
    draft_model: &'a mut D,
    /// KV cache for target model
    target_cache: &'a mut Vec<Option<C>>,
    /// KV cache for draft model
    draft_cache: &'a mut Vec<Option<C>>,
    /// Number of draft tokens to generate per step
    num_draft_tokens: usize,
    /// Temperature for sampling
    temperature: f32,
    // ...
}

/// Generate draft tokens using the draft model
fn generate_draft_tokens(&mut self, start_token: &Array) -> Result<Vec<Array>, Exception> {
    let mut tokens = Vec::with_capacity(self.num_draft_tokens);
    let mut current = start_token.clone();

    for _ in 0..self.num_draft_tokens {
        let input = current.index((.., NewAxis));
        let logits = self.draft_model.forward_speculative(&input, self.draft_cache)?;
        let token = self.sample(&logits)?;
        let _ = async_eval([&token]);
        tokens.push(token.clone());
        current = token;
    }

    Ok(tokens)
}

/// Verify draft tokens with target model
/// Returns (accepted_count, all_tokens, all_logprobs)
fn verify_draft_tokens(
    &mut self,
    input_tokens: &Array,
) -> Result<(usize, Vec<Array>, Vec<Array>), Exception> {
    // Forward all tokens through target model at once
    let logits = self.target_model.forward_speculative(input_tokens, self.target_cache)?;

    let seq_len = input_tokens.shape()[1] as usize;
    let mut tokens = Vec::with_capacity(seq_len);
    let mut logprobs = Vec::with_capacity(seq_len);

    // Sample from each position's logits
    for i in 0..seq_len {
        let pos_logits = logits.index((.., i as i32, ..));
        let token = self.sample(&pos_logits)?;

        // Compute log probabilities
        let log_sum_exp = logsumexp_axis(&pos_logits, -1, true)?;
        let lp = pos_logits.subtract(&log_sum_exp)?;

        tokens.push(token);
        logprobs.push(lp);
    }

    eval(&tokens)?;
    Ok((seq_len, tokens, logprobs))
}

// Compare draft tokens with target tokens to find acceptance
let mut accepted = 0;
for i in 0..draft_tokens.len() {
    let draft_id = draft_tokens[i].item::<u32>();
    let target_id = target_tokens[i].item::<u32>();

    if draft_id == target_id {
        accepted += 1;
        // Queue accepted token
        self.pending_tokens.push(SpeculativeToken {
            token: draft_tokens[i].clone(),
            logprobs: target_logprobs[i].clone(),
            from_draft: true,
        });
    } else {
        break;  // Stop at first mismatch
    }
}

// The token at position `accepted` is from target model (correction or next)
let final_token = target_tokens[accepted].clone();

/// Result of speculative decoding step
pub struct SpeculativeToken {
    /// The generated token
    pub token: Array,
    /// Log probabilities
    pub logprobs: Array,
    /// Whether this token was accepted from draft model (true) or from target model (false)
    pub from_draft: bool,
}

enum SpeculativeState<'a> {
    /// Initial state - need to process prompt
    Prefill { prompt: &'a Array },
    /// Main generation loop
    Generate { last_token: Array },
}

/// Token count
token_count: usize,
/// Pending tokens from verification (when draft tokens are accepted)
pending_tokens: Vec<SpeculativeToken>,