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