memra_engine/
round_stream.rs

1//! ROUND-STREAM: the model-generic device machinery for pre-issued M-round speculative
2//! bursts with zero per-round host readbacks (extracted from the qwen spec loop 2026-07-12
3//! so the gemma/next-model loops reuse it instead of re-growing their own).
4//!
5//! The pieces that live here are pure device-buffer plumbing — every model-specific thing
6//! (the draft-chain graph, the verify trunk, commit semantics) stays in the caller:
7//!   - `StreamBufs`: the per-burst device buffers (verify tokens, break, pending, ring,
8//!     accept counters, device position) sized from (k, m_rounds).
9//!   - `kv_len_ptr_table`: the per-layer `kvl.len_d` pointer table the device rollback
10//!     kernel walks (pointers are stable for the cache's lifetime — cache.rs note).
11//!   - `drain_ring`: the one host sync per M rounds — reads the ring and appends tokens.
12//!
13//! The kernels these feed (`spec_accept_greedy_dc`, `spec_ring_commit`,
14//! `spec_rollback_stream`, `spec_seed_gather`, `spec_assemble_verify`) are already
15//! model-generic in lib.rs; this module is the buffer/lifecycle half.
16
17use crate::cache::Cache;
18use crate::Engine;
19use cudarc::driver::CudaSlice;
20
21pub struct StreamBufs {
22    /// assembled verify tokens [k+1]
23    pub vtok_d: CudaSlice<u32>,
24    /// p-min break markers [2]
25    pub brk_d: CudaSlice<u32>,
26    /// pending (bonus-fold) token [1]
27    pub pend_d: CudaSlice<u32>,
28    /// last verify prediction [1]
29    pub last_pred_d: CudaSlice<u32>,
30    /// device position counter (rope/append base)
31    pub pos_ctr: CudaSlice<i32>,
32    /// round-start position (rollback anchor)
33    pub pos_start_d: CudaSlice<i32>,
34    /// committed-token ring [m*(k+1)+1] (slot 0 = count)
35    pub ring_d: CudaSlice<u32>,
36    /// device accept counters [2]
37    pub acc_d: CudaSlice<u32>,
38    pub m_rounds: usize,
39    pub k: usize,
40}
41
42impl StreamBufs {
43    pub fn new(e: &Engine, k: usize, m_rounds: usize) -> Result<Self, Box<dyn std::error::Error>> {
44        Ok(StreamBufs {
45            vtok_d: e.alloc_u32_zeroed(k + 1)?,
46            brk_d: e.alloc_u32_zeroed(2)?,
47            pend_d: e.alloc_u32_zeroed(1)?,
48            last_pred_d: e.alloc_u32_zeroed(1)?,
49            pos_ctr: e.htod_i32(&[0])?,
50            pos_start_d: e.htod_i32(&[0])?,
51            ring_d: e.alloc_u32_zeroed(m_rounds * (k + 1) + 1)?,
52            acc_d: e.alloc_u32_zeroed(2)?,
53            m_rounds,
54            k,
55        })
56    }
57
58    /// Drain the ring after a burst (THE one host sync per M rounds): returns the committed
59    /// tokens in order and resets nothing — the caller zeroes the ring count for the next
60    /// burst via `e.set_u32_one(&mut ring_d, 0)`.
61    pub fn drain_ring(&self, e: &Engine) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
62        let h = e.dtoh_u32(&self.ring_d)?;
63        let cnt = (h[0] as usize).min(self.ring_d.len() - 1);
64        Ok(h[1..1 + cnt].to_vec())
65    }
66}
67
68/// Per-layer `kvl.len_d` device-pointer table (+ the position counter appended when
69/// `pos_ctr` is given) for `spec_rollback_stream`. Pointers are stable for the cache's
70/// lifetime; 0 marks layers without KV (linear-attention / KV-shared).
71pub fn kv_len_ptr_table(
72    e: &Engine,
73    cache: &Cache,
74    pos_ctr: Option<&CudaSlice<i32>>,
75) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
76    use cudarc::driver::DevicePtr;
77    let mut ptrs: Vec<u64> = cache
78        .kv
79        .iter()
80        .map(|kv| match kv.as_ref() {
81            Some(kvl) => {
82                let __s_g = e.stream();
83                let (p, _g) = kvl.len_d.device_ptr(&__s_g);
84                p as u64
85            }
86            None => 0u64,
87        })
88        .collect();
89    if let Some(pc) = pos_ctr {
90        let __s_g = e.stream();
91        let (p, _g) = pc.device_ptr(&__s_g);
92        ptrs.push(p as u64);
93    }
94    Ok(e.htod_u64(&ptrs)?)
95}
memra_engine/round_stream.rs

memra_engine/
round_stream.rs
