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;
20use std::ops::Range;
21
22pub struct StreamBufs {
23    /// assembled verify tokens [k+1]
24    pub vtok_d: CudaSlice<u32>,
25    /// p-min break markers [2]
26    pub brk_d: CudaSlice<u32>,
27    /// pending (bonus-fold) token [1]
28    pub pend_d: CudaSlice<u32>,
29    /// last verify prediction [1]
30    pub last_pred_d: CudaSlice<u32>,
31    /// device position counter (rope/append base)
32    pub pos_ctr: CudaSlice<i32>,
33    /// round-start position (rollback anchor)
34    pub pos_start_d: CudaSlice<i32>,
35    /// committed-token ring [m*(k+1)+1] (slot 0 = count)
36    pub ring_d: CudaSlice<u32>,
37    /// device accept counters [2]
38    pub acc_d: CudaSlice<u32>,
39    pub m_rounds: usize,
40    pub k: usize,
41}
42
43impl StreamBufs {
44    pub fn new(e: &Engine, k: usize, m_rounds: usize) -> Result<Self, Box<dyn std::error::Error>> {
45        Ok(StreamBufs {
46            vtok_d: e.alloc_u32_zeroed(k + 1)?,
47            brk_d: e.alloc_u32_zeroed(2)?,
48            pend_d: e.alloc_u32_zeroed(1)?,
49            last_pred_d: e.alloc_u32_zeroed(1)?,
50            pos_ctr: e.htod_i32(&[0])?,
51            pos_start_d: e.htod_i32(&[0])?,
52            ring_d: e.alloc_u32_zeroed(m_rounds * (k + 1) + 1)?,
53            acc_d: e.alloc_u32_zeroed(2)?,
54            m_rounds,
55            k,
56        })
57    }
58
59    /// Drain the ring after a burst (THE one host sync per M rounds): returns the committed
60    /// tokens in order and resets nothing — the caller zeroes the ring count for the next
61    /// burst via `e.set_u32_one(&mut ring_d, 0)`.
62    pub fn drain_ring(&self, e: &Engine) -> Result<Vec<u32>, Box<dyn std::error::Error>> {
63        let h = e.dtoh_u32(&self.ring_d)?;
64        let cnt = (h[0] as usize).min(self.ring_d.len() - 1);
65        Ok(h[1..1 + cnt].to_vec())
66    }
67}
68
69/// Per-layer `kvl.len_d` device-pointer table (+ the position counter appended when
70/// `pos_ctr` is given) for `spec_rollback_stream`. Pointers are stable for the cache's
71/// lifetime; 0 marks layers without KV (linear-attention / KV-shared).
72pub fn kv_len_ptr_table(
73    e: &Engine,
74    cache: &Cache,
75    pos_ctr: Option<&CudaSlice<i32>>,
76) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
77    kv_len_ptr_table_range(e, cache, 0..cache.kv.len(), pos_ctr)
78}
79
80/// Stage-local twin of [`kv_len_ptr_table`]. Only pointers owned by `layers` enter the table,
81/// so a reconcile kernel launched through a PP stage's engine never dereferences another
82/// device's `len_d`. The returned table is dense over the requested range; callers pass
83/// `layers.len()` to the matching kernel.
84pub fn kv_len_ptr_table_range(
85    e: &Engine,
86    cache: &Cache,
87    layers: Range<usize>,
88    pos_ctr: Option<&CudaSlice<i32>>,
89) -> Result<CudaSlice<u64>, Box<dyn std::error::Error>> {
90    use cudarc::driver::DevicePtr;
91    assert!(layers.start <= layers.end && layers.end <= cache.kv.len());
92    let mut ptrs: Vec<u64> = cache
93        .kv
94        .get(layers)
95        .expect("validated KV layer range")
96        .iter()
97        .map(|kv| match kv.as_ref() {
98            Some(kvl) => {
99                let __s_g = e.stream();
100                let (p, _g) = kvl.len_d.device_ptr(&__s_g);
101                p as u64
102            }
103            None => 0u64,
104        })
105        .collect();
106    if let Some(pc) = pos_ctr {
107        let __s_g = e.stream();
108        let (p, _g) = pc.device_ptr(&__s_g);
109        ptrs.push(p as u64);
110    }
111    Ok(e.htod_u64(&ptrs)?)
112}
memra_engine/round_stream.rs

memra_engine/
round_stream.rs
