File size: 10,303 Bytes
ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 19886eb ab0f9c6 a31d7f2 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 | //! 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)) })
}
}
|