diff --git a/src/config.rs b/src/config.rs index dde95cce..24345202 100644 --- a/src/config.rs +++ b/src/config.rs @@ -6,6 +6,23 @@ use anyhow::{Context, bail}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum LayerKind { Delta, Attn } +/// Standalone EAGLE3 1-layer draft head parameters. +#[derive(Clone, Debug)] +pub struct Eagle3Config { + pub extract_layers: [usize; 3], + pub n_ff: usize, + pub n_head: usize, + pub n_head_kv: usize, + pub head_dim: usize, + pub n_rot: usize, + pub rope_base: f32, + pub eps: f32, + pub draft_vocab: usize, + pub norm_before_residual: bool, + pub rope_neox: bool, + pub d2t: Vec, +} + /// What differs between the supported architectures; everything else comes from GGUF metadata. struct Arch { /// DeltaNet layers with full attention every `full_attention_interval` layers; otherwise all attention. @@ -88,6 +105,7 @@ pub struct ModelConfig { pub q_gate: bool, pub k_interleave: bool, pub aliases: &'static [(&'static str, &'static str)], + pub eagle3: Option, } impl ModelConfig { @@ -175,6 +193,7 @@ impl ModelConfig { k_interleave: a.k_interleave, aliases: a.aliases, arch, + eagle3: None, }; // Kernels divide by these and index value heads by nv / nk key heads. let (nk, nv) = (c.ssm_n_group, c.ssm_dt_rank); diff --git a/src/cpu_worker.rs b/src/cpu_worker.rs index 142d6117..4c89e139 100644 --- a/src/cpu_worker.rs +++ b/src/cpu_worker.rs @@ -22,8 +22,8 @@ pub const VERDICT: usize = 8; pub const DONE: usize = 16; pub const IDS: usize = 32; -/// Most tokens one request may carry. -pub const TMAX: usize = 4; +/// Most tokens one request may carry (1 current + up to 16 speculative draft tokens). +pub const TMAX: usize = 17; /// Per layer: ids, weights and slots for up to TMAX tokens of `k` entries (token t at t * k; the three arrays are /// each NTOK * k long), then TMAX activations of `n` floats at `h` and TMAX results at `y`. diff --git a/src/expert_cache.rs b/src/expert_cache.rs index fcb559d7..de22cdb4 100644 --- a/src/expert_cache.rs +++ b/src/expert_cache.rs @@ -5,7 +5,6 @@ use crate::backend::{Backend, Buf}; use crate::model::Model; use crate::quant::GgmlType; -use std::collections::BTreeMap; use std::path::PathBuf; /// Share of each count kept per `adapt` (every 8 tokens by default): `counts` and `long`. @@ -63,8 +62,9 @@ impl ExpertCache { let n_layers = range.end; let c = &m.cfg; let n = c.n_expert; + const MAX_POOL_FULL_BYTES: usize = 8192 << 20; let mut layers: Vec> = Vec::new(); - let mut classes = BTreeMap::::new(); // slot bytes -> experts of that size + let mut pool_specs: Vec<(usize, usize)> = Vec::new(); // (slot_bytes, experts in pool) for il in 0..n_layers { if !c.is_moe_layer(il) || !range.contains(&il) { layers.push(None); continue; } let mut off = 0; @@ -74,10 +74,17 @@ impl ExpertCache { mats[i] = Mat { ty: t.ty, rows: t.rows, cols: t.cols, off, bytes: t.data.len() }; off = (off + t.data.len()).next_multiple_of(16); // GEMV kernels read 8-byte words } - *classes.entry(off).or_default() += n; - layers.push(Some((off, mats))); + let pid = pool_specs + .iter() + .position(|&(sb, cnt)| sb == off && (cnt + n) * off <= MAX_POOL_FULL_BYTES) + .unwrap_or_else(|| { + pool_specs.push((off, 0)); + pool_specs.len() - 1 + }); + pool_specs[pid].1 += n; + layers.push(Some((pid, mats))); } - let total: usize = classes.iter().map(|(b, k)| b * k).sum(); + let total: usize = pool_specs.iter().map(|(b, k)| b * k).sum(); let mut frac = (budget as f64 / total.max(1) as f64).min(1.0); let profile = std::env::var_os("STRATA_FIXED_EXPERTS").is_none().then(|| profile_path(m)); let counts = profile.as_ref().and_then(|p| std::fs::read(p).ok()) @@ -94,7 +101,7 @@ impl ExpertCache { // The budget comes from the driver's free-memory estimate, which can be more than it will allocate (AMD's // Windows driver): on a failed allocation, free the pools and try a tenth less. let pools = loop { - let r = classes.iter().map(|(&slot_bytes, &k)| { + let r = pool_specs.iter().map(|&(slot_bytes, k)| { let slots = (frac * k as f64) as usize; Ok(Pool { buf: gpu.zeros::(slots * slot_bytes)?, slot_bytes, owner: Vec::with_capacity(slots) }) }).collect::>>(); @@ -111,8 +118,8 @@ impl ExpertCache { for l in &layers { out.push(match l { None => None, - Some((sb, mats)) => Some(LayerExperts { - pool: classes.keys().position(|k| k == sb).unwrap(), mats: *mats, table: vec![-1; n], + Some((pid, mats)) => Some(LayerExperts { + pool: *pid, mats: *mats, table: vec![-1; n], table_dev: { let b = gpu.zeros::(n)?; gpu.put(&vec![-1i32; n], &b, 0)?; b }, }), }); diff --git a/src/gguf.rs b/src/gguf.rs index cc867493..69a60e39 100644 --- a/src/gguf.rs +++ b/src/gguf.rs @@ -153,4 +153,249 @@ impl Gguf { let o = data_offset + t.offset as usize; &mmap[o..o + n] } + + /// Raw bytes for auxiliary metadata tensors whose GGML type is not in `GgmlType` (e.g. EAGLE3 `d2t` I32/I64). + pub fn raw_bytes(&self, t: &TensorInfo, byte_len: usize) -> anyhow::Result<&[u8]> { + let (mmap, data_offset) = &self.files[t.file]; + let o = data_offset + t.offset as usize; + mmap.get(o..o + byte_len).context("GGUF tensor bytes out of range") + } +} + +/// Convert a HuggingFace EAGLE3 `.safetensors` checkpoint + `config.json` into a GGUF v3 draft model. +/// +/// Mirrors `tools/convert_eagle3_to_gguf.py` natively in Rust: 2D linear weights stay in BF16 (or convert from +/// F32/F16), 1D RMS norms convert to F32, NeoX RoPE is preserved (`eagle3.rope.neox = true`), and `d2t` relative +/// offsets are stored as absolute target vocabulary IDs (`I32`). +pub fn convert_safetensors_eagle3_to_gguf( + st_path: &Path, + cfg_path: &Path, + out_gguf: &Path, + target_layers_default: usize, +) -> anyhow::Result<()> { + use half::{bf16, f16}; + use std::borrow::Cow; + use std::io::{BufWriter, Write}; + + const GGUF_VERSION: u32 = 3; + const ALIGNMENT: usize = 32; + const GGUF_TYPE_UINT32: u32 = 4; + const GGUF_TYPE_INT32: u32 = 5; + const GGUF_TYPE_FLOAT32: u32 = 6; + const GGUF_TYPE_BOOL: u32 = 7; + const GGUF_TYPE_STRING: u32 = 8; + const GGUF_TYPE_ARRAY: u32 = 9; + const GGML_TYPE_F32: u32 = 0; + const GGML_TYPE_I32: u32 = 26; + const GGML_TYPE_BF16: u32 = 30; + + const TENSOR_MAP: [(&str, &str, u32); 13] = [ + ("fc.weight", "fc.weight", GGML_TYPE_BF16), + ("midlayer.input_layernorm.weight", "blk.0.attn_norm.weight", GGML_TYPE_F32), + ("midlayer.hidden_norm.weight", "blk.0.attn_norm_2.weight", GGML_TYPE_F32), + ("midlayer.self_attn.q_proj.weight", "blk.0.attn_q.weight", GGML_TYPE_BF16), + ("midlayer.self_attn.k_proj.weight", "blk.0.attn_k.weight", GGML_TYPE_BF16), + ("midlayer.self_attn.v_proj.weight", "blk.0.attn_v.weight", GGML_TYPE_BF16), + ("midlayer.self_attn.o_proj.weight", "blk.0.attn_output.weight", GGML_TYPE_BF16), + ("midlayer.post_attention_layernorm.weight", "blk.0.ffn_norm.weight", GGML_TYPE_F32), + ("midlayer.mlp.gate_proj.weight", "blk.0.ffn_gate.weight", GGML_TYPE_BF16), + ("midlayer.mlp.up_proj.weight", "blk.0.ffn_up.weight", GGML_TYPE_BF16), + ("midlayer.mlp.down_proj.weight", "blk.0.ffn_down.weight", GGML_TYPE_BF16), + ("norm.weight", "output_norm.weight", GGML_TYPE_F32), + ("lm_head.weight", "output.weight", GGML_TYPE_BF16), + ]; + + let cfg_str = std::fs::read_to_string(cfg_path).with_context(|| format!("read {}", cfg_path.display()))?; + let cfg: serde_json::Value = serde_json::from_str(&cfg_str).with_context(|| format!("parse {}", cfg_path.display()))?; + let layer_cfg = cfg.get("transformer_layer_config").unwrap_or(&cfg); + + let req_u64 = |k: &str| -> anyhow::Result { + layer_cfg.get(k).or_else(|| cfg.get(k)).and_then(|v| v.as_u64()) + .with_context(|| format!("{}: missing integer key {k}", cfg_path.display())) + }; + let n_embd = req_u64("hidden_size")? as u32; + let n_ff = req_u64("intermediate_size")? as u32; + let n_head = req_u64("num_attention_heads")? as u32; + let n_head_kv = req_u64("num_key_value_heads")? as u32; + let head_dim = layer_cfg.get("head_dim").or_else(|| cfg.get("head_dim")).and_then(|v| v.as_u64()).map(|v| v as u32).unwrap_or_else(|| n_embd / n_head.max(1)); + let eps = layer_cfg.get("rms_norm_eps").or_else(|| cfg.get("rms_norm_eps")).and_then(|v| v.as_f64()).unwrap_or(1e-6) as f32; + let rope_theta = layer_cfg.get("rope_theta").or_else(|| cfg.get("rope_theta")).and_then(|v| v.as_f64()).unwrap_or(1_000_000.0) as f32; + let norm_before_residual = cfg.get("norm_before_residual").and_then(|v| v.as_bool()).unwrap_or(false); + + let st_file = std::fs::File::open(st_path).with_context(|| format!("open {}", st_path.display()))?; + let mmap = unsafe { Mmap::map(&st_file)? }; + anyhow::ensure!(mmap.len() >= 8, "{}: safetensors file truncated", st_path.display()); + let hdr_len = usize::try_from(u64::from_le_bytes(mmap[0..8].try_into()?)).context("safetensors header length out of range")?; + let data_off = 8usize.checked_add(hdr_len).context("safetensors header overflow")?; + anyhow::ensure!(data_off <= mmap.len(), "{}: safetensors header past EOF", st_path.display()); + let hdr: serde_json::Map = + serde_json::from_slice(&mmap[8..data_off]).with_context(|| format!("parse safetensors header {}", st_path.display()))?; + + struct StTensor<'a> { + shape: Vec, + dtype: String, + raw: &'a [u8], + } + let mut tensors: BTreeMap> = BTreeMap::new(); + for (name, val) in hdr { + if name == "__metadata__" { continue; } + let obj = val.as_object().with_context(|| format!("tensor {name} entry is not an object"))?; + let dtype = obj.get("dtype").and_then(|v| v.as_str()).with_context(|| format!("tensor {name} missing dtype"))?.to_string(); + let shape = obj.get("shape").and_then(|v| v.as_array()).with_context(|| format!("tensor {name} missing shape"))? + .iter().map(|d| d.as_u64().context("bad tensor dim")).collect::>>()?; + let offs = obj.get("data_offsets").and_then(|v| v.as_array()).with_context(|| format!("tensor {name} missing data_offsets"))?; + anyhow::ensure!(offs.len() == 2, "tensor {name}: expected 2 data_offsets"); + let start = usize::try_from(offs[0].as_u64().context("bad start offset")?)?; + let end = usize::try_from(offs[1].as_u64().context("bad end offset")?)?; + let raw = mmap.get(data_off + start..data_off + end).with_context(|| format!("tensor {name} offsets out of range"))?; + tensors.insert(name, StTensor { shape, dtype, raw }); + } + + let draft_vocab = cfg.get("draft_vocab_size").and_then(|v| v.as_u64()).map(|v| v as usize) + .or_else(|| tensors.get("d2t").and_then(|t| t.shape.first().copied()).map(|v| v as usize)) + .unwrap_or(32000); + let extract_layers: Vec = match cfg.get("eagle_aux_hidden_state_layer_ids").and_then(|v| v.as_array()) { + Some(arr) if !arr.is_empty() => arr.iter().map(|v| v.as_i64().map(|x| x as i32).context("bad extract layer")).collect::>()?, + _ => { + let n = target_layers_default as i32; + vec![2, n / 2, (n - 3).max(0)] + } + }; + anyhow::ensure!(extract_layers.len() == 3, "expected 3 EAGLE3 extract_layers, got {extract_layers:?}"); + + let mut prepared: Vec<(String, Vec, u32, Cow<'_, [u8]>)> = Vec::with_capacity(TENSOR_MAP.len() + 1); + for &(hf_name, gguf_name, dst_ty) in &TENSOR_MAP { + let st = tensors.get(hf_name) + .or_else(|| hf_name.strip_prefix("midlayer.").and_then(|rest| tensors.get(&format!("layers.0.{rest}")))) + .with_context(|| format!("{}: missing EAGLE3 tensor {hf_name}", st_path.display()))?; + let gguf_dims: Vec = st.shape.iter().copied().rev().collect(); + let bytes: Cow<'_, [u8]> = match (dst_ty, st.dtype.as_str()) { + (GGML_TYPE_BF16, "BF16") | (GGML_TYPE_F32, "F32") => Cow::Borrowed(st.raw), + (GGML_TYPE_BF16, "F32") => Cow::Owned( + st.raw.chunks_exact(4).flat_map(|b| bf16::from_f32(f32::from_le_bytes(b.try_into().unwrap())).to_le_bytes()).collect(), + ), + (GGML_TYPE_BF16, "F16") => Cow::Owned( + st.raw.chunks_exact(2).flat_map(|b| bf16::from_f32(f16::from_le_bytes(b.try_into().unwrap()).to_f32()).to_le_bytes()).collect(), + ), + (GGML_TYPE_F32, "BF16") => Cow::Owned( + st.raw.chunks_exact(2).flat_map(|b| bf16::from_le_bytes(b.try_into().unwrap()).to_f32().to_le_bytes()).collect(), + ), + (GGML_TYPE_F32, "F16") => Cow::Owned( + st.raw.chunks_exact(2).flat_map(|b| f16::from_le_bytes(b.try_into().unwrap()).to_f32().to_le_bytes()).collect(), + ), + (_, dt) => bail!("{hf_name}: unsupported dtype {dt} for GGUF target type {dst_ty}"), + }; + prepared.push((gguf_name.to_string(), gguf_dims, dst_ty, bytes)); + } + + let d2t = tensors.get("d2t").with_context(|| format!("{}: missing EAGLE3 d2t tensor", st_path.display()))?; + anyhow::ensure!(d2t.shape.len() == 1 && d2t.shape[0] as usize == draft_vocab, "unexpected d2t shape {:?} (draft_vocab={draft_vocab})", d2t.shape); + let abs_d2t: Vec = match d2t.dtype.as_str() { + "I64" => d2t.raw.chunks_exact(8).enumerate().flat_map(|(i, b)| ((i as i64 + i64::from_le_bytes(b.try_into().unwrap())) as i32).to_le_bytes()).collect(), + "I32" => d2t.raw.chunks_exact(4).enumerate().flat_map(|(i, b)| (i as i32 + i32::from_le_bytes(b.try_into().unwrap())).to_le_bytes()).collect(), + dt => bail!("unexpected d2t dtype {dt}"), + }; + prepared.push(("d2t".to_string(), vec![draft_vocab as u64], GGML_TYPE_I32, Cow::Owned(abs_d2t))); + + let w_str = |buf: &mut Vec, s: &str| { + buf.extend_from_slice(&(s.len() as u64).to_le_bytes()); + buf.extend_from_slice(s.as_bytes()); + }; + let w_kv = |buf: &mut Vec, k: &str, ty: u32, payload: &[u8]| { + w_str(buf, k); + buf.extend_from_slice(&ty.to_le_bytes()); + buf.extend_from_slice(payload); + }; + + let mut kvs = Vec::new(); + let mut nkv = 0u64; + let push_str_kv = |buf: &mut Vec, n: &mut u64, k: &str, v: &str| { + let mut p = Vec::with_capacity(8 + v.len()); + w_str(&mut p, v); + w_kv(buf, k, GGUF_TYPE_STRING, &p); + *n += 1; + }; + let push_u32_kv = |buf: &mut Vec, n: &mut u64, k: &str, v: u32| { + w_kv(buf, k, GGUF_TYPE_UINT32, &v.to_le_bytes()); + *n += 1; + }; + let push_f32_kv = |buf: &mut Vec, n: &mut u64, k: &str, v: f32| { + w_kv(buf, k, GGUF_TYPE_FLOAT32, &v.to_le_bytes()); + *n += 1; + }; + let push_bool_kv = |buf: &mut Vec, n: &mut u64, k: &str, v: bool| { + w_kv(buf, k, GGUF_TYPE_BOOL, &[u8::from(v)]); + *n += 1; + }; + + push_str_kv(&mut kvs, &mut nkv, "general.architecture", "eagle3"); + push_u32_kv(&mut kvs, &mut nkv, "general.alignment", ALIGNMENT as u32); + push_u32_kv(&mut kvs, &mut nkv, "eagle3.block_count", 1); + push_u32_kv(&mut kvs, &mut nkv, "eagle3.embedding_length", n_embd); + push_u32_kv(&mut kvs, &mut nkv, "eagle3.feed_forward_length", n_ff); + push_u32_kv(&mut kvs, &mut nkv, "eagle3.attention.head_count", n_head); + push_u32_kv(&mut kvs, &mut nkv, "eagle3.attention.head_count_kv", n_head_kv); + push_u32_kv(&mut kvs, &mut nkv, "eagle3.attention.key_length", head_dim); + push_f32_kv(&mut kvs, &mut nkv, "eagle3.attention.layer_norm_rms_epsilon", eps); + push_u32_kv(&mut kvs, &mut nkv, "eagle3.rope.dimension_count", head_dim); + push_f32_kv(&mut kvs, &mut nkv, "eagle3.rope.freq_base", rope_theta); + push_bool_kv(&mut kvs, &mut nkv, "eagle3.rope.neox", true); + push_bool_kv(&mut kvs, &mut nkv, "eagle3.norm_before_residual", norm_before_residual); + let mut tl_payload = Vec::with_capacity(12 + extract_layers.len() * 4); + tl_payload.extend_from_slice(&GGUF_TYPE_INT32.to_le_bytes()); + tl_payload.extend_from_slice(&(extract_layers.len() as u64).to_le_bytes()); + for &v in &extract_layers { tl_payload.extend_from_slice(&v.to_le_bytes()); } + w_kv(&mut kvs, "eagle3.target_layers", GGUF_TYPE_ARRAY, &tl_payload); + nkv += 1; + + let mut tensor_infos = Vec::new(); + let mut cur_off = 0usize; + for (name, dims, ty, data) in &prepared { + cur_off = cur_off.next_multiple_of(ALIGNMENT); + w_str(&mut tensor_infos, name); + tensor_infos.extend_from_slice(&(dims.len() as u32).to_le_bytes()); + for &d in dims { tensor_infos.extend_from_slice(&d.to_le_bytes()); } + tensor_infos.extend_from_slice(&ty.to_le_bytes()); + tensor_infos.extend_from_slice(&(cur_off as u64).to_le_bytes()); + cur_off += data.len(); + } + + let mut meta_block = Vec::with_capacity(24 + kvs.len() + tensor_infos.len()); + meta_block.extend_from_slice(b"GGUF"); + meta_block.extend_from_slice(&GGUF_VERSION.to_le_bytes()); + meta_block.extend_from_slice(&(prepared.len() as u64).to_le_bytes()); + meta_block.extend_from_slice(&nkv.to_le_bytes()); + meta_block.extend_from_slice(&kvs); + meta_block.extend_from_slice(&tensor_infos); + let hdr_pad = (ALIGNMENT - (meta_block.len() % ALIGNMENT)) % ALIGNMENT; + + if let Some(parent) = out_gguf.parent().filter(|p| !p.as_os_str().is_empty()) { + std::fs::create_dir_all(parent)?; + } + let tmp_out = out_gguf.with_extension("gguf.tmp"); + { + let f = std::fs::File::create(&tmp_out).with_context(|| format!("create {}", tmp_out.display()))?; + let mut w = BufWriter::new(f); + w.write_all(&meta_block)?; + let zeros = [0u8; ALIGNMENT]; + if hdr_pad > 0 { w.write_all(&zeros[..hdr_pad])?; } + let mut written = 0usize; + for (_, _, _, data) in &prepared { + let tpad = (ALIGNMENT - (written % ALIGNMENT)) % ALIGNMENT; + if tpad > 0 { + w.write_all(&zeros[..tpad])?; + written += tpad; + } + w.write_all(data)?; + written += data.len(); + } + w.flush()?; + } + std::fs::rename(&tmp_out, out_gguf).with_context(|| format!("rename {} -> {}", tmp_out.display(), out_gguf.display()))?; + eprintln!( + "converted EAGLE3 safetensors -> {} ({:.1} MiB, extract_layers={extract_layers:?}, heads={n_head}x{head_dim}, kv={n_head_kv}, ff={n_ff}, draft_vocab={draft_vocab})", + out_gguf.display(), + std::fs::metadata(out_gguf).map(|m| m.len() as f64 / 1048576.0).unwrap_or(0.0), + ); + Ok(()) } diff --git a/src/gpu_model.rs b/src/gpu_model.rs index bb5d4a90..50958381 100644 --- a/src/gpu_model.rs +++ b/src/gpu_model.rs @@ -120,14 +120,37 @@ impl Parked { fn bytes(&self) -> usize { self.parts.iter().flat_map(|p| p.slot.iter().chain(&p.kv)).map(Vec::len).sum() } } +/// Standalone EAGLE3 1-layer GPU weights and scratch buffers. +struct Eagle3Gpu { + cfg: crate::config::Eagle3Config, + fc: GpuTensor, + output: GpuTensor, + /// `[TMAX, 3 * n_embd]` concatenated hidden states extracted before `extract_layers`. + feat: Buf, + /// `[TMAX, n_embd]` output of `fc` for the last target forward. + cur_g: Buf, + q: Buf, + k: Buf, + v: Buf, + att: Buf, + attp: Buf, + ge: Buf, + ue: Buf, +} + struct MtpHead { - layer: Layer, eh_proj: GpuTensor, enorm: GpuTensor, hnorm: GpuTensor, head_norm: GpuTensor, - worker: CpuWorker, step_buf: Buf, emb: Buf, eh: Buf, x: Buf, pos: usize, - /// Row 0 of the main model's `x` at `save_state` (the hidden state the next entry pairs with). + layer: Layer, eh_proj: Option, enorm: GpuTensor, hnorm: GpuTensor, head_norm: GpuTensor, + worker: Option, step_buf: Buf, emb: Buf, eh: Buf, x: Buf, pos: usize, + eagle: Option, + /// Row 0 of the main model's `x` (or `eagle.cur_g`) at `save_state`. saved_x: Buf, /// Accepted drafts of the last `spec_step` whose draft head entries (from the main model's hidden rows 0..) are /// still to be written; the next `spec_step` writes them in its first draft pass. pend: Vec, + /// 12-step rolling window of draft acceptance bits (LSB = latest step), used for adaptive speculative depth. + adapt_window: u16, + /// Remaining 1-token bypass steps before probing EAGLE3/MTP again on partially-resident MoE runs. + adapt_skip: u8, } /// What a slot keeps while another slot is current (`use_slot`): the `GpuModel` fields of the same names. @@ -137,6 +160,11 @@ struct Seq { pos: usize, saved: Vec<(Buf, Buf)>, saved_pos: usize, saved_ids: Ve /// See `GpuModel::sys`. type Stash = (Vec>, usize, [i64; 2], Vec); +/// Maximum number of intermediate DeltaNet + conv state snapshots written during a multi-token verify pass. +/// Any rollback with `keep <= SNAP_MAX` restores directly via `dtod`; `keep > SNAP_MAX` restores `snap[SNAP_MAX - 1]` +/// and replays the remaining `keep - SNAP_MAX` accepted tokens in one batch pass. +pub const SNAP_MAX: usize = 8; + pub struct GpuModel<'a> { m: &'a Model<'a>, pub gpu: Box, @@ -186,6 +214,8 @@ pub struct GpuModel<'a> { sys_owner: u64, /// Host memory for `parked`, in bytes (0 = keep none). pub cache_ram: usize, + /// Optional disk path for persisting `parked` / `saved` conversation states across restarts (Warm Start). + pub prompt_cache_path: Option, ctx: usize, pub pos: usize, pub cache: Option, @@ -223,6 +253,8 @@ pub struct GpuModel<'a> { /// Prefill scratch (None: prompts go through `step`). pf: Option, mtp: Option, + /// Speculative N-gram / lookup cache drafter (standalone or combined with `mtp` via `hybrid_spec_step`). + pub ngram: Option, /// The layers of this part are `first..first + layers.len()`. first: usize, /// The part on the next device of a model split by layers (`load_split`): this part's hidden state goes to it @@ -315,8 +347,9 @@ impl<'a> GpuModel<'a> { let ring = gpu.zeros::(if qsa_on { c.n_layer * QSA_RING } else { 1 })?; // ponytail: no MTP or batched prefill with hyper-connections yet (their state and row layouts assume one stream). - let k = if c.n_nextn > 0 && c.hc_n == 0 { k.min(TMAX - 1) } else { 0 }; - let mtp = k > 0; + let k = if c.hc_n == 0 { k.min(TMAX - 1) } else { 0 }; + let mtp = k > 0 && c.n_nextn > 0; + let nsnap = k.min(SNAP_MAX); let load_layer = |il: usize, kind: LayerKind| -> anyhow::Result { let mixer = if kind == LayerKind::Delta { Mixer::Delta { @@ -325,7 +358,7 @@ impl<'a> GpuModel<'a> { conv_w: lt(il, "ssm_conv1d.weight")?, norm: lt(il, "ssm_norm.weight")?, out: lt(il, "ssm_out.weight")?, conv: (0..ns).map(|_| gpu.zeros::((c.ssm_conv - 1) * channels)).collect::>()?, s: (0..ns).map(|_| gpu.zeros::(nv * c.ssm_dv() * dk)).collect::>()?, - snap: if mtp { Some((gpu.zeros::(k * (c.ssm_conv - 1) * channels)?, gpu.zeros::(k * c.ssm_d_inner * dk)?)) } else { None }, + snap: if nsnap > 0 { Some((gpu.zeros::(nsnap * (c.ssm_conv - 1) * channels)?, gpu.zeros::(nsnap * c.ssm_d_inner * dk)?)) } else { None }, } } else { Mixer::Attn { @@ -370,11 +403,86 @@ impl<'a> GpuModel<'a> { Ok(Layer { attn_norm, post_norm, hc, ple, mixer, ffn }) }; let layers = range.clone().map(|il| load_layer(il, c.layers[il])).collect::>>()?; - let mtp_parts = if mtp { + let mtp_head = if mtp { let il = c.n_layer; - Some((load_layer(il, LayerKind::Attn)?, lt(il, "nextn.eh_proj.weight")?, lt(il, "nextn.enorm.weight")?, - lt(il, "nextn.hnorm.weight")?, lt(il, "nextn.shared_head_norm.weight")?, - gpu.zeros::(2 * TMAX)?, gpu.zeros::(TMAX * c.n_embd)?, gpu.zeros::(TMAX * 2 * c.n_embd)?, gpu.zeros::(TMAX * c.n_embd)?)) + if let Some(ec) = c.eagle3.clone() { + let eg_kv = ec.n_head_kv * ec.head_dim; + let layer = Layer { + attn_norm: None, + post_norm: Some(lt(il, "post_attention_norm.weight")?), + hc: None, + ple: None, + mixer: Mixer::Attn { + q: lt(il, "attn_q.weight")?, + k: lt(il, "attn_k.weight")?, + v: lt(il, "attn_v.weight")?, + qn: lt(il, "nextn.enorm.weight")?, + kn: lt(il, "nextn.hnorm.weight")?, + o: lt(il, "attn_output.weight")?, + kc: (0..ns).map(|_| gpu.zeros::(gpu.kv_words(ctx * eg_kv))).collect::>()?, + vc: (0..ns).map(|_| gpu.zeros::(gpu.kv_words(ctx * eg_kv))).collect::>()?, + qsa: None, + sink: None, + }, + ffn: Ffn::Dense { + gate: lt(il, "ffn_gate.weight")?, + up: lt(il, "ffn_up.weight")?, + down: lt(il, "ffn_down.weight")?, + }, + }; + let attp_len = ec.n_head * ATT_SPLITS * (ec.head_dim + 2); + let eg = Eagle3Gpu { + fc: lt(il, "eagle3.fc.weight")?, + output: lt(il, "eagle3.output.weight")?, + feat: gpu.zeros::(TMAX * 3 * c.n_embd)?, + cur_g: gpu.zeros::(TMAX * c.n_embd)?, + q: gpu.zeros::(TMAX * ec.n_head * ec.head_dim)?, + k: gpu.zeros::(TMAX * eg_kv)?, + v: gpu.zeros::(TMAX * eg_kv)?, + att: gpu.zeros::(TMAX * ec.n_head * ec.head_dim)?, + attp: gpu.zeros::(TMAX * attp_len)?, + ge: gpu.zeros::(TMAX * ec.n_ff)?, + ue: gpu.zeros::(TMAX * ec.n_ff)?, + cfg: ec, + }; + Some(MtpHead { + layer, + eh_proj: None, + enorm: lt(il, "nextn.enorm.weight")?, + hnorm: lt(il, "nextn.hnorm.weight")?, + head_norm: lt(il, "nextn.shared_head_norm.weight")?, + worker: None, + step_buf: gpu.zeros::(2 * TMAX)?, + emb: gpu.zeros::(TMAX * c.n_embd)?, + eh: gpu.zeros::(TMAX * 2 * c.n_embd)?, + x: gpu.zeros::(TMAX * c.n_embd)?, + pos: 0, + eagle: Some(eg), + saved_x: gpu.zeros::(c.n_embd)?, + pend: vec![], + adapt_window: 0x0FFF, + adapt_skip: 0, + }) + } else { + Some(MtpHead { + layer: load_layer(il, LayerKind::Attn)?, + eh_proj: Some(lt(il, "nextn.eh_proj.weight")?), + enorm: lt(il, "nextn.enorm.weight")?, + hnorm: lt(il, "nextn.hnorm.weight")?, + head_norm: lt(il, "nextn.shared_head_norm.weight")?, + worker: Some(CpuWorker::new(m, &*gpu, vec![c.n_layer])?), + step_buf: gpu.zeros::(2 * TMAX)?, + emb: gpu.zeros::(TMAX * c.n_embd)?, + eh: gpu.zeros::(TMAX * 2 * c.n_embd)?, + x: gpu.zeros::(TMAX * c.n_embd)?, + pos: 0, + eagle: None, + saved_x: gpu.zeros::(c.n_embd)?, + pend: vec![], + adapt_window: 0x0FFF, + adapt_skip: 0, + }) + } } else { None }; let (output_norm, output_hc) = if !last { (None, None) } else if c.hc_n > 0 { (None, Some(hc_read("output_hc_")?)) } else { (Some(up("output_norm.weight")?), None) }; let ple_ring = if c.n_embd_ple > 0 { Some((gpu.zeros::(16 * hcn * c.n_embd)?, gpu.zeros::(16 * hcn * c.n_embd)?)) } else { None }; @@ -426,19 +534,29 @@ impl<'a> GpuModel<'a> { margin: lower --ctx or --parallel, use --kv q8 or a smaller --chunk (--vram-margin 0 skips this check)", free1 >> 20, floor >> 20); let cache = if c.is_moe() { - // STRATA_EXPERT_MB caps the expert slots, so that runs place the same experts whatever VRAM is free (a - // backend comparison, or other programs holding VRAM). - let cap = std::env::var("STRATA_EXPERT_MB").ok().and_then(|v| v.parse::().ok()).map_or(usize::MAX, |v| v << 20); - let budget = free1.saturating_sub(vram_margin).min(cap); - Some(ExpertCache::new(m, &*gpu, budget, range.start..range.end + mtp as usize)?) + // STRATA_EXPERT_MB sets the expert slot budget explicitly when > 0; otherwise use free VRAM minus margin. + let budget = match std::env::var("STRATA_EXPERT_MB") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&v| v > 0) + { + Some(mb) => { + eprintln!("note: STRATA_EXPERT_MB={mb} override active (reported free device heap: {} MB)", free1 >> 20); + mb << 20 + } + None => free1.saturating_sub(vram_margin), + }; + Some(ExpertCache::new(m, &*gpu, budget, range.start..range.end + (mtp && c.eagle3.is_none()) as usize)?) } else { None }; let moe_layers: Vec = range.clone().filter(|&il| c.is_moe_layer(il)).collect(); let worker = if !moe_layers.is_empty() { Some(CpuWorker::new(m, &*gpu, moe_layers)?) } else { None }; - // R3a: prefetch the next layer's experts predicted by the previous token's routing there. Default off; the GPU - // window of issue #34 measures it and accepts or rejects it. CUDA has the kernel, Vulkan's `moe_prefetch` is a - // no-op. - let prefetch = c.is_moe() && gpu.moe_prefetch_supported() - && std::env::var("STRATA_MOE_PREFETCH").is_ok_and(|v| v != "0"); + // R3a: prefetch the next layer's experts predicted by the previous token's routing there. Disabled on gfx11xx/iGPU + // where the small L2 cache causes prefetching to evict active lines and hurt decode throughput. + let want_prefetch = std::env::var("STRATA_MOE_PREFETCH").is_ok_and(|v| v != "0"); + if c.is_moe() && want_prefetch && !gpu.moe_prefetch_supported() { + eprintln!("note: STRATA_MOE_PREFETCH ignored on {} (gfx11xx/iGPU L2 cache too small for MoE expert prefetch)", gpu.device()); + } + let prefetch = c.is_moe() && gpu.moe_prefetch_supported() && want_prefetch; let pids = if prefetch { Some(gpu.zeros::((c.n_layer + 1) * c.n_expert_used)?) } else { None }; let mb = |b: usize| b as f64 / 1048576.0; eprintln!( @@ -457,17 +575,11 @@ impl<'a> GpuModel<'a> { mb(per), mb(ca.token_bytes(k)), ca.total_experts() / c.n_expert.max(1), if prefetch { ", next-layer L2 prefetch on (STRATA_MOE_PREFETCH)" } else { "" }); } - let mtp = match mtp_parts { - Some((layer, eh_proj, enorm, hnorm, head_norm, step_buf, emb, eh, x)) => Some(MtpHead { - layer, eh_proj, enorm, hnorm, head_norm, worker: CpuWorker::new(m, &*gpu, vec![c.n_layer])?, - step_buf, emb, eh, x, pos: 0, pend: vec![], saved_x: gpu.zeros::(c.n_embd)?, - }), - None => None, - }; + let mtp = mtp_head; let adapt_every = std::env::var("STRATA_ADAPT_EVERY").ok().and_then(|v| v.parse().ok()).unwrap_or(8); let adapt_swaps = std::env::var("STRATA_ADAPT_SWAPS").ok().and_then(|v| v.parse().ok()).unwrap_or(64); let reset_counts_after_prefill = std::env::var_os("STRATA_RESET_COUNTS").is_some(); - Ok(GpuModel { m, gpu, layers, output_norm, output_hc, output, ple_ring, qsa, ple_prev: [-1, -1], saved_ple_prev: [-1, -1], bufs, x, worker, logits, amax, ids, slots, step_buf, rope_buf, rope_offset: 0, ctx, pos: 0, saved, saved_pos: 0, saved_ids: vec![], sys: None, parked: vec![], owner: 0, saved_owner: 0, sys_owner: 0, cache_ram: 0, cache, adapt_every, adapt_swaps, reset_counts_after_prefill, hits: 0, misses: 0, tokens: 0, mailbox_writes: Default::default(), mailbox_reads: Default::default(), pids, cpu_time: Duration::ZERO, resumed: Default::default(), lora_applied: Default::default(), drafts: (0, 0), split_time: Default::default(), k, short: 32, pf, mtp, first: range.start, next: None, + Ok(GpuModel { m, gpu, layers, output_norm, output_hc, output, ple_ring, qsa, ple_prev: [-1, -1], saved_ple_prev: [-1, -1], bufs, x, worker, logits, amax, ids, slots, step_buf, rope_buf, rope_offset: 0, ctx, pos: 0, saved, saved_pos: 0, saved_ids: vec![], sys: None, parked: vec![], owner: 0, saved_owner: 0, sys_owner: 0, cache_ram: 0, prompt_cache_path: None, cache, adapt_every, adapt_swaps, reset_counts_after_prefill, hits: 0, misses: 0, tokens: 0, mailbox_writes: Default::default(), mailbox_reads: Default::default(), pids, cpu_time: Duration::ZERO, resumed: Default::default(), lora_applied: Default::default(), drafts: (0, 0), split_time: Default::default(), k, short: 32, pf, mtp, ngram: None, first: range.start, next: None, cur: 0, seqs, rows: vec![], row_tokens: 0 }) } @@ -475,7 +587,8 @@ impl<'a> GpuModel<'a> { pub fn reset(&mut self) -> anyhow::Result<()> { self.saved_ids.clear(); self.sys = None; - if let Some(h) = self.mtp.as_mut() { h.pos = 0; h.pend.clear(); } + if let Some(h) = self.mtp.as_mut() { h.pos = 0; h.pend.clear(); h.adapt_window = 0x0FFF; h.adapt_skip = 0; } + if let Some(ng) = self.ngram.as_mut() { ng.reset(); } for l in self.layers.iter().chain(self.mtp.as_ref().map(|h| &h.layer)) { match &l.mixer { Mixer::Delta { conv, s, .. } => { self.gpu.memset(&conv[self.cur])?; self.gpu.memset(&s[self.cur])?; } @@ -567,7 +680,7 @@ impl<'a> GpuModel<'a> { /// and attention regimes, so keep the rows in slot order. pub fn step_rows(&mut self, rows: &[(usize, u32, bool)]) -> anyhow::Result)>> { let (nt, nvoc) = (rows.len(), self.m.cfg.n_vocab); - ensure!((1..=TMAX).contains(&nt) && self.mtp.is_none(), "step_rows: {nt} rows"); + ensure!((1..=4).contains(&nt) && self.mtp.is_none(), "step_rows: {nt} rows"); let pos: Vec = rows.iter().map(|&(s, ..)| *self.slot_pos(s)).collect(); if pos.iter().any(|&p| p >= self.ctx) { bail!("context full ({} tokens; raise --ctx)", self.ctx); } self.put_embed(&rows.iter().map(|r| r.1).collect::>())?; @@ -610,7 +723,6 @@ impl<'a> GpuModel<'a> { self.gpu.put(&pos.iter().flat_map(|&p| [p as i32, seq as i32]).collect::>(), &self.step_buf, 0)?; self.gpu.put(&pos.iter().flat_map(|&p| [p as i32 + self.rope_offset, seq as i32]).collect::>(), &self.rope_buf, 0)?; self.rows = rows.iter().zip(pos).map(|(r, &p)| (r.0, p)).collect(); - const _: () = assert!(TMAX <= 4, "step_rows graph key"); let mut key = 1 << 63 | nt as u64; for (r, &(sl, p)) in self.rows.iter().enumerate() { key |= (sl as u64) << (8 + 8 * r) | ((gq_min(p, 1) == 0) as u64) << (48 + r); } if self.m.loras.is_empty() { self.gpu.graph(key, &mut || self.forward(nt, 0, None))?; } else { self.forward(nt, 0, None)?; } @@ -667,15 +779,25 @@ impl<'a> GpuModel<'a> { /// Run one token; returns logits over the vocabulary. pub fn step(&mut self, token: u32) -> anyhow::Result> { if self.pos >= self.ctx { bail!("context full ({} tokens; raise --ctx)", self.ctx); } + let prof = std::env::var_os("STRATA_VK_PROF").is_some(); + let t0 = prof.then(std::time::Instant::now); let seq = self.begin_seq(); self.put_embed(&[token])?; self.gpu.put(&[self.pos as i32, seq as i32], &self.step_buf, 0)?; self.gpu.put(&[self.pos as i32 + self.rope_offset, seq as i32], &self.rope_buf, 0)?; + let t1 = prof.then(std::time::Instant::now); self.pass(1)?; + let t2 = prof.then(std::time::Instant::now); let t = self.tail(); let mut logits = t.gpu.get::(&t.logits, 0, self.m.cfg.n_vocab)?; if t.finish(1, gq_min(self.pos, 1), None)? { logits = t.gpu.get::(&t.logits, 0, self.m.cfg.n_vocab)?; } + let t3 = prof.then(std::time::Instant::now); self.after_forward(1)?; + if let (Some(t0), Some(t1), Some(t2), Some(t3)) = (t0, t1, t2, t3) { + let t4 = std::time::Instant::now(); + eprintln!("step_prof pos={}: embed_htod={:.1}us pass_submit={:.1}us gpu_wait_dtoh={:.1}us after_experts={:.1}us total={:.1}us", + self.pos, (t1 - t0).as_secs_f64() * 1e6, (t2 - t1).as_secs_f64() * 1e6, (t3 - t2).as_secs_f64() * 1e6, (t4 - t3).as_secs_f64() * 1e6, (t4 - t0).as_secs_f64() * 1e6); + } Ok(logits) } @@ -699,7 +821,15 @@ impl<'a> GpuModel<'a> { /// `verify` for greedy speculation: `keep` gets the argmax of every row; returns its row n and that row's argmax. fn verify_greedy(&mut self, tokens: &[u32], keep: impl FnOnce(&[u32]) -> usize) -> anyhow::Result<(usize, u32)> { let (nt, nvoc) = (tokens.len(), self.m.cfg.n_vocab); - self.verify_run(tokens)?; + let prof = std::env::var_os("STRATA_VK_PROF").is_some(); + let t0 = prof.then(std::time::Instant::now); + let seq = self.begin_seq(); + self.put_embed(tokens)?; + let steps: Vec = (0..nt).flat_map(|t| [(self.pos + t) as i32, seq as i32]).collect(); + self.gpu.put(&steps, &self.step_buf, 0)?; + let t1 = prof.then(std::time::Instant::now); + self.pass(nt)?; + let t2 = prof.then(std::time::Instant::now); let t = self.tail(); let g = &t.gpu; g.argmax_rows(t.logits.addr(), nvoc, nt, t.amax.addr())?; @@ -709,7 +839,13 @@ impl<'a> GpuModel<'a> { am = g.get::(&t.amax, 0, nt)?.into_iter().map(|x| x as u32).collect(); } let n = keep(&am); + let t3 = prof.then(std::time::Instant::now); self.after_forward(nt)?; + if let (Some(t0), Some(t1), Some(t2), Some(t3)) = (t0, t1, t2, t3) { + let t4 = std::time::Instant::now(); + eprintln!("greedy_prof pos={}: embed_htod={:.1}us pass_submit={:.1}us gpu_wait_dtoh={:.1}us after_experts={:.1}us total={:.1}us", + self.pos, (t1 - t0).as_secs_f64() * 1e6, (t2 - t1).as_secs_f64() * 1e6, (t3 - t2).as_secs_f64() * 1e6, (t4 - t3).as_secs_f64() * 1e6, (t4 - t0).as_secs_f64() * 1e6); + } Ok((n, am[n])) } @@ -740,7 +876,10 @@ impl<'a> GpuModel<'a> { g.dtod(s, 0, bs, 0, s.len())?; } } - if let Some(h) = &self.mtp { g.dtod(&self.x, 0, &h.saved_x, 0, 4 * self.m.cfg.n_embd)?; } + if let Some(h) = &self.mtp { + let src = h.eagle.as_ref().map_or(&self.x, |e| &e.cur_g); + g.dtod(src, 0, &h.saved_x, 0, 4 * self.m.cfg.n_embd)?; + } if let Some((r, sr)) = &self.ple_ring { g.dtod(r, 0, sr, 0, r.len())?; } if let Some(q) = &self.qsa { g.dtod(&q.ring, 0, &q.saved, 0, q.ring.len())?; } self.saved_ple_prev = self.ple_prev; @@ -782,20 +921,19 @@ impl<'a> GpuModel<'a> { }).collect() } - /// Copy the saved state to host memory, so a later request that continues it can take it back (`unpark`); the - /// least recently parked states go when they exceed `cache_ram`. - fn park(&mut self) -> anyhow::Result<()> { + /// Snapshot the current saved GPU state into a host-memory `Parked` record if it fits in `cache_ram`. + fn snapshot_saved(&self) -> anyhow::Result> { let n = self.saved_ids.len(); - if n == 0 { return Ok(()) } + if n == 0 { return Ok(None) } let mut size = 0; - let mut p = Some(&*self); + let mut p = Some(self); while let Some(q) = p { size += q.kv_bufs().iter().map(|b| b.len() / q.ctx * n).sum::() + q.slot().iter().map(|b| b.len()).sum::(); p = q.next.as_deref(); } - if size > self.cache_ram { return Ok(()) } + if size > self.cache_ram { return Ok(None) } let mut parts = vec![]; - let mut p = Some(&*self); + let mut p = Some(self); while let Some(q) = p { let g = &q.gpu; let slot = q.slot().iter().map(|b| g.get::(b, 0, b.len())).collect::>()?; @@ -803,8 +941,189 @@ impl<'a> GpuModel<'a> { parts.push(PartCopy { ple_prev: q.saved_ple_prev, slot, kv }); p = q.next.as_deref(); } - self.parked.insert(0, Parked { ids: self.saved_ids.clone(), owner: self.saved_owner, parts }); + Ok(Some(Parked { ids: self.saved_ids.clone(), owner: self.saved_owner, parts })) + } + + fn insert_parked(&mut self, pk: Parked) { + self.parked.retain(|old| !(old.owner == pk.owner && (pk.ids.starts_with(&old.ids) || old.ids == pk.ids))); + self.parked.insert(0, pk); while self.parked.iter().map(Parked::bytes).sum::() > self.cache_ram { self.parked.pop(); } + } + + /// Copy the saved state to host memory, so a later request that continues it can take it back (`unpark`); the + /// least recently parked states go when they exceed `cache_ram`. + fn park(&mut self) -> anyhow::Result<()> { + if let Some(pk) = self.snapshot_saved()? { + self.insert_parked(pk); + if self.prompt_cache_path.is_some() { + if let Err(e) = self.save_prompt_cache() { + eprintln!("warning: could not save prompt cache: {e:#}"); + } + } + } + Ok(()) + } + + /// Snapshot the current saved state into `parked` and persist to `prompt_cache_path` (Warm Start). + pub(crate) fn sync_saved_to_prompt_cache(&mut self) -> anyhow::Result<()> { + if self.prompt_cache_path.is_none() { return Ok(()) } + if let Some(pk) = self.snapshot_saved()? { + self.insert_parked(pk); + self.save_prompt_cache()?; + } + Ok(()) + } + + fn prompt_cache_fingerprint(&self) -> u64 { + let mut h = std::collections::hash_map::DefaultHasher::new(); + self.m.cfg.n_layer.hash(&mut h); + self.m.cfg.n_embd.hash(&mut h); + self.m.cfg.n_vocab.hash(&mut h); + self.m.cfg.n_head.hash(&mut h); + self.m.cfg.n_head_kv.hash(&mut h); + self.m.cfg.head_dim.hash(&mut h); + self.m.cfg.ssm_d_inner.hash(&mut h); + self.m.cfg.ssm_d_state.hash(&mut h); + let mut p = Some(self); + while let Some(q) = p { + q.layers.len().hash(&mut h); + q.kv_bufs().len().hash(&mut h); + for b in q.kv_bufs() { (b.len() / q.ctx).hash(&mut h); } + p = q.next.as_deref(); + } + h.finish() + } + + /// Save up to 4 most recent `Parked` states to `self.prompt_cache_path` atomically. + pub fn save_prompt_cache(&self) -> anyhow::Result<()> { + use std::io::Write; + let Some(path) = &self.prompt_cache_path else { return Ok(()) }; + if self.parked.is_empty() { return Ok(()) } + if let Some(parent) = path.parent() && !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent)?; + } + let tmp = path.with_extension("tmp"); + let f = std::fs::File::create(&tmp)?; + let mut w = std::io::BufWriter::with_capacity(8 << 20, f); + w.write_all(b"STRATAPC")?; + w.write_all(&1u32.to_le_bytes())?; + w.write_all(&self.prompt_cache_fingerprint().to_le_bytes())?; + let entries: Vec<&Parked> = self.parked.iter().take(4).collect(); + w.write_all(&(entries.len() as u32).to_le_bytes())?; + for pk in entries { + w.write_all(&pk.owner.to_le_bytes())?; + w.write_all(&(pk.ids.len() as u32).to_le_bytes())?; + for &id in &pk.ids { + w.write_all(&id.to_le_bytes())?; + } + w.write_all(&(pk.parts.len() as u32).to_le_bytes())?; + for part in &pk.parts { + w.write_all(&part.ple_prev[0].to_le_bytes())?; + w.write_all(&part.ple_prev[1].to_le_bytes())?; + w.write_all(&(part.slot.len() as u32).to_le_bytes())?; + for s in &part.slot { + w.write_all(&(s.len() as u64).to_le_bytes())?; + w.write_all(s)?; + } + w.write_all(&(part.kv.len() as u32).to_le_bytes())?; + for k in &part.kv { + w.write_all(&(k.len() as u64).to_le_bytes())?; + w.write_all(k)?; + } + } + } + w.flush()?; + drop(w); + std::fs::rename(&tmp, path)?; + Ok(()) + } + + /// Enable Warm Start persistence at `path` and load any existing compatible states into `self.parked`. + pub fn load_prompt_cache(&mut self, path: &std::path::Path) -> anyhow::Result<()> { + use std::io::Read; + self.prompt_cache_path = Some(path.to_path_buf()); + if !path.exists() { + return Ok(()); + } + let f = std::fs::File::open(path)?; + let mut r = std::io::BufReader::with_capacity(8 << 20, f); + let mut magic = [0u8; 8]; + r.read_exact(&mut magic)?; + ensure!(&magic == b"STRATAPC", "invalid magic in prompt cache {}", path.display()); + let mut u4 = [0u8; 4]; + let mut u8b = [0u8; 8]; + r.read_exact(&mut u4)?; + let ver = u32::from_le_bytes(u4); + ensure!(ver == 1, "unsupported prompt cache version {ver}"); + r.read_exact(&mut u8b)?; + let fp = u64::from_le_bytes(u8b); + let expected_fp = self.prompt_cache_fingerprint(); + ensure!(fp == expected_fp, "prompt cache fingerprint mismatch (model or --kv changed)"); + r.read_exact(&mut u4)?; + let n_entries = u32::from_le_bytes(u4) as usize; + let n_parts_expected = self.parts().count(); + let mut loaded = Vec::new(); + for _ in 0..n_entries { + r.read_exact(&mut u8b)?; + let owner = u64::from_le_bytes(u8b); + r.read_exact(&mut u4)?; + let n_ids = u32::from_le_bytes(u4) as usize; + if n_ids == 0 || n_ids > self.ctx { + bail!("cached sequence length {n_ids} exceeds --ctx {}", self.ctx); + } + let mut ids = Vec::with_capacity(n_ids); + for _ in 0..n_ids { + r.read_exact(&mut u4)?; + ids.push(u32::from_le_bytes(u4)); + } + r.read_exact(&mut u4)?; + let n_parts = u32::from_le_bytes(u4) as usize; + ensure!(n_parts == n_parts_expected, "prompt cache part count {n_parts} != {n_parts_expected}"); + let mut parts = Vec::with_capacity(n_parts); + for _ in 0..n_parts { + r.read_exact(&mut u8b)?; + let p0 = i64::from_le_bytes(u8b); + r.read_exact(&mut u8b)?; + let p1 = i64::from_le_bytes(u8b); + r.read_exact(&mut u4)?; + let n_slot = u32::from_le_bytes(u4) as usize; + let mut slot = Vec::with_capacity(n_slot); + for _ in 0..n_slot { + r.read_exact(&mut u8b)?; + let len = u64::from_le_bytes(u8b) as usize; + let mut buf = vec![0u8; len]; + r.read_exact(&mut buf)?; + slot.push(buf); + } + r.read_exact(&mut u4)?; + let n_kv = u32::from_le_bytes(u4) as usize; + let mut kv = Vec::with_capacity(n_kv); + for _ in 0..n_kv { + r.read_exact(&mut u8b)?; + let len = u64::from_le_bytes(u8b) as usize; + let mut buf = vec![0u8; len]; + r.read_exact(&mut buf)?; + kv.push(buf); + } + parts.push(PartCopy { ple_prev: [p0, p1], slot, kv }); + } + loaded.push(Parked { ids, owner, parts }); + } + let mut total_toks = 0usize; + let mut total_bytes = 0usize; + for pk in loaded.into_iter().rev() { + total_toks += pk.ids.len(); + total_bytes += pk.bytes(); + self.insert_parked(pk); + } + if !self.parked.is_empty() { + eprintln!( + "warm start: loaded {} cached prefix state(s) ({total_toks} tokens, {:.1} MB) from {}", + self.parked.len(), + total_bytes as f64 / 1048576.0, + path.display() + ); + } Ok(()) } @@ -837,10 +1156,21 @@ impl<'a> GpuModel<'a> { self.restore_state() } + fn has_delta(&self) -> bool { + self.parts().any(|p| p.layers.iter().any(|l| matches!(l.mixer, Mixer::Delta { .. }))) + } + /// Keep only the first `keep` tokens of the last `verify`, which started at position `start` (1 <= keep <= `k`): /// recurrent states come from the snapshots it wrote, KV entries past them are dropped by rewinding `pos`. pub fn rollback(&mut self, start: usize, keep: usize) -> anyhow::Result<()> { - if !(1..=self.k).contains(&keep) || start + keep > self.pos { bail!("rollback: keep {keep} of verify at {start}"); } + if !self.has_delta() { + ensure!((1..=TMAX).contains(&keep) && start + keep <= self.pos, "rollback: keep {keep} of verify at {start}"); + self.pos = start + keep; + if let Some(n) = self.next.as_mut() { n.rollback(start, keep)?; } + return Ok(()); + } + let nsnap = self.k.min(SNAP_MAX); + if !(1..=nsnap).contains(&keep) || start + keep > self.pos { bail!("rollback: keep {keep} of verify at {start}"); } let g = &self.gpu; for l in &self.layers { if let Mixer::Delta { conv, s, snap: Some((sc, ss)), .. } = &l.mixer { @@ -852,9 +1182,29 @@ impl<'a> GpuModel<'a> { } } self.pos = start + keep; + if let Some(n) = self.next.as_mut() { n.rollback(start, keep)?; } Ok(()) } + /// Keep the first `keep` tokens (`1 <= keep < toks.len()`) of a `verify` that started at `start`. + /// Pure-attention models only rewind `pos`; DeltaNet models restore up to `SNAP_MAX` snapshots and replay any remainder. + pub fn rollback_tokens(&mut self, start: usize, keep: usize, toks: &[u32]) -> anyhow::Result<()> { + if !self.has_delta() { + return self.rollback(start, keep); + } + let nsnap = self.k.min(SNAP_MAX); + if keep <= nsnap { + return self.rollback(start, keep); + } + self.rollback(start, nsnap)?; + let rem = &toks[nsnap..keep]; + self.verify_run(rem)?; + let t = self.tail(); + let _ = t.finish(rem.len(), gq_min(self.pos, rem.len()), None)?; + t.gpu.sync()?; + self.after_forward(rem.len()) + } + /// Return to the state of the last `save_state`. pub fn restore_state(&mut self) -> anyhow::Result<()> { let g = &self.gpu; @@ -873,9 +1223,12 @@ impl<'a> GpuModel<'a> { self.ple_prev = self.saved_ple_prev; // The draft head's entries before the saved position are still the ones its prompt fill wrote (see `mtp_synced`). if let Some(h) = self.mtp.as_mut() { - g.dtod(&h.saved_x, 0, &self.x, 0, 4 * self.m.cfg.n_embd)?; + let dst = h.eagle.as_ref().map_or(&self.x, |e| &e.cur_g); + g.dtod(&h.saved_x, 0, dst, 0, 4 * self.m.cfg.n_embd)?; h.pos = self.pos.saturating_sub(1); h.pend.clear(); + h.adapt_window = 0x0FFF; + h.adapt_skip = 0; } if let Some(n) = self.next.as_mut() { n.restore_state()?; } Ok(()) @@ -905,6 +1258,16 @@ impl<'a> GpuModel<'a> { let mut res = 0; for (i, l) in self.layers.iter().enumerate().skip(from.map_or(0, |il| il - self.first)) { let il = self.first + i; + if from != Some(il) { + if let Some(eg) = self.mtp.as_ref().and_then(|h| h.eagle.as_ref()) { + if let Some(idx) = eg.cfg.extract_layers.iter().position(|&el| el == il) { + if res != 0 { g.add(x, res, nt * ne)?; res = 0; } + for r in 0..nt { + g.dtod(&self.x, 4 * r * ne, &eg.feat, 4 * (r * 3 + idx) * ne, 4 * ne)?; + } + } + } + } res = self.layer(il, l, x, step, self.worker.as_ref(), nt, res, gq, from == Some(il))?; } if self.next.is_some() { return if res != 0 { g.add(x, res, nt * ne) } else { Ok(()) }; } @@ -913,6 +1276,9 @@ impl<'a> GpuModel<'a> { (None, Some(hr)) => self.hc_read(hr, x, nt, &self.hc_at())?, _ => unreachable!(), } + if let Some(eg) = self.mtp.as_ref().and_then(|h| h.eagle.as_ref()) { + self.mv(&eg.fc, eg.feat.addr(), 3 * ne, eg.cur_g.addr(), ne, nt)?; + } g.gemvq(self.output.as_ref().unwrap(), p("h"), ne, p("hq"), self.logits.addr(), c.n_vocab, nt) } @@ -991,43 +1357,174 @@ impl<'a> GpuModel<'a> { pub fn spec_step(&mut self, t: u32, row: usize, smp: &mut Sampler) -> anyhow::Result<(Vec, u32, usize)> { let k = self.k; if self.pos + k + 1 > self.ctx { return Ok((vec![t], smp.pick(&self.step(t)?)? as u32, 0)); } + let partial_moe = self.cache.as_ref().is_some_and(|c| !c.fully_resident()) + && std::env::var_os("STRATA_MTP_ADAPT").is_none_or(|v| v != "0"); let h = self.mtp.as_mut().unwrap(); let mut first = std::mem::take(&mut h.pend); if first.len() != row { first.clear(); } first.push(t); let nf = first.len(); let hp = h.pos + nf - 1; + let hits12 = (h.adapt_window & 0x0FFF).count_ones(); + let hits6 = (h.adapt_window & 0x3F).count_ones(); + if partial_moe && h.adapt_skip > 0 { + h.adapt_skip -= 1; + let hid = self.row_ptr(row + 1 - nf); + self.mtp_run(&first, hid, false)?; + self.mtp.as_mut().unwrap().pos = hp + 1; + let next = if smp.temp <= 0.0 { + self.step_greedy(t)? + } else { + smp.pick(&self.step(t)?)? as u32 + }; + return Ok((vec![t], next, 0)); + } + let is_eagle = h.eagle.is_some(); + let k_eff = if partial_moe && k > 1 && hits12 < if is_eagle { 7 } else { 11 } { 1 } else { k }; let mut toks = vec![t]; let mut hid = self.row_ptr(row + 1 - nf); - for i in 0..k { + for i in 0..k_eff { let d = if i == 0 { self.mtp_run(&first, hid, true)? } else { self.mtp_run(&toks[i..=i], hid, true)? }; hid = self.mtp.as_ref().unwrap().x.addr() + 4 * ((if i == 0 { nf } else { 1 } - 1) * self.m.cfg.n_embd) as u64; toks.push(d); } let start = self.pos; let (n, next) = if smp.temp <= 0.0 { - self.verify_greedy(&toks, |am| (0..k).take_while(|&i| am[i] == toks[i + 1]).count())? + self.verify_greedy(&toks, |am| (0..k_eff).take_while(|&i| am[i] == toks[i + 1]).count())? } else { let l = self.verify(&toks)?; let mut n = 0; loop { - if n == k { break (n, smp.pick(&l[n])? as u32); } + if n == k_eff { break (n, smp.pick(&l[n])? as u32); } match smp.accept(&l[n], toks[n + 1] as usize)? { None => n += 1, Some(x) => break (n, x as u32) } } }; - self.drafts.0 += k; + self.drafts.0 += k_eff; self.drafts.1 += n; - if n < k { self.rollback(start, n + 1)?; } + if n < k_eff { self.rollback(start, n + 1)?; } // Draft head entries past t were made from draft hidden states: the next call redoes them from main rows. let h = self.mtp.as_mut().unwrap(); h.pos = hp + 1; h.pend = toks[1..=n].to_vec(); + h.adapt_window = ((h.adapt_window << 1) | u16::from(n > 0)) & 0x0FFF; + if partial_moe && n == 0 && (if is_eagle { hits12 <= 3 && hits6 <= 1 } else { hits12 <= 8 || hits6 <= 3 }) { + h.adapt_skip = if is_eagle { 2 } else { 3 }; + } + toks.truncate(n + 1); + Ok((toks, next, n)) + } + + /// Hybrid speculative decode combining context N-gram / lookup bursts (`NgramDrafter`) with MTP / EAGLE3 drafting: + /// uses zero-cost N-gram chains when matched, falls back to `spec_step` (EAGLE3/MTP) on a miss, keeps `MtpHead` + /// synchronized across N-gram steps, and feeds verified tokens back into `NgramDrafter`. + pub fn hybrid_spec_step(&mut self, t: u32, row: usize, smp: &mut Sampler) -> anyhow::Result<(Vec, u32, usize)> { + let max_draft = if self.has_delta() { + self.k.min(SNAP_MAX) + } else { + self.ngram.as_ref().map_or(self.k, |ng| ng.cfg.draft_max).min(TMAX - 1) + }; + let room = self.ctx.saturating_sub(self.pos + 1).min(max_draft); + let draft = match self.ngram.as_mut() { + Some(ng) if room > 0 => ng.draft(t, room), + _ => Vec::new(), + }; + if draft.is_empty() { + let (toks, next, rw) = self.spec_step(t, row, smp)?; + if let Some(ng) = self.ngram.as_mut() { ng.accept(&toks); } + return Ok((toks, next, rw)); + } + // Advance the draft head over [pend.., t] using previous main hidden states before `verify` overwrites them. + let h = self.mtp.as_mut().unwrap(); + let mut first = std::mem::take(&mut h.pend); + if first.len() != row { first.clear(); } + first.push(t); + let nf = first.len(); + let hp = h.pos + nf - 1; + let hid = self.row_ptr(row + 1 - nf); + self.mtp_run(&first, hid, false)?; + + let m = draft.len(); + let mut toks = Vec::with_capacity(1 + m); + toks.push(t); + toks.extend_from_slice(&draft); + let start = self.pos; + let (n, next) = if smp.temp <= 0.0 { + self.verify_greedy(&toks, |am| (0..m).take_while(|&i| am[i] == toks[i + 1]).count())? + } else { + let l = self.verify(&toks)?; + let mut n = 0; + loop { + if n == m { break (n, smp.pick(&l[n])? as u32); } + match smp.accept(&l[n], toks[n + 1] as usize)? { None => n += 1, Some(x) => break (n, x as u32) } + } + }; + self.drafts.0 += m; + self.drafts.1 += n; + if n < m { self.rollback(start, n + 1)?; } + let h = self.mtp.as_mut().unwrap(); + h.pos = hp + 1; + h.pend = toks[1..=n].to_vec(); + if n > 0 { + h.adapt_window = ((h.adapt_window << 1) | 1) & 0x0FFF; + h.adapt_skip = 0; + } toks.truncate(n + 1); + if let Some(ng) = self.ngram.as_mut() { ng.accept(&toks); } Ok((toks, next, n)) } + /// Speculative decode with context N-gram / lookup cache drafting (`NgramDrafter`). + /// When no N-gram match is found (`draft.is_empty()`), executes a normal 1-token decode step (`step_greedy` or `step`) + /// with zero GPU overhead. When `m > 0` draft tokens are proposed, verifies `[t, draft..]` in a single pass and + /// rolls back any rejected suffix via `rollback_tokens`. + pub fn ngram_spec_step(&mut self, t: u32, smp: &mut Sampler) -> anyhow::Result<(Vec, u32)> { + let max_draft = self.ngram.as_ref().map_or(self.k, |ng| ng.cfg.draft_max.max(self.k)).min(TMAX - 1); + let room = self.ctx.saturating_sub(self.pos + 1).min(max_draft); + let draft = match self.ngram.as_mut() { + Some(ng) if room > 0 => ng.draft(t, room), + _ => Vec::new(), + }; + if draft.is_empty() { + if let Some(ng) = self.ngram.as_mut() { ng.accept(&[t]); } + let next = if smp.temp <= 0.0 { + self.step_greedy(t)? + } else { + smp.pick(&self.step(t)?)? as u32 + }; + return Ok((vec![t], next)); + } + let m = draft.len(); + let mut toks = Vec::with_capacity(1 + m); + toks.push(t); + toks.extend_from_slice(&draft); + let start = self.pos; + let (n, next) = if smp.temp <= 0.0 { + self.verify_greedy(&toks, |am| (0..m).take_while(|&i| am[i] == toks[i + 1]).count())? + } else { + let l = self.verify(&toks)?; + let mut n = 0; + loop { + if n == m { break (n, smp.pick(&l[n])? as u32); } + match smp.accept(&l[n], toks[n + 1] as usize)? { None => n += 1, Some(x) => break (n, x as u32) } + } + }; + self.drafts.0 += m; + self.drafts.1 += n; + if n < m { + self.rollback_tokens(start, n + 1, &toks)?; + } + toks.truncate(n + 1); + if let Some(ng) = self.ngram.as_mut() { + ng.accept(&toks); + } + Ok((toks, next)) + } + /// Address of hidden row `row` of the last `step` or `verify`. - fn row_ptr(&self, row: usize) -> u64 { self.x.addr() + 4 * (row * self.m.cfg.n_embd) as u64 } + fn row_ptr(&self, row: usize) -> u64 { + let base = self.mtp.as_ref().and_then(|h| h.eagle.as_ref()).map_or_else(|| self.x.addr(), |e| e.cur_g.addr()); + base + 4 * (row * self.m.cfg.n_embd) as u64 + } /// Draft head on `tokens` with hidden states at `hid` (rows of n_embd), as one batch at the next draft positions; /// returns the argmax of the last token's logits (left in `logits` row 0), or 0 when `head` = false skips the head. @@ -1036,8 +1533,8 @@ impl<'a> GpuModel<'a> { let c = &self.m.cfg; let Some(h) = self.mtp.as_mut() else { bail!("MTP head not loaded") }; if h.pos + nt > self.ctx { bail!("context full ({} tokens; raise --ctx)", self.ctx); } - let fast = self.cache.as_ref().is_some_and(|c| c.fully_resident()) && self.gpu.moe_mailbox_optional(); - let seq = if fast { 0 } else { h.worker.begin() }; + let fast = h.eagle.is_some() || (self.cache.as_ref().is_some_and(|c| c.fully_resident()) && self.gpu.moe_mailbox_optional()); + let seq = if fast { 0 } else { h.worker.as_mut().unwrap().begin() }; let g = &self.gpu; let emb: Vec = tokens.iter().flat_map(|&t| self.m.embed(t)).collect(); g.put(&emb, &h.emb, 0)?; @@ -1045,19 +1542,24 @@ impl<'a> GpuModel<'a> { g.put(&steps, &h.step_buf, 0)?; // `run` keys are token counts (below 8); the top bit keeps these keys apart from them. let gq = gq_min(h.pos, nt); - let key = { let mut s = std::hash::DefaultHasher::new(); (hid, nt, head, gq).hash(&mut s); s.finish() | 1 << 63 }; + let key = { let mut s = std::hash::DefaultHasher::new(); (hid, nt, head, gq, self.cur).hash(&mut s); s.finish() | 1 << 63 }; if self.m.loras.is_empty() { g.graph(key, &mut || self.mtp_forward(hid, nt, head, gq, false))?; } else { self.mtp_forward(hid, nt, head, gq, false)?; } - let mut d = if head { g.get::(&self.amax, 0, 1)?[0] as u32 } else { g.sync()?; 0 }; + let is_eagle = self.mtp.as_ref().is_some_and(|h| h.eagle.is_some()); + let mut d = if head { g.get::(&self.amax, 0, 1)?[0] as u32 } else { if !is_eagle { g.sync()?; } 0 }; if self.finish(nt, gq, Some((hid, head)))? && head { d = self.gpu.get::(&self.amax, 0, 1)?[0] as u32; } let h = self.mtp.as_mut().unwrap(); h.pos += nt; + if let Some(eg) = &h.eagle { + if head { d = eg.cfg.d2t.get(d as usize).copied().unwrap_or(d as i32) as u32; } + return Ok(d); + } // Fast path: every expert of the draft layer is resident, so nothing can miss and there is nothing to record // (R2). The routing counters keep the head's share of the hit rate. if fast { self.hits += c.n_expert_used * nt; } else if let Some(cache) = self.cache.as_mut() { // One mailbox read per layer, as in `after_experts` (each is a bulk copy of the routing words). - let r = h.worker.routing(c.n_layer); + let r = h.worker.as_mut().unwrap().routing(c.n_layer); self.mailbox_reads.set(self.mailbox_reads.get() + 1); let (hits, misses) = tally(&r); self.hits += hits; @@ -1076,7 +1578,7 @@ impl<'a> GpuModel<'a> { let mut again = false; let mut last = None; while self.gpu.stalled() { - let w = match mtp { Some(_) => &self.mtp.as_ref().unwrap().worker, None => self.worker.as_ref().unwrap() }; + let w = match mtp { Some(_) => self.mtp.as_ref().unwrap().worker.as_ref().unwrap(), None => self.worker.as_ref().unwrap() }; let il = w.stalled_layer().context("a MoE combine gave up waiting, but no layer says so")?; // The worker had answered the last time: the GPU does not see the mailbox. ensure!(last != Some(il), "moe_combine: the GPU does not see the CPU expert worker's answers. On an AMD RX 9000 \ @@ -1097,14 +1599,58 @@ impl<'a> GpuModel<'a> { /// Enqueue the draft head (see `mtp_run`). fn mtp_forward(&self, hid: u64, nt: usize, head: bool, gq: i32, resume: bool) -> anyhow::Result<()> { let (c, g, h) = (&self.m.cfg, &self.gpu, self.mtp.as_ref().unwrap()); + if let Some(eg) = &h.eagle { + let ec = &eg.cfg; + let (eps, ne, hp) = (ec.eps, c.n_embd, self.bufs["h"].addr()); + let (eh, x) = (h.eh.addr(), h.x.addr()); + let (nh, nkv, hd) = (ec.n_head, ec.n_head_kv, ec.head_dim); + let kv = nkv * hd; + g.rms_norm(h.emb.addr(), ne, h.enorm.ptr, eh, 2 * ne, ne, nt, eps)?; + g.rms_norm(hid, ne, h.hnorm.ptr, eh + 4 * ne as u64, 2 * ne, ne, nt, eps)?; + if ec.norm_before_residual { + g.rms_norm(hid, ne, h.hnorm.ptr, x, ne, ne, nt, eps)?; + } else if hid != x { + g.memset(&h.x)?; + g.add(x, hid, nt * ne)?; + } + let Mixer::Attn { q, k, v, o, kc, vc, .. } = &h.layer.mixer else { unreachable!() }; + let Ffn::Dense { gate, up, down } = &h.layer.ffn else { unreachable!() }; + g.batch(&mut || { + self.mv(q, eh, 2 * ne, eg.q.addr(), nh * hd, nt)?; + self.mv(k, eh, 2 * ne, eg.k.addr(), kv, nt)?; + self.mv(v, eh, 2 * ne, eg.v.addr(), kv, nt) + })?; + g.batch(&mut || { + g.rope_neox(eg.q.addr(), hd, nt * nh, nh, ec.n_rot, h.step_buf.addr(), ec.rope_base)?; + g.rope_neox(eg.k.addr(), hd, nt * nkv, nkv, ec.n_rot, h.step_buf.addr(), ec.rope_base) + })?; + g.kv_store(eg.k.addr(), eg.v.addr(), kc[self.cur].addr(), vc[self.cur].addr(), kv, nt, h.step_buf.addr(), hd)?; + g.attn_decode_q(eg.q.addr(), kc[self.cur].addr(), vc[self.cur].addr(), 0, eg.att.addr(), + eg.attp.addr(), h.step_buf.addr(), nh, nkv, hd, 1.0 / (hd as f32).sqrt(), gq, nt, 0, 0, 0, 0, 0)?; + self.mv(o, eg.att.addr(), nh * hd, self.bufs["mix"].addr(), ne, nt)?; + g.rms_norm_q(x, ne, h.layer.post_norm.as_ref().unwrap().ptr, hp, ne, ne, nt, eps, 0, self.bufs["mix"].addr())?; + g.batch(&mut || { + self.mv(gate, hp, ne, eg.ge.addr(), ec.n_ff, nt)?; + self.mv(up, hp, ne, eg.ue.addr(), ec.n_ff, nt) + })?; + g.swiglu(eg.ge.addr(), eg.ue.addr(), nt * ec.n_ff)?; + self.mv(down, eg.ge.addr(), ec.n_ff, self.bufs["ye"].addr(), ne, nt)?; + g.add(x, self.bufs["ye"].addr(), nt * ne)?; + if head { + g.rms_norm(x + 4 * ((nt - 1) * ne) as u64, 0, h.head_norm.ptr, hp, 0, ne, 1, eps)?; + g.gemv(&eg.output, hp, self.logits.addr())?; + g.argmax_rows(self.logits.addr(), ec.draft_vocab, 1, self.amax.addr())?; + } + return Ok(()); + } let (eps, ne, hp) = (c.eps, c.n_embd, self.bufs["h"].addr()); let (eh, x) = (h.eh.addr(), h.x.addr()); if !resume { g.rms_norm(h.emb.addr(), ne, h.enorm.ptr, eh, 2 * ne, ne, nt, eps)?; g.rms_norm(hid, ne, h.hnorm.ptr, eh + 4 * ne as u64, 2 * ne, ne, nt, eps)?; - self.mv(&h.eh_proj, eh, 2 * ne, x, ne, nt)?; + self.mv(h.eh_proj.as_ref().unwrap(), eh, 2 * ne, x, ne, nt)?; } - let res = self.layer(c.n_layer, &h.layer, x, h.step_buf.addr(), Some(&h.worker), nt, 0, gq, resume)?; + let res = self.layer(c.n_layer, &h.layer, x, h.step_buf.addr(), h.worker.as_ref(), nt, 0, gq, resume)?; if res != 0 { g.add(x, res, nt * ne)?; } if head { g.rms_norm(x + 4 * ((nt - 1) * ne) as u64, 0, h.head_norm.ptr, hp, 0, ne, 1, eps)?; @@ -1149,28 +1695,26 @@ impl<'a> GpuModel<'a> { } match &l.mixer { Mixer::Delta { qkv, gate, beta, alpha, dt_bias, a, conv_w, norm, out, conv, s, snap } => { - let (sc, ss, ns) = snap.as_ref().map_or((0, 0, 0), |(a, b)| (a.addr(), b.addr(), self.k)); - g.batch(&mut || { - mv2(qkv, gate, "qkv", "z")?; - mv2(beta, alpha, "beta", "g") - })?; + let (sc, ss, ns) = snap.as_ref().map_or((0, 0, 0), |(a, b)| (a.addr(), b.addr(), self.k.min(SNAP_MAX))); + mv(qkv, "h", ne, "qkv", channels)?; // One token with 128-wide heads: conv_silu also does the q and k l2 norms. + // Overlap conv_silu/conv_seq (reads "qkv", writes "y") with gate/beta/alpha GEMVs (read "hq", write "z"/"beta"/"g"). let rows = !self.rows.is_empty(); let nn = if (nt == 1 || rows) && dk == 128 { 2 * nk } else { 0 }; let at = |n: &str, r: usize, w: usize| p(n) + (4 * r * w) as u64; - if rows { - // One token per slot (`step_rows`): each row has its own conv and recurrent state. - g.batch(&mut || { + g.batch(&mut || { + if rows { for (r, &(sl, _)) in self.rows.iter().enumerate() { g.conv_silu(at("qkv", r, channels), conv[sl].addr(), conv_w.ptr, at("y", r, channels), channels, c.ssm_conv, nn, eps)?; } - Ok(()) - })?; - } else if nt == 1 { - g.conv_silu(p("qkv"), conv[self.cur].addr(), conv_w.ptr, p("y"), channels, c.ssm_conv, nn, eps)?; - } else { - g.conv_seq(p("qkv"), conv[self.cur].addr(), conv_w.ptr, p("y"), channels, c.ssm_conv, nt, sc, ns)?; - } + } else if nt == 1 { + g.conv_silu(p("qkv"), conv[self.cur].addr(), conv_w.ptr, p("y"), channels, c.ssm_conv, nn, eps)?; + } else { + g.conv_seq(p("qkv"), conv[self.cur].addr(), conv_w.ptr, p("y"), channels, c.ssm_conv, nt, sc, ns)?; + } + mv(gate, "h", ne, "z", c.ssm_d_inner)?; + mv2(beta, alpha, "beta", "g") + })?; if nn == 0 { g.l2_norm(p("y"), dk, 2 * nk, channels, nt * 2 * nk, eps)?; } let scale = 1.0 / (dk as f32).sqrt(); if rows { diff --git a/src/gpu_model/prefill.rs b/src/gpu_model/prefill.rs index 8a39a9c1..4341a3f4 100644 --- a/src/gpu_model/prefill.rs +++ b/src/gpu_model/prefill.rs @@ -20,6 +20,8 @@ pub(super) struct Scratch { pub chunk: usize, lay: Lay, slab: Buf, + /// `[chunk, 3 * n_embd]` hidden states captured before EAGLE3 `extract_layers`. + feat: Option, /// ids | xrow | inv, `chunk * k` each. ints: Buf, tiles: Buf, @@ -63,19 +65,37 @@ impl Scratch { let n_tiles = 3 * (chunk * k / GEMM_BM + 2 * c.n_expert); Ok(Scratch { chunk, lay: Lay { h, a, b, c: c2, d, e, rl, wts, ge, ye, shg, shy, sg, xn, lo, hg, pk, pq, pg, pn, pv, inj, mix, mo, emb }, - slab: gpu.zeros::(total)?, ints: gpu.zeros::(3 * chunk * k)?, tiles: gpu.zeros::(2 * n_tiles)?, + slab: gpu.zeros::(total)?, + feat: if c.eagle3.is_some() { Some(gpu.zeros::(3 * ne * chunk)?) } else { None }, + ints: gpu.zeros::(3 * chunk * k)?, tiles: gpu.zeros::(2 * n_tiles)?, staging: gpu.zeros::(if c.is_moe() { 3 * STAGE * stage_bytes } else { 0 })?, }) } /// Bytes of VRAM this scratch takes. - pub fn bytes(&self) -> usize { self.slab.len() + self.ints.len() + self.tiles.len() + self.staging.len() } + pub fn bytes(&self) -> usize { self.slab.len() + self.feat.as_ref().map_or(0, Buf::len) + self.ints.len() + self.tiles.len() + self.staging.len() } } impl GpuModel<'_> { /// Run `tokens` from the current position; returns the logits after the last one. Short prompts (and models /// loaded without a prefill chunk) go through the decode step. - pub fn prefill(&mut self, tokens: &[u32]) -> anyhow::Result> { self.prefill_with(tokens, &mut |_| Ok(())) } + pub fn prefill(&mut self, tokens: &[u32]) -> anyhow::Result> { + if self.prompt_cache_path.is_some() && self.pos == 0 && tokens.len() > 1 { + let t0 = std::time::Instant::now(); + let (logits, reused) = self.prefill_cached(tokens, 0, tokens.len() - 1)?; + if reused > 0 { + eprintln!( + "warm start: reused {reused}/{} prompt tokens from cache (TTFT {:.1} ms)", + tokens.len(), + t0.elapsed().as_secs_f64() * 1e3 + ); + } + return Ok(logits); + } + let r = self.prefill_with(tokens, &mut |_| Ok(()))?; + if let Some(ng) = self.ngram.as_mut() { ng.begin(tokens); } + Ok(r) + } /// `prefill`, calling `between` between two chunks (`serve --parallel` decodes the other slots there; it must /// leave the current slot's state alone). @@ -172,8 +192,13 @@ impl GpuModel<'_> { self.prefill_with(&tokens[from..keep], between)?; self.save_state()?; (self.saved_ids, self.saved_owner) = (tokens[..keep].to_vec(), own); + if let Err(e) = self.sync_saved_to_prompt_cache() { + eprintln!("warning: could not save prompt cache: {e:#}"); + } } - Ok((self.prefill_with(&tokens[keep..], between)?, reused)) + let logits = self.prefill_with(&tokens[keep..], between)?; + if let Some(ng) = self.ngram.as_mut() { ng.begin(tokens); } + Ok((logits, reused)) } fn step_all(&mut self, tokens: &[u32]) -> anyhow::Result> { @@ -229,7 +254,8 @@ impl GpuModel<'_> { let kv = nkv * hd; // The draft head runs as one more layer on the chunk while it is in step with the main model. let fill = self.mtp_synced(); - let mtp = self.mtp.as_ref().filter(|_| fill); + let is_eagle = self.mtp.as_ref().is_some_and(|mh| mh.eagle.is_some()); + let mtp = self.mtp.as_ref().filter(|mh| fill && mh.eagle.is_none()); // STRATA_NAN_TRACE[=N]: from position N (default 0), one line per layer with the largest magnitude of each intermediate buffer, and the first // row of one holding a non-finite value. Each probe waits for the GPU: diagnostic runs only. let trace = std::env::var("STRATA_NAN_TRACE").is_ok_and(|v| self.pos >= v.parse::().unwrap_or(0)); @@ -243,6 +269,15 @@ impl GpuModel<'_> { Ok(()) }; for (il, ly) in self.layers.iter().enumerate().map(|(i, ly)| (self.first + i, ly)).chain(mtp.map(|mh| (c.n_layer, &mh.layer))) { + if il < c.n_layer { + if let (Some(pff), Some(eg)) = (&pf.feat, self.mtp.as_ref().and_then(|mh| mh.eagle.as_ref()).filter(|_| fill)) { + if let Some(idx) = eg.cfg.extract_layers.iter().position(|&el| el == il) { + for r in 0..t { + g.dtod(&pf.slab, 4 * r * ne, pff, 4 * (r * 3 + idx) * ne, 4 * ne)?; + } + } + } + } if let Some(mh) = mtp.filter(|_| il == c.n_layer) { // Its entry for token i pairs the token with the hidden state before it (`mtp_run`): h gets the last // hidden state of the previous chunk, then rows 0..t - 1; the first token of the sequence has none. @@ -255,7 +290,7 @@ impl GpuModel<'_> { t -= skip; g.rms_norm(x, ne, mh.enorm.ptr, at(l.a), 2 * ne, ne, t, eps)?; g.rms_norm(h, ne, mh.hnorm.ptr, at(l.a + ne), 2 * ne, ne, t, eps)?; - g.gemm_dense(&mh.eh_proj, t, at(l.a), 2 * ne, x, ne, 0)?; + g.gemm_dense(mh.eh_proj.as_ref().unwrap(), t, at(l.a), 2 * ne, x, ne, 0)?; g.put(&[mh.pos as i32, 0], &self.step_buf, 2)?; g.put(&[mh.pos as i32 + self.rope_offset, 0], &self.rope_buf, 2)?; step += 8; @@ -517,10 +552,29 @@ impl GpuModel<'_> { if trace { eprintln!("nan-trace pos {} layer {il}:{tr}", self.pos); tr.clear(); } g.kick()?; } - match self.mtp.as_mut() { - Some(mh) if fill => mh.pos += t, - // Row 0 of `x` gets the last hidden state (with the draft head, the copy above did it). - _ => g.dtod(&pf.slab, 4 * (t - 1) * d, &self.x, 0, 4 * d)?, + if fill && is_eagle { + g.dtod(&pf.slab, 4 * (t - 1) * d, &self.x, 0, 4 * d)?; + let (lh, slab_h, old_cur_g) = { + let pf = self.pf.as_ref().unwrap(); + let eg = self.mtp.as_ref().unwrap().eagle.as_ref().unwrap(); + self.gpu.gemm_dense(&eg.fc, t, pf.feat.as_ref().unwrap().addr(), 3 * ne, pf.slab.addr() + 4 * l.h as u64, ne, 0)?; + (l.h, pf.slab.addr() + 4 * l.h as u64, eg.cur_g.addr()) + }; + if self.pos > 0 { + self.mtp_run(&toks[..1], old_cur_g, false)?; + } + for (off, ch) in toks[1..].chunks(crate::cpu_worker::TMAX).enumerate() { + self.mtp_run(ch, slab_h + 4 * (off * crate::cpu_worker::TMAX * ne) as u64, false)?; + } + let pf = self.pf.as_ref().unwrap(); + let eg = self.mtp.as_ref().unwrap().eagle.as_ref().unwrap(); + self.gpu.dtod(&pf.slab, 4 * (lh + (t - 1) * ne), &eg.cur_g, 0, 4 * ne)?; + } else { + match self.mtp.as_mut() { + Some(mh) if fill => mh.pos += t, + // Row 0 of `x` gets the last hidden state (with the draft head, the copy above did it). + _ => g.dtod(&pf.slab, 4 * (t - 1) * d, &self.x, 0, 4 * d)?, + } } self.pos += toks.len(); let Some(nx) = self.next.as_deref_mut() else { return Ok(()) }; diff --git a/src/lib.rs b/src/lib.rs index ebb33e24..96f17186 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,3 +20,4 @@ pub mod hf; pub mod deepseek; pub mod gemma4; pub mod glm; +pub mod ngram; diff --git a/src/main.rs b/src/main.rs index 8973856d..d084ad71 100644 --- a/src/main.rs +++ b/src/main.rs @@ -38,6 +38,37 @@ struct Engine { /// [default: calibrated, else 0]. #[arg(long, value_name = "N", default_missing_value = "1", num_args = 0..=1)] mtp: Option, + /// Standalone MTP draft head GGUF (e.g. ~/models-local/mtp-Q8_0.gguf). + #[arg(long)] + mtp_model: Option, + /// External speculative draft model GGUF socket (e.g. future DFlash/EAGLE draft). + #[arg(long)] + draft_model: Option, + /// Speculative context-lookup drafter type(s), comma-separated: none, ngram-simple, ngram-map-k, ngram-map-k4v, + /// ngram-mod, ngram-cache (combines with --mtp/--mtp-model in hybrid cascade). + #[arg(long, default_value = "none")] + spec_type: String, + /// Maximum number of speculative draft tokens verified per step (1..=16). + #[arg(long, default_value_t = 16)] + spec_draft_n_max: usize, + /// N-gram lookup key length in tokens (for ngram-simple, ngram-map-k, ngram-map-k4v, ngram-mod). + #[arg(long, default_value_t = 12)] + spec_ngram_size_n: usize, + /// N-gram draft window length in tokens (for ngram-simple, ngram-map-k, ngram-map-k4v). + #[arg(long, default_value_t = 48)] + spec_ngram_size_m: usize, + /// Minimum hit count before drafting (for ngram-map-k, ngram-map-k4v). + #[arg(long, default_value_t = 1)] + spec_ngram_min_hits: u16, + /// Static lookup cache file (llama.cpp binary ngram cache or UTF-8 text scaffold corpus). + #[arg(long)] + lookup_cache_static: Option, + /// Dynamic lookup cache file loaded at startup and updated/saved across requests. + #[arg(long)] + lookup_cache_dynamic: Option, + /// Persistent disk cache file for KV and recurrent prefix states (Warm Start across requests and restarts). + #[arg(long)] + prompt_cache: Option, /// Sampling temperature; 0 is greedy. #[arg(long, default_value_t = 0.0)] temp: f32, @@ -231,7 +262,10 @@ impl Runner<'_> { if let Runner::Gpu(g) = self { let sum = |f: &dyn Fn(&GpuModel) -> usize| g.parts().map(f).sum::(); let st = self.expert_stats().expect("Runner::Gpu has expert stats"); - if g.drafts.0 > 0 { eprintln!("MTP drafts accepted {}/{} ({:.1}%)", g.drafts.1, g.drafts.0, 100.0 * g.drafts.1 as f64 / g.drafts.0 as f64); } + if g.drafts.0 > 0 { + let kind = if g.ngram.is_some() { "speculative ngram" } else { "MTP" }; + eprintln!("{kind} drafts accepted {}/{} ({:.1}%)", g.drafts.1, g.drafts.0, 100.0 * g.drafts.1 as f64 / g.drafts.0 as f64); + } eprintln!("expert hit rate {:.1}% ({} misses), cache swaps {}", 100.0 * st.hits as f64 / (st.hits + st.misses).max(1) as f64, st.misses, st.swaps); // Issue #34 R1: printed on every run, so a fully resident run reads as a measured 0 in the log instead of // an absent line. `tokens` counts the token forwards the worker could have been asked about. @@ -281,8 +315,20 @@ impl Engine { let v: serde_json::Value = std::fs::read(&path).ok().and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default(); if !v.is_null() { eprintln!("calibration: {}", path.display()); } let get = |k: &str| v[k].as_u64().map(|x| x as usize); + if let Some(dp) = &self.draft_model { + anyhow::ensure!(dp.exists(), "draft model path does not exist: {}", dp.display()); + eprintln!("draft model socket: registered candidate {}", dp.display()); + } + let mut ngram_cfg = self.ngram_config()?; + ngram_cfg.n_vocab = m.cfg.n_vocab; self.chunk.get_or_insert(get("chunk").unwrap_or(4096)); - self.mtp.get_or_insert(get("mtp").unwrap_or(0)); + if self.mtp_model.is_some() && self.mtp.is_none() { + self.mtp = Some(1); + } else if ngram_cfg.is_enabled() && self.mtp.is_none() { + self.mtp = Some(0); + } else { + self.mtp.get_or_insert(get("mtp").unwrap_or(0)); + } self.short = get("short").unwrap_or(32); // The global pool can be built once; a model loaded later by `serve` keeps the first one's threads. // GPU expert work is memory-bound, so SMT siblings only add contention (6 threads beat 12 on the reference @@ -307,15 +353,45 @@ impl Engine { let types = rusting_strata::gpu_model::gpu_weight_types(&m.w.g); for b in &bk { b.preflight(&types, false)?; } let margins = vram_margins(&self.vram_margin, &self.device)?.into_iter().map(|mb| mb << 20).collect::>(); + let spec_k = if mtp > 0 { mtp } else if ngram_cfg.is_enabled() && self.parallel <= 1 { ngram_cfg.draft_max } else { 0 }; let mut g = if self.parallel > 1 { GpuModel::load_split_margins(m, bk, self.ctx, self.chunk(), &margins, 0, self.parallel).context(LOAD_HINT)? } else { - GpuModel::load_split_margins(m, bk, self.ctx, self.chunk(), &margins, self.mtp(), 1).context(LOAD_HINT)? + GpuModel::load_split_margins(m, bk, self.ctx, self.chunk(), &margins, spec_k, 1).context(LOAD_HINT)? }; + if ngram_cfg.is_enabled() && self.parallel <= 1 { + let tok = Tokenizer::from_gguf(&m.w.g).ok(); + g.ngram = Some(rusting_strata::ngram::NgramDrafter::new(ngram_cfg, tok.as_ref())?); + } g.short = self.short; g.cache_ram = self.cache_ram << 20; + if let Some(p) = &self.prompt_cache { + if g.cache_ram == 0 { + g.cache_ram = 2048 << 20; + } + if let Err(e) = g.load_prompt_cache(p) { + eprintln!("warning: could not load prompt cache {}: {e:#}", p.display()); + } + } Ok(Runner::Gpu(Box::new(g))) } + fn ngram_config(&self) -> anyhow::Result { + let types = rusting_strata::ngram::SpecType::parse_list(&self.spec_type)?; + anyhow::ensure!((1..=rusting_strata::cpu_worker::TMAX - 1).contains(&self.spec_draft_n_max), + "--spec-draft-n-max {}: must be between 1 and {}", self.spec_draft_n_max, rusting_strata::cpu_worker::TMAX - 1); + Ok(rusting_strata::ngram::NgramConfig { + types, + draft_max: self.spec_draft_n_max, + size_n: self.spec_ngram_size_n.max(1), + size_m: self.spec_ngram_size_m.max(1), + min_hits: self.spec_ngram_min_hits.max(1), + mod_n_match: self.spec_ngram_size_n.max(1), + mod_n_min: 2, + lookup_cache_static: self.lookup_cache_static.clone(), + lookup_cache_dynamic: self.lookup_cache_dynamic.clone(), + n_vocab: 0, + }) + } fn sampler(&self) -> Sampler { Sampler { temp: self.temp, top_k: self.top_k, top_p: self.top_p, rng: seed(self.seed) } } @@ -820,8 +896,8 @@ fn run(cmd: Cmd) -> anyhow::Result<()> { let (prompt, cont) = (ids("prompt")?, ids("continuation")?); anyhow::ensure!(!prompt.is_empty() && !cont.is_empty(), "empty prompt or continuation"); let g = Gguf::open(&model)?; - let mut m = Model::load(&g)?; - if !engine.lora.is_empty() { m.apply_lora(&engine.lora_specs())?; } + let aux_g = open_mtp_gguf(engine.mtp_model.as_deref())?; + let m = Model::load_with_aux(&g, aux_g.as_ref(), &engine.lora_specs())?; if let Some(&t) = prompt.iter().chain(&cont).find(|&&t| t as usize >= m.cfg.n_vocab) { anyhow::bail!("token id {t} out of range (vocabulary has {} tokens)", m.cfg.n_vocab); } @@ -860,8 +936,8 @@ fn run(cmd: Cmd) -> anyhow::Result<()> { let prompt = if prompt == "-" { let mut s = String::new(); std::io::stdin().read_to_string(&mut s)?; s } else { prompt }; return rusting_strata::deepseek::chat::generate(&g, &prompt, n); } - let mut m = Model::load(&g)?; - if !engine.lora.is_empty() { m.apply_lora(&engine.lora_specs())?; } + let aux_g = open_mtp_gguf(engine.mtp_model.as_deref())?; + let m = Model::load_with_aux(&g, aux_g.as_ref(), &engine.lora_specs())?; let tok = Tokenizer::from_gguf(&g)?; let prompt = if prompt == "-" { let mut s = String::new(); @@ -901,8 +977,8 @@ fn run(cmd: Cmd) -> anyhow::Result<()> { } Cmd::Perplexity { model, file, ctx_len: n, chunks, kld_base, mut engine } => { let g = Gguf::open(&model)?; - let mut m = Model::load(&g)?; - if !engine.lora.is_empty() { m.apply_lora(&engine.lora_specs())?; } + let aux_g = open_mtp_gguf(engine.mtp_model.as_deref())?; + let m = Model::load_with_aux(&g, aux_g.as_ref(), &engine.lora_specs())?; let tok = Tokenizer::from_gguf(&g)?; // ponytail: no BOS at chunk starts (llama-perplexity adds one where the vocabulary asks for it; Qwen does not) let ids = tok.encode(&std::fs::read_to_string(&file)?)?; @@ -948,8 +1024,8 @@ fn run(cmd: Cmd) -> anyhow::Result<()> { }; } } - let mut m = Model::load(&g)?; - if !engine.lora.is_empty() { m.apply_lora(&engine.lora_specs())?; } + let aux_g = open_mtp_gguf(engine.mtp_model.as_deref())?; + let m = Model::load_with_aux(&g, aux_g.as_ref(), &engine.lora_specs())?; let tok = Tokenizer::from_gguf(&g)?; let mut r = engine.runner(&m, &model)?; let mut results = Vec::with_capacity(pp.len()); @@ -1056,8 +1132,8 @@ fn run(cmd: Cmd) -> anyhow::Result<()> { anyhow::ensure!(image.is_empty(), "--image is not supported for deepseek41"); return rusting_strata::deepseek::chat::chat(&g, system, n, think); } - let mut m = Model::load(&g)?; - if !engine.lora.is_empty() { m.apply_lora(&engine.lora_specs())?; } + let aux_g = open_mtp_gguf(engine.mtp_model.as_deref())?; + let m = Model::load_with_aux(&g, aux_g.as_ref(), &engine.lora_specs())?; let tok = Tokenizer::from_gguf(&g)?; let image_input = if image.is_empty() { None } else { let mmproj = engine.mmproj.clone().or_else(|| vision::find_sibling(&model)); @@ -1524,33 +1600,46 @@ fn decode(r: &mut Runner, tok: &Tokenizer, mtp: bool, smp: &mut Sampler, mut log let eos = tok.eos_ids(); // `ahead`: the next token when `spec_step` or `step_greedy` already picked it. // `start`: the first token whose text is not emitted yet (the text before it ends on a whole character). - let (mut out, mut start, mut row, mut ahead) = (Vec::new(), 0, 0, None); - while out.len() < n { + let (mut out, mut start, mut row, mut ahead, mut stopped) = (Vec::new(), 0, 0, None, false); + 'outer: while out.len() < n { let next = match ahead.take() { Some(t) => t, None => smp.pick(&logits)? as u32 }; - if eos.contains(&next) { return Ok((out, true)); } + if eos.contains(&next) { stopped = true; break; } let toks = match r { + Runner::Gpu(g) if mtp && g.ngram.is_some() => { + let (toks, nx, rw) = g.hybrid_spec_step(next, row, smp)?; + (ahead, row) = (Some(nx), rw); + toks + } Runner::Gpu(g) if mtp => { let (toks, nx, rw) = g.spec_step(next, row, smp)?; (ahead, row) = (Some(nx), rw); toks } + Runner::Gpu(g) if g.ngram.is_some() => { + let (toks, nx) = g.ngram_spec_step(next, smp)?; + ahead = Some(nx); + toks + } // Greedy without MTP: the argmax on the GPU, so only the token id is copied back, not the logits. Runner::Gpu(g) if smp.temp <= 0.0 => { ahead = Some(g.step_greedy(next)?); vec![next] } Runner::Slot(c) if smp.temp <= 0.0 => { ahead = Some(c.step(next, false)?.0); vec![next] } _ => { logits = r.step(next)?; vec![next] } }; for t in toks { - if eos.contains(&t) { return Ok((out, true)); } + if eos.contains(&t) { stopped = true; break 'outer; } if out.len() == n { break; } out.push(t); let s = tok.decode(&out[start..])?; if !s.ends_with('\u{fffd}') { - if emit(&s)? { return Ok((out, true)); } + if emit(&s)? { stopped = true; break 'outer; } start = out.len(); } } } - Ok((out, false)) + if let Runner::Gpu(g) = r && let Some(ng) = g.ngram.as_mut() { + let _ = ng.save_dynamic(); + } + Ok((out, stopped)) } /// Tokens left for a reply: `n`, capped by the context after a `used`-token prompt. @@ -1746,7 +1835,7 @@ impl Sched { /// One forward over the oldest steps (with more than `TMAX` slots decoding, each waits its turn), in slot order /// (fewer graphs to capture). fn run(&mut self, g: &mut GpuModel) { - let mut batch: Vec<_> = self.steps.drain(..self.steps.len().min(rusting_strata::cpu_worker::TMAX)).collect(); + let mut batch: Vec<_> = self.steps.drain(..self.steps.len().min(4)).collect(); batch.sort_by_key(|s| s.0); self.since = if self.steps.is_empty() { None } else { Some(Instant::now() - Self::GATHER) }; match g.step_rows(&batch.iter().map(|b| (b.0, b.1, b.2)).collect::>()) { @@ -1841,6 +1930,37 @@ impl Serving { } } +/// Resolve `--mtp-model` (`.gguf` file, HF `.safetensors` file, or directory containing either `model.gguf` or +/// `model.safetensors` + `config.json`) and open it as a companion GGUF. If `.gguf` is missing (or +/// `STRATA_EAGLE3_RECONVERT` is set) and `.safetensors` + `config.json` are present, converts natively in Rust. +fn open_mtp_gguf(path: Option<&Path>) -> anyhow::Result> { + let Some(p) = path else { return Ok(None) }; + let reconvert = std::env::var_os("STRATA_EAGLE3_RECONVERT").is_some(); + let (gguf, st_cfg) = if p.is_dir() { + (p.join("model.gguf"), Some((p.join("model.safetensors"), p.join("config.json")))) + } else if p.extension().and_then(|e| e.to_str()) == Some("safetensors") { + (p.with_extension("gguf"), Some((p.to_path_buf(), p.with_file_name("config.json")))) + } else if p.extension().and_then(|e| e.to_str()) == Some("bin") { + (p.with_file_name("model.gguf"), None) + } else { + (p.to_path_buf(), None) + }; + if let Some((st, cfg)) = st_cfg + && (!gguf.exists() || reconvert) + && st.exists() + && cfg.exists() + { + eprintln!("note: native Rust EAGLE3 conversion {} -> {}", st.display(), gguf.display()); + rusting_strata::gguf::convert_safetensors_eagle3_to_gguf(&st, &cfg, &gguf, 48)?; + } + anyhow::ensure!( + gguf.exists(), + "{}: not found (for PyTorch .bin EAGLE3 checkpoints, convert first via tools/convert_eagle3_to_gguf.py)", + gguf.display() + ); + Ok(Some(Gguf::open(&gguf)?)) +} + /// Loads `model`, answers `w` (the /api/load request that asked for it) and serves until another model is asked /// for. The model and its GPU memory are freed on return. fn serve_model(model: &Path, mut engine: Engine, s: &Serving, w: &mut Option) -> anyhow::Result> { @@ -1855,7 +1975,8 @@ fn serve_model(model: &Path, mut engine: Engine, s: &Serving, w: &mut Option Vec { impl<'a> Model<'a> { pub fn load(g: &'a Gguf) -> anyhow::Result> { Self::load_with_lora(g, &[]) } pub fn load_with_lora(g: &'a Gguf, loras: &[(String, f32)]) -> anyhow::Result> { + Self::load_with_aux(g, None, loras) + } + pub fn load_with_aux(g: &'a Gguf, aux: Option<&'a Gguf>, loras: &[(String, f32)]) -> anyhow::Result> { let mut cfg = ModelConfig::from_gguf(g)?; - let w = Weights::new(g)?; + let w = Weights::new_with_aux(g, aux)?; + if cfg.n_nextn == 0 { + if let Some(ag) = aux.filter(|_| w.layer(cfg.n_layer, "eagle3.fc.weight").is_ok()) { + let il = cfg.n_layer; + let wq = w.layer(il, "attn_q.weight")?; + let wk = w.layer(il, "attn_k.weight")?; + let gate = w.layer(il, "ffn_gate.weight")?; + let out = w.layer(il, "eagle3.output.weight")?; + let n_head = ag.meta.get("eagle3.attention.head_count").and_then(|v| v.as_u64()).map(|v| v as usize).unwrap_or(cfg.n_head); + let head_dim = wq.rows / n_head.max(1); + let n_head_kv = (wk.rows / head_dim.max(1)).max(1); + let n_ff = gate.rows; + let draft_vocab = out.rows; + let n_rot = ag.meta.get("eagle3.rope.dimension_count").and_then(|v| v.as_u64()).map(|v| v as usize).unwrap_or(head_dim); + let rope_base = ag.meta.get("eagle3.rope.freq_base").and_then(|v| v.as_f32()).unwrap_or(cfg.rope_base); + let eps = ag.meta.get("eagle3.attention.layer_norm_rms_epsilon").and_then(|v| v.as_f32()).unwrap_or(cfg.eps); + let norm_before_residual = matches!(ag.meta.get("eagle3.norm_before_residual"), Some(crate::gguf::Value::Bool(true))); + let rope_neox = !matches!(ag.meta.get("eagle3.rope.neox"), Some(crate::gguf::Value::Bool(false))); + let extract_layers = match ag.meta.get("eagle3.target_layers").or_else(|| ag.meta.get("eagle3.extract_layers")).and_then(|v| v.as_arr()) { + Some(arr) if arr.len() == 3 => { + let mut ls = [0usize; 3]; + for (i, v) in arr.iter().enumerate() { + ls[i] = v.as_u64().context("bad eagle3 extract layer")? as usize; + } + ls + } + _ => [2, cfg.n_layer / 2, cfg.n_layer.saturating_sub(3)], + }; + let d2t = if let Some(t) = ag.tensors.get("d2t") { + let n = t.dims[0] as usize; + anyhow::ensure!(n == draft_vocab, "d2t len {n} != draft_vocab {draft_vocab}"); + let mut ids: Vec = match t.ty { + 26 => ag.raw_bytes(t, 4 * n)?.chunks_exact(4).map(|b| i32::from_le_bytes(b.try_into().unwrap())).collect(), + 27 => ag.raw_bytes(t, 8 * n)?.chunks_exact(8).map(|b| i64::from_le_bytes(b.try_into().unwrap()) as i32).collect(), + ty => anyhow::bail!("unsupported EAGLE3 d2t tensor type {ty}"), + }; + if ids.windows(2).any(|w| w[0] >= w[1]) { + for (i, v) in ids.iter_mut().enumerate() { *v += i as i32; } + } + ids + } else { + (0..draft_vocab as i32).collect() + }; + cfg.n_nextn = 1; + cfg.eagle3 = Some(Eagle3Config { + extract_layers, n_ff, n_head, n_head_kv, head_dim, n_rot, rope_base, eps, + draft_vocab, norm_before_residual, rope_neox, d2t, + }); + eprintln!( + "note: loaded standalone EAGLE3 draft head (extract_layers={extract_layers:?}, heads={n_head}x{head_dim}, kv={n_head_kv}, ff={n_ff}, draft_vocab={draft_vocab})" + ); + } else if aux.is_some_and(|ag| ag.tensors.contains_key(&format!("blk.{}.nextn.eh_proj.weight", cfg.n_layer))) { + cfg.n_nextn = 1; + eprintln!("note: loaded standalone MTP draft head from companion GGUF (layer {})", cfg.n_layer); + } + } if !w.validate(cfg.n_layer)? && cfg.n_nextn > 0 { eprintln!("note: the MTP (nextn) layer has a quant type this engine does not support; loading without MTP"); cfg.n_nextn = 0; } - w.get("token_embd.weight")?; + let emb_t = w.get("token_embd.weight")?; + // Pre-touch mmap pages of token_embd.weight so decode step put_embed() never stalls on SSD page faults. + let mut _touch = 0u8; + for off in (0..emb_t.data.len()).step_by(4096) { + _touch ^= unsafe { std::ptr::read_volatile(&emb_t.data[off]) }; + } + std::hint::black_box(_touch); let ple = if cfg.n_embd_ple > 0 { let p = |k: &str| format!("{}.ple.{k}", cfg.arch); let arr = |k: &str| -> anyhow::Result> { @@ -175,7 +239,8 @@ impl<'a> Model<'a> { pub fn embed(&self, token: u32) -> Vec { let mut x = vec![0f32; self.cfg.n_embd]; - self.w.get("token_embd.weight").unwrap().row(token as usize, &mut x); + let idx = (token as usize).min(self.cfg.n_vocab.saturating_sub(1)); + self.w.get("token_embd.weight").unwrap().row(idx, &mut x); x } @@ -665,4 +730,215 @@ mod tests { assert_eq!(e.a[0], (0..8).map(|i| i as f32).collect::>()); assert_eq!(e.b[0], (0..6).map(|i| 10.0 + i as f32).collect::>()); } + + #[test] + fn load_eagle3_aux_config() { + let (mut hdr, mut off) = (Vec::new(), 0u64); + let push_tensor = |h: &mut Vec, body: &mut Vec, off: &mut u64, name: &str, dims: &[u64], ty: u32, data: &[u8]| { + w_str(h, name); + h.extend((dims.len() as u32).to_le_bytes()); + for &d in dims { h.extend(d.to_le_bytes()); } + h.extend(ty.to_le_bytes()); + h.extend((*off).to_le_bytes()); + body.extend_from_slice(data); + let pad = (32 - data.len() % 32) % 32; + body.extend(std::iter::repeat_n(0u8, pad)); + *off += (data.len() + pad) as u64; + }; + let mut body = Vec::new(); + let mut kvs = Vec::new(); + let add_u32 = |b: &mut Vec, k: &str, v: u32| { w_str(b, k); b.extend(4u32.to_le_bytes()); b.extend(v.to_le_bytes()); }; + let add_f32 = |b: &mut Vec, k: &str, v: f32| { w_str(b, k); b.extend(6u32.to_le_bytes()); b.extend(v.to_le_bytes()); }; + let add_str = |b: &mut Vec, k: &str, v: &str| { w_str(b, k); b.extend(8u32.to_le_bytes()); w_str(b, v); }; + // Build tiny target GGUF (`qwen3moe` with 1 layer, embd=4, n_vocab=8, nh=2, nkv=1, hd=2) + add_str(&mut kvs, "general.architecture", "qwen3moe"); + add_u32(&mut kvs, "qwen3moe.block_count", 1); + add_u32(&mut kvs, "qwen3moe.embedding_length", 4); + add_u32(&mut kvs, "qwen3moe.attention.head_count", 2); + add_u32(&mut kvs, "qwen3moe.attention.head_count_kv", 1); + add_f32(&mut kvs, "qwen3moe.attention.layer_norm_rms_epsilon", 1e-6); + let z4 = vec![0u8; 4 * 4]; + let z4x8 = vec![0u8; 32 * 4]; + push_tensor(&mut hdr, &mut body, &mut off, "token_embd.weight", &[4, 8], 0, &z4x8); + push_tensor(&mut hdr, &mut body, &mut off, "output_norm.weight", &[4], 0, &z4); + push_tensor(&mut hdr, &mut body, &mut off, "blk.0.attn_norm.weight", &[4], 0, &z4); + let mut tgt = Vec::new(); + tgt.extend(b"GGUF"); + tgt.extend(3u32.to_le_bytes()); + tgt.extend(3u64.to_le_bytes()); + tgt.extend(6u64.to_le_bytes()); + tgt.extend(kvs); + tgt.extend(hdr); + while tgt.len() % 32 != 0 { tgt.push(0); } + tgt.extend(body); + + // Build tiny EAGLE3 aux GGUF (`draft_vocab=3`, draft `heads=2x4`, `kv=1`, `ff=6`, `d2t` relative `[0, 0, 5]` -> abs `[0, 1, 7]`) + let (mut ahdr, mut abody, mut aoff) = (Vec::new(), Vec::new(), 0u64); + let mut akvs = Vec::new(); + add_str(&mut akvs, "general.architecture", "eagle3"); + add_u32(&mut akvs, "eagle3.attention.head_count", 2); + add_f32(&mut akvs, "eagle3.rope.freq_base", 1000000.0); + w_str(&mut akvs, "eagle3.target_layers"); + akvs.extend(9u32.to_le_bytes()); + akvs.extend(5u32.to_le_bytes()); + akvs.extend(3u64.to_le_bytes()); + for v in [0i32, 0, 0] { akvs.extend(v.to_le_bytes()); } + let fc_d = vec![0u8; 12 * 4 * 4]; + let norm_d = vec![0u8; 4 * 4]; + let qk_d = vec![0u8; 8 * 8 * 4]; + let kv_d = vec![0u8; 8 * 4 * 4]; + let o_d = vec![0u8; 8 * 4 * 4]; + let ff_d = vec![0u8; 4 * 6 * 4]; + let out_d = vec![0u8; 4 * 3 * 4]; + let d2t_d: Vec = [0i64, 0, 5].into_iter().flat_map(|v| v.to_le_bytes()).collect(); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "fc.weight", &[12, 4], 0, &fc_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "output.weight", &[4, 3], 0, &out_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "output_norm.weight", &[4], 0, &norm_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.attn_norm.weight", &[4], 0, &norm_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.attn_norm_2.weight", &[4], 0, &norm_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.ffn_norm.weight", &[4], 0, &norm_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.attn_q.weight", &[8, 8], 0, &qk_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.attn_k.weight", &[8, 4], 0, &kv_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.attn_v.weight", &[8, 4], 0, &kv_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.attn_output.weight", &[8, 4], 0, &o_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.ffn_gate.weight", &[4, 6], 0, &ff_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.ffn_up.weight", &[4, 6], 0, &ff_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "blk.0.ffn_down.weight", &[6, 4], 0, &ff_d); + push_tensor(&mut ahdr, &mut abody, &mut aoff, "d2t", &[3], 27, &d2t_d); + let mut aux = Vec::new(); + aux.extend(b"GGUF"); + aux.extend(3u32.to_le_bytes()); + aux.extend(14u64.to_le_bytes()); + aux.extend(4u64.to_le_bytes()); + aux.extend(akvs); + aux.extend(ahdr); + while aux.len() % 32 != 0 { aux.push(0); } + aux.extend(abody); + + let pt = std::env::temp_dir().join(format!("strata-e3-tgt-{}.gguf", std::process::id())); + let pa = std::env::temp_dir().join(format!("strata-e3-aux-{}.gguf", std::process::id())); + std::fs::write(&pt, &tgt).unwrap(); + std::fs::write(&pa, &aux).unwrap(); + let gt = Gguf::open(&pt).unwrap(); + let ga = Gguf::open(&pa).unwrap(); + let m = Model::load_with_aux(>, Some(&ga), &[]).unwrap(); + std::fs::remove_file(&pt).ok(); + std::fs::remove_file(&pa).ok(); + let ec = m.cfg.eagle3.as_ref().expect("eagle3 config loaded"); + assert_eq!(m.cfg.n_nextn, 1); + assert_eq!((ec.n_head, ec.head_dim, ec.n_head_kv, ec.n_ff, ec.draft_vocab), (2, 4, 1, 6, 3)); + assert_eq!(ec.d2t, vec![0, 1, 7]); + } + + #[test] + fn eagle3_safetensors_auto_converts_and_loads() { + let dir = std::env::temp_dir().join(format!("strata-e3-st-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + + // Build tiny target GGUF (`qwen3moe` 1 layer, embd=4, n_vocab=8, nh=2, nkv=1, hd=2) + let (mut hdr, mut body, mut off) = (Vec::new(), Vec::new(), 0u64); + let push_tensor = |h: &mut Vec, b: &mut Vec, o: &mut u64, name: &str, dims: &[u64], ty: u32, data: &[u8]| { + w_str(h, name); + h.extend((dims.len() as u32).to_le_bytes()); + for &d in dims { h.extend(d.to_le_bytes()); } + h.extend(ty.to_le_bytes()); + h.extend((*o).to_le_bytes()); + b.extend_from_slice(data); + let pad = (32 - data.len() % 32) % 32; + b.extend(std::iter::repeat_n(0u8, pad)); + *o += (data.len() + pad) as u64; + }; + let mut kvs = Vec::new(); + let add_u32 = |b: &mut Vec, k: &str, v: u32| { w_str(b, k); b.extend(4u32.to_le_bytes()); b.extend(v.to_le_bytes()); }; + let add_f32 = |b: &mut Vec, k: &str, v: f32| { w_str(b, k); b.extend(6u32.to_le_bytes()); b.extend(v.to_le_bytes()); }; + let add_str = |b: &mut Vec, k: &str, v: &str| { w_str(b, k); b.extend(8u32.to_le_bytes()); w_str(b, v); }; + add_str(&mut kvs, "general.architecture", "qwen3moe"); + add_u32(&mut kvs, "qwen3moe.block_count", 1); + add_u32(&mut kvs, "qwen3moe.embedding_length", 4); + add_u32(&mut kvs, "qwen3moe.attention.head_count", 2); + add_u32(&mut kvs, "qwen3moe.attention.head_count_kv", 1); + add_f32(&mut kvs, "qwen3moe.attention.layer_norm_rms_epsilon", 1e-6); + let z4 = vec![0u8; 4 * 4]; + let z4x8 = vec![0u8; 32 * 4]; + push_tensor(&mut hdr, &mut body, &mut off, "token_embd.weight", &[4, 8], 0, &z4x8); + push_tensor(&mut hdr, &mut body, &mut off, "output_norm.weight", &[4], 0, &z4); + push_tensor(&mut hdr, &mut body, &mut off, "blk.0.attn_norm.weight", &[4], 0, &z4); + let mut tgt = Vec::new(); + tgt.extend(b"GGUF"); + tgt.extend(3u32.to_le_bytes()); + tgt.extend(3u64.to_le_bytes()); + tgt.extend(6u64.to_le_bytes()); + tgt.extend(kvs); + tgt.extend(hdr); + while tgt.len() % 32 != 0 { tgt.push(0); } + tgt.extend(body); + let pt = dir.join("target.gguf"); + std::fs::write(&pt, &tgt).unwrap(); + + // Build synthetic HuggingFace EAGLE3 `config.json` + `model.safetensors` + let cfg_path = dir.join("config.json"); + std::fs::write( + &cfg_path, + r#"{ + "hidden_size": 4, + "intermediate_size": 6, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "rms_norm_eps": 1e-6, + "rope_theta": 1000000.0, + "norm_before_residual": false, + "draft_vocab_size": 3, + "eagle_aux_hidden_state_layer_ids": [0, 0, 0] + }"#, + ).unwrap(); + + let d2t_bytes: Vec = [0i64, 0, 5].into_iter().flat_map(|v| v.to_le_bytes()).collect(); + let st_specs: Vec<(&str, &str, Vec, Vec)> = vec![ + ("fc.weight", "BF16", vec![4, 12], vec![0u8; 4 * 12 * 2]), + ("midlayer.input_layernorm.weight", "BF16", vec![4], vec![0u8; 4 * 2]), + ("midlayer.hidden_norm.weight", "F32", vec![4], vec![0u8; 4 * 4]), + ("midlayer.self_attn.q_proj.weight", "BF16", vec![8, 8], vec![0u8; 8 * 8 * 2]), + ("midlayer.self_attn.k_proj.weight", "BF16", vec![4, 8], vec![0u8; 4 * 8 * 2]), + ("midlayer.self_attn.v_proj.weight", "BF16", vec![4, 8], vec![0u8; 4 * 8 * 2]), + ("midlayer.self_attn.o_proj.weight", "BF16", vec![4, 8], vec![0u8; 4 * 8 * 2]), + ("midlayer.post_attention_layernorm.weight", "F32", vec![4], vec![0u8; 4 * 4]), + ("midlayer.mlp.gate_proj.weight", "BF16", vec![6, 4], vec![0u8; 6 * 4 * 2]), + ("midlayer.mlp.up_proj.weight", "BF16", vec![6, 4], vec![0u8; 6 * 4 * 2]), + ("midlayer.mlp.down_proj.weight", "BF16", vec![4, 6], vec![0u8; 4 * 6 * 2]), + ("norm.weight", "F32", vec![4], vec![0u8; 4 * 4]), + ("lm_head.weight", "BF16", vec![3, 4], vec![0u8; 3 * 4 * 2]), + ("d2t", "I64", vec![3], d2t_bytes), + ]; + let mut st_hdr = serde_json::Map::new(); + let mut st_body = Vec::new(); + for (name, dt, shape, data) in st_specs { + let start = st_body.len(); + st_body.extend_from_slice(&data); + let end = st_body.len(); + st_hdr.insert( + name.to_string(), + serde_json::json!({"dtype": dt, "shape": shape, "data_offsets": [start, end]}), + ); + } + let hdr_bytes = serde_json::to_vec(&st_hdr).unwrap(); + let mut st_file = Vec::with_capacity(8 + hdr_bytes.len() + st_body.len()); + st_file.extend_from_slice(&(hdr_bytes.len() as u64).to_le_bytes()); + st_file.extend_from_slice(&hdr_bytes); + st_file.extend_from_slice(&st_body); + let st_path = dir.join("model.safetensors"); + std::fs::write(&st_path, &st_file).unwrap(); + + let aux_gguf = dir.join("model.gguf"); + crate::gguf::convert_safetensors_eagle3_to_gguf(&st_path, &cfg_path, &aux_gguf, 48).unwrap(); + + let gt = Gguf::open(&pt).unwrap(); + let ga = Gguf::open(&aux_gguf).unwrap(); + let m = Model::load_with_aux(>, Some(&ga), &[]).unwrap(); + std::fs::remove_dir_all(&dir).ok(); + let ec = m.cfg.eagle3.as_ref().expect("eagle3 config loaded from converted safetensors"); + assert_eq!(m.cfg.n_nextn, 1); + assert_eq!((ec.n_head, ec.head_dim, ec.n_head_kv, ec.n_ff, ec.draft_vocab), (2, 4, 1, 6, 3)); + assert_eq!(ec.d2t, vec![0, 1, 7]); + } } diff --git a/src/ngram.rs b/src/ngram.rs new file mode 100644 index 00000000..6d0707a2 --- /dev/null +++ b/src/ngram.rs @@ -0,0 +1,1247 @@ +//! Context-lookup speculative decoding (port of `llama.cpp` `ngram-simple`, `ngram-map-k`, `ngram-map-k4v`, +//! `ngram-mod`, and `ngram-cache` with static/dynamic lookup caches). +//! +//! All drafting runs on the CPU in microseconds with zero VRAM overhead, producing draft token sequences that +//! `GpuModel::ngram_spec_step` verifies in a single batched forward pass on the GPU. + +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufWriter, Read, Write}; +use std::path::{Path, PathBuf}; + +pub const LLAMA_TOKEN_NULL: i32 = -1; +pub const LLAMA_NGRAM_MIN: usize = 1; +pub const LLAMA_NGRAM_MAX: usize = 4; +pub const LLAMA_NGRAM_STATIC: usize = 2; + +const LCG_FACTOR: u32 = 2654435761; +const COMMON_NGRAM_MAX_VALUES: usize = 4; +const COMMON_NGRAM_HASH_MAP_SIZE: usize = 262_144; +const COMMON_NGRAM_MAX_VALUE_COUNT: u16 = 16_380; + +const DRAFT_MIN_SAMPLE_SIZE_LAX: [i32; LLAMA_NGRAM_MAX] = [2, 2, 1, 1]; +const DRAFT_MIN_PERCENT_LAX: [i32; LLAMA_NGRAM_MAX] = [66, 50, 50, 50]; +const DRAFT_MIN_SAMPLE_SIZE_STRICT: [i32; LLAMA_NGRAM_MAX] = [4, 3, 2, 2]; +const DRAFT_MIN_PERCENT_STRICT: [i32; LLAMA_NGRAM_MAX] = [75, 66, 66, 66]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SpecType { + None, + NgramSimple, + NgramMapK, + NgramMapK4v, + NgramMod, + NgramCache, +} + +impl SpecType { + pub fn parse_list(s: &str) -> anyhow::Result> { + let mut out = Vec::new(); + for part in s.split(',').map(str::trim).filter(|p| !p.is_empty()) { + let ty = match part { + "none" => return Ok(vec![]), + "ngram-simple" => SpecType::NgramSimple, + "ngram-map-k" => SpecType::NgramMapK, + "ngram-map-k4v" => SpecType::NgramMapK4v, + "ngram-mod" => SpecType::NgramMod, + "ngram-cache" => SpecType::NgramCache, + other => anyhow::bail!( + "unknown --spec-type {other:?} (expected: none, ngram-simple, ngram-map-k, ngram-map-k4v, ngram-mod, ngram-cache)" + ), + }; + if !out.contains(&ty) { + out.push(ty); + } + } + Ok(out) + } + + pub fn as_str(self) -> &'static str { + match self { + SpecType::None => "none", + SpecType::NgramSimple => "ngram-simple", + SpecType::NgramMapK => "ngram-map-k", + SpecType::NgramMapK4v => "ngram-map-k4v", + SpecType::NgramMod => "ngram-mod", + SpecType::NgramCache => "ngram-cache", + } + } +} + +#[derive(Debug, Clone)] +pub struct NgramConfig { + pub types: Vec, + pub draft_max: usize, + pub size_n: usize, + pub size_m: usize, + pub min_hits: u16, + pub mod_n_match: usize, + pub mod_n_min: usize, + pub lookup_cache_static: Option, + pub lookup_cache_dynamic: Option, + /// Target model vocabulary size (0 = unchecked). Tokens >= `n_vocab` are filtered from caches and drafts. + pub n_vocab: usize, +} + +impl Default for NgramConfig { + fn default() -> Self { + Self { + types: Vec::new(), + draft_max: 16, + size_n: 12, + size_m: 48, + min_hits: 1, + mod_n_match: 24, + mod_n_min: 4, + lookup_cache_static: None, + lookup_cache_dynamic: None, + n_vocab: 0, + } + } +} + +impl NgramConfig { + pub fn is_enabled(&self) -> bool { + !self.types.is_empty() && self.draft_max > 0 + } +} + +// --------------------------------------------------------------------------- +// 1. N-gram Simple (`common_ngram_simple_draft` in `llama.cpp`) +// --------------------------------------------------------------------------- + +pub fn ngram_simple_draft(tokens: &[u32], sampled: u32, size_n: usize, size_m: usize, draft_max: usize) -> Vec { + ngram_simple_draft_eos(tokens, sampled, size_n, size_m, draft_max, &[]) +} + +pub fn ngram_simple_draft_eos( + tokens: &[u32], + sampled: u32, + size_n: usize, + size_m: usize, + draft_max: usize, + eos: &[u32], +) -> Vec { + let cur_len = tokens.len(); + let n = size_n.max(1); + let m = size_m.max(1); + if cur_len <= n + 1 || draft_max == 0 { + return Vec::new(); + } + let pat_prefix = &tokens[cur_len + 1 - n..cur_len]; + let first_tok = if n == 1 { sampled } else { pat_prefix[0] }; + + // Search backwards from `cur_len - n - 1` down to `1` (matching llama.cpp's `j > 0` loop). + // If a candidate window crosses an EOS/turn token (such as `<|im_end|>` at the end of a user turn), + // remember it as a fallback and keep searching for an earlier clean match. + let Some(start_j) = cur_len.checked_sub(n + 1) else { + return Vec::new(); + }; + let mut match_pos = 0usize; + let mut fallback_pos = 0usize; + for j in (1..=start_j).rev() { + if tokens[j + n - 1] == sampled && tokens[j] == first_tok && tokens[j..j + n - 1] == *pat_prefix { + let avail = cur_len.saturating_sub(j + n); + let take = m.min(avail).min(draft_max); + if !eos.is_empty() && tokens[j + n..j + n + take].iter().any(|t| eos.contains(t)) { + if fallback_pos == 0 { + fallback_pos = j; + } + continue; + } + match_pos = j; + break; + } + } + if match_pos == 0 { + match_pos = fallback_pos; + } + if match_pos == 0 { + return Vec::new(); + } + let avail = cur_len.saturating_sub(match_pos + n); + let copy_max = m.min(avail); + if copy_max < n.min(draft_max) { + return Vec::new(); + } + let take = copy_max.min(draft_max); + tokens[match_pos + n..match_pos + n + take].to_vec() +} + +// --------------------------------------------------------------------------- +// 2. N-gram Map (`common_ngram_map` in `llama.cpp`: `ngram-map-k` & `ngram-map-k4v`) +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Copy, Default)] +struct NgramMapValue { + value_idx: usize, + value_num: u16, + n_accepted: i16, +} + +#[derive(Debug, Clone)] +struct NgramMapKey { + key_idx: usize, + stat_idx: usize, + key_num: u16, + values: [NgramMapValue; COMMON_NGRAM_MAX_VALUES], +} + +#[derive(Debug, Clone)] +pub struct NgramMap { + size_key: usize, + size_value: usize, + key_only: bool, + min_hits: u16, + keys: Vec, + size_last_begin: usize, + last_draft_created: bool, + last_draft_key_idx: usize, + last_draft_value_idx: usize, + idx_last_check: usize, + key_map: Vec, + key_map_last_idx: usize, +} + +fn ngram_map_hash(tokens: &[u32]) -> u32 { + let mut hash = 0u32; + for &t in tokens { + hash = hash.wrapping_mul(LCG_FACTOR).wrapping_add(t); + } + hash +} + +impl NgramMap { + pub fn new(size_key: usize, size_value: usize, key_only: bool, min_hits: u16) -> Self { + Self { + size_key: size_key.max(1), + size_value: size_value.max(1), + key_only, + min_hits: min_hits.max(1), + keys: Vec::new(), + size_last_begin: 0, + last_draft_created: false, + last_draft_key_idx: 0, + last_draft_value_idx: 0, + idx_last_check: 0, + key_map: vec![0; COMMON_NGRAM_HASH_MAP_SIZE], + key_map_last_idx: 0, + } + } + + pub fn begin(&mut self, tokens: &[u32]) { + let size_begin = tokens.len(); + let mut idx_begin_cleanup = self.size_last_begin; + if idx_begin_cleanup > size_begin { + idx_begin_cleanup = size_begin.saturating_sub(self.size_key + self.size_value); + } + if !self.key_map.is_empty() && size_begin < self.idx_last_check { + for entry in &mut self.key_map { + if *entry != 0 && (*entry as usize) >= idx_begin_cleanup { + *entry = 0; + } + } + self.key_map_last_idx = idx_begin_cleanup.saturating_sub(1); + } + if size_begin < self.idx_last_check && !self.keys.is_empty() { + let key_only = self.key_only; + self.keys.retain_mut(|key| { + if key.key_idx >= idx_begin_cleanup { + return false; + } + if key_only { + return true; + } + for j in (0..COMMON_NGRAM_MAX_VALUES).rev() { + if key.values[j].value_idx != 0 && key.values[j].value_idx >= idx_begin_cleanup { + for k in j..COMMON_NGRAM_MAX_VALUES - 1 { + key.values[k] = key.values[k + 1]; + } + key.values[COMMON_NGRAM_MAX_VALUES - 1] = NgramMapValue::default(); + } + } + key.values[0].value_idx != 0 + }); + } + self.idx_last_check = size_begin; + self.size_last_begin = size_begin; + } + + pub fn draft(&mut self, inp: &[u32], sampled: u32, draft_max: usize) -> Vec { + self.last_draft_created = false; + self.last_draft_key_idx = 0; + self.last_draft_value_idx = 0; + + let cur_len = inp.len(); + let n = self.size_key; + let m = self.size_value; + if cur_len < 2 * n + m || draft_max == 0 { + return Vec::new(); + } + if self.idx_last_check > cur_len || self.size_last_begin > cur_len { + self.begin(inp); + } + self.idx_last_check = cur_len; + + let mut key_tokens = Vec::with_capacity(n); + key_tokens.extend_from_slice(&inp[cur_len + 1 - n..cur_len]); + key_tokens.push(sampled); + + let mut match_pos = 0usize; + if !self.key_map.is_empty() { + let idx_hash = (ngram_map_hash(&key_tokens) as usize) % self.key_map.len(); + let idx_key = self.key_map[idx_hash] as usize; + if idx_key != 0 && idx_key + n + m + 1 < cur_len && inp[idx_key..idx_key + n] == key_tokens[..] { + match_pos = idx_key; + } + } + if match_pos == 0 && self.size_last_begin > n + m + 1 { + let hi = self.size_last_begin - n - m - 1; + for j in (self.key_map_last_idx + 1..=hi).rev() { + if inp[j..j + n] == key_tokens[..] { + match_pos = j; + break; + } + } + } + if match_pos == 0 && cur_len > n + m + 1 { + let hi = cur_len - n - m - 1; + let lo = self.size_last_begin.max(self.key_map_last_idx); + if hi > lo { + for j in (lo + 1..=hi).rev() { + if inp[j..j + n] == key_tokens[..] { + match_pos = j; + break; + } + } + } + } + + if !self.key_map.is_empty() { + if self.size_last_begin > n + m + 1 { + let hi = self.size_last_begin - n - m - 1; + for j in (self.key_map_last_idx + 1..=hi).rev() { + let idx_hash = (ngram_map_hash(&inp[j..j + n]) as usize) % self.key_map.len(); + if self.key_map[idx_hash] == 0 { + self.key_map[idx_hash] = j as u32; + } + } + } + if cur_len > n + m + 1 { + let hi = cur_len - n - m - 1; + let lo = self.size_last_begin.max(self.key_map_last_idx); + if hi > lo { + for j in (lo + 1..=hi).rev() { + let idx_hash = (ngram_map_hash(&inp[j..j + n]) as usize) % self.key_map.len(); + if self.key_map[idx_hash] == 0 { + self.key_map[idx_hash] = j as u32; + } + } + } + self.key_map_last_idx = self.key_map_last_idx.max(hi); + } + } + + if match_pos == 0 { + return Vec::new(); + } + + let mut key_offset = self.keys.len(); + for (i, k) in self.keys.iter().enumerate() { + if k.key_idx + n <= inp.len() && inp[k.key_idx..k.key_idx + n] == key_tokens[..] { + key_offset = i; + break; + } + } + if key_offset == self.keys.len() { + let mut values = [NgramMapValue::default(); COMMON_NGRAM_MAX_VALUES]; + for v in &mut values { + v.n_accepted = m as i16; + } + self.keys.push(NgramMapKey { + key_idx: match_pos, + stat_idx: 0, + key_num: 0, + values, + }); + } + + let curr_key = &mut self.keys[key_offset]; + curr_key.key_num = (curr_key.key_num + 1).min(COMMON_NGRAM_MAX_VALUE_COUNT); + + if self.key_only { + let n_draft = m + .min(curr_key.values[0].n_accepted.max(0) as usize) + .min(draft_max) + .min(cur_len.saturating_sub(match_pos + n)); + self.last_draft_created = n_draft > 0; + self.last_draft_key_idx = key_offset; + self.last_draft_value_idx = 0; + return inp[match_pos + n..match_pos + n + n_draft].to_vec(); + } + + if curr_key.key_num < self.min_hits { + return Vec::new(); + } + + for i in curr_key.stat_idx..=match_pos { + if i + n + m > cur_len || inp[i..i + n] != key_tokens[..] { + continue; + } + let idx_begin_value_key = i + n; + let mut idx_value = None; + for v in 0..COMMON_NGRAM_MAX_VALUES { + let idx_v = curr_key.values[v].value_idx; + if idx_v == 0 { + curr_key.values[v].value_idx = idx_begin_value_key; + curr_key.values[v].value_num = 0; + curr_key.values[v].n_accepted = m as i16; + idx_value = Some(v); + break; + } + if idx_v + m <= cur_len && inp[idx_begin_value_key..idx_begin_value_key + m] == inp[idx_v..idx_v + m] { + idx_value = Some(v); + break; + } + } + if let Some(v) = idx_value { + curr_key.values[v].value_num = (curr_key.values[v].value_num + 1).min(COMMON_NGRAM_MAX_VALUE_COUNT); + } + } + curr_key.stat_idx = match_pos; + + let mut max_occur = 0u16; + let mut slot_max = 0usize; + for v in 0..COMMON_NGRAM_MAX_VALUES { + if curr_key.values[v].value_num > max_occur { + max_occur = curr_key.values[v].value_num; + slot_max = v; + } + } + let sum_occur: u32 = (0..COMMON_NGRAM_MAX_VALUES) + .filter(|&v| v != slot_max) + .map(|v| curr_key.values[v].value_num as u32) + .sum(); + if sum_occur > 0 && (max_occur as u32) < 2 * sum_occur { + return Vec::new(); + } + + let n_draft = m + .min(curr_key.values[slot_max].n_accepted.max(0) as usize) + .min(draft_max) + .min(cur_len.saturating_sub(match_pos + n)); + self.last_draft_created = n_draft > 0; + self.last_draft_key_idx = key_offset; + self.last_draft_value_idx = slot_max; + inp[match_pos + n..match_pos + n + n_draft].to_vec() + } + + pub fn accept(&mut self, n_accepted: usize) { + if !self.last_draft_created { + return; + } + if let Some(key) = self.keys.get_mut(self.last_draft_key_idx) { + if let Some(val) = key.values.get_mut(self.last_draft_value_idx) { + val.n_accepted = n_accepted as i16; + } + } + } +} + +// --------------------------------------------------------------------------- +// 3. N-gram Mod (`common_ngram_mod` in `llama.cpp`) +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +pub struct NgramMod { + n: usize, + n_min: usize, + n_max: usize, + entries: Vec, + used: usize, + i_last: usize, + n_draft_last: usize, + n_low: usize, +} + +impl NgramMod { + pub const EMPTY: i32 = -1; + pub const DEFAULT_SIZE: usize = 4 * 1024 * 1024; + + pub fn new(n: usize, n_min: usize, n_max: usize) -> Self { + Self { + n: n.max(1), + n_min, + n_max, + entries: vec![Self::EMPTY; Self::DEFAULT_SIZE], + used: 0, + i_last: 0, + n_draft_last: 0, + n_low: 0, + } + } + + fn idx(&self, tokens: &[u32]) -> usize { + let mut res = 0u64; + for &t in &tokens[..self.n] { + res = res.wrapping_mul(6364136223846793005u64).wrapping_add(t as u64); + } + (res as usize) % self.entries.len() + } + + fn add(&mut self, window: &[u32]) { + let i = self.idx(window); + if self.entries[i] == Self::EMPTY { + self.used += 1; + } + self.entries[i] = window[self.n] as i32; + } + + fn get(&self, window: &[u32]) -> i32 { + self.entries[self.idx(window)] + } + + pub fn reset_table(&mut self) { + self.entries.fill(Self::EMPTY); + self.used = 0; + } + + pub fn begin(&mut self, prompt: &[u32]) { + self.i_last = 0; + self.n_draft_last = 0; + let n = self.n; + if prompt.len() <= n { + return; + } + for i in 0..prompt.len() - n { + self.add(&prompt[i..=i + n]); + } + self.i_last = prompt.len() - n; + if (self.used as f64) / (self.entries.len() as f64) > 0.25 { + self.reset_table(); + } + } + + pub fn draft(&mut self, prompt: &[u32], sampled: u32, draft_max: usize) -> Vec { + self.n_draft_last = 0; + let cur_len = prompt.len(); + let n = self.n; + let max_len = self.n_max.min(draft_max); + if cur_len < n || max_len == 0 { + return Vec::new(); + } + if self.i_last + 32 < cur_len && cur_len > n { + for i in self.i_last..cur_len - n { + self.add(&prompt[i..=i + n]); + } + self.i_last = cur_len - n; + } + let mut buf = Vec::with_capacity(n + max_len); + buf.extend_from_slice(&prompt[cur_len + 1 - n..cur_len]); + buf.push(sampled); + for i in 0..max_len { + let tok = self.get(&buf[i..i + n]); + if tok == Self::EMPTY { + if i < self.n_min.min(max_len) { + return Vec::new(); + } + break; + } + buf.push(tok as u32); + } + let out = buf[n..].to_vec(); + self.n_draft_last = out.len(); + out + } + + pub fn accept(&mut self, n_accepted: usize) { + if self.n_draft_last > 0 { + let f_acc = (n_accepted as f64) / (self.n_draft_last as f64); + if f_acc < 0.25 { + self.n_low += 1; + if self.n_low >= 5 { + self.reset_table(); + self.n_low = 0; + self.i_last = 0; + } + } else { + self.n_low = 0; + } + } + } +} + +// --------------------------------------------------------------------------- +// 4. N-gram Cache (`common_ngram_cache` in `llama.cpp`) +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct NgramKey(pub [i32; LLAMA_NGRAM_MAX]); + +impl Default for NgramKey { + fn default() -> Self { + Self([LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX]) + } +} + +impl NgramKey { + pub fn from_slice(tokens: &[u32]) -> Self { + let mut k = [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX]; + for (i, &t) in tokens.iter().take(LLAMA_NGRAM_MAX).enumerate() { + k[i] = t as i32; + } + Self(k) + } +} + +impl std::hash::Hash for NgramKey { + fn hash(&self, state: &mut H) { + let mut h = (self.0[0] as u64).wrapping_mul(11400714819323198485u64); + for i in 1..LLAMA_NGRAM_MAX { + h ^= (self.0[i] as u64).wrapping_mul(11400714819323198485u64); + } + state.write_u64(h); + } +} + +pub type NgramCachePart = HashMap; +pub type NgramCacheMap = HashMap; + +pub fn ngram_cache_update(cache: &mut NgramCacheMap, ngram_min: usize, ngram_max: usize, inp: &[u32], nnew: usize) { + let inp_size = inp.len(); + if inp_size < 2 || nnew == 0 { + return; + } + let nmin = ngram_min.clamp(1, LLAMA_NGRAM_MAX); + let nmax = ngram_max.clamp(nmin, LLAMA_NGRAM_MAX); + for ngram_size in nmin..=nmax { + let i_start = inp_size.saturating_sub(nnew).max(ngram_size); + for i in i_start..inp_size { + let ngram = NgramKey::from_slice(&inp[i - ngram_size..i]); + let token = inp[i]; + let part = cache.entry(ngram).or_default(); + *part.entry(token).or_insert(0) += 1; + } + } +} + +fn get_seq_token(inp: &[u32], draft: &[u32], i: usize) -> u32 { + if i < inp.len() { + inp[i] + } else { + draft[1 + i - inp.len()] + } +} + +fn try_draft_static(nc_static: &NgramCacheMap, ngram_static: NgramKey) -> Option { + let part = nc_static.get(&ngram_static)?; + let mut max_count = 0i32; + let mut sum_count = 0i32; + let mut max_token = None; + for (&tok, &cnt) in part { + if cnt > max_count { + max_count = cnt; + max_token = Some(tok); + } + sum_count += cnt; + } + if sum_count < DRAFT_MIN_SAMPLE_SIZE_LAX[LLAMA_NGRAM_STATIC - 1] { + return None; + } + if 100 * max_count < DRAFT_MIN_PERCENT_LAX[LLAMA_NGRAM_STATIC - 1] * sum_count { + return None; + } + max_token +} + +fn try_draft_primary( + nc_primary: &NgramCacheMap, + ngrams_primary: &[NgramKey], + part_static: Option<&NgramCachePart>, + min_sample_size: &[i32; LLAMA_NGRAM_MAX], + min_percent: &[i32; LLAMA_NGRAM_MAX], +) -> Option { + for &ngram in ngrams_primary.iter().rev() { + let Some(part_primary) = nc_primary.get(&ngram) else { + continue; + }; + let idx = ngram + .0 + .iter() + .take_while(|&&t| t != LLAMA_TOKEN_NULL) + .count() + .saturating_sub(1) + .min(LLAMA_NGRAM_MAX - 1); + let mut max_count_primary = 0i64; + let mut max_count_static = 0i64; + let mut sum_count_primary = 0i32; + let mut max_token = None; + + for (&token, &count_primary) in part_primary { + let count_static = part_static + .and_then(|ps| ps.get(&token)) + .map_or(1i64, |&c| 100 * (c as i64)); + let cp = count_primary as i64; + if cp * count_static > max_count_primary * max_count_static { + max_token = Some(token); + max_count_primary = cp; + max_count_static = count_static; + } + sum_count_primary += count_primary; + } + if sum_count_primary < min_sample_size[idx] { + continue; + } + if 100 * (max_count_primary as i32) < min_percent[idx] * sum_count_primary { + continue; + } + if max_token.is_some() { + return max_token; + } + } + None +} + +pub fn ngram_cache_draft( + inp: &[u32], + sampled: u32, + n_draft: usize, + ngram_min: usize, + ngram_max: usize, + nc_context: &NgramCacheMap, + nc_dynamic: &NgramCacheMap, + nc_static: &NgramCacheMap, +) -> Vec { + if inp.len() + 1 < LLAMA_NGRAM_STATIC || n_draft == 0 { + return Vec::new(); + } + let mut full_inp = Vec::with_capacity(inp.len() + 1); + full_inp.extend_from_slice(inp); + full_inp.push(sampled); + let inp_size = full_inp.len(); + + let mut draft = vec![sampled]; + let nmin = ngram_min.clamp(1, LLAMA_NGRAM_MAX); + let nmax = ngram_max.clamp(nmin, LLAMA_NGRAM_MAX); + + while draft.len() - 1 < n_draft { + let total_len = inp_size + draft.len() - 1; + let ngram_static = if total_len >= LLAMA_NGRAM_STATIC { + let ngram_start_static = total_len - LLAMA_NGRAM_STATIC; + let mut k = NgramKey::default(); + for j in 0..LLAMA_NGRAM_STATIC { + k.0[j] = get_seq_token(&full_inp, &draft, ngram_start_static + j) as i32; + } + Some(k) + } else { + None + }; + let part_static = ngram_static.as_ref().and_then(|k| nc_static.get(k)); + + let mut ngrams_cd = Vec::with_capacity(nmax - nmin + 1); + for ngram_size_cd in nmin..=nmax { + if total_len < ngram_size_cd { + break; + } + let start_cd = total_len - ngram_size_cd; + let mut key = NgramKey::default(); + for j in 0..ngram_size_cd { + key.0[j] = get_seq_token(&full_inp, &draft, start_cd + j) as i32; + } + ngrams_cd.push(key); + } + + let drafted = try_draft_primary( + nc_context, + &ngrams_cd, + part_static, + &DRAFT_MIN_SAMPLE_SIZE_LAX, + &DRAFT_MIN_PERCENT_LAX, + ) + .or_else(|| { + try_draft_primary( + nc_dynamic, + &ngrams_cd, + part_static, + &DRAFT_MIN_SAMPLE_SIZE_STRICT, + &DRAFT_MIN_PERCENT_STRICT, + ) + }) + .or_else(|| ngram_static.and_then(|k| try_draft_static(nc_static, k))); + + match drafted { + Some(tok) => draft.push(tok), + None => break, + } + } + draft.into_iter().skip(1).collect() +} + +pub fn ngram_cache_merge(target: &mut NgramCacheMap, add: &NgramCacheMap) { + for (&ngram, part) in add { + let dst = target.entry(ngram).or_default(); + for (&tok, &cnt) in part { + *dst.entry(tok).or_insert(0) += cnt; + } + } +} + +/// Save `cache` in `llama.cpp`'s `common_ngram_cache_save` binary format. +pub fn ngram_cache_save(cache: &NgramCacheMap, path: &Path) -> anyhow::Result<()> { + if let Some(parent) = path.parent() && !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent)?; + } + let mut out = BufWriter::new(File::create(path)?); + for (ngram, part) in cache { + if part.is_empty() { + continue; + } + for &t in &ngram.0 { + out.write_all(&t.to_le_bytes())?; + } + let ntokens = part.len() as i32; + out.write_all(&ntokens.to_le_bytes())?; + for (&tok, &cnt) in part { + out.write_all(&(tok as i32).to_le_bytes())?; + out.write_all(&cnt.to_le_bytes())?; + } + } + out.flush()?; + Ok(()) +} + +/// Load a lookup cache in `llama.cpp`'s binary format (`common_ngram_cache_load`). +pub fn ngram_cache_load_binary(bytes: &[u8]) -> anyhow::Result { + let mut r = std::io::Cursor::new(bytes); + let mut cache = NgramCacheMap::new(); + let mut buf4 = [0u8; 4]; + while (r.position() as usize) < bytes.len() { + let mut key = [LLAMA_TOKEN_NULL; LLAMA_NGRAM_MAX]; + for k in &mut key { + r.read_exact(&mut buf4)?; + *k = i32::from_le_bytes(buf4); + } + r.read_exact(&mut buf4)?; + let ntokens = i32::from_le_bytes(buf4); + anyhow::ensure!(ntokens > 0 && ntokens <= 1_000_000, "invalid ntokens {ntokens} in binary ngram cache"); + let mut part = NgramCachePart::with_capacity(ntokens as usize); + for _ in 0..ntokens { + r.read_exact(&mut buf4)?; + let tok = i32::from_le_bytes(buf4); + r.read_exact(&mut buf4)?; + let cnt = i32::from_le_bytes(buf4); + anyhow::ensure!(tok >= 0 && cnt > 0, "invalid token/count ({tok}, {cnt}) in binary ngram cache"); + part.insert(tok as u32, cnt); + } + cache.insert(NgramKey(key), part); + } + Ok(cache) +} + +/// Load a lookup cache from `path`: tries `llama.cpp` binary format first, or if the file is valid UTF-8 text and a +/// tokenizer is supplied, tokenizes the text corpus and extracts 1..=4 n-grams into the cache. +pub fn ngram_cache_load(path: &Path, tok: Option<&crate::tokenizer::Tokenizer>) -> anyhow::Result { + let bytes = std::fs::read(path)?; + if bytes.is_empty() { + return Ok(NgramCacheMap::new()); + } + if let Ok(cache) = ngram_cache_load_binary(&bytes) && !cache.is_empty() { + return Ok(cache); + } + if let (Ok(text), Some(tok)) = (std::str::from_utf8(&bytes), tok) { + let ids = tok.encode(text)?; + let mut cache = NgramCacheMap::new(); + ngram_cache_update(&mut cache, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, &ids, ids.len()); + return Ok(cache); + } + ngram_cache_load_binary(&bytes) +} + +/// Drop cache entries whose key or candidate tokens exceed `n_vocab` (e.g. a cache saved by a larger-vocab model). +pub fn ngram_cache_filter_vocab(cache: &mut NgramCacheMap, n_vocab: usize) -> usize { + if n_vocab == 0 { + return 0; + } + let mut pruned = 0usize; + cache.retain(|key, part| { + let key_ok = key.0.iter().all(|&t| t == LLAMA_TOKEN_NULL || (t >= 0 && (t as usize) < n_vocab)); + if !key_ok { + pruned += part.len().max(1); + return false; + } + let before = part.len(); + part.retain(|&tok, _| (tok as usize) < n_vocab); + pruned += before - part.len(); + !part.is_empty() + }); + pruned +} + +// --------------------------------------------------------------------------- +// Unified Drafter (`NgramDrafter`) owned by `GpuModel` +// --------------------------------------------------------------------------- + +pub struct NgramDrafter { + pub cfg: NgramConfig, + pub history: Vec, + eos: Vec, + map_k: Option, + map_k4v: Option, + mod_state: Option, + cache_size: usize, + nc_context: NgramCacheMap, + nc_dynamic: NgramCacheMap, + nc_static: NgramCacheMap, + last_used: Option, + last_draft_len: usize, + short_cooldown: usize, +} + +impl NgramDrafter { + pub fn new(cfg: NgramConfig, tok: Option<&crate::tokenizer::Tokenizer>) -> anyhow::Result { + let eos = tok.map(|t| t.eos_ids()).unwrap_or_default(); + let mut nc_static = match &cfg.lookup_cache_static { + Some(p) => ngram_cache_load(p, tok)?, + None => NgramCacheMap::new(), + }; + let mut nc_dynamic = match &cfg.lookup_cache_dynamic { + Some(p) if p.exists() => ngram_cache_load(p, tok).unwrap_or_default(), + _ => NgramCacheMap::new(), + }; + let pruned = ngram_cache_filter_vocab(&mut nc_static, cfg.n_vocab) + + ngram_cache_filter_vocab(&mut nc_dynamic, cfg.n_vocab); + if pruned > 0 { + eprintln!( + "note: pruned {pruned} out-of-vocabulary entries (n_vocab={}) from ngram lookup cache", + cfg.n_vocab + ); + } + let map_k = cfg + .types + .contains(&SpecType::NgramMapK) + .then(|| NgramMap::new(cfg.size_n, cfg.size_m, true, cfg.min_hits)); + let map_k4v = cfg + .types + .contains(&SpecType::NgramMapK4v) + .then(|| NgramMap::new(cfg.size_n, cfg.size_m, false, cfg.min_hits)); + let mod_state = cfg + .types + .contains(&SpecType::NgramMod) + .then(|| NgramMod::new(cfg.mod_n_match, cfg.mod_n_min, cfg.draft_max)); + + Ok(Self { + cfg, + history: Vec::new(), + eos, + map_k, + map_k4v, + mod_state, + cache_size: 0, + nc_context: NgramCacheMap::new(), + nc_dynamic, + nc_static, + last_used: None, + last_draft_len: 0, + short_cooldown: 0, + }) + } + + fn flush_dynamic_context(&mut self) { + if self.cfg.lookup_cache_dynamic.is_some() { + if self.cache_size < self.history.len() { + let nnew = self.history.len() - self.cache_size; + ngram_cache_update(&mut self.nc_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, &self.history, nnew); + self.cache_size = self.history.len(); + } + if !self.nc_context.is_empty() { + ngram_cache_merge(&mut self.nc_dynamic, &self.nc_context); + } + } + } + + pub fn reset(&mut self) { + self.flush_dynamic_context(); + self.history.clear(); + self.nc_context.clear(); + self.cache_size = 0; + self.last_used = None; + self.last_draft_len = 0; + self.short_cooldown = 0; + if let Some(m) = &mut self.map_k { + m.begin(&[]); + } + if let Some(m) = &mut self.map_k4v { + m.begin(&[]); + } + } + + /// Initialize or replace history with `prompt` (called by `prefill` and `prefill_cached`). + pub fn begin(&mut self, prompt: &[u32]) { + if !prompt.starts_with(&self.history) { + self.flush_dynamic_context(); + self.nc_context.clear(); + self.cache_size = 0; + } + self.history.clear(); + self.history.extend_from_slice(prompt); + self.last_used = None; + self.last_draft_len = 0; + self.short_cooldown = 0; + if let Some(m) = &mut self.map_k { + m.begin(&self.history); + } + if let Some(m) = &mut self.map_k4v { + m.begin(&self.history); + } + if let Some(m) = &mut self.mod_state { + m.begin(&self.history); + } + } + + /// Propose up to `max_draft` tokens following `sampled` (which is not yet in `self.history`). + pub fn draft(&mut self, sampled: u32, max_draft: usize) -> Vec { + let cap = max_draft.min(self.cfg.draft_max); + self.last_used = None; + self.last_draft_len = 0; + if cap == 0 { + return Vec::new(); + } + let mut out = Vec::new(); + for &ty in &self.cfg.types { + let res = match ty { + SpecType::None => Vec::new(), + SpecType::NgramSimple => { + let mut r = ngram_simple_draft_eos(&self.history, sampled, self.cfg.size_n, self.cfg.size_m, cap, &self.eos); + if r.is_empty() && self.short_cooldown == 0 { + if self.cfg.size_n > 6 { + r = ngram_simple_draft_eos(&self.history, sampled, 6, self.cfg.size_m, cap, &self.eos); + } + if r.is_empty() && self.cfg.size_n > 4 { + r = ngram_simple_draft_eos(&self.history, sampled, 4, self.cfg.size_m, cap.min(8), &self.eos); + } + } + r + } + SpecType::NgramMapK => self.map_k.as_mut().map_or_else(Vec::new, |m| m.draft(&self.history, sampled, cap)), + SpecType::NgramMapK4v => self.map_k4v.as_mut().map_or_else(Vec::new, |m| m.draft(&self.history, sampled, cap)), + SpecType::NgramMod => self.mod_state.as_mut().map_or_else(Vec::new, |m| m.draft(&self.history, sampled, cap)), + SpecType::NgramCache => { + let mut inp = self.history.clone(); + inp.push(sampled); + if self.cache_size < inp.len() { + let nnew = inp.len() - self.cache_size; + ngram_cache_update(&mut self.nc_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, &inp, nnew); + self.cache_size = inp.len(); + } + let exact = if self.short_cooldown == 0 { + ngram_simple_draft_eos(&self.history, sampled, 6, self.cfg.size_m, cap, &self.eos) + } else { + Vec::new() + }; + if !exact.is_empty() { + exact + } else if self.short_cooldown == 0 || !self.nc_static.is_empty() || !self.nc_dynamic.is_empty() { + let cache_cap = if self.nc_static.is_empty() && self.nc_dynamic.is_empty() { + cap.min(6) + } else { + cap + }; + ngram_cache_draft( + &self.history, + sampled, + cache_cap, + 2, + LLAMA_NGRAM_MAX, + &self.nc_context, + &self.nc_dynamic, + &self.nc_static, + ) + } else { + Vec::new() + } + } + }; + if !res.is_empty() { + self.last_used = Some(ty); + out = res; + break; + } + } + if out.is_empty() + && !self.cfg.types.contains(&SpecType::NgramCache) + && (!self.nc_static.is_empty() || !self.nc_dynamic.is_empty()) + { + let res = ngram_cache_draft( + &self.history, + sampled, + cap, + 3, + LLAMA_NGRAM_MAX, + &NgramCacheMap::new(), + &self.nc_dynamic, + &self.nc_static, + ); + if !res.is_empty() { + self.last_used = Some(SpecType::NgramCache); + out = res; + } + } + if let Some(pos) = out.iter().position(|t| self.eos.contains(t)) { + out.truncate(pos + 1); + } + if self.cfg.n_vocab > 0 { + if let Some(pos) = out.iter().position(|&t| (t as usize) >= self.cfg.n_vocab) { + out.truncate(pos); + } + } + self.last_draft_len = out.len(); + out + } + + /// Record the tokens accepted in this step (`[sampled, accepted_drafts..]`) and notify the active drafter. + pub fn accept(&mut self, accepted_tokens: &[u32]) { + let n_accepted_drafts = accepted_tokens.len().saturating_sub(1); + if self.short_cooldown > 0 { + self.short_cooldown -= 1; + } + if self.last_draft_len > 0 && n_accepted_drafts == 0 { + self.short_cooldown = 8; + } + match self.last_used { + Some(SpecType::NgramMapK) => { + if let Some(m) = &mut self.map_k { + m.accept(n_accepted_drafts); + } + } + Some(SpecType::NgramMapK4v) => { + if let Some(m) = &mut self.map_k4v { + m.accept(n_accepted_drafts); + } + } + Some(SpecType::NgramMod) => { + if let Some(m) = &mut self.mod_state { + m.accept(n_accepted_drafts); + } + } + _ => {} + } + self.history.extend_from_slice(accepted_tokens); + } + + /// Persist the dynamic lookup cache to disk if configured. + pub fn save_dynamic(&mut self) -> anyhow::Result<()> { + self.flush_dynamic_context(); + let Some(path) = &self.cfg.lookup_cache_dynamic else { + return Ok(()); + }; + if !self.nc_dynamic.is_empty() { + ngram_cache_save(&self.nc_dynamic, path)?; + } + Ok(()) + } +} + +impl Drop for NgramDrafter { + fn drop(&mut self) { + let _ = self.save_dynamic(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ngram_simple_finds_exact_match_and_respects_draft_max() { + // Build a 100-token sequence with a repeated 30-token block. + let mut tokens: Vec = (1..=80).collect(); + // Repeat tokens 20..32 (12 tokens) at the end: 20..31 in history, 31 as `sampled`. + tokens.extend(20..31); + let draft = ngram_simple_draft(&tokens, 31, 12, 48, 16); + assert_eq!(draft, (32..48).collect::>()); + } + + #[test] + fn ngram_map_k_and_k4v_draft_and_adapt_acceptance() { + let mut tokens: Vec = (1..=100).collect(); + tokens.extend(10..21); + let mut map = NgramMap::new(12, 16, true, 1); + map.begin(&tokens); + let d1 = map.draft(&tokens, 21, 16); + assert_eq!(d1, (22..38).collect::>()); + map.accept(5); + let d2 = map.draft(&tokens, 21, 16); + assert_eq!(d2.len(), 5); + } + + #[test] + fn ngram_mod_hashes_and_drafts_long_sequences() { + let mut mod_drafter = NgramMod::new(8, 2, 16); + let prompt: Vec = (100..200).collect(); + mod_drafter.begin(&prompt); + let mut ctx = prompt.clone(); + ctx.extend_from_slice(&prompt[10..17]); + let draft = mod_drafter.draft(&ctx, prompt[17], 12); + assert_eq!(draft, prompt[18..30].to_vec()); + } + + #[test] + fn ngram_cache_roundtrip_and_draft() { + let mut cache = NgramCacheMap::new(); + let seq = vec![10u32, 20, 30, 40, 50, 10, 20, 30, 40, 50, 10, 20, 30, 40, 50]; + ngram_cache_update(&mut cache, 1, 4, &seq, seq.len()); + let tmp = std::env::temp_dir().join("strata_test_ngram_cache.bin"); + ngram_cache_save(&cache, &tmp).unwrap(); + let loaded = ngram_cache_load(&tmp, None).unwrap(); + let _ = std::fs::remove_file(&tmp); + assert_eq!(cache.len(), loaded.len()); + let draft = ngram_cache_draft(&[10, 20], 30, 4, 1, 4, &loaded, &NgramCacheMap::new(), &NgramCacheMap::new()); + assert_eq!(draft, vec![40, 50, 10, 20]); + } + + #[test] + fn test_ngram_cache_draft_short_history() { + let empty = NgramCacheMap::new(); + let draft = ngram_cache_draft(&[], 42, 4, 1, 4, &empty, &empty, &empty); + assert!(draft.is_empty()); + let mut c = NgramCacheMap::new(); + ngram_cache_update(&mut c, 1, 4, &[], 1); + ngram_cache_update(&mut c, 1, 4, &[42], 1); + assert!(c.is_empty()); + } + + #[test] + fn test_ngram_cache_vocab_filter() { + let mut cache = NgramCacheMap::new(); + // Token ID 231622 comes from a 248320-vocab model; target is 151936-vocab Qwen3-30B-A3B. + let cross_seq = vec![10u32, 20, 231_622, 10, 20, 231_622, 10, 20, 231_622]; + ngram_cache_update(&mut cache, 1, 4, &cross_seq, cross_seq.len()); + let valid_seq = vec![10u32, 20, 30, 10, 20, 30, 10, 20, 30]; + ngram_cache_update(&mut cache, 1, 4, &valid_seq, valid_seq.len()); + let tmp = std::env::temp_dir().join("strata_test_ngram_vocab_filter.bin"); + ngram_cache_save(&cache, &tmp).unwrap(); + let cfg = NgramConfig { + types: vec![SpecType::NgramCache], + n_vocab: 151_936, + lookup_cache_static: Some(tmp.clone()), + ..Default::default() + }; + let mut drafter = NgramDrafter::new(cfg, None).unwrap(); + let _ = std::fs::remove_file(&tmp); + drafter.begin(&[10]); + let drafted = drafter.draft(20, 8); + assert!(drafted.iter().all(|&t| (t as usize) < 151_936)); + assert_eq!(drafted.first().copied(), Some(30)); + } + + #[test] + fn test_ngram_cache_try_draft_primary_indexes_by_ngram_len() { + let mut cache = NgramCacheMap::new(); + // Single occurrence of [10, 20, 30] -> 40. With nmin=2, a 3-token context key must consult + // DRAFT_MIN_SAMPLE_SIZE_LAX[2] (= 1), not index [1] (= 2). + let seq = vec![10u32, 20, 30, 40]; + ngram_cache_update(&mut cache, 2, 4, &seq, seq.len()); + let draft = ngram_cache_draft(&[10, 20], 30, 2, 2, 4, &cache, &NgramCacheMap::new(), &NgramCacheMap::new()); + assert_eq!(draft, vec![40]); + } +} diff --git a/src/vulkan/mod.rs b/src/vulkan/mod.rs index 010ecdbf..f7790733 100644 --- a/src/vulkan/mod.rs +++ b/src/vulkan/mod.rs @@ -141,6 +141,10 @@ struct Inner { tsem: vk::Semaphore, mem: vk::PhysicalDeviceMemoryProperties, budget: bool, + /// Integrated GPU (APU sharing unified system LPDDR memory between device and host heaps). + igpu: bool, + /// AMD RDNA3 / RDNA3.5 (gfx11xx) or iGPU whose small L2 cache hurts from next-layer MoE prefetch. + gfx11: bool, /// The device runs fixed 64-lane subgroups, split into logical 32-lane warps. wave64: bool, /// VK_KHR_cooperative_matrix with the 16 x 16 x 16 f16 / f32 shape is enabled (`gemm_f`, `attn_prefill`). @@ -164,6 +168,8 @@ struct Inner { /// `maxComputeSharedMemorySize` of the device (or STRATA_VK_SMEM_LIMIT, a test hook): no pipeline may need more /// `Workgroup` bytes than this (`Vk::pipeline`), or the driver accepts it and faults at dispatch (Xid 13 on Volta). smem_limit: u32, + /// Nanoseconds per GPU timestamp tick (`VkPhysicalDeviceLimits::timestampPeriod`). + ts_period: f32, /// STRATA_VK_STATS: print the driver's statistics (registers, ...) of each pipeline as it is created. stats: bool, /// The one pipeline layout: a 128-byte push constant range (the Vulkan minimum), no descriptor sets. @@ -177,12 +183,16 @@ struct State { spare: Vec, /// Submitted since the last wait. used: Vec, + /// Closed command buffers and graph replays queued in execution order for the next `submit`. + queued: Vec, /// The open command buffer. cur: Option, /// The command buffer with conditional rendering on (`cond_on`), ended before the buffer is. cond: Option, /// Conditional rendering failed the self-test in `Vk::new`: not used. no_cond: bool, + /// Set by `skip_rest` when a non-fast-path MoE combine (`mb != 0`) uses conditional rendering. + need_cond: bool, /// A MoE combine gave up waiting for the CPU worker since the last `stalled` (see `Backend::stalled`). stalled: bool, /// Work was submitted since the last wait. @@ -260,9 +270,14 @@ impl Drop for Inner { fn drop(&mut self) { let st = self.st.get_mut().unwrap_or_else(|e| e.into_inner()); if let Some((_, _, tot)) = &st.prof { + let scale = self.ts_period as f64; let mut v: Vec<_> = tot.iter().collect(); v.sort_by_key(|x| std::cmp::Reverse(x.1.0)); - let all: u64 = v.iter().map(|x| x.1.0).sum(); - for (n, (ns, c)) in v { eprintln!("prof {n:16} {:9.1} ms {:5.1}% {c:7} calls {:8.2} us", *ns as f64 / 1e6, *ns as f64 * 100.0 / all as f64, *ns as f64 / 1e3 / *c as f64); } + let all: f64 = v.iter().filter(|(k, _)| **k != "_barrier_gap").map(|x| x.1.0 as f64 * scale).sum(); + eprintln!("vulkan prof (timestampPeriod = {:.2} ns, active kernel sum = {:.2} ms):", self.ts_period, all / 1e6); + for (n, (ticks, c)) in v { + let ns = *ticks as f64 * scale; + eprintln!("prof {n:24} {:9.2} ms {:5.1}% {c:7} calls {:8.2} us", ns / 1e6, ns * 100.0 / all.max(1.0), ns / 1e3 / (*c).max(1) as f64); + } } unsafe { let _ = self.dev.device_wait_idle(); @@ -369,7 +384,13 @@ impl Vk { && std::env::var_os("STRATA_VK_NO_COOP").is_none(); // STRATA_VK_NO_COOP: run the kernels for devices without them let coop_i8 = coop && shape(32, vk::ComponentTypeKHR::SINT8, vk::ComponentTypeKHR::SINT32); // The tiled fallback needs a 1024-thread workgroup and 32 KB of shared memory (see attn_prefill_tile.comp). - let limits = unsafe { instance.get_physical_device_properties(pd) }.limits; + let dev_props = unsafe { instance.get_physical_device_properties(pd) }; + let limits = dev_props.limits; + let igpu = dev_props.device_type == vk::PhysicalDeviceType::INTEGRATED_GPU; + let name_up = name.to_ascii_uppercase(); + let gfx11 = dev_props.vendor_id == 0x1002 && (igpu || (coop && !coop_i8) + || ["GFX11", "STRIX", "PHOENIX", "KRACKAN", "NAVI3", "RADEON 8", "RADEON 7", "RX 7"] + .iter().any(|k| name_up.contains(k))); let dp_tile = !coop && limits.max_compute_work_group_invocations >= 1024 && limits.max_compute_shared_memory_size >= 32 << 10 && std::env::var_os("STRATA_VK_DP_OLD").is_none(); if coop { exts.push(ash::khr::cooperative_matrix::NAME.as_ptr()) } @@ -437,12 +458,13 @@ impl Vk { let fence = unsafe { dev.create_fence(&vk::FenceCreateInfo::default(), None) }?; let pc = [vk::PushConstantRange::default().stage_flags(vk::ShaderStageFlags::COMPUTE).size(128)]; let layout = unsafe { dev.create_pipeline_layout(&vk::PipelineLayoutCreateInfo::default().push_constant_ranges(&pc), None) }?; - let st = State { pool, spare: vec![], used: vec![], cur: None, cond: None, no_cond: false, stalled: false, busy: false, gpool, graphs: HashMap::new(), rec: None, batch: false, beside: false, tpool, tcur: None, + let st = State { pool, spare: vec![], used: vec![], queued: vec![], cur: None, cond: None, no_cond: false, need_cond: false, stalled: false, busy: false, gpool, graphs: HashMap::new(), rec: None, batch: false, beside: false, tpool, tcur: None, tspare: vec![], tused: vec![], tbytes: 0, tstart: 0, tring: Default::default(), cstaged: false, cval: 0, tval: 0, bwait: 0, twait: 0, tlag: 0, pipes: HashMap::new(), fence, staging: Default::default(), map: 0, staged: 0, ops: 0, copied: (0, 0), bufs: BTreeMap::new(), err: Default::default(), prof: if std::env::var_os("STRATA_VK_PROF").is_some() { Some((unsafe { dev.create_query_pool(&vk::QueryPoolCreateInfo::default().query_type(vk::QueryType::TIMESTAMP).query_count(65536), None) }?, vec![], HashMap::new())) } else { None }, step_pos: None }; let cr = cond.then(|| ash::ext::conditional_rendering::Device::new(&instance, &dev)); - let inner = Inner { _entry: entry, instance, pd, dev, queue, tq, qfs: [qf, tf.unwrap_or(qf)], csem, tsem, mem, budget, wave64, coop, dp_tile, coop_i8, amd_coherent, amd_plain, cr, llama: crate::cuda::llama_numerics() as u32, kv_q8: crate::cuda::kv_q8(), smem_limit, stats, layout, st: Mutex::new(st) }; + let ts_period = limits.timestamp_period; + let inner = Inner { _entry: entry, instance, pd, dev, queue, tq, qfs: [qf, tf.unwrap_or(qf)], csem, tsem, mem, budget, igpu, gfx11, wave64, coop, dp_tile, coop_i8, amd_coherent, amd_plain, cr, llama: crate::cuda::llama_numerics() as u32, kv_q8: crate::cuda::kv_q8(), smem_limit, ts_period, stats, layout, st: Mutex::new(st) }; let staging = inner.buffer(STAGING, HOST)?; let map = unsafe { inner.dev.map_memory(staging.1, 0, vk::WHOLE_SIZE, vk::MemoryMapFlags::empty()) }; let mut st = inner.st(); @@ -511,13 +533,31 @@ impl Inner { let req = unsafe { self.dev.get_buffer_memory_requirements(buf) }; // Memory the GPU does not cache (AMD flags 0xc0) only when asked for: it would be slow for anything else. let ok = |f: vk::MemoryPropertyFlags, want| f.contains(want) && (f & vk::MemoryPropertyFlags::from_raw(0xc0) & !want).is_empty(); - let ty = wants.iter().find_map(|&want| (0..self.mem.memory_type_count).find(|&i| req.memory_type_bits & (1 << i) != 0 - && ok(self.mem.memory_types[i as usize].property_flags, want))); + // On an iGPU, device-local and GTT heaps live in the same LPDDR: fall back to HOST if DEVICE heap is full. + let candidates: Vec = if wants == DEVICE && self.igpu { + wants.iter().chain(HOST.iter()).copied().collect() + } else { + wants.to_vec() + }; let mut flags = vk::MemoryAllocateFlagsInfo::default().flags(vk::MemoryAllocateFlags::DEVICE_ADDRESS); - let r = ty.context("no memory type").and_then(|ty| Ok(unsafe { self.dev.allocate_memory(&vk::MemoryAllocateInfo::default() - .allocation_size(req.size).memory_type_index(ty).push_next(&mut flags), None) } - .with_context(|| format!("allocate {bytes} bytes"))?)); - let mem = match r { Ok(m) => m, Err(e) => { unsafe { self.dev.destroy_buffer(buf, None) }; return Err(e) } }; + let mut last_err = anyhow::anyhow!("no matching memory type for {bytes} bytes"); + let mut mem_ok = None; + 'outer: for want in candidates { + for i in 0..self.mem.memory_type_count { + if req.memory_type_bits & (1 << i) != 0 && ok(self.mem.memory_types[i as usize].property_flags, want) { + match unsafe { + self.dev.allocate_memory( + &vk::MemoryAllocateInfo::default().allocation_size(req.size).memory_type_index(i).push_next(&mut flags), + None, + ) + } { + Ok(m) => { mem_ok = Some(m); break 'outer; } + Err(e) => last_err = anyhow::Error::from(e).context(format!("allocate {bytes} bytes")), + } + } + } + } + let mem = match mem_ok { Some(m) => m, None => { unsafe { self.dev.destroy_buffer(buf, None) }; return Err(last_err); } }; unsafe { self.dev.bind_buffer_memory(buf, mem, 0) }?; Ok((buf, mem)) } @@ -536,7 +576,7 @@ impl Inner { unsafe { self.dev.begin_command_buffer(c, &vk::CommandBufferBeginInfo::default() .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT)) }?; st.cur = Some(c); - self.cond_on(st, c); + if st.need_cond { self.cond_on(st, c); } Ok(c) } @@ -555,6 +595,7 @@ impl Inner { /// while the error word is 0, so nothing after a combine that gave up runs until the host has seen it (no-op /// without conditional rendering). Every compute command buffer starts this way too, for the submissions after. fn skip_rest(&self, st: &mut State) -> anyhow::Result<()> { + st.need_cond = true; let cb = self.cb(st)?; self.end_cond(st, cb); self.cond_on(st, cb); @@ -581,17 +622,23 @@ impl Inner { } } - /// Submit the open command buffer, then `graph`; with `wait`, wait for all submitted work. - fn submit(&self, st: &mut State, graph: Option, wait: bool) -> anyhow::Result<()> { - let mut cbs = vec![]; + /// Close `st.cur` if open and append it (plus `graph` if given) to `st.queued` in order. + fn enqueue_cb(&self, st: &mut State, graph: Option) -> anyhow::Result<()> { if let Some(c) = st.cur.take() { st.copied.1 = 0; self.end_cond(st, c); unsafe { self.dev.end_command_buffer(c) }?; st.used.push(c); - cbs.push(c); + st.queued.push(c); } - cbs.extend(graph); + st.queued.extend(graph); + Ok(()) + } + + /// Submit the open command buffer, then `graph`; with `wait`, wait for all submitted work. + fn submit(&self, st: &mut State, graph: Option, wait: bool) -> anyhow::Result<()> { + self.enqueue_cb(st, graph)?; + let cbs = std::mem::take(&mut st.queued); if cbs.is_empty() && !(wait && st.busy) { return Ok(()); } let fence = if wait { st.fence } else { vk::Fence::null() }; st.cval += 1; @@ -634,11 +681,21 @@ impl Inner { if let Some((qp, names, tot)) = &mut st.prof { if !names.is_empty() { let mut t = vec![0u64; 2 * names.len()]; unsafe { self.dev.get_query_pool_results(*qp, 0, &mut t, vk::QueryResultFlags::TYPE_64 | vk::QueryResultFlags::WAIT) }?; - for (i, n) in names.drain(..).enumerate() { let e = tot.entry(n).or_default(); e.0 += t[2 * i + 1] - t[2 * i]; e.1 += 1; } + for (i, n) in names.drain(..).enumerate() { + let e = tot.entry(n).or_default(); + e.0 += t[2 * i + 1].saturating_sub(t[2 * i]); + e.1 += 1; + if i > 0 && t[2 * i] >= t[2 * i - 1] { + let g = tot.entry("_barrier_gap").or_default(); + g.0 += t[2 * i] - t[2 * i - 1]; + g.1 += 1; + } + } } } let used = std::mem::take(&mut st.used); st.spare.extend(used); st.busy = false; + st.need_cond = false; st.staged = 0; if !e.is_null() { unsafe { e.add(1).write_volatile(0) }; } if !e.is_null() && unsafe { e.read_volatile() } != 0 { @@ -797,7 +854,19 @@ impl Vk { } if !st.batch || name == "quant_xq" { self.0.barrier(cb, false); } if st.rec.is_none() { if let Some((qp, names, _)) = &mut st.prof { if names.len() < 32767 { - unsafe { d.cmd_write_timestamp(cb, vk::PipelineStageFlags::BOTTOM_OF_PIPE, *qp, 2 * names.len() as u32 + 1) }; names.push(name); } } } + let label = match (name, spec.get(1).copied()) { + ("gemv_q", Some(11)) => "gemv_q:Q3_K", + ("gemv_q", Some(13)) => "gemv_q:Q5_K", + ("gemv_q", Some(14)) => "gemv_q:Q6_K", + ("gemv_q", Some(8)) => "gemv_q:Q8_0", + ("gemv_q", Some(18)) => "gemv_q:IQ3_XXS", + ("moe_gemv_q", Some(18)) => "moe_gemv_q:IQ3_XXS", + ("moe_gemv_q", Some(11)) => "moe_gemv_q:Q3_K", + ("moe_down_combine", Some(18)) => "moe_down_combine:IQ3_XXS", + ("moe_down_combine", Some(11)) => "moe_down_combine:Q3_K", + _ => name, + }; + unsafe { d.cmd_write_timestamp(cb, vk::PipelineStageFlags::BOTTOM_OF_PIPE, *qp, 2 * names.len() as u32 + 1) }; names.push(label); } } } Ok(()) } @@ -1073,14 +1142,14 @@ impl Backend for Vk { } fn graph(&self, key: u64, f: &mut dyn FnMut() -> anyhow::Result<()>) -> anyhow::Result<()> { let mut st = self.0.st(); - if let Some(&g) = st.graphs.get(&key) { return self.0.submit(&mut st, Some(g), false); } + if let Some(&g) = st.graphs.get(&key) { return self.0.enqueue_cb(&mut st, Some(g)); } if std::env::var_os("STRATA_NO_GRAPH").is_some() { drop(st); return f(); } let c = unsafe { self.0.dev.allocate_command_buffers(&vk::CommandBufferAllocateInfo::default().command_pool(st.gpool) .command_buffer_count(1)) }?[0]; unsafe { self.0.dev.begin_command_buffer(c, &vk::CommandBufferBeginInfo::default() .flags(vk::CommandBufferUsageFlags::SIMULTANEOUS_USE)) }?; st.rec = Some(c); - self.0.cond_on(&mut st, c); + if st.need_cond { self.0.cond_on(&mut st, c); } drop(st); let r = f(); let mut st = self.0.st(); @@ -1089,9 +1158,11 @@ impl Backend for Vk { unsafe { self.0.dev.end_command_buffer(c) }?; if let Err(e) = r { unsafe { self.0.dev.free_command_buffers(st.gpool, &[c]) }; return Err(e); } st.graphs.insert(key, c); - self.0.submit(&mut st, Some(c), false) + self.0.enqueue_cb(&mut st, Some(c)) } + fn moe_mailbox_optional(&self) -> bool { true } + /// Device memory mapped through the BAR (`BAR`, or `BAR_AMD`). fn host_mailbox(&self, words: usize) -> anyhow::Result<(Buf, *mut u32)> { let (b, p) = self.buf(4 * words.max(1), if self.0.amd_coherent { BAR_AMD } else { BAR })?; @@ -1118,7 +1189,7 @@ impl Backend for Vk { self.gemvq(w, x, x_stride, 0, y, y_stride, nt) } fn gemvq(&self, w: &GpuTensor, x: u64, x_stride: usize, xq: u64, y: u64, y_stride: usize, nt: usize) -> anyhow::Result<()> { - ensure!((1..=4).contains(&nt), "gemvn: {nt} inputs (1 to 4)"); + ensure!((1..=crate::cpu_worker::TMAX).contains(&nt), "gemvn: {nt} inputs (1 to {})", crate::cpu_worker::TMAX); let (nc, spec) = (gemv_cols(self.0.llama, w.ty, w.cols)?, [&self.wspec(w.ty)[..], &[nt as u32]].concat()); if matches!(w.ty, GgmlType::F32 | GgmlType::BF16) { // Rows per workgroup as CUDA gemv_cfg: 8 for tall matrices, else 1 (hc_attn_up 10240 x 320). @@ -1315,6 +1386,7 @@ impl Backend for Vk { ensure!(n <= 512 && k <= 32, "moe_route: {n} experts, top {k}"); self.run("moe_route", &[], [1, 1, 1], &[&p, &table, &ids, &wts, &slots, &wh, &h, &mb, &step, &(n as u32), &(k as u32), &extra_scale, &(nh as u32), &(h_off as u32), &(t as u32), &(sigmoid as u32)]) } + fn moe_prefetch_supported(&self) -> bool { !self.0.gfx11 && !self.0.igpu } fn moe_combine(&self, x: u64, ye: u64, w: u64, k: usize, sh: u64, sg: u64, mb: u64, y_off: usize, step: u64, slots: u64, n: usize, t: usize) -> anyhow::Result<()> { let err = self.0.st().err.3; self.run("moe_combine", &[], [n.div_ceil(256), t, 1], &[&x, &ye, &w, &sh, &sg, &mb, &step, &slots, &err, &(k as u32), &(y_off as u32), &(n as u32)])?; @@ -1323,7 +1395,7 @@ impl Backend for Vk { } fn moe_down_combine(&self, ty: GgmlType, rows: usize, cols: usize, w: u64, stride: usize, slots: u64, k: usize, h: u64, ye: u64, xq: u64, x: u64, wh: u64, sh: u64, sg: u64, mb: u64, y_off: usize, step: u64, t: usize) -> anyhow::Result<()> { - if xq == 0 || cols % 256 != 0 || k > 32 { // the shader keeps up to 32 slots in shared memory + if t > 1 || xq == 0 || cols % 256 != 0 || k > 32 { // t > 1 deduplicates shared experts across tokens in moe_gemv self.moe_gemv(ty, rows, cols, w, stride, slots, t * k, h, cols, 1, ye, None, xq)?; return self.moe_combine(x, ye, wh, k, sh, sg, mb, y_off, step, slots, rows, t); } diff --git a/src/vulkan/preflight.rs b/src/vulkan/preflight.rs index 4feeac84..43e69a6e 100644 --- a/src/vulkan/preflight.rs +++ b/src/vulkan/preflight.rs @@ -37,6 +37,31 @@ impl Vk { } eprintln!("vulkan: shared-memory preflight: {n} (kernel, specialization) pairs checked, largest {} bytes \ ({}) of {} allowed, {:.1} ms", worst.0, worst.1, self.0.smem_limit, t0.elapsed().as_secs_f64() * 1e3); + if !verbose { + let mut st = self.0.st(); + for &ty in types { + let ws = self.wspec(ty); + if !matches!(ty, GgmlType::F32 | GgmlType::BF16) { + for nt in [1u32, 3, 5, 9, 17] { + let _ = self.0.pipeline(&mut st, "gemv_q", &[&ws[..], &[nt]].concat()); + } + let _ = self.0.pipeline(&mut st, "moe_gemv_q", &ws); + let _ = self.0.pipeline(&mut st, "moe_down_combine", &ws); + } else { + for nt in [1u32, 3, 5, 9, 17] { + for wpb in [1u32, 8] { + let _ = self.0.pipeline(&mut st, "gemv_f", &[&ws[..], &[nt, 32 * wpb]].concat()); + } + } + } + } + for name in ["moe_combine", "deltanet", "conv_silu", "conv_seq", "conv_state", "l2_norm", "norm_gate", "moe_route", "argmax_rows", "attn_merge"] { + let _ = self.0.pipeline(&mut st, name, &[]); + } + for name in ["rms_norm", "swiglu_q", "quant_xq"] { + let _ = self.0.pipeline(&mut st, name, &[self.0.llama]); + } + } Ok(()) } diff --git a/src/vulkan/shaders/attn_merge.comp b/src/vulkan/shaders/attn_merge.comp index c98bfd1f..f40baa46 100644 --- a/src/vulkan/shaders/attn_merge.comp +++ b/src/vulkan/shaders/attn_merge.comp @@ -29,10 +29,10 @@ void main() { float v = fdiv(o, L) * gate; p.out_.v[(z * nh + h) * hd + d] = v; if (uint64_t(p.xq) != 0) { - float mx = warp_max(abs(v)), id = mx > 0.0 ? fdiv(127.0, mx) : 0.0; + float mx = warp_max(abs(v)), id = mx > 0.0 ? rcp_d(127.0, mx) : 0.0; uint n = nh * hd, i = z * n + h * hd + d; p.xq.v[i] = uint8_t(int(roundEven(v * id)) & 255); - if ((tid & 31) == 0) F32(uint64_t(p.xq) + gl_NumWorkGroups.z * n).v[i / 32] = fdiv(mx, 127.0); + if ((tid & 31) == 0) F32(uint64_t(p.xq) + gl_NumWorkGroups.z * n).v[i / 32] = rcp_d(mx, 127.0); } } } diff --git a/src/vulkan/shaders/common.glsl b/src/vulkan/shaders/common.glsl index 6ee5af76..af0e68a8 100644 --- a/src/vulkan/shaders/common.glsl +++ b/src/vulkan/shaders/common.glsl @@ -29,9 +29,9 @@ float fdiv(float a, float b) { // An infinite operand or b = 0 makes the steps inf * 0 = NaN (sigmoid(-100): 1 / inf); a / b is exact there. return isinf(a) || isinf(b) || b == 0.0 ? a / b : r; } -// Quantizer reciprocals on wave64: a correctly rounded double division keeps values on a roundEven() tie -// from moving by one int8 step. Wave32 keeps the fdiv path that its kernels were tuned with. -float rcp_d(float a, float b) { return W64 ? float(double(a) / double(b)) : fdiv(a, b); } +// Quantizer reciprocals: a correctly rounded double division keeps values on a roundEven() tie +// from moving by one int8 step across both wave64 and wave32 devices. +float rcp_d(float a, float b) { return float(double(a) / double(b)); } float sqrt_(float x) { float s = sqrt(x); precise float r = fma(fma(-s, s, x), fdiv(0.5, s), s); return x > 0.0 && !isinf(x) ? r : s; } // expf as CUDA compiles it (its PTX): x = j ln2 + r with j from a rounded-down fma, then ex2.approx(r) * 2^j. diff --git a/src/vulkan/shaders/deq.glsl b/src/vulkan/shaders/deq.glsl index 3b8c52ea..acd55f8c 100644 --- a/src/vulkan/shaders/deq.glsl +++ b/src/vulkan/shaders/deq.glsl @@ -9,14 +9,16 @@ layout(constant_id = 2) const uint SBB = 1024; #ifdef SHARED_TABLES // Sized by T (1 entry for the other types), so only the table in use takes shared memory. shared int s_kv[T == 20 || T == 23 ? 16 : 1]; -shared uint s_g3[T == 21 ? 512 : 1]; +shared uint s_g3[T == 21 ? 512 : (T == 18 ? 256 : 1)]; shared uint64_t s_g2[T == 22 ? 1024 : 1]; #define KV_IQ4NL s_kv #define IQ3S_GRID s_g3 +#define IQ3XXS_GRID s_g3 #define IQ2S_GRID s_g2 #else #define KV_IQ4NL kvalues_iq4nl #define IQ3S_GRID iq3s_grid +#define IQ3XXS_GRID iq3xxs_grid #define IQ2S_GRID iq2s_grid void load_tables() {} // the tables come from the buffer #endif @@ -72,6 +74,7 @@ uvec2 ld8(U8 w, uint o) { void load_tables() { uint li = gl_LocalInvocationIndex, n = gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z; if (T == 20 || T == 23) { if (li < 16) s_kv[li] = kvalues_iq4nl[li]; barrier(); } + if (T == 18) { for (uint i = li; i < 256; i += n) s_g3[i] = iq3xxs_grid[i]; barrier(); } if (T == 21) { for (uint i = li; i < 512; i += n) s_g3[i] = iq3s_grid[i]; barrier(); } if (T == 22) { for (uint i = li; i < 1024; i += n) s_g2[i] = iq2s_grid[i]; barrier(); } } @@ -99,7 +102,7 @@ void qw_k(U8 w, uint b, uint sub, out uint lo, out uint hi, out float d, out flo lo = (v.x >> (2 * j)) & 0x03030303u; hi = (v.y >> (2 * j)) & 0x03030303u; } else if (T == 11) { // Q3_K: value = 2 low bits + 4 * high bit - 4 uint e = sub * 8, n = e >> 7, r = e & 127, j = r >> 5, i0 = r & 31, is = n * 8 + j * 2 + (i0 >> 4), sc = b + 96; - uint low = is < 8 ? u8_(w, sc + is) & 15 : u8_(w, sc + is - 8) >> 4; + uint low = (u8_(w, sc + (is & 7)) >> ((is & 8) >> 1)) & 15; uint hs = (u8_(w, sc + 8 + (is & 3)) >> (2 * (is >> 2))) & 3, hb = 4 * n + j; d = ldh(w, b + 108) * float(int(low | hs << 4) - 32); uvec2 h = ld8(w, b + i0); @@ -181,7 +184,7 @@ void iq_grid(U8 w, uint b, uint ib, uint l, out uint g0, out uint g1, out uint s uint aux = ld32(w, b + 66 + 4 * ib); d = ldh(w, b) * (0.5 + float(aux >> 28)) * 0.5; sg = ksign((aux >> (7 * l)) & 127); - g0 = iq3xxs_grid[qq & 255]; g1 = iq3xxs_grid[qq >> 8]; + g0 = IQ3XXS_GRID[qq & 255]; g1 = IQ3XXS_GRID[qq >> 8]; } else { uint hb = u8_(w, b + 66 + ib); sg = u8_(w, b + 74 + 4 * ib + l); @@ -274,6 +277,71 @@ void qw(U8 w, uint b, uint sub, out uint lo, out uint hi, out float d, out float lo = pack4(q[0], q[1], q[2], q[3]); hi = pack4(q[4], q[5], q[6], q[7]); } +// Shuffle within a logical 32-lane warp (works on both wave32 and wave64). +uint shfl32(uint v, uint src) { return subgroupShuffle(v, W64 ? ((gl_SubgroupInvocationID & ~31u) | src) : src); } + +// Warp-cooperative super-block dequantization for GEMV (all 32 lanes of the logical warp pass the same w and b, with sub = LANE). +void qw_sb(U8 w, uint b, uint sub, out uint lo, out uint hi, out float d, out float m) { + if (T == 11) { // Q3_K (110 bytes = 28 uint32 words from 4-aligned base): 1 coalesced U32 load per lane + register shuffles + m = 0.0; + uint64_t addr = uint64_t(w) + b, base4 = addr & ~3ul; + uint align_sh = uint(addr & 2ul) * 8u; + uint raw = U32(base4).v[min(sub, 27u)]; + uint next_raw = shfl32(raw, min(sub + 1u, 27u)); + uint w32 = align_sh == 0u ? raw : ((raw >> 16) | (next_raw << 16)); + uint n = sub >> 4, j = (sub >> 2) & 3u, hw = (sub & 3u) * 2u, is = sub >> 1, hb = 4u * n + j; + uint sc_lo_w = shfl32(w32, 24u + ((is >> 2) & 1u)), sc_hi_w = shfl32(w32, 26u); + uint low = (sc_lo_w >> (8u * (is & 3u) + ((is & 8u) >> 1))) & 15u; + uint hs = (sc_hi_w >> (8u * (is & 3u) + 2u * (is >> 2))) & 3u; + d = h2f(shfl32(w32, 27u) & 0xFFFFu) * float(int(low | (hs << 4)) - 32); + uvec2 h = uvec2(shfl32(w32, hw), shfl32(w32, hw + 1u)); + uint vw = 8u + 8u * n + hw; + uvec2 v = uvec2(shfl32(w32, vw), shfl32(w32, vw + 1u)); + lo = vsub4(((v.x >> (2u * j)) & 0x03030303u) | ((h.x >> hb) & 0x01010101u) << 2, 0x04040404u); + hi = vsub4(((v.y >> (2u * j)) & 0x03030303u) | ((h.y >> hb) & 0x01010101u) << 2, 0x04040404u); + return; + } + if (T == 14) { // Q6_K (210 bytes = 53 uint32 words from 4-aligned base): 2 coalesced U32 loads per lane + register shuffles + m = 0.0; + uint64_t addr = uint64_t(w) + b, base4 = addr & ~3ul; + uint align_sh = uint(addr & 2ul) * 8u; + uint raw0 = U32(base4).v[sub], raw1 = U32(base4).v[min(sub + 32u, 52u)]; + uint s0 = shfl32(raw0, min(sub + 1u, 31u)), s1 = shfl32(raw1, 0u); + uint next0 = sub < 31u ? s0 : s1; + uint next1 = shfl32(raw1, min(sub + 1u, 20u)); + uint w0 = align_sh == 0u ? raw0 : ((raw0 >> 16) | (next0 << 16)); + uint w1 = align_sh == 0u ? raw1 : ((raw1 >> 16) | (next1 << 16)); + uint n = sub >> 4, qt = (sub >> 2) & 3u, kw = (sub & 3u) * 2u, ns = 4u * (qt >> 1), hs = 2u * qt; + uint sc_idx = 8u * n + ((sub >> 1) & 1u) + 2u * qt; + uint sc_w = shfl32(w1, 16u + (sc_idx >> 2)); + d = h2f(shfl32(w1, 20u) & 0xFFFFu) * float(sb(sc_w >> (8u * (sc_idx & 3u)))); + uint lw = 16u * n + (qt & 1u) * 8u + kw, hw = 8u * n + kw; + uvec2 l = uvec2(shfl32(w0, lw), shfl32(w0, lw + 1u)); + uvec2 h = uvec2(shfl32(w1, hw), shfl32(w1, hw + 1u)); + lo = vsub4(((l.x >> ns) & 0x0F0F0F0Fu) | ((h.x >> hs) & 0x03030303u) << 4, 0x20202020u); + hi = vsub4(((l.y >> ns) & 0x0F0F0F0Fu) | ((h.y >> hs) & 0x03030303u) << 4, 0x20202020u); + return; + } + if (T == 18) { // IQ3_XXS (98 bytes = 25 uint32 words from 4-aligned base): 1 coalesced U32 load per lane + LDS grid + m = 0.0; + uint64_t addr = uint64_t(w) + b, base4 = addr & ~3ul; + uint align_sh = uint(addr & 2ul) * 8u; + uint raw = U32(base4).v[min(sub, 24u)]; + uint next_raw = shfl32(raw, min(sub + 1u, 24u)); + float d0 = h2f((shfl32(raw, 0u) >> align_sh) & 0xFFFFu); + uint w32 = align_sh == 0u ? ((raw >> 16) | (next_raw << 16)) : next_raw; + uint ib = sub >> 2, l = sub & 3u; + uint qq = (shfl32(w32, sub >> 1) >> ((sub & 1u) * 16u)) & 0xFFFFu; + uint aux = shfl32(w32, 16u + ib); + d = d0 * (0.5 + float(aux >> 28)) * 0.5; + uint sg = ksign((aux >> (7u * l)) & 127u); + uint g0 = IQ3XXS_GRID[qq & 255u], g1 = IQ3XXS_GRID[qq >> 8]; + lo = apply_signs(g0, sg & 15u); hi = apply_signs(g1, sg >> 4); + return; + } + qw(w, b, sub, lo, hi, d, m); +} + // Values [8 sub, 8 sub + 8) of the super-block at byte b (deq8 in kernels.cu; d * q - m is the same float). void deq8(U8 w, uint b, uint sub, out float o[8]) { if (T == 0) { // F32 diff --git a/src/vulkan/shaders/gemv_q.comp b/src/vulkan/shaders/gemv_q.comp index 10842a00..6d6a28e5 100644 --- a/src/vulkan/shaders/gemv_q.comp +++ b/src/vulkan/shaders/gemv_q.comp @@ -15,7 +15,7 @@ layout(constant_id = 5) const uint NT = 1; void sb_step(F32 xd, uint row, uint lane, uint s, inout precise float acc[NT]) { uint lo, hi; float d, m; - qw(p.w, row * rowb(p.cols) + s * SBB, lane, lo, hi, d, m); + qw_sb(p.w, row * rowb(p.cols) + s * SBB, lane, lo, hi, d, m); [[unroll]] for (uint t = 0; t < NT; t++) { uint o = (t * p.cols + s * 256) / 4 + lane * 2, a = p.xq.v[o], c = p.xq.v[o + 1]; int dot = dotPacked4x8EXT(int(lo), int(a)) + dotPacked4x8EXT(int(hi), int(c)); @@ -36,7 +36,10 @@ void main() { [[unroll]] for (uint t = 0; t < NT; t++) acc[t] = 0.0; // Four super-blocks per iteration so their loads overlap (kernels.cu: unroll 4); the sums keep their order. uint s = 0; - for (; s + 4 <= nsb; s += 4) [[unroll]] for (uint j = 0; j < 4; j++) sb_step(xd, row, lane, s + j, acc); + if (NT <= 4) { + for (; s + 4 <= nsb; s += 4) [[unroll]] for (uint j = 0; j < 4; j++) sb_step(xd, row, lane, s + j, acc); + } + for (; s + 2 <= nsb; s += 2) [[unroll]] for (uint j = 0; j < 2; j++) sb_step(xd, row, lane, s + j, acc); for (; s < nsb; s++) sb_step(xd, row, lane, s, acc); if ((p.cols & 255) > lane * 8) sb_step(xd, row, lane, nsb, acc); [[unroll]] for (uint t = 0; t < NT; t++) { diff --git a/src/vulkan/shaders/moe_down_combine.comp b/src/vulkan/shaders/moe_down_combine.comp index 155c8853..8a7e8b46 100644 --- a/src/vulkan/shaders/moe_down_combine.comp +++ b/src/vulkan/shaders/moe_down_combine.comp @@ -6,6 +6,7 @@ // sums in shared memory, then invocation i combines row r0 + i as moe_combine does (same order, same result). layout(local_size_x = 256) in; #include "common.glsl" +#define SHARED_TABLES #include "deq.glsl" #extension GL_EXT_integer_dot_product : require #extension GL_EXT_control_flow_attributes : require @@ -15,7 +16,7 @@ layout(push_constant) uniform P { U8 w; I32 slots; U32 xq; uint64_t stride; F32 void sb_step(U8 w, F32 xd, uint row, uint lane, uint t, uint s, inout precise float acc) { uint lo, hi; float d, m; - qw(w, (row * p.nsb + s) * SBB, lane, lo, hi, d, m); + qw_sb(w, (row * p.nsb + s) * SBB, lane, lo, hi, d, m); uint o = (t * p.nsb + s) * 64 + lane * 2, a = p.xq.v[o], c = p.xq.v[o + 1]; int dot = dotPacked4x8EXT(int(lo), int(a)) + dotPacked4x8EXT(int(hi), int(c)); int sx = dotPacked4x8EXT(0x01010101, int(a)) + dotPacked4x8EXT(0x01010101, int(c)); @@ -24,11 +25,29 @@ void sb_step(U8 w, F32 xd, uint row, uint lane, uint t, uint s, inout precise fl if (MIN) acc = fma(-(m * dx), float(sx), acc); // types with a min } +void sb_step2(U8 w, F32 xd, uint row0, uint row1, uint lane, uint t, uint s, inout precise float acc0, inout precise float acc1) { + uint lo0, hi0, lo1, hi1; float d0, m0, d1, m1; + qw_sb(w, (row0 * p.nsb + s) * SBB, lane, lo0, hi0, d0, m0); + qw_sb(w, (row1 * p.nsb + s) * SBB, lane, lo1, hi1, d1, m1); + uint o = (t * p.nsb + s) * 64 + lane * 2, a = p.xq.v[o], c = p.xq.v[o + 1]; + float dx = xd.v[(t * p.nsb + s) * 8 + (lane >> 2)]; + int dot0 = dotPacked4x8EXT(int(lo0), int(a)) + dotPacked4x8EXT(int(hi0), int(c)); + int dot1 = dotPacked4x8EXT(int(lo1), int(a)) + dotPacked4x8EXT(int(hi1), int(c)); + acc0 = fma(d0 * dx, float(dot0), acc0); + acc1 = fma(d1 * dx, float(dot1), acc1); + if (MIN) { + int sx = dotPacked4x8EXT(0x01010101, int(a)) + dotPacked4x8EXT(0x01010101, int(c)); + acc0 = fma(-(m0 * dx), float(sx), acc0); + acc1 = fma(-(m1 * dx), float(sx), acc1); + } +} + #include "mb_wait.glsl" const uint RPW = 8; // rows per workgroup (rpw) at most shared float part[32 * RPW]; shared bool miss, ok; void main() { + load_tables(); uint tk = gl_WorkGroupID.y, k = p.k, lane = LANE, n = p.rows; uint rpw = p.rpw, r0 = gl_WorkGroupID.x * rpw; F32 xd = F32(uint64_t(p.xq) + p.ntok * p.nsb * 256); @@ -37,10 +56,19 @@ void main() { int e = p.slots.v[t]; if (e < 0) continue; // a miss: weight 0, the CPU's result is added below U8 w = U8(uint64_t(p.w) + uint64_t(e) * p.stride); - for (uint i = 0; i < rpw && r0 + i < n; i++) { + uint i = 0; + for (; i + 1 < rpw && r0 + i + 1 < n; i += 2) { + precise float acc0 = 0.0, acc1 = 0.0; + uint s = 0; + for (; s + 2 <= p.nsb; s += 2) [[unroll]] for (uint q = 0; q < 2; q++) sb_step2(w, xd, r0 + i, r0 + i + 1, lane, t, s + q, acc0, acc1); + for (; s < p.nsb; s++) sb_step2(w, xd, r0 + i, r0 + i + 1, lane, t, s, acc0, acc1); + acc0 = warp_sum(acc0); acc1 = warp_sum(acc1); + if (lane == 0) { part[j * RPW + i] = acc0; part[j * RPW + i + 1] = acc1; } + } + for (; i < rpw && r0 + i < n; i++) { precise float acc = 0.0; uint s = 0; - for (; s + 4 <= p.nsb; s += 4) [[unroll]] for (uint q = 0; q < 4; q++) sb_step(w, xd, r0 + i, lane, t, s + q, acc); + for (; s + 2 <= p.nsb; s += 2) [[unroll]] for (uint q = 0; q < 2; q++) sb_step(w, xd, r0 + i, lane, t, s + q, acc); for (; s < p.nsb; s++) sb_step(w, xd, r0 + i, lane, t, s, acc); acc = warp_sum(acc); if (lane == 0) part[j * RPW + i] = acc; diff --git a/src/vulkan/shaders/moe_gemv_q.comp b/src/vulkan/shaders/moe_gemv_q.comp index 958f0f2a..ed0b0826 100644 --- a/src/vulkan/shaders/moe_gemv_q.comp +++ b/src/vulkan/shaders/moe_gemv_q.comp @@ -12,9 +12,13 @@ layout(local_size_x = 256) in; layout(push_constant) uniform P { U8 w; I32 ids; U32 xq; F32 y; uint64_t stride; uint64_t w2; uint64_t y2; uint rows; uint cols; uint xg; uint ntok; uint rpw; } p; +shared uint s_slots[18]; +shared uint s_nm; +shared bool s_dup; + void sb_step(U8 w, F32 xd, uint row, uint lane, uint t, uint s, inout precise float acc) { uint lo, hi; float d, m; - qw(w, row * rowb(p.cols) + s * SBB, lane, lo, hi, d, m); + qw_sb(w, row * rowb(p.cols) + s * SBB, lane, lo, hi, d, m); uint o = (t * p.cols + s * 256) / 4 + lane * 2, a = p.xq.v[o], c = p.xq.v[o + 1]; int dot = dotPacked4x8EXT(int(lo), int(a)) + dotPacked4x8EXT(int(hi), int(c)); int sx = dotPacked4x8EXT(0x01010101, int(a)) + dotPacked4x8EXT(0x01010101, int(c)); @@ -23,23 +27,70 @@ void sb_step(U8 w, F32 xd, uint row, uint lane, uint t, uint s, inout precise fl if (MIN) acc = fma(-(m * dx), float(sx), acc); // types with a min } +void sb_step_multi(U8 w, F32 xd, uint row, uint lane, uint nm, uint s, inout precise float acc[17]) { + uint lo, hi; float d, m; + qw_sb(w, row * rowb(p.cols) + s * SBB, lane, lo, hi, d, m); + for (uint u = 0; u < nm; u++) { + uint t = s_slots[u] / p.xg; + uint o = (t * p.cols + s * 256) / 4 + lane * 2, a = p.xq.v[o], c = p.xq.v[o + 1]; + int dot = dotPacked4x8EXT(int(lo), int(a)) + dotPacked4x8EXT(int(hi), int(c)); + int sx = dotPacked4x8EXT(0x01010101, int(a)) + dotPacked4x8EXT(0x01010101, int(c)); + float dx = xd.v[(t * p.cols + s * 256) / 32 + (lane >> 2)]; + acc[u] = fma(d * dx, float(dot), acc[u]); + if (MIN) acc[u] = fma(-(m * dx), float(sx), acc[u]); + } +} + void main() { - uint k = gl_WorkGroupID.y, lane = LANE; + uint k = gl_WorkGroupID.y, lane = LANE, nslots = gl_NumWorkGroups.y; int e = p.ids.v[k]; if (e < 0) return; // the whole workgroup + uint nm = 1; + if (nslots > 8 && nslots <= 256) { + uint tid = gl_LocalInvocationID.x; + if (tid == 0) { s_nm = 0; s_dup = false; } + barrier(); + if (tid < nslots && p.ids.v[tid] == e) { + if (tid < k) { + s_dup = true; + } else { + uint idx = atomicAdd(s_nm, 1u); + if (idx < 17u) s_slots[idx] = tid; + } + } + barrier(); + if (s_dup) return; + nm = min(s_nm, 17u); + } load_tables(); U8 w = U8(uint64_t(p.w) + uint64_t(e) * p.stride + gl_WorkGroupID.z * p.w2); F32 y = F32(uint64_t(p.y) + gl_WorkGroupID.z * p.y2); F32 xd = F32(uint64_t(p.xq) + p.ntok * p.cols); uint t = k / p.xg, nsb = p.cols >> 8, r0 = gl_WorkGroupID.x * 8 * p.rpw, r1 = min(p.rows, r0 + 8 * p.rpw); - for (uint row = r0 + WARP; row < r1; row += 8) { - precise float acc = 0.0; - // Four super-blocks per iteration so their loads overlap (kernels.cu: unroll 4); the sums keep their order. - uint s = 0; - for (; s + 4 <= nsb; s += 4) [[unroll]] for (uint j = 0; j < 4; j++) sb_step(w, xd, row, lane, t, s + j, acc); - for (; s < nsb; s++) sb_step(w, xd, row, lane, t, s, acc); - if ((p.cols & 255) > lane * 8) sb_step(w, xd, row, lane, t, nsb, acc); // partial last super-block - acc = warp_sum(acc); - if (lane == 0) y.v[k * p.rows + row] = acc; + if (nm <= 1) { + for (uint row = r0 + WARP; row < r1; row += 8) { + precise float acc = 0.0; + // Four super-blocks per iteration so their loads overlap (kernels.cu: unroll 4); the sums keep their order. + uint s = 0; + for (; s + 4 <= nsb; s += 4) [[unroll]] for (uint j = 0; j < 4; j++) sb_step(w, xd, row, lane, t, s + j, acc); + for (; s + 2 <= nsb; s += 2) [[unroll]] for (uint j = 0; j < 2; j++) sb_step(w, xd, row, lane, t, s + j, acc); + for (; s < nsb; s++) sb_step(w, xd, row, lane, t, s, acc); + if ((p.cols & 255) > lane * 8) sb_step(w, xd, row, lane, t, nsb, acc); // partial last super-block + acc = warp_sum(acc); + if (lane == 0) y.v[k * p.rows + row] = acc; + } + } else { + for (uint row = r0 + WARP; row < r1; row += 8) { + precise float acc[17]; + for (uint u = 0; u < nm; u++) acc[u] = 0.0; + uint s = 0; + for (; s + 2 <= nsb; s += 2) [[unroll]] for (uint j = 0; j < 2; j++) sb_step_multi(w, xd, row, lane, nm, s + j, acc); + for (; s < nsb; s++) sb_step_multi(w, xd, row, lane, nm, s, acc); + if ((p.cols & 255) > lane * 8) sb_step_multi(w, xd, row, lane, nm, nsb, acc); + for (uint u = 0; u < nm; u++) { + float a = warp_sum(acc[u]); + if (lane == 0) y.v[s_slots[u] * p.rows + row] = a; + } + } } } diff --git a/src/vulkan/spv/attn_merge.spv b/src/vulkan/spv/attn_merge.spv index 42adf215..9e7c9f80 100644 Binary files a/src/vulkan/spv/attn_merge.spv and b/src/vulkan/spv/attn_merge.spv differ diff --git a/src/vulkan/spv/dequant.spv b/src/vulkan/spv/dequant.spv index d2391723..a207e30e 100644 Binary files a/src/vulkan/spv/dequant.spv and b/src/vulkan/spv/dequant.spv differ diff --git a/src/vulkan/spv/gemm_f.spv b/src/vulkan/spv/gemm_f.spv index 207b2363..981e103d 100644 Binary files a/src/vulkan/spv/gemm_f.spv and b/src/vulkan/spv/gemm_f.spv differ diff --git a/src/vulkan/spv/gemm_f_dp.spv b/src/vulkan/spv/gemm_f_dp.spv index a0c876b4..7a03d767 100644 Binary files a/src/vulkan/spv/gemm_f_dp.spv and b/src/vulkan/spv/gemm_f_dp.spv differ diff --git a/src/vulkan/spv/gemm_fs.spv b/src/vulkan/spv/gemm_fs.spv index 6f68c505..dddf75d8 100644 Binary files a/src/vulkan/spv/gemm_fs.spv and b/src/vulkan/spv/gemm_fs.spv differ diff --git a/src/vulkan/spv/gemmi.spv b/src/vulkan/spv/gemmi.spv index c5710b2b..aa102a54 100644 Binary files a/src/vulkan/spv/gemmi.spv and b/src/vulkan/spv/gemmi.spv differ diff --git a/src/vulkan/spv/gemmi_dp.spv b/src/vulkan/spv/gemmi_dp.spv index dede9ce9..962b6e70 100644 Binary files a/src/vulkan/spv/gemmi_dp.spv and b/src/vulkan/spv/gemmi_dp.spv differ diff --git a/src/vulkan/spv/gemmi_s.spv b/src/vulkan/spv/gemmi_s.spv index 95d599ed..8df0d35c 100644 Binary files a/src/vulkan/spv/gemmi_s.spv and b/src/vulkan/spv/gemmi_s.spv differ diff --git a/src/vulkan/spv/gemv_f.spv b/src/vulkan/spv/gemv_f.spv index 3c40d9f4..3a118005 100644 Binary files a/src/vulkan/spv/gemv_f.spv and b/src/vulkan/spv/gemv_f.spv differ diff --git a/src/vulkan/spv/gemv_q.spv b/src/vulkan/spv/gemv_q.spv index 6ef82e08..9cf8eb25 100644 Binary files a/src/vulkan/spv/gemv_q.spv and b/src/vulkan/spv/gemv_q.spv differ diff --git a/src/vulkan/spv/kv_store.spv b/src/vulkan/spv/kv_store.spv index ae1fa79e..ff935661 100644 Binary files a/src/vulkan/spv/kv_store.spv and b/src/vulkan/spv/kv_store.spv differ diff --git a/src/vulkan/spv/moe_down_combine.spv b/src/vulkan/spv/moe_down_combine.spv index 3c0c5358..8069f065 100644 Binary files a/src/vulkan/spv/moe_down_combine.spv and b/src/vulkan/spv/moe_down_combine.spv differ diff --git a/src/vulkan/spv/moe_gemv_f.spv b/src/vulkan/spv/moe_gemv_f.spv index cdf44a15..d341d825 100644 Binary files a/src/vulkan/spv/moe_gemv_f.spv and b/src/vulkan/spv/moe_gemv_f.spv differ diff --git a/src/vulkan/spv/moe_gemv_q.spv b/src/vulkan/spv/moe_gemv_q.spv index 2407383a..c77b1bb0 100644 Binary files a/src/vulkan/spv/moe_gemv_q.spv and b/src/vulkan/spv/moe_gemv_q.spv differ diff --git a/src/vulkan/spv/norm_gate.spv b/src/vulkan/spv/norm_gate.spv index 69aa1017..5ef4ad1b 100644 Binary files a/src/vulkan/spv/norm_gate.spv and b/src/vulkan/spv/norm_gate.spv differ diff --git a/src/vulkan/spv/quant_q8.spv b/src/vulkan/spv/quant_q8.spv index e6739177..64b3c786 100644 Binary files a/src/vulkan/spv/quant_q8.spv and b/src/vulkan/spv/quant_q8.spv differ diff --git a/src/vulkan/spv/quant_xq.spv b/src/vulkan/spv/quant_xq.spv index 746c8347..eed4f206 100644 Binary files a/src/vulkan/spv/quant_xq.spv and b/src/vulkan/spv/quant_xq.spv differ diff --git a/src/vulkan/spv/rms_norm.spv b/src/vulkan/spv/rms_norm.spv index 2fa66760..14a56e63 100644 Binary files a/src/vulkan/spv/rms_norm.spv and b/src/vulkan/spv/rms_norm.spv differ diff --git a/src/weights.rs b/src/weights.rs index 70653b97..04125eb0 100644 --- a/src/weights.rs +++ b/src/weights.rs @@ -64,6 +64,7 @@ impl<'a> Tensor<'a> { pub struct Weights<'a> { pub g: &'a Gguf, + pub aux_g: Option<&'a Gguf>, aliases: &'static [(&'static str, &'static str)], /// Loader-owned overrides: synthesized `ssm_beta`/`ssm_alpha` and LoRA-baked tensors. pub built: HashMap>, @@ -75,8 +76,12 @@ pub struct Weights<'a> { impl<'a> Weights<'a> { pub fn new(g: &'a Gguf) -> anyhow::Result> { + Self::new_with_aux(g, None) + } + + pub fn new_with_aux(g: &'a Gguf, aux_g: Option<&'a Gguf>) -> anyhow::Result> { let c = ModelConfig::from_gguf(g)?; - let mut w = Weights { g, aliases: c.aliases, built: HashMap::new(), baked: HashMap::new(), lora_baked: std::collections::HashSet::new() }; + let mut w = Weights { g, aux_g, aliases: c.aliases, built: HashMap::new(), baked: HashMap::new(), lora_baked: std::collections::HashSet::new() }; // llama.cpp src/models/qwen3next.cpp: per key head group, r = nv / nk beta rows, then r alpha rows. for il in 0..c.n_layer { let Ok(ba) = w.layer(il, "ssm_ba.weight") else { continue }; @@ -95,21 +100,54 @@ impl<'a> Weights<'a> { w.built.insert(format!("blk.{il}.ssm_beta.weight"), Tensor { data: leak(b), ..half }); w.built.insert(format!("blk.{il}.ssm_alpha.weight"), Tensor { data: leak(a), ..half }); } + // Standalone EAGLE3 draft GGUF: remap layer-0 draft weights onto blk.{n_layer}.* so they never shadow layer 0. + if let Some(ag) = aux_g.filter(|ag| ag.tensors.contains_key("fc.weight")) { + let nl = c.n_layer; + for (name, t) in &ag.tensors { + if matches!(name.as_str(), "d2t" | "t2d" | "token_embd.weight") { continue; } + let Ok(ty) = GgmlType::from_u32(t.ty) else { continue }; + let dst = match name.as_str() { + "fc.weight" => format!("blk.{nl}.eagle3.fc.weight"), + "output.weight" => format!("blk.{nl}.eagle3.output.weight"), + "output_norm.weight" => format!("blk.{nl}.nextn.shared_head_norm.weight"), + "blk.0.attn_norm.weight" => format!("blk.{nl}.nextn.enorm.weight"), + "blk.0.attn_norm_2.weight" => format!("blk.{nl}.nextn.hnorm.weight"), + "blk.0.ffn_norm.weight" => format!("blk.{nl}.post_attention_norm.weight"), + _ if name.starts_with("blk.0.") => format!("blk.{nl}.{}", &name["blk.0.".len()..]), + _ => continue, + }; + w.built.insert(dst, Tensor { + ty, + cols: t.dims[0] as usize, + rows: t.dims[1..].iter().product::() as usize, + data: ag.bytes(t), + }); + } + } Ok(w) } /// `name` with the architecture's alias for a canonical suffix tried second. pub fn get(&self, name: &str) -> anyhow::Result> { if let Some(t) = self.built.get(name) { return Ok(*t); } - let t = self.g.tensors.get(name) - .or_else(|| self.aliases.iter().find_map(|&(c, a)| self.g.tensors.get(&format!("{}{a}", name.strip_suffix(c)?)))) + let find_t = |g: &'a Gguf| -> Option<(&'a crate::gguf::TensorInfo, &'a Gguf)> { + if let Some(t) = g.tensors.get(name) { + return Some((t, g)); + } + if let Some(t) = self.aliases.iter().find_map(|&(c, a)| g.tensors.get(&format!("{}{a}", name.strip_suffix(c)?))) { + return Some((t, g)); + } + None + }; + let (t, owner_g) = find_t(self.g) + .or_else(|| self.aux_g.and_then(find_t)) .with_context(|| format!("missing tensor {name}"))?; let ty = GgmlType::from_u32(t.ty).with_context(|| format!("tensor {name}"))?; Ok(Tensor { ty, cols: t.dims[0] as usize, rows: t.dims[1..].iter().product::() as usize, - data: self.g.bytes(t), + data: owner_g.bytes(t), }) } @@ -161,7 +199,9 @@ impl<'a> Weights<'a> { } pub fn validate(&self, n_layer: usize) -> anyhow::Result { let mut nextn = true; - for t in self.g.tensors.values() { + let iter = self.g.tensors.values().chain(self.aux_g.into_iter().flat_map(|ag| ag.tensors.values())); + for t in iter { + if matches!(t.name.as_str(), "d2t" | "t2d") { continue; } if let Err(e) = GgmlType::from_u32(t.ty) { let il = t.name.strip_prefix("blk.").and_then(|r| r.split('.').next()).and_then(|i| i.parse::().ok()); if il.is_some_and(|il| il >= n_layer) { nextn = false; continue; } diff --git a/tests/gpu_model.rs b/tests/gpu_model.rs index c142db31..e7572178 100644 --- a/tests/gpu_model.rs +++ b/tests/gpu_model.rs @@ -216,6 +216,59 @@ fn mtp_spec_matches_greedy() { } } +/// Hybrid speculative decoding (`hybrid_spec_step`, combining N-gram bursts and MTP/EAGLE3 drafting) +/// matches plain greedy token sequences bit-for-bit across both N-gram hits and MTP fallbacks. +#[test] +fn hybrid_spec_matches_greedy() { + use rusting_strata::ngram::{NgramConfig, NgramDrafter, SpecType}; + let _gpu = common::gpu_lock(); + let Some(p) = common::model_path() else { return }; + let g = Gguf::open(&p).unwrap(); + let m = Model::load(&g).unwrap(); + if m.cfg.n_nextn == 0 { eprintln!("SKIP: no nextn layer"); return } + let tok = Tokenizer::from_gguf(&g).unwrap(); + for (bk, gpu) in common::backends() { + eprintln!("backend {bk}"); + let mut gm = GpuModel::load(&m, gpu, 512, 0, 512 << 20, 1).unwrap(); + gm.adapt_every = 0; + let ids = tok.encode( + "Repeat verbatim: fn alpha(x: i32) -> i32 { x + 1 }\nRepeat verbatim: fn alpha(x: i32) -> i32 {" + ).unwrap(); + let n = 64; + gm.reset().unwrap(); + let mut t = ops::argmax(&gm.prefill(&ids).unwrap()) as u32; + let mut plain = vec![]; + while plain.len() < n { + plain.push(t); + t = ops::argmax(&gm.step(t).unwrap()) as u32; + } + plain.truncate(n); + + gm.ngram = Some(NgramDrafter::new( + NgramConfig { + types: vec![SpecType::NgramSimple], + draft_max: 8, + size_n: 4, + size_m: 16, + n_vocab: m.cfg.n_vocab, + ..NgramConfig::default() + }, + Some(&tok), + ).unwrap()); + gm.drafts = (0, 0); + gm.reset().unwrap(); + let mut t = ops::argmax(&gm.prefill(&ids).unwrap()) as u32; + let (mut hybrid, mut row) = (vec![], 0); + while hybrid.len() < n { + let (toks, nx, r) = gm.hybrid_spec_step(t, row, &mut ops::Sampler::greedy()).unwrap(); + (t, row) = (nx, r); + hybrid.extend(toks); + } + hybrid.truncate(n); + assert_eq!(hybrid, plain, "hybrid_spec_step greedy parity ({bk})"); + } +} + /// `rollback` to the first token of a 2-token verify leaves the same state as stepping that token alone. #[test] fn rollback_matches_step() { diff --git a/tests/kernels.rs b/tests/kernels.rs index dbe219c8..80ce657e 100644 --- a/tests/kernels.rs +++ b/tests/kernels.rs @@ -752,7 +752,9 @@ fn attn_prefill_qsa_mask_all_ones_matches_unmasked() { unsafe { std::env::set_var("STRATA_VK_ATTN_BUDGET", "1000") }; gpu.attn_prefill_m(dq.addr(), dk.addr(), dv.addr(), 0, got_buf.addr(), step.addr(), t, nh, nkv, hd, scale, dpat.addr(), mw, 0, 0).unwrap(); - assert_eq!(down::(gpu, &got_buf), want, "{bk} attn_prefill sliced mask"); + let got = down::(gpu, &got_buf); + let e = rel_l2(&got, &want); + assert!(e <= 1e-6, "{bk} attn_prefill sliced mask: rel err {e}"); unsafe { std::env::remove_var("STRATA_VK_ATTN_BUDGET") }; } } diff --git a/tests/serve.rs b/tests/serve.rs index 3354f27f..57991465 100644 --- a/tests/serve.rs +++ b/tests/serve.rs @@ -15,7 +15,7 @@ fn prompts_match_llama_cpp() { for (case, body) in &reqs { let q = chat_request(body.dump().as_bytes()); match want[case].as_str() { - Some(w) => { let q = q.unwrap(); assert_eq!(tpl.render(&q.messages, &q.tools, q.think, false, true), w, "{name} {case}") } + Some(w) => { let q = q.unwrap(); assert_eq!(tpl.render(&q.messages, &q.tools, q.think.unwrap_or(false), false, true), w, "{name} {case}") } None => assert!(q.is_err(), "{name} {case} should be rejected"), } } @@ -134,7 +134,7 @@ fn parses_request() { assert_eq!((r.method.as_str(), r.path.as_str(), r.host.as_str(), r.json), ("POST", "/v1/chat/completions", "x", true)); let q = chat_request(&r.body).unwrap(); assert_eq!(q.messages, vec![Msg::new("user", "a\nb")]); - assert!(q.stream && !q.think); + assert!(q.stream && !q.think.unwrap_or(false)); assert_eq!(q.max_tokens, Some(7)); assert_eq!((q.temperature, q.top_k, q.top_p, q.seed), (Some(0.7), None, Some(0.8), Some(5))); assert!(chat_request(br#"{"messages":[]}"#).is_err());