engine/crates/runtime/src/think.rs

//! Thought decoding behind a narrow trait: one `step` advances every active branch by one cycle and returns the ACCEPTED tokens per branch.
//! `GreedyDecoder` = plain greedy, one token per branch per step (all branches of all requests in ONE decode forward, weights read once).
//! `SpecDecoder` = MTP draft k tokens + one verify forward (`dllm_model::spec`): it may return several tokens per branch per cycle; the scheduler
//! stops consuming at the first stop token / token budget and ignores surplus. A decoder only writes KV at positions `pos ..= pos + lookahead()`
//! of the pages it is given (the scheduler guarantees those pages exist); everything beyond the accepted tokens is overwritten later.
//! Number answers go through the same decoder (`answer_step`): the accepted tokens come with the target's log-probabilities from the full logits.
use crate::model::{ModelRunner, RowBatch};
use anyhow::Result;
use dllm_model::spec::greedy_logprob;

/// one active thinking branch: feed `tok` at position `pos` (its KV slot is page `pages[pos / 16]`)
pub struct ThinkReq<'a> {
    pub pages: &'a [i32],
    pub pos: usize,
    pub tok: i32,
    /// stable id of the branch (scheduler branch slot); lets a decoder remember per-branch state between steps
    pub uid: usize,
}

pub trait ThinkDecoder: Send {
    /// extra positions beyond `pos` that one step may write
    fn lookahead(&self) -> usize {
        0
    }
    fn step(&mut self, m: &mut dyn ModelRunner, reqs: &[ThinkReq]) -> Result<Vec<Vec<i32>>>;
    /// number answers: like `step`, but every returned token is the greedy token of the target's full logits, with its log-probability
    fn answer_step(&mut self, m: &mut dyn ModelRunner, reqs: &[ThinkReq]) -> Result<Vec<Vec<(i32, f64)>>> {
        greedy_answer_step(m, reqs)
    }
    /// branch `uid` continues with `tok` at `pos`; row `row` of the last forward produced `tok` (lets a decoder seed its per-branch state)
    fn prime(&mut self, _m: &mut dyn ModelRunner, _uid: usize, _row: usize, _pos: usize, _tok: i32) -> Result<()> {
        Ok(())
    }
}

/// one decode row per answering branch (all in one forward per 16), greedy token + log-probability from the full logits
fn greedy_answer_step(m: &mut dyn ModelRunner, reqs: &[ThinkReq]) -> Result<Vec<Vec<(i32, f64)>>> {
    let mut out = Vec::with_capacity(reqs.len());
    for chunk in reqs.chunks(m.max_rows().min(16).max(1)) {
        let toks: Vec<i32> = chunk.iter().map(|r| r.tok).collect();
        let seqs: Vec<(&[i32], usize)> = chunk.iter().map(|r| (r.pages, r.pos)).collect();
        m.forward(&RowBatch::decode(&toks, &seqs), None)?;
        let lg = m.logits_rows(&(0..chunk.len()).collect::<Vec<_>>())?;
        out.extend(lg.iter().map(|z| vec![greedy_logprob(z)]));
    }
    Ok(out)
}

#[derive(Default)]
pub struct GreedyDecoder;

impl ThinkDecoder for GreedyDecoder {
    fn step(&mut self, m: &mut dyn ModelRunner, reqs: &[ThinkReq]) -> Result<Vec<Vec<i32>>> {
        let mut out = Vec::with_capacity(reqs.len());
        let cap = m.max_rows().max(1);
        for chunk in reqs.chunks(cap) {
            let toks: Vec<i32> = chunk.iter().map(|r| r.tok).collect();
            let seqs: Vec<(&[i32], usize)> = chunk.iter().map(|r| (r.pages, r.pos)).collect();
            let b = RowBatch::decode(&toks, &seqs);
            m.forward(&b, None)?;
            let rows: Vec<usize> = (0..chunk.len()).collect();
            let a = m.argmax_rows(&rows, &[])?;
            out.extend(a.into_iter().map(|t| vec![t]));
        }
        Ok(out)
    }
}

/// speculative decoder on the shared paged pool: MTP draft chain + one batched verify forward per step (see `dllm_model::spec`)
pub struct SpecDecoder {
    pub spec: dllm_model::spec::Spec,
    fallback: GreedyDecoder,
    /// above this many simultaneously thinking branches the plain batched greedy step (CUDA-graph buckets) is faster than a verify batch: delegate to it
    pub max_branches: usize,
}
impl SpecDecoder {
    pub fn new(spec: dllm_model::spec::Spec, max_branches: usize) -> Self {
        SpecDecoder { spec, fallback: GreedyDecoder, max_branches }
    }
}
impl ThinkDecoder for SpecDecoder {
    fn lookahead(&self) -> usize {
        self.spec.max_lookahead()
    }
    fn step(&mut self, m: &mut dyn ModelRunner, reqs: &[ThinkReq]) -> Result<Vec<Vec<i32>>> {
        match m.as_gpu().filter(|_| reqs.len() <= self.max_branches) {
            Some(g) => {
                let rs: Vec<dllm_model::spec::SpecReq> = reqs.iter().map(|r| dllm_model::spec::SpecReq { uid: r.uid, pages: r.pages, pos: r.pos, tok: r.tok }).collect();
                self.spec.step(g, &rs)
            }
            None => self.fallback.step(m, reqs),
        }
    }
    fn answer_step(&mut self, m: &mut dyn ModelRunner, reqs: &[ThinkReq]) -> Result<Vec<Vec<(i32, f64)>>> {
        match m.as_gpu().filter(|_| reqs.len() <= self.max_branches) {
            Some(g) => {
                let rs: Vec<dllm_model::spec::SpecReq> = reqs.iter().map(|r| dllm_model::spec::SpecReq { uid: r.uid, pages: r.pages, pos: r.pos, tok: r.tok }).collect();
                self.spec.answer_step(g, &rs)
            }
            None => greedy_answer_step(m, reqs),
        }
    }
    fn prime(&mut self, m: &mut dyn ModelRunner, uid: usize, row: usize, pos: usize, tok: i32) -> Result<()> {
        match m.as_gpu() {
            Some(g) => self.spec.prime(g, uid, row, pos, tok),
            None => Ok(()),
        }
    }
}