File size: 4,714 Bytes
19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | //! 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)
}
}
|