logicaffeine_compile/codegen_cuda/lane/
weights.rs

1//! Where each constant of a compiled step comes from.
2//!
3//! A projection or a table the author declares in Q8_0 is read straight out of the checkpoint: the rows its binding's
4//! view picks, in the order it picks them, copied byte for byte — the program maps the checkpoint and uploads those
5//! rows, so a weight is never rewritten or quantized again. Every other constant — a norm weight, a rotary table, a
6//! small `f32` projection — the compiler computes once, as [`bind::values`] reads it, and ships beside the program.
7
8use crate::model::bind::{self, Source, View};
9use crate::model::gguf::Gguf;
10use crate::model::quant::GgmlType;
11use crate::model::spec::{ModelSpec, RopeScaling};
12
13use super::LaneError;
14
15/// `rows` consecutive checkpoint rows, the first at byte `offset` of the file.
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub struct Run {
18    pub offset: u64,
19    pub rows: usize,
20}
21
22/// A quantized matrix read straight from the checkpoint: `rows` rows of `cols` weights in `ty`, the device copy's rows
23/// the runs' rows in order. `transposed` when the program declares the matrix the other way round — a projection's
24/// rows are its outputs, the columns of the `in × out` matrix the program multiplies by.
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct Rows {
27    pub tensor: String,
28    pub ty: GgmlType,
29    pub rows: usize,
30    pub cols: usize,
31    pub transposed: bool,
32    pub runs: Vec<Run>,
33}
34
35/// How a constant reaches the device.
36#[derive(Clone, Debug, PartialEq)]
37pub enum Placement {
38    /// Quantized rows straight from the checkpoint.
39    Rows(Rows),
40    /// Values the compiler computes, in the layout the program declares (`rows × cols`, row-major).
41    Values(Vec<f32>),
42}
43
44/// The rows of a tensor of `n` rows a view picks, in order, and whether it transposes them — `None` for a view that
45/// reorders within a row.
46fn picked(view: &View, n: usize) -> Option<(Vec<usize>, bool)> {
47    match *view {
48        View::Same => Some(((0..n).collect(), false)),
49        View::Transposed => Some(((0..n).collect(), true)),
50        View::NeoxPairs { head_dim, from, dims, transpose: true } => Some((
51            (0..n).map(|r| (r / head_dim) * head_dim + bind::neox_source_from(r % head_dim, from, dims)).collect(),
52            true,
53        )),
54        View::NeoxPairs { transpose: false, .. } => None,
55        View::RowRange { first, count, transpose } => Some(((first..first + count).collect(), transpose)),
56        View::HeadPart { stride, offset, width, transpose } => {
57            Some(((0..n).filter(|r| (offset..offset + width).contains(&(r % stride))).collect(), transpose))
58        }
59        // Each head's block of rows in place: the rows are the head's outputs, read transposed head by head.
60        View::FusedPart { first, count, head_dim, dims, bias: false } => Some((
61            (0..count).map(|r| first + (r / head_dim) * head_dim + bind::neox_source_from(r % head_dim, 0, dims)).collect(),
62            true,
63        )),
64        View::FusedPart { bias: true, .. } => None,
65        View::HeadBlocks { stride, offset, width } => {
66            Some(((0..n).filter(|r| (offset..offset + width).contains(&(r % stride))).collect(), true))
67        }
68    }
69}
70
71/// Where constant `var` (of `len` elements) of a step authored from `spec` comes from. `q8` when the author declared
72/// it a quantized matrix — rows read in place, or values when the checkpoint keeps it in a float format; `positions`
73/// the rows of the rotary caches.
74pub fn place(
75    spec: &ModelSpec,
76    g: &mut Gguf,
77    var: &str,
78    len: usize,
79    q8: bool,
80    positions: usize,
81) -> Result<Placement, LaneError> {
82    if let Some(table) = rotary_table(spec, g, var, positions)? {
83        if table.len() != len {
84            return Err(LaneError::Shape { var: var.to_string(), declared: len, bound: table.len() });
85        }
86        return Ok(Placement::Values(table));
87    }
88    let binding = match bind::bind(spec, var) {
89        Ok(Source::Weight(b)) => b,
90        Ok(Source::State) => return Err(LaneError::Unbound(format!("{var} is carried state, not a constant"))),
91        Err(bind::Unbound(v)) => return Err(LaneError::Unbound(v)),
92    };
93    let info = g.info(&binding.tensor).ok_or_else(|| LaneError::Unbound(binding.tensor.clone()))?.clone();
94    let ty = g.tensor_type(&binding.tensor).map_err(|e| LaneError::Io(e.to_string()))?;
95    // A matrix the checkpoint keeps in a float format is read as its values (the f32 kernels take it); every other
96    // format is read as rows in place.
97    let float = matches!(ty, GgmlType::F32 | GgmlType::F16 | GgmlType::Bf16);
98    if q8 && !float {
99        if !(ty == GgmlType::Q8_0 || super::kernels::decodes(ty)) {
100            return Err(LaneError::Format {
101                tensor: binding.tensor,
102                ty: ty.name(),
103                why: "a matrix the lane reads in place is Q8_0 or a format the block library decodes; this one has no decoder",
104            });
105        }
106        let cols = info.dims.first().copied().unwrap_or(1) as usize;
107        let n = info.elements() / cols.max(1);
108        let Some((picked, transposed)) = picked(&binding.view, n) else {
109            return Err(LaneError::Format {
110                tensor: binding.tensor,
111                ty: ty.name(),
112                why: "a view that reorders within a row cannot be read as whole quantized rows",
113            });
114        };
115        if picked.len() * cols != len {
116            return Err(LaneError::Shape { var: var.to_string(), declared: len, bound: picked.len() * cols });
117        }
118        let row_bytes = ty.bytes_for(cols) as u64;
119        let base = g.data_start() + info.offset;
120        let mut runs: Vec<Run> = Vec::new();
121        let mut last: Option<usize> = None;
122        for &r in &picked {
123            match (runs.last_mut(), last) {
124                (Some(run), Some(prev)) if r == prev + 1 => run.rows += 1,
125                _ => runs.push(Run { offset: base + r as u64 * row_bytes, rows: 1 }),
126            }
127            last = Some(r);
128        }
129        return Ok(Placement::Rows(Rows {
130            tensor: binding.tensor,
131            ty,
132            rows: picked.len(),
133            cols,
134            transposed,
135            runs,
136        }));
137    }
138    let values = bind::values(&binding, g).map_err(|e| LaneError::Io(e.to_string()))?;
139    if values.len() != len {
140        return Err(LaneError::Shape { var: var.to_string(), declared: len, bound: values.len() });
141    }
142    Ok(Placement::Values(values.into_iter().map(|v| v as f32).collect()))
143}
144
145/// A rotary table's values when `var` names one ([`bind::RotaryTable`]): a decode step's caches of `positions` rows, or
146/// a forward's one row (position 0), of the rotation the name gives.
147fn rotary_table(spec: &ModelSpec, g: &mut Gguf, var: &str, positions: usize) -> Result<Option<Vec<f32>>, LaneError> {
148    let Some(table) = bind::RotaryTable::named(var) else { return Ok(None) };
149    let rows = if table.cached { positions } else { 1 };
150    let rope = *bind::rotations(spec)
151        .get(table.rotation)
152        .ok_or_else(|| LaneError::Unbound(format!("{var}: the model has no rotation {}", table.rotation)))?;
153    let factors = match rope.scaling {
154        RopeScaling::Factors => Some(g.float_tensor("rope_freqs.weight").map_err(|e| LaneError::Io(e.to_string()))?),
155        _ => None,
156    };
157    let factors: Option<Vec<f64>> = factors.map(|f| f.into_iter().map(f64::from).collect());
158    let (c, s) = bind::rope_tables(rope, rows, factors.as_deref());
159    Ok(Some((if table.cos { c } else { s }).into_iter().map(|v| v as f32).collect()))
160}
logicaffeine_compile/codegen_cuda/lane/weights.rs

logicaffeine_compile/codegen_cuda/lane/
weights.rs
