rlx_flow/blocks/
llama_decoder.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::op::MaskKind;
18
19use super::BlockStage;
20use crate::context::FlowCtx;
21use crate::value::FlowValue;
22#[derive(Debug, Clone)]
23pub struct LlamaDecoderStage {
24    pub layer_prefix: String,
25    pub num_heads: usize,
26    pub head_dim: usize,
27    pub num_kv_heads: usize,
28    pub eps: f32,
29    pub mask: MaskKind,
30    pub hidden_shape: rlx_ir::Shape,
31    pub rope_style: rlx_ir::RopeStyle,
32}
33
34impl LlamaDecoderStage {
35    pub fn layer(layer_idx: usize, spec: LlamaDecoderSpec) -> Self {
36        Self {
37            layer_prefix: format!("model.layers.{layer_idx}"),
38            num_heads: spec.num_heads,
39            head_dim: spec.head_dim,
40            num_kv_heads: spec.num_kv_heads,
41            eps: spec.eps,
42            mask: spec.mask,
43            hidden_shape: spec.hidden_shape,
44            rope_style: spec.rope_style,
45        }
46    }
47}
48
49#[derive(Debug, Clone)]
50pub struct LlamaDecoderSpec {
51    pub num_heads: usize,
52    pub head_dim: usize,
53    pub num_kv_heads: usize,
54    pub eps: f32,
55    pub mask: MaskKind,
56    pub hidden_shape: rlx_ir::Shape,
57    /// RoPE pairing flavor (GGUF Llama → [`rlx_ir::RopeStyle::GptJ`]).
58    pub rope_style: rlx_ir::RopeStyle,
59}
60
61impl BlockStage for LlamaDecoderStage {
62    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
63        let lp = &self.layer_prefix;
64        let zero_beta = ctx
65            .state
66            .zero_beta
67            .ok_or_else(|| anyhow::anyhow!("LlamaDecoder requires ZeroBeta stage"))?;
68        let cos = ctx
69            .state
70            .rope_cos
71            .ok_or_else(|| anyhow::anyhow!("LlamaDecoder requires RopeTables stage"))?;
72        let sin = ctx
73            .state
74            .rope_sin
75            .ok_or_else(|| anyhow::anyhow!("LlamaDecoder requires RopeTables stage"))?;
76
77        let in_ln_g = ctx.load_param(&format!("{lp}.input_layernorm.weight"), false)?;
78        let q_w = ctx.load_param(&format!("{lp}.self_attn.q_proj.weight"), true)?;
79        let k_w = ctx.load_param(&format!("{lp}.self_attn.k_proj.weight"), true)?;
80        let v_w = ctx.load_param(&format!("{lp}.self_attn.v_proj.weight"), true)?;
81        let o_w = ctx.load_param(&format!("{lp}.self_attn.o_proj.weight"), true)?;
82        let post_ln_g = ctx.load_param(&format!("{lp}.post_attention_layernorm.weight"), false)?;
83        let gate_w = ctx.load_param(&format!("{lp}.mlp.gate_proj.weight"), true)?;
84        let up_w = ctx.load_param(&format!("{lp}.mlp.up_proj.weight"), true)?;
85        let down_w = ctx.load_param(&format!("{lp}.mlp.down_proj.weight"), true)?;
86
87        let id = ctx.hir().llama_decoder_block(
88            input.id,
89            in_ln_g,
90            zero_beta,
91            q_w,
92            k_w,
93            v_w,
94            o_w,
95            post_ln_g,
96            zero_beta,
97            gate_w,
98            up_w,
99            down_w,
100            cos,
101            sin,
102            None,
103            self.num_heads,
104            self.head_dim,
105            self.num_kv_heads,
106            self.eps,
107            self.mask,
108            self.rope_style,
109            self.hidden_shape.clone(),
110        );
111
112        Ok(Some(ctx.wrap(id, self.hidden_shape.clone())))
113    }
114}
rlx_flow/blocks/llama_decoder.rs

rlx_flow/blocks/
llama_decoder.rs
