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    /// Leading per-head dims that get rotary-rotated (`head_dim` for plain
54    /// RoPE; smaller for Phi-3 partial RoPE).
55    pub n_rot: usize,
56    pub num_kv_heads: usize,
57    pub eps: f32,
58    pub mask: MaskKind,
59    pub hidden_shape: rlx_ir::Shape,
60    /// RoPE pairing flavor (GGUF Llama → [`rlx_ir::RopeStyle::GptJ`]).
61    pub rope_style: rlx_ir::RopeStyle,
62}
63
64impl BlockStage for LlamaDecoderStage {
65    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
66        let lp = &self.layer_prefix;
67        let zero_beta = ctx
68            .state
69            .zero_beta
70            .ok_or_else(|| anyhow::anyhow!("LlamaDecoder requires ZeroBeta stage"))?;
71        let cos = ctx
72            .state
73            .rope_cos
74            .ok_or_else(|| anyhow::anyhow!("LlamaDecoder requires RopeTables stage"))?;
75        let sin = ctx
76            .state
77            .rope_sin
78            .ok_or_else(|| anyhow::anyhow!("LlamaDecoder requires RopeTables stage"))?;
79
80        let in_ln_g = ctx.load_param(&format!("{lp}.input_layernorm.weight"), false)?;
81        let q_w = ctx.load_param(&format!("{lp}.self_attn.q_proj.weight"), true)?;
82        let k_w = ctx.load_param(&format!("{lp}.self_attn.k_proj.weight"), true)?;
83        let v_w = ctx.load_param(&format!("{lp}.self_attn.v_proj.weight"), true)?;
84        let o_w = ctx.load_param(&format!("{lp}.self_attn.o_proj.weight"), true)?;
85        let post_ln_g = ctx.load_param(&format!("{lp}.post_attention_layernorm.weight"), false)?;
86        let gate_w = ctx.load_param(&format!("{lp}.mlp.gate_proj.weight"), true)?;
87        let up_w = ctx.load_param(&format!("{lp}.mlp.up_proj.weight"), true)?;
88        let down_w = ctx.load_param(&format!("{lp}.mlp.down_proj.weight"), true)?;
89
90        let id = ctx.hir().llama_decoder_block(
91            input.id,
92            in_ln_g,
93            zero_beta,
94            q_w,
95            k_w,
96            v_w,
97            o_w,
98            post_ln_g,
99            zero_beta,
100            gate_w,
101            up_w,
102            down_w,
103            cos,
104            sin,
105            None,
106            self.num_heads,
107            self.head_dim,
108            self.num_kv_heads,
109            self.eps,
110            self.mask,
111            self.rope_style,
112            self.hidden_shape.clone(),
113        );
114
115        Ok(Some(ctx.wrap(id, self.hidden_shape.clone())))
116    }
117}
rlx_flow/blocks/llama_decoder.rs

rlx_flow/blocks/
llama_decoder.rs
