Something went wrong. Try again.
a tiny pet with an even tinier brain, living in your browser
Something went wrong. Try again.
38 kB · 1040 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041//! One mutable ONNX session per worker. No concurrent calls into a session.use anyhow::{Context, Result, ensure};use ort::{ memory::{AllocationDevice, Allocator, AllocatorType, MemoryInfo, MemoryType}, session::Session, value::{DynTensor, DynValue, Tensor, TensorElementType, ValueType},};use serde::{Deserialize, Serialize};use std::{path::Path, time::Instant};use tokenizers::Tokenizer;/// Requests, the stop rule and context fitting are shared with the browser's worker.pub use tokenling_core::turn::{Generation, TOP_K, TOP_P};use tokenling_core::world;
#[derive(Clone, Debug, Deserialize, Serialize)]#[serde(rename_all = "camelCase")]pub struct Generated { pub text: String, pub prompt: String, pub reason: String, pub prompt_tokens: usize, pub prefill_ms: f64, /// Post-first-token samples, including the terminal EOS compute when present. pub decode_tokens: usize, /// First-to-last sampled token interval, excluding final text decoding. pub decode_ms: f64, pub tps: f64, pub generated_tokens: usize, pub total_ms: f64,}/// Export weight precision, independent of the graph's cache element type.#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]#[serde(rename_all = "lowercase")]pub enum ModelDtype { #[default] Q4, Q8, Fp32,}impl std::str::FromStr for ModelDtype { type Err = anyhow::Error; fn from_str(value: &str) -> Result<Self> { match value { "q4" => Ok(Self::Q4), "q8" => Ok(Self::Q8), "fp32" => Ok(Self::Fp32), _ => anyhow::bail!("model dtype must be q4, q8 or fp32"), } }}impl ModelDtype { pub fn filename(self) -> &'static str { match self { Self::Q4 => "model_q4.onnx", Self::Q8 => "model_quantized.onnx", Self::Fp32 => "model.onnx", } }}
#[derive(Debug)]struct Dimensions { layers: usize, kv_heads: usize, width: usize, vocab: usize, context: usize, eos: Vec<u32>,}impl Dimensions { fn from_config(cfg: &serde_json::Value) -> Result<Self> { let model = cfg["model_type"] .as_str() .context("config requires model_type")?; ensure!( matches!(model, "gpt2" | "gpt_neox" | "llama"), "unsupported model_type: {model}" ); // Accept either spelling, but never silently choose between conflicting aliases. let number = |keys: &[&str]| -> Result<usize> { let mut found = None; for key in keys { if let Some(value) = cfg.get(*key) { let n = value .as_u64() .filter(|n| *n > 0) .with_context(|| format!("config {key} must be a positive integer"))?; let n = usize::try_from(n)?; ensure!( found.is_none_or(|old| old == n), "conflicting config dimensions: {keys:?}" ); found = Some(n); } } found.with_context(|| format!("config requires a positive dimension: {keys:?}")) }; let heads = number(&["num_attention_heads", "n_head"])?; let hidden = number(&["hidden_size", "n_embd"])?; let width = if cfg.get("head_dim").is_some_and(|v| !v.is_null()) { number(&["head_dim"])? } else { ensure!( hidden % heads == 0, "hidden width must divide attention heads exactly" ); hidden / heads }; let kv_heads = if cfg.get("num_key_value_heads").is_some_and(|v| !v.is_null()) { number(&["num_key_value_heads"])? } else { heads }; ensure!( heads % kv_heads == 0, "KV heads must divide attention heads" ); ensure!( model == "llama" || (kv_heads == heads && width == hidden / heads && hidden % heads == 0), "GPT2/NeoX require ordinary attention dimensions" ); let vocab = number(&["vocab_size"])?; ensure!( vocab <= u32::MAX as usize, "vocabulary exceeds token ID range" ); let raw = cfg .get("eos_token_id") .context("config requires eos_token_id")?; let ids: Vec<&serde_json::Value> = match raw { serde_json::Value::Array(ids) => ids.iter().collect(), _ => vec![raw], }; ensure!(!ids.is_empty(), "EOS list must not be empty"); let mut eos = Vec::with_capacity(ids.len()); for id in ids { let id = id .as_u64() .context("EOS must be an integer or list of integers")?; ensure!(id < vocab as u64, "EOS token is outside vocabulary"); ensure!(!eos.contains(&(id as u32)), "duplicate EOS token"); eos.push(id as u32); } Ok(Self { layers: number(&["num_hidden_layers", "n_layer"])?, kv_heads, width, vocab, context: number(&["max_position_embeddings", "n_positions"])?, eos, }) } fn fit( &self, g: &Generation, count: &mut dyn FnMut(&str) -> std::result::Result<usize, String>, ) -> Result<String> { let ctx = g.ctx.min(self.context).min(1024); let prompt = world::fit(&g.header, &g.lines, &g.prefix, g.max_new, ctx, count) .map_err(anyhow::Error::msg)?; let tokens = count(&prompt).map_err(anyhow::Error::msg)?; ensure!( tokens <= ctx.saturating_sub(g.max_new), "tokenized prompt exceeds model positional capacity" ); Ok(prompt) }}
#[derive(Debug)]struct Input { name: String, ty: TensorElementType, rank: usize,}#[derive(Debug)]struct Cache { input: String, output: String, ty: TensorElementType,}#[derive(Debug)]struct Contract { dimensions: Dimensions, inputs: Vec<Input>, caches: Vec<Cache>,}fn tensor_meta<'a>(name: &str, value: &'a ValueType) -> Result<(TensorElementType, &'a [i64])> { match value { ValueType::Tensor { ty, shape, .. } => Ok((*ty, shape)), _ => anyhow::bail!("{name} must be a tensor"), }}fn cache_key(name: &str, output: bool) -> Option<(usize, &str)> { let mut parts = name.split('.'); let prefix = parts.next()?; if if output { prefix != "present" && prefix != "present_key_values" } else { prefix != "past_key_values" } { return None; } let layer = parts.next()?; let index = layer.parse::<usize>().ok()?; // Reject aliases such as .00.key; each layer must have one canonical key/value. if index.to_string() != layer { return None; } let kind = parts.next()?; (matches!(kind, "key" | "value") && parts.next().is_none()).then_some((index, kind))}fn cache_meta(name: &str, value: &ValueType, d: &Dimensions) -> Result<TensorElementType> { let (ty, shape) = tensor_meta(name, value)?; ensure!( matches!( ty, TensorElementType::Float16 | TensorElementType::Float32 | TensorElementType::Float64 ), "{name} cache must be float16/float32/float64" ); // Fixed head axes prove the layout. Dynamic head axes are ambiguous, even // when a config happens to suggest the intended interpretation. ensure!( shape.len() == 4 && (shape[0] == 1 || shape[0] == -1) && shape[1] == d.kv_heads as i64 && shape[2] == -1 && shape[3] == d.width as i64, "{name} requires [batch=1, KV heads={}, dynamic sequence, head width={}], got {shape:?}", d.kv_heads, d.width ); if let ValueType::Tensor { dimension_symbols, .. } = value { ensure!( !dimension_symbols[0].to_lowercase().contains("seq"), "{name} has a sequence axis in the batch position" ); } Ok(ty)}impl Contract { fn new( d: Dimensions, inputs: &[(String, ValueType)], outputs: &[(String, ValueType)], ) -> Result<Self> { let unique = |values: &[(String, ValueType)]| -> Result<()> { let mut names = std::collections::HashSet::new(); ensure!( values.iter().all(|(name, _)| names.insert(name)), "duplicate graph tensor name" ); Ok(()) }; unique(inputs)?; unique(outputs)?; let logits = outputs .iter() .find(|(name, _)| name == "logits") .context("graph requires logits")?; let (ty, shape) = tensor_meta("logits", &logits.1)?; ensure!(ty == TensorElementType::Float32, "logits must be float32"); ensure!( shape.len() == 3 && (shape[0] == -1 || shape[0] == 1) && (shape[1] == -1 || shape[1] == 1) && (shape[2] == -1 || shape[2] == d.vocab as i64), "logits must be [batch=1, sequence, config vocabulary={}], got {shape:?}", d.vocab ); let mut presents = std::collections::BTreeMap::new(); for (name, value) in outputs { if name == "logits" { continue; } let key = cache_key(name, true) .with_context(|| format!("unsupported causal output: {name}"))?; ensure!( presents.insert(key, (name, value)).is_none(), "ambiguous present cache: {name}" ); } let mut caches = Vec::new(); let mut ordinary = Vec::new(); let mut seen = std::collections::BTreeSet::new(); for (name, value) in inputs { if let Some(key) = cache_key(name, false) { let (output, output_meta) = presents .get(&key) .with_context(|| format!("missing present cache for {name}"))?; let ty = cache_meta(name, value, &d)?; ensure!( cache_meta(output, output_meta, &d)? == ty, "cache dtype changes for {name}" ); ensure!(seen.insert(key), "ambiguous past cache: {name}"); caches.push(Cache { input: name.clone(), output: (*output).clone(), ty, }); } else { ensure!( matches!( name.as_str(), "input_ids" | "attention_mask" | "position_ids" | "token_type_ids" | "cache_position" ), "unsupported causal input: {name}; use an unmerged with-past export" ); let (ty, shape) = tensor_meta(name, value)?; ensure!( matches!(ty, TensorElementType::Int64 | TensorElementType::Int32), "{name} must be int64/int32" ); let rank = shape.len(); ensure!( (rank == 2 && (shape[0] == -1 || shape[0] == 1) && shape[1] == -1) || (name == "cache_position" && rank == 1 && shape[0] == -1), "unsupported {name} shape: {shape:?}" ); ordinary.push(Input { name: name.clone(), ty, rank, }); } } ensure!( ordinary.iter().any(|i| i.name == "input_ids"), "graph requires input_ids" ); let expected: std::collections::BTreeSet<_> = (0..d.layers) .flat_map(|i| [(i, "key"), (i, "value")]) .collect(); ensure!( seen == expected && presents .keys() .copied() .collect::<std::collections::BTreeSet<_>>() == expected, "graph cache layers/key/value tensors disagree with config" ); Ok(Self { dimensions: d, inputs: ordinary, caches, }) } fn feed(&self, ids: Vec<i64>, past_length: usize) -> Result<Vec<(String, DynValue)>> { let tokens = ids.len(); let total = past_length + tokens; let mut ids = Some(ids); self.inputs .iter() .map(|input| { let values: Vec<i64> = match input.name.as_str() { "input_ids" => ids.take().context("duplicate input_ids")?, "attention_mask" => vec![1; total], "position_ids" | "cache_position" => { (past_length as i64..total as i64).collect() } "token_type_ids" => vec![0; tokens], _ => unreachable!("validated causal input"), }; let shape = if input.rank == 1 { vec![values.len()] } else { vec![1, values.len()] }; let tensor = match input.ty { TensorElementType::Int64 => Tensor::from_array((shape, values))?.into_dyn(), TensorElementType::Int32 => { let values = values .into_iter() .map(i32::try_from) .collect::<std::result::Result<Vec<_>, _>>()?; Tensor::from_array((shape, values))?.into_dyn() } _ => unreachable!("validated integer input"), }; Ok((input.name.clone(), tensor)) }) .collect() } fn validate_logits(&self, shape: &[i64], tokens: usize) -> Result<usize> { ensure!( shape.len() == 3 && shape[0] == 1 && (shape[1] == 1 || shape[1] == tokens as i64) && shape[2] == self.dimensions.vocab as i64, "invalid logits shape: {shape:?}" ); Ok(shape[2] as usize) } fn validate_cache(&self, cache: &Cache, value: &ValueType, length: usize) -> Result<()> { let (ty, shape) = tensor_meta(&cache.output, value)?; let d = &self.dimensions; ensure!( ty == cache.ty && shape == [1, d.kv_heads as i64, length as i64, d.width as i64], "invalid {} cache dtype/shape: {value:?}", cache.output ); Ok(()) }}pub struct Brain { session: Session, tokenizer: Tokenizer, contract: Contract, gpu: Option<i32>,}impl Brain { pub fn load( model_dir: &Path, dtype: ModelDtype, threads: usize, gpu: Option<i32>, ) -> Result<Self> { ensure!(threads > 0, "inference threads must be positive"); ensure!( gpu.is_none_or(|device| device >= 0), "CUDA device must be nonnegative" ); let tokenizer = Tokenizer::from_file(model_dir.join("tokenizer.json")) .map_err(|e| anyhow::anyhow!(e.to_string()))?; let cfg: serde_json::Value = serde_json::from_slice(&std::fs::read(model_dir.join("config.json"))?)?; let dimensions = Dimensions::from_config(&cfg)?; ensure!( tokenizer .get_vocab(true) .values() .all(|id| (*id as usize) < dimensions.vocab), "tokenizer emits IDs outside model vocabulary" ); let mut builder = Session::builder()?; builder = builder .with_intra_threads(threads) .map_err(|e| anyhow::anyhow!(e.to_string()))?; builder = builder .with_inter_threads(1) .map_err(|e| anyhow::anyhow!(e.to_string()))?; for (key, value) in [ ("session.disable_quant_qdq", "1"), ("session.disable_prepacking", "1"), ("session.intra_op.allow_spinning", "0"), ] { builder = builder .with_config_entry(key, value) .map_err(|e| anyhow::anyhow!(e.to_string()))?; } if let Some(device) = gpu { builder = builder .with_execution_providers([ort::ep::CUDA::default() .with_device_id(device) .with_tf32(false) .build() .error_on_failure()]) .map_err(|e| anyhow::anyhow!(e.to_string()))?; } let session = builder .commit_from_file(model_dir.join("onnx").join(dtype.filename())) .map_err(|e| anyhow::anyhow!(e.to_string()))?; let inputs = session .inputs() .iter() .map(|i| (i.name().to_owned(), i.dtype().clone())) .collect::<Vec<_>>(); let outputs = session .outputs() .iter() .map(|o| (o.name().to_owned(), o.dtype().clone())) .collect::<Vec<_>>(); let contract = Contract::new(dimensions, &inputs, &outputs)?; let mut brain = Self { session, tokenizer, contract, gpu, }; brain.generate(&Generation { header: String::new(), lines: vec![], prefix: "hi".into(), mode: "line".into(), max_new: 1, temperature: 0.0, top_k: TOP_K, top_p: TOP_P, ctx: 1024, })?; Ok(brain) } pub fn count(&self, text: &str) -> Result<usize> { Ok(self .tokenizer .encode(text, false) .map_err(|e| anyhow::anyhow!(e.to_string()))? .len()) } pub fn fit(&self, g: &Generation) -> Result<String> { self.contract .dimensions .fit(g, &mut |text| self.count(text).map_err(|e| e.to_string())) } pub fn generate(&mut self, g: &Generation) -> Result<Generated> { self.generate_cancellable(g, || false) } pub fn generate_cancellable( &mut self, g: &Generation, cancelled: impl Fn() -> bool, ) -> Result<Generated> { self.generate_inner(g, cancelled, false) } /// Diagnostic compute only: fixed greedy samples, including EOS, with no /// wire stop. Production callers must use generate/generate_cancellable. pub fn benchmark(&mut self, g: &Generation, tokens: usize) -> Result<Generated> { ensure!((1..=220).contains(&g.max_new), "invalid output budget"); ensure!( (1..=220).contains(&tokens), "invalid benchmark output budget" ); ensure!( g.temperature == 0.0, "benchmark requires greedy temperature 0" ); let mut request = g.clone(); request.max_new = tokens; self.generate_inner(&request, || false, true) } fn generate_inner( &mut self, g: &Generation, cancelled: impl Fn() -> bool, benchmark: bool, ) -> Result<Generated> { ensure!((1..=220).contains(&g.max_new), "invalid output budget"); ensure!( matches!(g.mode.as_str(), "line" | "diary"), "invalid generation mode" ); ensure!( g.temperature.is_finite() && (0.0..=2.0).contains(&g.temperature), "invalid temperature" ); ensure!( g.top_k > 0 && g.top_p > 0.0 && g.top_p <= 1.0, "invalid sampling settings" ); let prompt = self.fit(g)?; let encoded = self .tokenizer .encode(prompt.clone(), false) .map_err(|e| anyhow::anyhow!(e.to_string()))?; let ids: Vec<i64> = encoded.get_ids().iter().map(|i| *i as i64).collect(); let prompt_tokens = ids.len(); ensure!( prompt_tokens > 0, "generation requires a nonempty tokenized prompt" ); ensure!( ids.iter() .all(|id| *id >= 0 && (*id as usize) < self.contract.dimensions.vocab), "input token is outside model vocabulary" ); let mut current = ids; let mut generated = Vec::new(); let d = &self.contract.dimensions; let mut past: Vec<(String, DynValue)> = self .contract .caches .iter() .map(|cache| { Ok(( cache.input.clone(), DynTensor::new(&Allocator::default(), cache.ty, [1, d.kv_heads, 0, d.width])? .into_dyn(), )) }) .collect::<Result<_>>()?; let started = Instant::now(); let mut prefill_ms = 0.0; let mut sampled_tokens = 0usize; let mut last_token_ms = 0.0; let mut reason = if benchmark { "fixed_length" } else { "length" }; let mut text = String::new(); let mut binding = self .gpu .map(|_| self.session.create_binding()) .transpose()?; let gpu_memory = self .gpu .map(|device| { MemoryInfo::new( AllocationDevice::CUDA, device, AllocatorType::Device, MemoryType::Default, ) }) .transpose()?; let cpu_memory = MemoryInfo::new( AllocationDevice::CPU, 0, AllocatorType::Device, MemoryType::CPUOutput, )?; for step in 0..g.max_new { ensure!(!cancelled(), "generation cancelled"); let total = prompt_tokens + step; let input_tokens = current.len(); let mut input = self.contract.feed(current, total - input_tokens)?; input.append(&mut past); let mut output = if let Some(binding) = &mut binding { binding.clear(); for (name, value) in &input { binding.bind_input(name, value)?; } // Sampling needs only logits on the CPU. Keep all KV outputs // on CUDA and feed their Values directly into the next step. // Fresh output bindings avoid aliasing a previous cache buffer. binding.bind_output_to_device("logits", &cpu_memory)?; for cache in &self.contract.caches { binding.bind_output_to_device(&cache.output, gpu_memory.as_ref().unwrap())?; } let output = self.session.run_binding(binding)?; binding.synchronize_outputs()?; output } else { self.session.run(input)? }; // Read only metadata for KV tensors, including CUDA-resident values. // This proves cache growth without transferring the cache to the CPU. for cache in &self.contract.caches { self.contract.validate_cache( cache, output[cache.output.as_str()].dtype(), total, )?; } let (shape, logits) = output["logits"].try_extract_tensor::<f32>()?; let vocab = self.contract.validate_logits(shape, input_tokens)?; let token = sample(&logits[logits.len() - vocab..], g)?; last_token_ms = started.elapsed().as_secs_f64() * 1000.0; sampled_tokens += 1; if step == 0 { prefill_ms = last_token_ms; } if !benchmark && self.contract.dimensions.eos.contains(&token) { reason = "eos"; break; } generated.push(token); if !benchmark { text = self .tokenizer .decode(&generated, true) .map_err(|e| anyhow::anyhow!(e.to_string()))?; if let Some((why, kept)) = world::stop(&g.mode, &text) { text = kept; reason = why; break; } } for cache in &self.contract.caches { let value = output .remove(cache.output.as_str()) .context("missing KV output")?; past.push((cache.input.clone(), value)); } current = vec![token as i64]; } if benchmark { text = self .tokenizer .decode(&generated, true) .map_err(|e| anyhow::anyhow!(e.to_string()))?; } let decode_tokens = sampled_tokens.saturating_sub(1); let decode_ms = (last_token_ms - prefill_ms).max(0.0); Ok(Generated { text, prompt, reason: reason.into(), prompt_tokens, prefill_ms, decode_tokens, decode_ms, tps: if decode_ms > 0.0 { decode_tokens as f64 * 1000.0 / decode_ms } else { 0.0 }, generated_tokens: generated.len(), total_ms: started.elapsed().as_secs_f64() * 1000.0, }) }}
fn sample(logits: &[f32], g: &Generation) -> Result<u32> { let temperature = g.temperature; ensure!( logits.iter().all(|v| v.is_finite()), "non-finite model logits" ); if temperature == 0.0 { return Ok(logits .iter() .enumerate() .max_by(|a, b| a.1.total_cmp(b.1)) .unwrap() .0 as u32); } let mut ranked: Vec<_> = logits.iter().copied().enumerate().collect(); let k = g.top_k.clamp(1, ranked.len()); ranked.select_nth_unstable_by(k - 1, |a, b| b.1.total_cmp(&a.1)); ranked.truncate(k); ranked.sort_unstable_by(|a, b| b.1.total_cmp(&a.1)); let high = ranked[0].1; let probabilities: Vec<_> = ranked .iter() .map(|(_, v)| ((*v - high) / temperature).exp()) .collect(); let total: f32 = probabilities.iter().sum(); let mut kept = 0.0; let mut n = 0; for p in &probabilities { kept += p; n += 1; if kept / total >= g.top_p { break; } } let mut draw = rand::random::<f32>() * kept; for (i, p) in probabilities.iter().take(n).enumerate() { draw -= p; if draw <= 0.0 { return Ok(ranked[i].0 as u32); } } Ok(ranked[n - 1].0 as u32)}
#[cfg(test)]mod tests { use super::*; use ort::value::{Shape, SymbolicDimensions}; use serde_json::{Value, json};
fn config(model: &str) -> Value { if model == "gpt2" { json!({"model_type":model,"n_head":4,"n_embd":32,"n_layer":1,"n_positions":1024,"vocab_size":97,"eos_token_id":96}) } else { json!({"model_type":model,"num_attention_heads":4,"hidden_size":32, "num_hidden_layers":1,"max_position_embeddings":1024,"vocab_size":97,"eos_token_id":0}) } } fn meta(ty: TensorElementType, shape: &[i64]) -> ValueType { ValueType::Tensor { ty, shape: Shape::new(shape.iter().copied()), dimension_symbols: SymbolicDimensions::empty(shape.len()), } } fn graph(d: &Dimensions) -> (Vec<(String, ValueType)>, Vec<(String, ValueType)>) { let mut inputs = vec![( "input_ids".into(), meta(TensorElementType::Int64, &[-1, -1]), )]; let mut outputs = vec![( "logits".into(), meta(TensorElementType::Float32, &[-1, 1, d.vocab as i64]), )]; for layer in 0..d.layers { for kind in ["key", "value"] { let value = meta( TensorElementType::Float32, &[-1, d.kv_heads as i64, -1, d.width as i64], ); inputs.push((format!("past_key_values.{layer}.{kind}"), value.clone())); outputs.push((format!("present.{layer}.{kind}"), value)); } } (inputs, outputs) } #[test] fn supported_configs_and_gqa_do_not_assume_gpt2_dimensions() { for model in ["gpt2", "gpt_neox", "llama"] { let mut cfg = config(model); if model == "llama" { cfg["num_key_value_heads"] = json!(2); cfg["head_dim"] = json!(16); } let d = Dimensions::from_config(&cfg).unwrap(); assert_eq!(d.kv_heads, if model == "llama" { 2 } else { 4 }); assert_eq!(d.width, if model == "llama" { 16 } else { 8 }); let (inputs, outputs) = graph(&d); Contract::new(d, &inputs, &outputs).unwrap(); } } #[test] fn rejects_unsupported_and_ambiguous_configs() { for (key, value) in [ ("model_type", json!("mistral")), ("n_head", json!(0)), ("num_attention_heads", json!(3)), ("n_embd", json!(33)), ("vocab_size", json!(0)), ("n_layer", json!(null)), ("num_key_value_heads", json!(2)), ] { let mut cfg = config("gpt2"); cfg[key] = value; assert!(Dimensions::from_config(&cfg).is_err(), "accepted {key}"); } let mut cfg = config("llama"); cfg["num_key_value_heads"] = json!(3); assert!(Dimensions::from_config(&cfg).is_err()); } #[test] fn positional_capacity_is_required_positive_and_unambiguous() { for value in [json!(0), json!(-1), json!(null), json!(2.5)] { let mut cfg = config("llama"); cfg["max_position_embeddings"] = value; assert!(Dimensions::from_config(&cfg).is_err()); } let mut cfg = config("gpt2"); cfg.as_object_mut().unwrap().remove("n_positions"); assert!(Dimensions::from_config(&cfg).is_err()); cfg["n_positions"] = json!(32); cfg["max_position_embeddings"] = json!(64); assert!(Dimensions::from_config(&cfg).is_err()); cfg["max_position_embeddings"] = json!(32); assert_eq!(Dimensions::from_config(&cfg).unwrap().context, 32); } #[test] fn smaller_model_context_caps_history_and_rejects_unfit_header_or_prefix() { let mut cfg = config("llama"); cfg["max_position_embeddings"] = json!(32); let d = Dimensions::from_config(&cfg).unwrap(); let mut g = Generation { header: "head\n".into(), lines: vec!["older line".into(), "newer line".into()], prefix: "pet:".into(), mode: "line".into(), max_new: 4, temperature: 0.0, top_k: TOP_K, top_p: TOP_P, ctx: 1024, }; // The supplied counter isolates the positional fit contract from tokenization. assert_eq!( d.fit(&g, &mut |text| Ok(text.len())).unwrap(), "head\nnewer line\npet:" ); g.ctx = 16; assert!(d.fit(&g, &mut |text| Ok(text.len())).is_err()); g.ctx = 1024; g.header = "x".repeat(25); assert!(d.fit(&g, &mut |text| Ok(text.len())).is_err()); g.header.clear(); g.prefix = "x".repeat(25); assert!(d.fit(&g, &mut |text| Ok(text.len())).is_err()); g.prefix = "pet:".into(); g.max_new = 32; assert!(d.fit(&g, &mut |text| Ok(text.len())).is_err()); } #[test] fn eos_accepts_scalar_and_list_including_zero_not_baseline_id() { let mut cfg = config("llama"); cfg["eos_token_id"] = json!([0, 96]); assert_eq!(Dimensions::from_config(&cfg).unwrap().eos, [0, 96]); for value in [ json!([]), json!([0, 0]), json!(97), json!(-1), json!(null), json!([0, "1"]), json!(1.5), ] { cfg["eos_token_id"] = value; assert!(Dimensions::from_config(&cfg).is_err()); } } #[test] fn logits_vocabulary_and_runtime_shape_are_checked() { let d = Dimensions::from_config(&config("gpt_neox")).unwrap(); let (inputs, mut outputs) = graph(&d); outputs[0].1 = meta(TensorElementType::Float32, &[-1, 1, 98]); assert!(Contract::new(d, &inputs, &outputs).is_err()); let d = Dimensions::from_config(&config("gpt_neox")).unwrap(); outputs[0].1 = meta(TensorElementType::Float32, &[-1, -1, -1]); let contract = Contract::new(d, &inputs, &outputs).unwrap(); assert_eq!(contract.validate_logits(&[1, 3, 97], 3).unwrap(), 97); assert_eq!(contract.validate_logits(&[1, 1, 97], 3).unwrap(), 97); for shape in [&[1, 1, 50257][..], &[2, 1, 97], &[1, 2, 97], &[1, 97]] { assert!(contract.validate_logits(shape, 3).is_err()); } } #[test] fn separate_cache_contract_rejects_bad_layouts_and_missing_pairs() { for shape in [ vec![-1, -1, -1, 8], vec![-1, -1, 4, 8], vec![-1, 4, 0, 8], vec![-1, 4, -1, 9], vec![2, 4, -1, 8], vec![-1, 2, 4, -1, 8], ] { let d = Dimensions::from_config(&config("gpt2")).unwrap(); let (mut inputs, outputs) = graph(&d); inputs[1].1 = meta(TensorElementType::Float32, &shape); assert!( Contract::new(d, &inputs, &outputs).is_err(), "accepted {shape:?}" ); } for missing_input in [true, false] { let d = Dimensions::from_config(&config("gpt2")).unwrap(); let (mut inputs, mut outputs) = graph(&d); if missing_input { inputs.pop(); } else { outputs.pop(); } assert!(Contract::new(d, &inputs, &outputs).is_err()); } let d = Dimensions::from_config(&config("gpt2")).unwrap(); let (inputs, mut outputs) = graph(&d); outputs.push(("present_key_values.0.key".into(), outputs[1].1.clone())); assert!(Contract::new(d, &inputs, &outputs).is_err()); } #[test] fn cache_metadata_types_and_runtime_growth_are_checked() { for ty in [ TensorElementType::Float16, TensorElementType::Float32, TensorElementType::Float64, ] { let d = Dimensions::from_config(&config("gpt2")).unwrap(); let (mut inputs, mut outputs) = graph(&d); for (_, value) in inputs.iter_mut().skip(1).chain(outputs.iter_mut().skip(1)) { *value = meta(ty, &[-1, 4, -1, 8]); } let contract = Contract::new(d, &inputs, &outputs).unwrap(); let cache = &contract.caches[0]; assert!( contract .validate_cache(cache, &meta(ty, &[1, 4, 9, 8]), 9) .is_ok() ); assert!( contract .validate_cache(cache, &meta(ty, &[1, 4, 8, 8]), 9) .is_err() ); let other = if ty == TensorElementType::Float64 { TensorElementType::Float32 } else { TensorElementType::Float64 }; outputs[1].1 = meta(other, &[-1, 4, -1, 8]); let d = Dimensions::from_config(&config("gpt2")).unwrap(); assert!(Contract::new(d, &inputs, &outputs).is_err()); outputs[1].1 = meta(TensorElementType::Int64, &[-1, 4, -1, 8]); let d = Dimensions::from_config(&config("gpt2")).unwrap(); assert!(Contract::new(d, &inputs, &outputs).is_err()); } } #[test] fn integer_inputs_require_supported_ranks_and_types() { let d = Dimensions::from_config(&config("llama")).unwrap(); let (mut inputs, outputs) = graph(&d); inputs.push(( "cache_position".into(), meta(TensorElementType::Int32, &[-1]), )); inputs.push(( "position_ids".into(), meta(TensorElementType::Int64, &[-1, -1]), )); Contract::new(d, &inputs, &outputs).unwrap(); for (name, ty, shape) in [ ("input_ids", TensorElementType::Int64, vec![-1]), ("attention_mask", TensorElementType::Float32, vec![-1, -1]), ("cache_position", TensorElementType::Int64, vec![-1, -1, -1]), ("position_ids", TensorElementType::Int64, vec![-1, 8]), ] { let d = Dimensions::from_config(&config("llama")).unwrap(); let (mut inputs, outputs) = graph(&d); if name == "input_ids" { inputs[0].1 = meta(ty, &shape); } else { inputs.push((name.into(), meta(ty, &shape))); } assert!( Contract::new(d, &inputs, &outputs).is_err(), "accepted {name} {shape:?}" ); } } #[test] fn rejects_merged_control_inputs_and_noncausal_outputs() { for input in [true, false] { let d = Dimensions::from_config(&config("gpt2")).unwrap(); let (mut inputs, mut outputs) = graph(&d); if input { inputs.push(( "use_cache_branch".into(), meta(TensorElementType::Bool, &[]), )); } else { outputs.push(( "hidden_states".into(), meta(TensorElementType::Float32, &[-1, -1, 32]), )); } assert!(Contract::new(d, &inputs, &outputs).is_err()); } } #[test] fn unsupported_export_precision_is_rejected() { for dtype in ["fp16", "", "Q4", "q4/../fp32"] { assert!(dtype.parse::<ModelDtype>().is_err()); } }}