vrr/ruvector
2
0
Fork 0
mirror of https://github.com/ruvnet/RuVector.git synced 2026-07-09 17:28:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

perf(ruvllm): KV cache pre-allocation via scatter_set + greedy fast-path

Browse source 
Two decode-path optimizations:

1. KV cache pre-allocation (O(N²)→O(N) bandwidth across N decode steps)

Add KvLayerCache::GqaPrealloc { k, v: Tensor[b,kv_heads,max_seq,head_dim],
seq_len, max_seq }. When the cache holds a pre-allocated buffer, append uses
Tensor::scatter_set (candle 0.9 in-place op) instead of Tensor:🐱
  - Old: cat([past_k, k_cur], dim=2)  →  new [b,kv,N+1,hd] allocation + full copy
  - New: scatter_set(k_cur at pos N)  →  in-place write, O(kv_heads*head_dim)

MythosCache::with_prealloc(cfg, b, device, dtype) creates a cache with GQA
pre-allocated buffers. reset() resets seq_len (reuses the buffer).

2. Greedy fast-path in generate_sampled / generate_stream_sampled

When temperature=0 and no rep penalty, bypass sort_last_dim + topk transfer
(320B) and use last_argmax directly (4-byte scalar). Eliminates GPU sort for
the common greedy inference case.

Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruvnet 2026-06-18 14:19:47 -04:00
parent 7c6108bb03
commit 37fe37e5be
2 changed files with 129 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

53
crates/ruvllm/src/models/openmythos/attention.rs
View file

@ -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,
))
}
}

102
crates/ruvllm/src/models/openmythos/mod.rs
View file

@ -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<Self> {
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<Option<KvLayerCache>> {
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::<candle_core::Result<Vec<_>>>()?
} 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::<candle_core::Result<Vec<_>>>()?
} 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<u32> = 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<u32> = 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(())
}