engine/crates/runtime/src/model.rs

//! Model-side traits of the runtime: `ModelRunner` (what the scheduler needs from the GPU model), `VisionEncoder`, `ImagePrep`, soft-token injection.
use anyhow::Result;
pub use dllm_model::gpu_model::RowBatch;
use std::any::Any;

/// encoded soft tokens of one image: `n` x hidden bf16 rows on the device (`dev`), or host identities for the fake model.
pub struct SoftTokens {
    pub n: usize,
    pub dev: u64,
    /// fake model only: one identity per soft token
    pub host: Option<Vec<u64>>,
    /// keeps the device allocation alive
    pub keep: Option<Box<dyn Any + Send + Sync>>,
}

pub trait VisionEncoder: Send {
    /// `patches` = ImagePatches.patches, (gh, gw) = patch grid; returns n_soft = gh*gw/9 soft tokens
    fn encode(&mut self, patches: &[u8], gh: usize, gw: usize) -> Result<SoftTokens>;
}

/// one injected soft token: batch `row` receives token `idx` of image `img` (index into `SoftSel::imgs`)
#[derive(Debug, Clone, Copy)]
pub struct SoftRow {
    pub img: usize,
    pub row: i32,
    pub idx: i32,
}

pub struct SoftSel<'a> {
    pub imgs: Vec<&'a SoftTokens>,
    pub rows: Vec<SoftRow>,
}

/// preprocessed image (CPU side): identity hash + soft token count + patches for the tower
#[derive(Clone)]
pub struct PreparedImage {
    pub n_soft: usize,
    pub gh: usize,
    pub gw: usize,
    pub patches: std::sync::Arc<Vec<u8>>,
    pub hash: u64,
}

pub trait ImagePrep: Send + Sync {
    /// decode + preprocess raw encoded bytes. Errors are reported to the client as "cannot decode image ...".
    fn prepare(&self, bytes: &[u8]) -> Result<PreparedImage>;
}

pub trait ModelRunner: Send {
    fn num_pages(&self) -> usize;
    fn max_rows(&self) -> usize;
    /// run the decoder over `b` (K/V of all rows are appended to the pool first); `soft` replaces the embeddings of the listed rows
    fn forward(&mut self, b: &RowBatch, soft: Option<&SoftSel>) -> Result<()>;
    /// letter logits of the given rows of the LAST forward
    fn letter_logits(&mut self, rows: &[usize], label_ids: &[i32]) -> Result<Vec<Vec<f32>>>;
    /// greedy token of the rows of the last forward, never returning an id in `exclude` (<= 3 ids)
    fn argmax_rows(&mut self, rows: &[usize], exclude: &[i32]) -> Result<Vec<i32>>;
    /// softcapped full-vocab logits of the rows of the last forward (number answers: greedy token + its probability)
    fn logits_rows(&mut self, rows: &[usize]) -> Result<Vec<Vec<f32>>>;
    /// wait for the device
    fn sync(&mut self) -> Result<()> {
        Ok(())
    }
    /// the concrete GPU model (speculative decoding needs its KV pools / hidden rows); None for fakes
    fn as_gpu(&mut self) -> Option<&mut dllm_model::gpu_model::GpuModel> {
        None
    }
    /// capture CUDA graphs for the decode row buckets (call once before serving); Ok(false) = not supported
    fn enable_graphs(&mut self) -> Result<bool> {
        Ok(false)
    }
    /// sliding-window KV pool: pages of the separate sliding-layer pool (ids 1..=n; 0 = the model has no sliding layers to manage)
    fn sliding_pages(&self) -> usize {
        0
    }
    /// attention window of the sliding layers in positions (incl. self)
    fn sliding_window(&self) -> usize {
        0
    }
    /// sliding page map update before the next forward: smap[first + i] = vals[i] (full page id -> sliding page id, 0 = none)
    fn set_sliding_map(&mut self, _first: usize, _vals: &[i32]) -> Result<()> {
        Ok(())
    }
}

// ---------------------------------------------------------------------------------------------------------------- GpuModel
impl ModelRunner for dllm_model::gpu_model::GpuModel {
    fn num_pages(&self) -> usize {
        self.opts.num_pages
    }
    fn max_rows(&self) -> usize {
        self.opts.max_rows
    }
    fn forward(&mut self, b: &RowBatch, soft: Option<&SoftSel>) -> Result<()> {
        gpu_forward(self, b, soft)
    }
    fn letter_logits(&mut self, rows: &[usize], label_ids: &[i32]) -> Result<Vec<Vec<f32>>> {
        dllm_model::gpu_model::GpuModel::letter_logits(self, rows, label_ids)
    }
    fn argmax_rows(&mut self, rows: &[usize], exclude: &[i32]) -> Result<Vec<i32>> {
        let e0 = exclude.first().copied().unwrap_or(-1);
        let e1 = exclude.get(1).copied().unwrap_or(-1);
        if exclude.len() <= 2 && rows.len() <= 16 && !rows.is_empty() && rows.iter().enumerate().all(|(i, &r)| r == i) {
            // fused FP8 lm_head scan + argmax over the first `rows.len()` hidden rows (decode steps)
            return self.greedy_first_rows(rows.len(), e0, e1);
        }
        let mut r = dllm_model::gpu_model::GpuModel::argmax_rows(self, rows, e0, e1)?;
        if exclude.len() > 2 {
            // rare third exclusion: recheck on the host with full logits for rows that landed on an excluded id
            let bad: Vec<usize> = (0..rows.len()).filter(|&i| exclude.contains(&r[i])).collect();
            if !bad.is_empty() {
                let rr: Vec<usize> = bad.iter().map(|&i| rows[i]).collect();
                let lg = dllm_model::gpu_model::GpuModel::logits_rows(self, &rr)?;
                for (k, &i) in bad.iter().enumerate() {
                    let mut best = 0usize;
                    for (v, &x) in lg[k].iter().enumerate() {
                        if exclude.contains(&(v as i32)) {
                            continue;
                        }
                        if exclude.contains(&(best as i32)) || x > lg[k][best] {
                            best = v;
                        }
                    }
                    r[i] = best as i32;
                }
            }
        }
        Ok(r)
    }
    fn logits_rows(&mut self, rows: &[usize]) -> Result<Vec<Vec<f32>>> {
        dllm_model::gpu_model::GpuModel::logits_rows(self, rows)
    }
    fn sync(&mut self) -> Result<()> {
        self.gpu.be.sync()
    }
    fn as_gpu(&mut self) -> Option<&mut dllm_model::gpu_model::GpuModel> {
        Some(self)
    }
    fn enable_graphs(&mut self) -> Result<bool> {
        dllm_model::gpu_model::GpuModel::enable_graphs(self)?;
        Ok(self.graphs_enabled())
    }
    fn sliding_pages(&self) -> usize {
        dllm_model::gpu_model::GpuModel::sliding_pages(self)
    }
    fn sliding_window(&self) -> usize {
        dllm_model::gpu_model::GpuModel::sliding_window(self)
    }
    fn set_sliding_map(&mut self, first: usize, vals: &[i32]) -> Result<()> {
        dllm_model::gpu_model::GpuModel::set_sliding_map(self, first, vals)
    }
}

fn gpu_forward(m: &mut dllm_model::gpu_model::GpuModel, b: &RowBatch, soft: Option<&SoftSel>) -> Result<()> {
    let Some(sel) = soft else {
        if b.decode && m.graphs_enabled() {
            if let Some(bk) = dllm_model::gpu_model::graph_bucket(b.rows()) {
                if bk > b.rows() {
                    // pad to the graph bucket with duplicates of row 0 (identical metadata => identical, idempotent KV writes)
                    let mut p = b.clone();
                    for _ in b.rows()..bk {
                        p.tokens.push(b.tokens[0]);
                        p.pos.push(b.pos[0]);
                        p.slot.push(b.slot[0]);
                        p.bt_off.push(b.bt_off[0]);
                        p.kv_hi.push(b.kv_hi[0]);
                        p.kv_hi_sliding.push(b.kv_hi_sliding[0]);
                    }
                    return m.forward(&p);
                }
            }
        }
        return m.forward(b);
    };
    use dllm_kernels::DBuf;
    let h2 = m.cfg.hidden_size * 2;
    // all images referenced by this batch are concatenated on the device in the order of `sel.imgs` (single image: used in place)
    let mut base = vec![0i32; sel.imgs.len()];
    let mut total = 0usize;
    for (i, t) in sel.imgs.iter().enumerate() {
        base[i] = total as i32;
        total += t.n;
    }
    let concat;
    let src = if sel.imgs.len() == 1 {
        sel.imgs[0].dev
    } else {
        concat = DBuf::zeros(&m.gpu.be, total * h2)?;
        for (i, t) in sel.imgs.iter().enumerate() {
            m.gpu.be.copy_d2d(concat.at(base[i] as usize * h2), t.dev, t.n * h2)?;
        }
        concat.ptr
    };
    let mut bb = b.clone();
    bb.soft = Some(dllm_model::gpu_model::SoftInject { src, row_idx: sel.rows.iter().map(|r| r.row).collect(), src_idx: sel.rows.iter().map(|r| base[r.img] + r.idx).collect() });
    m.forward(&bb)
}

/// image preprocessing (CPU decode + resize + patchify, crate dllm-vision)
pub struct VisionPrep(pub dllm_vision::PreprocConfig);
impl ImagePrep for VisionPrep {
    fn prepare(&self, bytes: &[u8]) -> Result<PreparedImage> {
        // concurrency: the first image in preparation uses the global rayon pool (all cores); images prepared at the same time run on a
        // single-thread pool of their own (per calling thread). Many requests sharing the global pool were starved unfairly by rayon's
        // work stealing (16 concurrent images: p90 217 ms, max 308 ms per image vs 29 / 43 ms this way, 4x the throughput); the
        // results are identical (tests/prep_timing.rs).
        static ACTIVE: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
        thread_local!(static SOLO: Option<rayon::ThreadPool> = rayon::ThreadPoolBuilder::new().num_threads(1).build().ok());
        struct Dec;
        impl Drop for Dec {
            fn drop(&mut self) {
                ACTIVE.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
            }
        }
        let alone = ACTIVE.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0;
        let _d = Dec;
        let run = || dllm_vision::preprocess(bytes, &self.0);
        let p = if alone { run() } else { SOLO.with(|s| match s { Some(pool) => pool.install(run), None => run() }) };
        let p = p.map_err(|e| anyhow::anyhow!("{e}"))?;
        Ok(PreparedImage { n_soft: p.n_soft, gh: p.gh, gw: p.gw, patches: std::sync::Arc::new(p.patches), hash: p.hash })
    }
}

/// the vision tower on the GPU
pub struct TowerEncoder(pub dllm_model::vision::VisionTower);
impl VisionEncoder for TowerEncoder {
    fn encode(&mut self, patches: &[u8], gh: usize, gw: usize) -> Result<SoftTokens> {
        let buf = self.0.encode(patches, gh, gw)?;
        Ok(SoftTokens { n: gh * gw / 9, dev: buf.ptr, host: None, keep: Some(Box::new(buf)) })
    }
}