cram_codec/
kvgpu.rs

1//! Host side of the GPU KV decoder: packs GPU-layout TKV streams
2//! ([`crate::kvcodec`], `gpu = true`) into the flat tables the device kernel
3//! reads (`tang-compute` `tkv_decode`, MSL source in
4//! `tang-compute/src/kernels/kvdec_msl.rs`). No GPU dependency here.
5//!
6//! One kernel thread per lane = one (stream, block, head) arithmetic-coded
7//! substream. Tables (all `u32`):
8//! - `stream`: every stream's bytes, concatenated, little-endian words;
9//! - `lanes`: [`LANE_WORDS`] words per lane: byte offset, byte length, first
10//!   token, tokens, channel-table base, output base (elements), output token
11//!   stride, scratch base (words), unit (`f32` bits), sink, sink divisor,
12//!   lane width (channels), RoPE flag, RoPE theta (`f32` bits), first RoPE
13//!   pair of the lane, head dim;
14//! - `chans`: 4 words per channel, in lane-local order per (stream, head,
15//!   channel group): mean (`i32`), rho, class, dense class index within the
16//!   group;
17//! - `prior`: [`crate::kvcodec::gpu_prior_table`];
18//! - scratch: `scratch_words` words of adaptive-model state.
19
20use crate::bitio::DecodeError;
21use crate::kvcodec::{gpu_prior_table, lane_chan, KvStream, GPU_CTX_MODELS};
22
23/// Words per lane descriptor.
24pub const LANE_WORDS: usize = 16;
25
26/// Flat decoder tables for a set of streams.
27#[derive(Default)]
28pub struct Batch {
29    /// Concatenated stream bytes as words.
30    pub stream: Vec<u32>,
31    /// Lane descriptors, [`LANE_WORDS`] each.
32    pub lanes: Vec<u32>,
33    /// Channel table, 4 words per channel.
34    pub chans: Vec<u32>,
35    /// Model scratch size in words.
36    pub scratch_words: usize,
37    bytes: Vec<u8>,
38}
39
40/// Where a stream's output goes: element index of (head, token) row start
41/// and the token stride, plus optional RoPE (K given pre-RoPE).
42pub struct Placement<'f> {
43    /// `(head, token) -> element offset` of that row's channel 0.
44    pub row: &'f dyn Fn(usize, usize) -> usize,
45    /// Apply half-split RoPE with this theta to each decoded token row.
46    pub rope_theta: Option<f32>,
47    /// Decode without writing the output (benchmark: the cost a fused
48    /// decode + attention kernel cannot avoid).
49    pub decode_only: bool,
50}
51
52impl Batch {
53    /// Empty batch.
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    /// Add one stream; every lane writes `row(head, t) + c` for its tokens.
59    /// Lanes covering `[t0, t1)` only (whole blocks) are added.
60    pub fn add(
61        &mut self,
62        bytes: &[u8],
63        place: &Placement,
64        t0: usize,
65        t1: usize,
66    ) -> Result<(), DecodeError> {
67        let s = KvStream::parse(bytes)?;
68        let tb = s
69            .gpu_tables()
70            .ok_or(DecodeError::Invalid("not a gpu-layout stream"))?;
71        let p = &s.params;
72        let base = self.bytes.len();
73        self.bytes.extend_from_slice(bytes);
74        while !self.bytes.len().is_multiple_of(4) {
75            self.bytes.push(0);
76        }
77        let chan0 = self.chans.len() / 4;
78        let (w, ng) = (p.lane_width(), p.lane_groups());
79        if w > 128 {
80            return Err(DecodeError::Invalid("gpu lane wider than 128 channels"));
81        }
82        // dense class index per (head, channel group), lane-local channel order
83        let mut ncls = vec![0usize; p.heads * ng];
84        for (hg, nc) in ncls.iter_mut().enumerate() {
85            let (hh, g) = (hg / ng, hg % ng);
86            let mut seen: Vec<u32> = Vec::new();
87            for i in 0..w {
88                let j = hh * p.hd + lane_chan(p.hd, w, g, i);
89                let cl = tb.class[j];
90                let di = seen.iter().position(|&v| v == cl).unwrap_or_else(|| {
91                    seen.push(cl);
92                    seen.len() - 1
93                });
94                self.chans
95                    .extend_from_slice(&[tb.mu[j] as u32, tb.rho[j], cl, di as u32]);
96            }
97            *nc = seen.len();
98        }
99        let ctx_words = 3 * GPU_CTX_MODELS;
100        for l in &tb.lanes {
101            let [off, len, pos0, n, hh, g] = *l;
102            let (pos0, n, hh, g) = (pos0 as usize, n as usize, hh as usize, g as usize);
103            let hg = hh * ng + g;
104            if pos0 + n <= t0 || pos0 >= t1 {
105                continue;
106            }
107            let row0 = (place.row)(hh, pos0);
108            let stride = if n > 1 {
109                (place.row)(hh, pos0 + 1) - row0
110            } else {
111                p.hd
112            };
113            self.lanes.extend_from_slice(&[
114                (base + off as usize) as u32,
115                len,
116                pos0 as u32,
117                n as u32,
118                (chan0 + hg * w) as u32,
119                row0 as u32,
120                stride as u32,
121                self.scratch_words as u32,
122                p.unit().to_bits(),
123                p.sink as u32,
124                p.sink_div,
125                w as u32,
126                place.rope_theta.is_some() as u32 | (place.decode_only as u32) << 1,
127                place.rope_theta.unwrap_or(0.0).to_bits(),
128                (g * w / 2) as u32,
129                p.hd as u32,
130            ]);
131            self.scratch_words += ncls[hg] * ctx_words;
132        }
133        self.stream = self
134            .bytes
135            .chunks(4)
136            .map(|w| u32::from_le_bytes([w[0], w[1], w[2], w[3]]))
137            .collect();
138        Ok(())
139    }
140
141    /// Number of lanes.
142    pub fn n_lanes(&self) -> usize {
143        self.lanes.len() / LANE_WORDS
144    }
145
146    /// The shared prior table.
147    pub fn prior() -> Vec<u32> {
148        gpu_prior_table()
149    }
150}
cram_codec/kvgpu.rs

cram_codec/
kvgpu.rs
