kime_model/
laya.rs

1//! The laya compat family: Laya 0.3.7's `DecisionModel` over a ModernBERT encoder.
2//!
3//! The graph, from spec/05-model.md, with `d` the hidden size:
4//!
5//! ```text
6//! h   = LN(E[ids])                                     encoder.embeddings
7//! for each encoder layer i:
8//!     h = h + Wo(Attn(RoPE(q), RoPE(k), v))            [q k v] = Wqkv LN(h), no LN on layer 0
9//!     h = h + Wo(GELU(a) * g)                          [a g] = Wi LN(h)
10//! h   = LN(h)                                          encoder.final_norm
11//! h   = h + type_emb[qtype]
12//! for each head layer:                                 PyTorch TransformerEncoderLayer, norm first
13//!     h = h + out_proj(Attn(in_proj LN1(h)))           full attention, biases, no RoPE
14//!     h = h + linear2(ReLU(linear1(LN2(h))))
15//! m   = h[markers]
16//! z   = scorer.3(GELU(scorer.1(LN(m))))                one logit per option
17//! act = act_head.2(GELU(act_head.0([h[0], top1, top1 - top2, entropy, k / 255])))
18//! ```
19//!
20//! Encoder layers have no biases and LayerNorms without bias. Attention is global on every third
21//! layer and a 128 token sliding window on the rest, and GELU is the exact erf form throughout.
22//! [`LayaGraph`] binds each weight in that description to a tensor in a checkpoint, and binding is
23//! also how a checkpoint is checked: every missing tensor, extra tensor and wrong shape is named.
24
25use serde_json::Value;
26
27use kime_tensor::plan::{Epilogue, Graph, Op, Rows};
28
29use crate::error::{Error, Result};
30use crate::tensors::Tensors;
31
32/// Head dimension of every attention in the family.
33pub const HEAD_DIM: usize = 64;
34
35/// PyTorch's LayerNorm default, which the decision head and the scorer use.
36pub const TORCH_EPS: f64 = 1e-5;
37
38/// Width of the act head's hidden layer.
39pub const ACT_HIDDEN: usize = 256;
40
41/// The encoder's shape, from `encoder/config.json` (a HF ModernBERT config).
42#[derive(Debug, Clone, PartialEq)]
43pub struct EncoderConfig {
44    /// Hidden size.
45    pub d: usize,
46    /// GeGLU inner size. `mlp.Wi` has twice as many rows.
47    pub inter: usize,
48    /// Attention heads.
49    pub heads: usize,
50    /// Encoder layers.
51    pub layers: usize,
52    /// Vocabulary size, the rows of the embedding table.
53    pub vocab: usize,
54    /// Whether each layer attends globally, false meaning the sliding window.
55    pub global: Vec<bool>,
56    /// Sliding window width in tokens, 128 for both published encoders.
57    pub window: usize,
58    /// RoPE base on global layers.
59    pub rope_global: f64,
60    /// RoPE base on sliding window layers.
61    pub rope_local: f64,
62    /// LayerNorm epsilon.
63    pub norm_eps: f64,
64}
65
66/// Laya's own config, `rl_agent_config.json`.
67#[derive(Debug, Clone, PartialEq)]
68pub struct AgentConfig {
69    /// The HF id of the encoder the checkpoint was trained from.
70    pub encoder: String,
71    /// Transformer layers in the decision head.
72    pub head_layers: usize,
73    /// Longest sequence.
74    pub max_len: usize,
75    /// Budget for the question head and its options.
76    pub head_max_len: usize,
77    /// Per type temperatures, choice, score, noul.
78    pub temperature: [f64; 3],
79    /// Temperatures per type and option count bucket, such as `choice:3-5`.
80    pub temperature_by_options: Vec<(String, f64)>,
81}
82
83/// Both configs of a compat checkpoint.
84#[derive(Debug, Clone, PartialEq)]
85pub struct LayaSpec {
86    /// The name kime serves it under, `laya` or `laya-multilingual` for the published ones.
87    pub id: String,
88    /// The encoder.
89    pub encoder: EncoderConfig,
90    /// The decision head.
91    pub agent: AgentConfig,
92}
93
94fn get<'a>(v: &'a Value, key: &str, file: &str) -> Result<&'a Value> {
95    v.get(key).ok_or_else(|| Error::format(format!("{file}: missing {key:?}")))
96}
97
98fn usize_of(v: &Value, key: &str, file: &str) -> Result<usize> {
99    get(v, key, file)?
100        .as_u64()
101        .and_then(|n| usize::try_from(n).ok())
102        .ok_or_else(|| Error::format(format!("{file}: {key:?} is not a non negative integer")))
103}
104
105fn f64_of(v: &Value, key: &str, file: &str) -> Result<f64> {
106    get(v, key, file)?
107        .as_f64()
108        .ok_or_else(|| Error::format(format!("{file}: {key:?} is not a number")))
109}
110
111impl EncoderConfig {
112    /// Reads a HF ModernBERT config. Both the transformers 5 layout (`layer_types`,
113    /// `rope_parameters`) and the older one (`global_attn_every_n_layers`, `global_rope_theta`,
114    /// `local_rope_theta`) are accepted, since checkpoints in the wild have either.
115    ///
116    /// # Errors
117    ///
118    /// [`Error::Format`] naming the missing or bad field.
119    pub fn from_json(v: &Value) -> Result<Self> {
120        const F: &str = "encoder/config.json";
121        let d = usize_of(v, "hidden_size", F)?;
122        let heads = usize_of(v, "num_attention_heads", F)?;
123        let layers = usize_of(v, "num_hidden_layers", F)?;
124        if heads == 0 || d != heads * HEAD_DIM {
125            return Err(Error::format(format!(
126                "{F}: hidden_size {d} is not {heads} heads of {HEAD_DIM}"
127            )));
128        }
129        let global = match v.get("layer_types").and_then(Value::as_array) {
130            Some(types) => types
131                .iter()
132                .map(|t| match t.as_str() {
133                    Some("full_attention") => Ok(true),
134                    Some("sliding_attention") => Ok(false),
135                    _ => Err(Error::format(format!("{F}: unknown layer type {t}"))),
136                })
137                .collect::<Result<Vec<_>>>()?,
138            None => {
139                let every = usize_of(v, "global_attn_every_n_layers", F)?.max(1);
140                (0..layers).map(|i| i % every == 0).collect()
141            }
142        };
143        if global.len() != layers {
144            return Err(Error::format(format!(
145                "{F}: {} layer types for {layers} layers",
146                global.len()
147            )));
148        }
149        let (rope_global, rope_local) = match v.get("rope_parameters") {
150            Some(p) => (
151                f64_of(get(p, "full_attention", F)?, "rope_theta", F)?,
152                f64_of(get(p, "sliding_attention", F)?, "rope_theta", F)?,
153            ),
154            None => (f64_of(v, "global_rope_theta", F)?, f64_of(v, "local_rope_theta", F)?),
155        };
156        Ok(Self {
157            d,
158            inter: usize_of(v, "intermediate_size", F)?,
159            heads,
160            layers,
161            vocab: usize_of(v, "vocab_size", F)?,
162            global,
163            window: usize_of(v, "local_attention", F)?,
164            rope_global,
165            rope_local,
166            norm_eps: v.get("norm_eps").and_then(Value::as_f64).unwrap_or(1e-5),
167        })
168    }
169}
170
171impl AgentConfig {
172    /// Reads `rl_agent_config.json`.
173    ///
174    /// # Errors
175    ///
176    /// [`Error::Format`] naming the missing or bad field.
177    pub fn from_json(v: &Value) -> Result<Self> {
178        const F: &str = "rl_agent_config.json";
179        let t = get(v, "temperature", F)?
180            .as_array()
181            .filter(|a| a.len() == 3)
182            .and_then(|a| Some([a[0].as_f64()?, a[1].as_f64()?, a[2].as_f64()?]))
183            .ok_or_else(|| Error::format(format!("{F}: temperature must be three numbers")))?;
184        let by_options = match v.get("temperature_by_options") {
185            None | Some(Value::Null) => Vec::new(),
186            Some(Value::Object(m)) => m
187                .iter()
188                .map(|(k, t)| {
189                    t.as_f64().map(|t| (k.clone(), t)).ok_or_else(|| {
190                        Error::format(format!("{F}: temperature_by_options[{k:?}] is not a number"))
191                    })
192                })
193                .collect::<Result<_>>()?,
194            Some(_) => {
195                return Err(Error::format(format!("{F}: temperature_by_options is not an object")));
196            }
197        };
198        Ok(Self {
199            encoder: get(v, "encoder", F)?.as_str().unwrap_or_default().to_string(),
200            head_layers: usize_of(v, "head_layers", F)?,
201            max_len: usize_of(v, "max_len", F)?,
202            head_max_len: usize_of(v, "head_max_len", F)?,
203            temperature: t,
204            temperature_by_options: by_options,
205        })
206    }
207}
208
209impl LayaSpec {
210    /// Builds the spec from the two config files. The id follows the encoder: the two published
211    /// checkpoints get their Laya names and anything else is `laya-custom`.
212    ///
213    /// # Errors
214    ///
215    /// [`Error::Format`] from either config.
216    pub fn from_json(agent: &Value, encoder: &Value) -> Result<Self> {
217        let agent = AgentConfig::from_json(agent)?;
218        let encoder = EncoderConfig::from_json(encoder)?;
219        let id = match agent.encoder.as_str() {
220            "answerdotai/ModernBERT-large" => "laya",
221            "jhu-clsp/mmBERT-base" => "laya-multilingual",
222            _ => "laya-custom",
223        };
224        Ok(Self { id: id.to_string(), encoder, agent })
225    }
226
227    /// Every tensor the checkpoint must hold, with its shape, in Laya's module order.
228    #[must_use]
229    pub fn expected(&self) -> Vec<(String, Vec<usize>)> {
230        let e = &self.encoder;
231        let d = e.d;
232        let mut out: Vec<(String, Vec<usize>)> = Vec::new();
233        let mut put = |name: String, shape: &[usize]| out.push((name, shape.to_vec()));
234        put("encoder.embeddings.tok_embeddings.weight".into(), &[e.vocab, d]);
235        put("encoder.embeddings.norm.weight".into(), &[d]);
236        for i in 0..e.layers {
237            let p = format!("encoder.layers.{i}");
238            if i > 0 {
239                put(format!("{p}.attn_norm.weight"), &[d]);
240            }
241            put(format!("{p}.attn.Wqkv.weight"), &[3 * d, d]);
242            put(format!("{p}.attn.Wo.weight"), &[d, d]);
243            put(format!("{p}.mlp_norm.weight"), &[d]);
244            put(format!("{p}.mlp.Wi.weight"), &[2 * e.inter, d]);
245            put(format!("{p}.mlp.Wo.weight"), &[d, e.inter]);
246        }
247        put("encoder.final_norm.weight".into(), &[d]);
248        put("type_emb.weight".into(), &[3, d]);
249        for l in 0..self.agent.head_layers {
250            let p = format!("head.layers.{l}");
251            put(format!("{p}.self_attn.in_proj_weight"), &[3 * d, d]);
252            put(format!("{p}.self_attn.in_proj_bias"), &[3 * d]);
253            put(format!("{p}.self_attn.out_proj.weight"), &[d, d]);
254            put(format!("{p}.self_attn.out_proj.bias"), &[d]);
255            put(format!("{p}.linear1.weight"), &[4 * d, d]);
256            put(format!("{p}.linear1.bias"), &[4 * d]);
257            put(format!("{p}.linear2.weight"), &[d, 4 * d]);
258            put(format!("{p}.linear2.bias"), &[d]);
259            for n in ["norm1", "norm2"] {
260                put(format!("{p}.{n}.weight"), &[d]);
261                put(format!("{p}.{n}.bias"), &[d]);
262            }
263        }
264        put("scorer.0.weight".into(), &[d]);
265        put("scorer.0.bias".into(), &[d]);
266        put("scorer.1.weight".into(), &[d, d]);
267        put("scorer.1.bias".into(), &[d]);
268        put("scorer.3.weight".into(), &[1, d]);
269        put("scorer.3.bias".into(), &[1]);
270        put("act_head.0.weight".into(), &[256, d + 4]);
271        put("act_head.0.bias".into(), &[256]);
272        put("act_head.2.weight".into(), &[2, 256]);
273        put("act_head.2.bias".into(), &[2]);
274        put("temperature".into(), &[3]);
275        out
276    }
277}
278
279/// A tensor bound into the graph, as its index in [`Tensors`].
280pub type W = usize;
281
282/// One encoder layer.
283#[derive(Debug, Clone, Copy, PartialEq)]
284pub struct EncoderLayer {
285    /// None on layer 0, where ModernBERT skips the norm.
286    pub attn_norm: Option<W>,
287    /// `[3d, d]`, rows are q then k then v.
288    pub wqkv: W,
289    /// `[d, d]`.
290    pub wo: W,
291    /// `[d]`.
292    pub mlp_norm: W,
293    /// `[2I, d]`, rows are the GELU input then the gate.
294    pub wi: W,
295    /// `[d, I]`.
296    pub mlp_wo: W,
297    /// Global attention, or the sliding window.
298    pub global: bool,
299    /// RoPE base for this layer.
300    pub rope_theta: f64,
301}
302
303/// One layer of the decision head, a PyTorch `TransformerEncoderLayer` with `norm_first`.
304#[derive(Debug, Clone, Copy, PartialEq, Eq)]
305#[allow(missing_docs)]
306pub struct HeadLayer {
307    pub in_proj_w: W,
308    pub in_proj_b: W,
309    pub out_proj_w: W,
310    pub out_proj_b: W,
311    pub norm1_w: W,
312    pub norm1_b: W,
313    pub norm2_w: W,
314    pub norm2_b: W,
315    pub linear1_w: W,
316    pub linear1_b: W,
317    pub linear2_w: W,
318    pub linear2_b: W,
319}
320
321/// A linear layer or LayerNorm with a bias.
322#[derive(Debug, Clone, Copy, PartialEq, Eq)]
323pub struct Affine {
324    /// The weight.
325    pub w: W,
326    /// The bias.
327    pub b: W,
328}
329
330/// The whole compat graph with every weight bound.
331#[derive(Debug, Clone, PartialEq)]
332pub struct LayaGraph {
333    /// `[vocab, d]`.
334    pub tok_embeddings: W,
335    /// `[d]`.
336    pub embed_norm: W,
337    /// Encoder layers in order.
338    pub layers: Vec<EncoderLayer>,
339    /// `[d]`.
340    pub final_norm: W,
341    /// `[3, d]`, one row per question type.
342    pub type_emb: W,
343    /// Decision head layers in order.
344    pub head: Vec<HeadLayer>,
345    /// LayerNorm, `scorer.0`.
346    pub scorer_norm: Affine,
347    /// `scorer.1`, `[d, d]`.
348    pub scorer_in: Affine,
349    /// `scorer.3`, `[1, d]`.
350    pub scorer_out: Affine,
351    /// `act_head.0`, `[256, d + 4]`.
352    pub act_in: Affine,
353    /// `act_head.2`, `[2, 256]`.
354    pub act_out: Affine,
355    /// `[3]`, the calibrated temperatures stored in the checkpoint.
356    pub temperature: W,
357}
358
359impl LayaGraph {
360    /// Binds every weight of the graph to `tensors`, checking names, shapes and dtypes.
361    ///
362    /// # Errors
363    ///
364    /// [`Error::Mismatch`] listing every missing tensor, extra tensor, wrong shape and non float
365    /// dtype, in that order.
366    ///
367    /// # Panics
368    ///
369    /// Never. Binding looks up only names the check above has found.
370    pub fn bind(spec: &LayaSpec, tensors: &Tensors) -> Result<Self> {
371        let expected = spec.expected();
372        let mut problems = Vec::new();
373        for (name, shape) in &expected {
374            match tensors.get(name) {
375                None => problems.push(format!("missing {name} {shape:?}")),
376                Some(v) if v.shape != shape.as_slice() => {
377                    problems.push(format!("{name} has shape {:?}, expected {shape:?}", v.shape));
378                }
379                Some(v) if !v.dtype.is_float() => {
380                    problems.push(format!("{name} is {}, expected a float type", v.dtype));
381                }
382                Some(_) => {}
383            }
384        }
385        let known: std::collections::HashSet<&str> =
386            expected.iter().map(|(n, _)| n.as_str()).collect();
387        for e in tensors.entries() {
388            if !known.contains(e.name.as_str()) {
389                problems.push(format!("unexpected {} {:?}", e.name, e.shape));
390            }
391        }
392        if !problems.is_empty() {
393            return Err(Error::Mismatch(problems));
394        }
395        let w = |name: &str| tensors.index(name).expect("checked above");
396        let aff = |p: &str| Affine { w: w(&format!("{p}.weight")), b: w(&format!("{p}.bias")) };
397        let enc = &spec.encoder;
398        let layers = (0..enc.layers)
399            .map(|i| {
400                let p = format!("encoder.layers.{i}");
401                EncoderLayer {
402                    attn_norm: (i > 0).then(|| w(&format!("{p}.attn_norm.weight"))),
403                    wqkv: w(&format!("{p}.attn.Wqkv.weight")),
404                    wo: w(&format!("{p}.attn.Wo.weight")),
405                    mlp_norm: w(&format!("{p}.mlp_norm.weight")),
406                    wi: w(&format!("{p}.mlp.Wi.weight")),
407                    mlp_wo: w(&format!("{p}.mlp.Wo.weight")),
408                    global: enc.global[i],
409                    rope_theta: if enc.global[i] { enc.rope_global } else { enc.rope_local },
410                }
411            })
412            .collect();
413        let head = (0..spec.agent.head_layers)
414            .map(|l| {
415                let p = format!("head.layers.{l}");
416                HeadLayer {
417                    in_proj_w: w(&format!("{p}.self_attn.in_proj_weight")),
418                    in_proj_b: w(&format!("{p}.self_attn.in_proj_bias")),
419                    out_proj_w: w(&format!("{p}.self_attn.out_proj.weight")),
420                    out_proj_b: w(&format!("{p}.self_attn.out_proj.bias")),
421                    norm1_w: w(&format!("{p}.norm1.weight")),
422                    norm1_b: w(&format!("{p}.norm1.bias")),
423                    norm2_w: w(&format!("{p}.norm2.weight")),
424                    norm2_b: w(&format!("{p}.norm2.bias")),
425                    linear1_w: w(&format!("{p}.linear1.weight")),
426                    linear1_b: w(&format!("{p}.linear1.bias")),
427                    linear2_w: w(&format!("{p}.linear2.weight")),
428                    linear2_b: w(&format!("{p}.linear2.bias")),
429                }
430            })
431            .collect();
432        Ok(Self {
433            tok_embeddings: w("encoder.embeddings.tok_embeddings.weight"),
434            embed_norm: w("encoder.embeddings.norm.weight"),
435            layers,
436            final_norm: w("encoder.final_norm.weight"),
437            type_emb: w("type_emb.weight"),
438            head,
439            scorer_norm: aff("scorer.0"),
440            scorer_in: aff("scorer.1"),
441            scorer_out: aff("scorer.3"),
442            act_in: aff("act_head.0"),
443            act_out: aff("act_head.2"),
444            temperature: w("temperature"),
445        })
446    }
447
448    /// The forward pass as ops for a backend to lower, the graph in the module comment. Values are
449    /// token rows until the markers are gathered, and a backend runs each op on the rows the batch
450    /// has.
451    #[must_use]
452    pub fn plan(&self, spec: &LayaSpec) -> Graph {
453        let e = &spec.encoder;
454        let d = e.d;
455        let mut g = Graph::default();
456        let tok = |g: &mut Graph, w| g.val(Rows::Tokens, w);
457        let emb = tok(&mut g, d);
458        let h = tok(&mut g, d);
459        g.push(Op::Embed { table: self.tok_embeddings, out: emb });
460        g.push(Op::LayerNorm { x: emb, w: self.embed_norm, b: None, eps: e.norm_eps, out: h });
461        let gemm = |g: &mut Graph, a, w, b, epilogue, out| {
462            g.push(Op::Gemm { a, w, b, epilogue, out });
463        };
464        for l in &self.layers {
465            let x = match l.attn_norm {
466                Some(n) => {
467                    let x = tok(&mut g, d);
468                    g.push(Op::LayerNorm { x: h, w: n, b: None, eps: e.norm_eps, out: x });
469                    x
470                }
471                None => h,
472            };
473            let qkv = tok(&mut g, 3 * d);
474            gemm(&mut g, x, l.wqkv, None, Epilogue::None, qkv);
475            g.push(Op::Rope { qkv, theta: l.rope_theta });
476            let att = tok(&mut g, d);
477            let window = (!l.global).then_some(e.window / 2);
478            g.push(Op::Attention { qkv, window, out: att });
479            gemm(&mut g, att, l.wo, None, Epilogue::Accumulate, h);
480            let x = tok(&mut g, d);
481            g.push(Op::LayerNorm { x: h, w: l.mlp_norm, b: None, eps: e.norm_eps, out: x });
482            let u = tok(&mut g, 2 * e.inter);
483            gemm(&mut g, x, l.wi, None, Epilogue::None, u);
484            let a = tok(&mut g, e.inter);
485            g.push(Op::GeGlu { x: u, out: a });
486            gemm(&mut g, a, l.mlp_wo, None, Epilogue::Accumulate, h);
487        }
488        let h2 = tok(&mut g, d);
489        g.push(Op::LayerNorm { x: h, w: self.final_norm, b: None, eps: e.norm_eps, out: h2 });
490        let h = h2;
491        g.push(Op::AddType { h, table: self.type_emb });
492        for l in &self.head {
493            let x = tok(&mut g, d);
494            g.push(Op::LayerNorm {
495                x: h,
496                w: l.norm1_w,
497                b: Some(l.norm1_b),
498                eps: TORCH_EPS,
499                out: x,
500            });
501            let qkv = tok(&mut g, 3 * d);
502            gemm(&mut g, x, l.in_proj_w, Some(l.in_proj_b), Epilogue::None, qkv);
503            let att = tok(&mut g, d);
504            g.push(Op::Attention { qkv, window: None, out: att });
505            gemm(&mut g, att, l.out_proj_w, Some(l.out_proj_b), Epilogue::Accumulate, h);
506            let x = tok(&mut g, d);
507            g.push(Op::LayerNorm {
508                x: h,
509                w: l.norm2_w,
510                b: Some(l.norm2_b),
511                eps: TORCH_EPS,
512                out: x,
513            });
514            let f = tok(&mut g, 4 * d);
515            gemm(&mut g, x, l.linear1_w, Some(l.linear1_b), Epilogue::Relu, f);
516            gemm(&mut g, f, l.linear2_w, Some(l.linear2_b), Epilogue::Accumulate, h);
517        }
518        let m = g.val(Rows::Markers, d);
519        g.push(Op::GatherMarkers { h, out: m });
520        let mn = g.val(Rows::Markers, d);
521        let sn = self.scorer_norm;
522        g.push(Op::LayerNorm { x: m, w: sn.w, b: Some(sn.b), eps: TORCH_EPS, out: mn });
523        let z = g.val(Rows::Markers, d);
524        gemm(&mut g, mn, self.scorer_in.w, Some(self.scorer_in.b), Epilogue::Gelu, z);
525        let logits = g.val(Rows::Markers, 1);
526        gemm(&mut g, z, self.scorer_out.w, Some(self.scorer_out.b), Epilogue::None, logits);
527        let f = g.val(Rows::Seqs, d + 4);
528        g.push(Op::ActFeatures { h, logits, out: f });
529        let a = g.val(Rows::Seqs, ACT_HIDDEN);
530        gemm(&mut g, f, self.act_in.w, Some(self.act_in.b), Epilogue::Gelu, a);
531        let act = g.val(Rows::Seqs, 2);
532        gemm(&mut g, a, self.act_out.w, Some(self.act_out.b), Epilogue::None, act);
533        g.logits = Some(logits);
534        g.act = Some(act);
535        g
536    }
537}
kime_model/laya.rs

kime_model/
laya.rs
