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

kime_model/
laya.rs
