rlx_flow/blocks/
llama_decode_layer.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16use anyhow::Result;
17use rlx_ir::HirGraphExt;
18use rlx_ir::hir::HirMut;
19use rlx_ir::op::MaskKind;
20use rlx_ir::shape;
21
22use super::BlockStage;
23use crate::context::FlowCtx;
24use crate::value::FlowValue;
25#[derive(Debug, Clone)]
26pub struct LlamaDecodeLayerSpec {
27    pub num_heads: usize,
28    pub head_dim: usize,
29    /// Leading per-head dims that get rotary-rotated (`head_dim` unless
30    /// partial RoPE — Phi-3 / long-context variants).
31    pub n_rot: usize,
32    pub num_kv_heads: usize,
33    pub kv_group_size: usize,
34    pub eps: f32,
35    pub use_custom_mask: bool,
36    pub hidden_shape: rlx_ir::Shape,
37    /// RoPE pairing flavor. GGUF Llama weights need [`rlx_ir::RopeStyle::GptJ`];
38    /// HF-safetensors checkpoints use [`rlx_ir::RopeStyle::NeoX`] (default).
39    pub rope_style: rlx_ir::RopeStyle,
40}
41
42#[derive(Debug, Clone)]
43pub struct LlamaDecodeLayerStage {
44    pub layer_prefix: String,
45    pub spec: LlamaDecodeLayerSpec,
46    pub layer_idx: usize,
47    pub kv_out: std::sync::Arc<std::sync::Mutex<Vec<rlx_ir::HirNodeId>>>,
48    /// Optional EAGLE3-style tap for the pre-attention-norm layer
49    /// input. Mirrors the field on
50    /// [`crate::blocks::GemmaDecodeLayerStage`]; see that doc for
51    /// semantics and push-order guarantees.
52    pub aux_in_out: Option<std::sync::Arc<std::sync::Mutex<Vec<rlx_ir::HirNodeId>>>>,
53}
54
55impl LlamaDecodeLayerStage {
56    pub fn layer(
57        layer_idx: usize,
58        spec: LlamaDecodeLayerSpec,
59        kv_out: std::sync::Arc<std::sync::Mutex<Vec<rlx_ir::HirNodeId>>>,
60    ) -> Self {
61        Self {
62            layer_prefix: format!("model.layers.{layer_idx}"),
63            spec,
64            layer_idx,
65            kv_out,
66            aux_in_out: None,
67        }
68    }
69
70    pub fn with_aux_input_tap(
71        mut self,
72        sink: std::sync::Arc<std::sync::Mutex<Vec<rlx_ir::HirNodeId>>>,
73    ) -> Self {
74        self.aux_in_out = Some(sink);
75        self
76    }
77}
78
79impl BlockStage for LlamaDecodeLayerStage {
80    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
81        if let Some(sink) = self.aux_in_out.as_ref() {
82            sink.lock().expect("aux in out").push(input.id);
83        }
84
85        let decode = ctx
86            .state
87            .decode
88            .clone()
89            .ok_or_else(|| anyhow::anyhow!("LlamaDecodeLayer requires BindDecodeInputs"))?;
90        let zero_beta = ctx
91            .state
92            .zero_beta
93            .ok_or_else(|| anyhow::anyhow!("LlamaDecodeLayer requires ZeroBeta"))?;
94
95        let lp = &self.layer_prefix;
96        let spec = &self.spec;
97        let in_ln_g = ctx.load_param(&format!("{lp}.input_layernorm.weight"), false)?;
98        let q_w = ctx.load_param(&format!("{lp}.self_attn.q_proj.weight"), true)?;
99        let k_w = ctx.load_param(&format!("{lp}.self_attn.k_proj.weight"), true)?;
100        let v_w = ctx.load_param(&format!("{lp}.self_attn.v_proj.weight"), true)?;
101        let o_w = ctx.load_param(&format!("{lp}.self_attn.o_proj.weight"), true)?;
102        let post_ln_g = ctx.load_param(&format!("{lp}.post_attention_layernorm.weight"), false)?;
103        let gate_w = ctx.load_param(&format!("{lp}.mlp.gate_proj.weight"), true)?;
104        let up_w = ctx.load_param(&format!("{lp}.mlp.up_proj.weight"), true)?;
105        let down_w = ctx.load_param(&format!("{lp}.mlp.down_proj.weight"), true)?;
106
107        let past_k = decode.past_k.get(self.layer_idx);
108        let past_v = decode.past_v.get(self.layer_idx);
109
110        let mut gb = HirMut::new(ctx.hir());
111        let normed_in = gb.rms_norm(input.id, in_ln_g, zero_beta, spec.eps);
112        let q = gb.mm(normed_in, q_w);
113        let k = gb.mm(normed_in, k_w);
114        let v = gb.mm(normed_in, v_w);
115
116        let q_rope = gb.rope_n_styled(
117            q,
118            decode.cos,
119            decode.sin,
120            spec.head_dim,
121            spec.n_rot,
122            spec.rope_style,
123        );
124        let k_rope = gb.rope_n_styled(
125            k,
126            decode.cos,
127            decode.sin,
128            spec.head_dim,
129            spec.n_rot,
130            spec.rope_style,
131        );
132
133        let (new_k, new_v) = match (past_k, past_v) {
134            (Some(past_k), Some(past_v)) => (
135                gb.concat_(vec![*past_k, k_rope], 1),
136                gb.concat_(vec![*past_v, v], 1),
137            ),
138            _ => (k_rope, v),
139        };
140        self.kv_out.lock().expect("kv out").push(new_k);
141        self.kv_out.lock().expect("kv out").push(new_v);
142
143        let k_rep = super::self_attn::repeat_kv(
144            &mut gb,
145            new_k,
146            spec.num_kv_heads,
147            spec.head_dim,
148            spec.kv_group_size,
149        );
150        let v_rep = super::self_attn::repeat_kv(
151            &mut gb,
152            new_v,
153            spec.num_kv_heads,
154            spec.head_dim,
155            spec.kv_group_size,
156        );
157
158        let attn_shape = shape::attention_shape(gb.shape(q_rope));
159        let attn = if spec.use_custom_mask {
160            let mask = decode
161                .mask
162                .ok_or_else(|| anyhow::anyhow!("custom mask requested but not bound"))?;
163            gb.attention(
164                q_rope,
165                k_rep,
166                v_rep,
167                mask,
168                spec.num_heads,
169                spec.head_dim,
170                attn_shape,
171            )
172        } else {
173            gb.attention_kind(
174                q_rope,
175                k_rep,
176                v_rep,
177                spec.num_heads,
178                spec.head_dim,
179                MaskKind::Causal,
180                attn_shape,
181            )
182        };
183
184        let attn_out = gb.mm(attn, o_w);
185        let post_attn = gb.add(input.id, attn_out);
186        let normed_post = gb.rms_norm(post_attn, post_ln_g, zero_beta, spec.eps);
187        let gate = gb.mm(normed_post, gate_w);
188        let up = gb.mm(normed_post, up_w);
189        let gate_act = gb.silu(gate);
190        let swiglu = gb.mul(gate_act, up);
191        let ffn_out = gb.mm(swiglu, down_w);
192        let h_id = gb.add(post_attn, ffn_out);
193
194        Ok(Some(ctx.wrap(h_id, spec.hidden_shape.clone())))
195    }
196}
rlx_flow/blocks/llama_decode_layer.rs

rlx_flow/blocks/
llama_decode_layer.rs
