logicaffeine_compile/spatial/
handoff.rs

1//! What one chip hands another when they share a model's work.
2//!
3//! A prompt prefilled on one chip and decoded on another — a GPU's prefill handed to an RDU's decode, as SambaNova serves
4//! MiniMax, or the reverse — and a model whose layers are split over two chips meet at the same place: the state a
5//! decode step carries from one token to the next. Every attention layer's key and value cache, every gated-delta
6//! layer's memory and conv rings: the tensors the step reads and writes back in place, less those it assembles whole
7//! ([`carried_state`]).
8//!
9//! ★ THE CONTRACT IS THE AUTHOR'S. Every backend lowers the one program `model::author` writes, and that program
10//! declares each carried tensor once, by name and shape. The prefill program declares the same tensors the decode
11//! program reads, so the hand-off is those tensors' values, row-major as declared — whatever tiles, pages or arenas a
12//! backend lays them out in, it reads and writes this logical order at its edge. [`contract`] reads the declarations
13//! off the decode step's layer programs; a bundle's manifest states it, so a host on either side knows what to send;
14//! [`Contract::check`] refuses a hand-off that does not match it, by name.
15
16use std::collections::BTreeMap;
17
18use crate::model::program::{written, Cut};
19use crate::model::spec::ModelSpec;
20use crate::tensor_graph::front::carried_state;
21use crate::tensor_graph::graph::{Storage, TensorGraph};
22
23/// One tensor a decode step carries: its name in the authored programs, the layer it belongs to, and its declared shape.
24#[derive(Clone, Debug, PartialEq, Eq)]
25pub struct Carried {
26    pub name: String,
27    pub layer: usize,
28    pub rows: usize,
29    pub cols: usize,
30}
31
32/// The state a model's decode step carries over a cache of `ctx` positions, layer by layer.
33#[derive(Clone, Debug, PartialEq, Eq)]
34pub struct Contract {
35    pub ctx: usize,
36    /// The width of the hidden state a layer hands the next — what a model split over chips sends at its cut.
37    pub hidden: usize,
38    pub carried: Vec<Carried>,
39}
40
41/// A prompt's state after `position` tokens: each carried tensor's values by name, row-major as declared.
42#[derive(Clone, Debug, PartialEq)]
43pub struct Handoff {
44    pub position: usize,
45    pub state: BTreeMap<String, Vec<f64>>,
46}
47
48/// Why a hand-off does not meet its contract, or why a model has none.
49#[derive(Clone, Debug, PartialEq, Eq)]
50pub enum HandoffError {
51    /// A layer's decode program does not walk.
52    Program { layer: usize, why: String },
53    /// A carried tensor the hand-off leaves out.
54    Missing(String),
55    /// A tensor the contract does not carry.
56    Unknown(String),
57    /// A carried tensor of the wrong size.
58    Length { name: String, want: usize, got: usize },
59    /// A position whose next token's row the cache does not have: a decode after `position` tokens writes row
60    /// `position + 1` of `ctx`.
61    Position { position: usize, ctx: usize },
62}
63
64impl std::fmt::Display for HandoffError {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        match self {
67            HandoffError::Program { layer, why } => write!(f, "layer {layer}'s decode program: {why}"),
68            HandoffError::Missing(n) => write!(f, "the hand-off leaves out the carried `{n}`"),
69            HandoffError::Unknown(n) => write!(f, "the hand-off carries `{n}`, which no decode step reads"),
70            HandoffError::Length { name, want, got } => write!(f, "`{name}` is {want} values in the contract and {got} in the hand-off"),
71            HandoffError::Position { position, ctx } => {
72                write!(f, "after {position} tokens the next decode writes row {} of a cache of {ctx}", position + 1)
73            }
74        }
75    }
76}
77
78/// The carried state of `spec`'s decode step over `ctx` positions: every layer's, in layer order and, within a layer, in
79/// the order its program first touches it.
80pub fn contract(spec: &ModelSpec, ctx: usize) -> Result<Contract, HandoffError> {
81    let mut carried = Vec::new();
82    for layer in 0..spec.layers.len() {
83        let ((src, _, _), _) = written(spec, Cut::DecodeLayer { layer, ctx });
84        let lifted = TensorGraph::lift_source(&src, &|_| Storage::Values).map_err(|e| HandoffError::Program { layer, why: format!("{e:?}") })?;
85        let rp = &lifted.program;
86        for name in carried_state(rp) {
87            let (rows, cols) = rp.shapes.get(&name).copied().ok_or_else(|| HandoffError::Program {
88                layer,
89                why: format!("carries `{name}`, which it does not declare"),
90            })?;
91            carried.push(Carried { name, layer, rows, cols });
92        }
93    }
94    Ok(Contract { ctx, hidden: spec.d, carried })
95}
96
97impl Contract {
98    /// The carried tensors of layer `layer`.
99    pub fn of_layer(&self, layer: usize) -> impl Iterator<Item = &Carried> {
100        self.carried.iter().filter(move |c| c.layer == layer)
101    }
102
103    /// Whether `h` is this contract's: every carried tensor, at its size, nothing else, and a next row in the cache.
104    pub fn check(&self, h: &Handoff) -> Result<(), HandoffError> {
105        if let Some(name) = h.state.keys().find(|n| !self.carried.iter().any(|c| &c.name == *n)) {
106            return Err(HandoffError::Unknown(name.clone()));
107        }
108        for c in &self.carried {
109            let values = h.state.get(&c.name).ok_or_else(|| HandoffError::Missing(c.name.clone()))?;
110            if values.len() != c.rows * c.cols {
111                return Err(HandoffError::Length { name: c.name.clone(), want: c.rows * c.cols, got: values.len() });
112            }
113        }
114        if h.position >= self.ctx {
115            return Err(HandoffError::Position { position: h.position, ctx: self.ctx });
116        }
117        Ok(())
118    }
119
120    /// The hand-off a run leaves after `position` tokens: the carried tensors out of `tensors` (a run's every tensor, by
121    /// name), checked.
122    pub fn take(&self, position: usize, tensors: &BTreeMap<String, Vec<f64>>) -> Result<Handoff, HandoffError> {
123        let state = self
124            .carried
125            .iter()
126            .map(|c| tensors.get(&c.name).map(|v| (c.name.clone(), v.clone())).ok_or_else(|| HandoffError::Missing(c.name.clone())))
127            .collect::<Result<BTreeMap<_, _>, _>>()?;
128        let h = Handoff { position, state };
129        self.check(&h)?;
130        Ok(h)
131    }
132
133    /// The contract as a manifest states it.
134    pub fn to_json(&self) -> serde_json::Value {
135        serde_json::json!({
136            "ctx": self.ctx,
137            "hidden": self.hidden,
138            "carried": self.carried.iter().map(|c| serde_json::json!({ "name": c.name, "layer": c.layer, "rows": c.rows, "cols": c.cols })).collect::<Vec<_>>(),
139        })
140    }
141}
logicaffeine_compile/spatial/handoff.rs

logicaffeine_compile/spatial/
handoff.rs
