engine/crates/model/src/gpu_model/graphs.rs

//! CUDA-graph replay of the decode forward. Decode batches are padded by the caller to a row bucket (duplicates of row 0, idempotent), the number of KV splits is
//! quantised to {1,2,4,8}; per (rows, splits, metadata buffers) key the first use runs the forward eagerly, captures the identical launch sequence and later calls replay it
//! (metadata contents are uploaded to the fixed `ws` buffers before the replay; device pointers never change - a block-table reallocation drops all graphs).
use super::*;
use dllm_kernels::GraphExec;
use std::collections::HashMap;
use std::sync::atomic::Ordering;

pub(super) struct GraphStore {
    pub buckets: Vec<usize>,
    map: HashMap<(usize, usize, u64, u64, u64), Box<dyn GraphExec>>,
}

const GRAPH_BUCKETS: [usize; 7] = [1, 2, 4, 8, 16, 24, 32];
const GRAPH_SPLITS: [usize; 4] = [1, 2, 4, 8];

pub fn graph_bucket(r: usize) -> Option<usize> {
    GRAPH_BUCKETS.iter().copied().find(|&b| b >= r)
}
pub(super) fn quant_splits(s: usize) -> usize {
    GRAPH_SPLITS.iter().copied().find(|&b| b >= s).unwrap_or(8)
}

impl GpuModel {
    pub fn graphs_enabled(&self) -> bool {
        self.graphs.lock().unwrap().is_some()
    }

    pub(super) fn graphs_clear(&self) {
        if let Some(g) = self.graphs.lock().unwrap().as_mut() {
            g.map.clear();
        }
    }

    /// enable graph replay for decode batches of at most 32 rows and capture every (bucket, splits) variant now (eager run + capture each).
    /// Call before serving: it overwrites position 0 of KV page 0.
    pub fn enable_graphs(&mut self) -> Result<()> {
        // fixed block-table capacity: a reallocation would invalidate the captured pointers
        let want = 64 * 1024;
        if self.ws.bt_cap < want {
            self.ws.bt_cap = want;
            self.ws.bt = DBuf::zeros(&self.gpu.be, want * 4)?;
        }
        *self.graphs.lock().unwrap() = Some(GraphStore { buckets: GRAPH_BUCKETS.to_vec(), map: HashMap::new() });
        let res = self.warm_graphs();
        if res.is_err() || !self.graphs_enabled() {
            *self.graphs.lock().unwrap() = None;
        }
        res
    }

    fn warm_graphs(&mut self) -> Result<()> {
        for &r in &GRAPH_BUCKETS {
            if r > self.ws.rows {
                continue;
            }
            for &s in &GRAPH_SPLITS {
                let mut b = RowBatch { decode: true, ..Default::default() };
                for _ in 0..r {
                    b.tokens.push(0);
                    b.pos.push(0);
                    b.slot.push(0);
                    b.bt_off.push(0);
                    b.kv_hi.push(0);
                    b.kv_hi_sliding.push(0);
                }
                b.block_table = vec![0];
                self.force_splits.store(s, Ordering::Relaxed);
                let res = self.forward(&b);
                self.force_splits.store(0, Ordering::Relaxed);
                res?;
            }
        }
        self.gpu.be.sync()?;
        Ok(())
    }

    /// called by `forward_dev` for decode batches: replay / capture, or None when this batch is not graph-eligible
    pub(super) fn try_graph(&self, r: usize, splits: usize, m: &DevMeta) -> Option<Result<()>> {
        if self.graph_busy.load(Ordering::Relaxed) {
            return None;
        }
        let mut g = self.graphs.lock().unwrap();
        let st = g.as_mut()?;
        if !st.buckets.contains(&r) || !GRAPH_SPLITS.contains(&splits) {
            return None;
        }
        let key = (r, splits, m.tokens, m.pos, m.bt);
        if let Some(ex) = st.map.get(&key) {
            return Some(ex.launch());
        }
        self.graph_busy.store(true, Ordering::Relaxed);
        // eager run first: loads lazily-initialised kernels and produces this call's result
        let eager = self.forward_rows(r, true, splits, *m).and_then(|_| self.gpu.be.sync());
        let res = match eager {
            Err(e) => Err(e),
            Ok(()) => {
                let cap = (|| -> Result<Box<dyn GraphExec>> {
                    self.gpu.be.graph_begin()?;
                    let body = self.forward_rows(r, true, splits, *m);
                    let ex = self.gpu.be.graph_end();
                    body?;
                    ex
                })();
                match cap {
                    Ok(ex) => {
                        st.map.insert(key, ex);
                    }
                    Err(e) => {
                        eprintln!("CUDA graph capture failed ({e:#}); graphs disabled");
                        *g = None;
                    }
                }
                Ok(())
            }
        };
        self.graph_busy.store(false, Ordering::Relaxed);
        Some(res)
    }
}