diff --git a/crates/ruvllm/src/models/openmythos/attention.rs b/crates/ruvllm/src/models/openmythos/attention.rs index 5141da7fa..ed2c7fdbd 100644 --- a/crates/ruvllm/src/models/openmythos/attention.rs +++ b/crates/ruvllm/src/models/openmythos/attention.rs @@ -12,7 +12,16 @@ use crate::error::Result; #[derive(Clone)] pub enum KvLayerCache { /// GQA: rotated keys and values `[b, kv_heads, len, head_dim]`. + /// Grows via `Tensor::cat` on each decode step (legacy path). Gqa { k: Tensor, v: Tensor }, + /// GQA with pre-allocated buffers — uses `scatter_set` for O(1) per-step + /// appends instead of O(N) cat copies. Allocated up to `max_seq` positions. + GqaPrealloc { + k: Tensor, // [b, kv_heads, max_seq, head_dim] — full buffer + v: Tensor, + seq_len: usize, // positions 0..seq_len are valid + max_seq: usize, + }, /// MLA: compressed latent `[b, len, kv_lora_rank]` and rotated shared /// rope keys `[b, len, qk_rope_head_dim]`. Mla { c_kv: Tensor, k_rope: Tensor }, @@ -23,6 +32,7 @@ impl KvLayerCache { pub fn len(&self) -> usize { match self { KvLayerCache::Gqa { k, .. } => k.dim(2).unwrap_or(0), + KvLayerCache::GqaPrealloc { seq_len, .. } => *seq_len, KvLayerCache::Mla { c_kv, .. } => c_kv.dim(1).unwrap_or(0), } } @@ -132,13 +142,37 @@ impl GqaAttention { let q = apply_rope(&q, cos, sin)?; let k_cur = apply_rope(&k, cos, sin)?; - // Concatenate with past (already-rotated) keys/values. - let (k_full, v_full) = match past { - Some(KvLayerCache::Gqa { k: pk, v: pv }) => ( - Tensor::cat(&[pk, &k_cur], 2).map_err(cand)?, - Tensor::cat(&[pv, &v], 2).map_err(cand)?, - ), - _ => (k_cur, v), + // Accumulate KV: two paths depending on cache variant. + let (k_full, v_full, new_cache) = match past { + // Pre-allocated: scatter_set is O(new_data) not O(total); no new tensor. + Some(KvLayerCache::GqaPrealloc { k: buf_k, v: buf_v, seq_len, max_seq }) => { + // Index tensor: all positions write to `seq_len` along dim 2. + let idx = Tensor::full(*seq_len as u32, k_cur.shape(), k_cur.device()) + .map_err(cand)?; + buf_k.scatter_set(&idx, &k_cur, 2).map_err(cand)?; + buf_v.scatter_set(&idx, &v, 2).map_err(cand)?; + let new_seq = seq_len + seq; + let k_view = buf_k.narrow(2, 0, new_seq).map_err(cand)?; + let v_view = buf_v.narrow(2, 0, new_seq).map_err(cand)?; + let cache = KvLayerCache::GqaPrealloc { + k: buf_k.clone(), + v: buf_v.clone(), + seq_len: new_seq, + max_seq: *max_seq, + }; + (k_view, v_view, cache) + } + // Legacy cat path (first call or non-preallocated cache). + Some(KvLayerCache::Gqa { k: pk, v: pv }) => { + let k_f = Tensor::cat(&[pk, &k_cur], 2).map_err(cand)?; + let v_f = Tensor::cat(&[pv, &v], 2).map_err(cand)?; + let cache = KvLayerCache::Gqa { k: k_f.clone(), v: v_f.clone() }; + (k_f, v_f, cache) + } + _ => { + let cache = KvLayerCache::Gqa { k: k_cur.clone(), v: v.clone() }; + (k_cur, v, cache) + } }; let n_rep = self.n_heads / self.n_kv_heads; @@ -164,10 +198,7 @@ impl GqaAttention { let out = self.o_proj.forward(&ctx).map_err(cand)?; Ok(( out, - KvLayerCache::Gqa { - k: k_full, - v: v_full, - }, + new_cache, )) } } diff --git a/crates/ruvllm/src/models/openmythos/mod.rs b/crates/ruvllm/src/models/openmythos/mod.rs index 07e3f3773..c800c81cc 100644 --- a/crates/ruvllm/src/models/openmythos/mod.rs +++ b/crates/ruvllm/src/models/openmythos/mod.rs @@ -84,14 +84,63 @@ impl MythosCache { self.seq_len == 0 } + /// Create a cache with pre-allocated GQA KV buffers to avoid per-step + /// `Tensor::cat` growth (O(N²) → O(N) bandwidth across N decode steps). + /// + /// Pre-allocates `[b, kv_heads, max_seq, head_dim]` for every GQA layer. + /// The first forward call fills positions 0..prompt_len; subsequent single- + /// token decode steps use `scatter_set` to append at O(head_dim) cost. + pub fn with_prealloc( + cfg: &MythosConfig, + b: usize, + device: &candle_core::Device, + dtype: candle_core::DType, + ) -> candle_core::Result { + let kv_heads = cfg.n_kv_heads; + let head_dim = cfg.head_dim(); + let max_seq = cfg.max_seq_len; + let mk_buf = |_| -> candle_core::Result> { + let k = candle_core::Tensor::zeros((b, kv_heads, max_seq, head_dim), dtype, device)?; + let v = candle_core::Tensor::zeros((b, kv_heads, max_seq, head_dim), dtype, device)?; + Ok(Some(KvLayerCache::GqaPrealloc { k, v, seq_len: 0, max_seq })) + }; + // MLA layers share the same pre-alloc approach but use a different shape; + // for now only pre-alloc for GQA (AttnType::Gqa). + let prelude = if cfg.attn_type == AttnType::Gqa { + (0..cfg.prelude_layers).map(|_| mk_buf(())).collect::>>()? + } else { + vec![None; cfg.prelude_layers] + }; + let recurrent = if cfg.attn_type == AttnType::Gqa { mk_buf(())? } else { None }; + let coda = if cfg.attn_type == AttnType::Gqa { + (0..cfg.coda_layers).map(|_| mk_buf(())).collect::>>()? + } else { + vec![None; cfg.coda_layers] + }; + Ok(Self { prelude, recurrent, coda, seq_len: 0 }) + } + /// Clear all cached state. pub fn reset(&mut self) { for c in &mut self.prelude { - *c = None; + // For GqaPrealloc, reset seq_len (the buffer is reused). + if let Some(KvLayerCache::GqaPrealloc { seq_len, .. }) = c { + *seq_len = 0; + } else { + *c = None; + } + } + if let Some(KvLayerCache::GqaPrealloc { seq_len, .. }) = &mut self.recurrent { + *seq_len = 0; + } else { + self.recurrent = None; } - self.recurrent = None; for c in &mut self.coda { - *c = None; + if let Some(KvLayerCache::GqaPrealloc { seq_len, .. }) = c { + *seq_len = 0; + } else { + *c = None; + } } self.seq_len = 0; } @@ -354,9 +403,12 @@ impl OpenMythos { if prompt_ids.is_empty() { return Err(RuvLLMError::Generation("empty prompt".into())); } - // top_k_transfer: how many candidates to sort on GPU and transfer. - // Using top_k when set (40-200 typical); capped at 512 for top-p-only. - let top_k_transfer = if sampling.top_k > 0 { sampling.top_k } else { 512.min(self.cfg.vocab_size) }; + // Greedy with no rep penalty: bypass sort/transfer entirely — use on-device argmax. + let is_greedy = sampling.temperature <= 0.0 + && ((sampling.repetition_penalty - 1.0).abs() <= f32::EPSILON + || sampling.repetition_window == 0); + let top_k_transfer = + if sampling.top_k > 0 { sampling.top_k } else { 512.min(self.cfg.vocab_size) }; let mut sampler = Sampler::new(sampling); let mut cache = MythosCache::new(&self.cfg); let mut history: Vec = prompt_ids.to_vec(); @@ -365,8 +417,12 @@ impl OpenMythos { Tensor::from_slice(prompt_ids, (1, prompt_ids.len()), &self.device) .map_err(cand)?; let logits = self.forward_cached(&prompt, &mut cache, n_loops)?; - let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; - let mut next = sampler.sample_topk(&vals, &idxs, &history); + let mut next = if is_greedy { + self.last_argmax(&logits)? + } else { + let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; + sampler.sample_topk(&vals, &idxs, &history) + }; let mut out = Vec::with_capacity(max_new_tokens); for _ in 0..max_new_tokens { @@ -377,8 +433,12 @@ impl OpenMythos { } let step = Tensor::from_slice(&[next], (1, 1), &self.device).map_err(cand)?; let logits = self.forward_cached(&step, &mut cache, n_loops)?; - let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; - next = sampler.sample_topk(&vals, &idxs, &history); + next = if is_greedy { + self.last_argmax(&logits)? + } else { + let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; + sampler.sample_topk(&vals, &idxs, &history) + }; } Ok(out) } @@ -401,7 +461,11 @@ impl OpenMythos { if prompt_ids.is_empty() { return Err(RuvLLMError::Generation("empty prompt".into())); } - let top_k_transfer = if sampling.top_k > 0 { sampling.top_k } else { 512.min(self.cfg.vocab_size) }; + let is_greedy = sampling.temperature <= 0.0 + && ((sampling.repetition_penalty - 1.0).abs() <= f32::EPSILON + || sampling.repetition_window == 0); + let top_k_transfer = + if sampling.top_k > 0 { sampling.top_k } else { 512.min(self.cfg.vocab_size) }; let mut sampler = Sampler::new(sampling); let mut cache = MythosCache::new(&self.cfg); let mut history: Vec = prompt_ids.to_vec(); @@ -410,8 +474,12 @@ impl OpenMythos { Tensor::from_slice(prompt_ids, (1, prompt_ids.len()), &self.device) .map_err(cand)?; let logits = self.forward_cached(&prompt, &mut cache, n_loops)?; - let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; - let mut next = sampler.sample_topk(&vals, &idxs, &history); + let mut next = if is_greedy { + self.last_argmax(&logits)? + } else { + let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; + sampler.sample_topk(&vals, &idxs, &history) + }; for _ in 0..max_new_tokens { if !on_token(next) { @@ -423,8 +491,12 @@ impl OpenMythos { } let step = Tensor::from_slice(&[next], (1, 1), &self.device).map_err(cand)?; let logits = self.forward_cached(&step, &mut cache, n_loops)?; - let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; - next = sampler.sample_topk(&vals, &idxs, &history); + next = if is_greedy { + self.last_argmax(&logits)? + } else { + let (vals, idxs) = self.last_logits_topk(&logits, top_k_transfer)?; + sampler.sample_topk(&vals, &idxs, &history) + }; } Ok(()) }