rlx_flow/blocks/
qwen3_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 std::sync::{Arc, Mutex};
23
24use super::BlockStage;
25use super::qwen3_decoder::per_head_rms;
26use super::self_attn::repeat_kv;
27use crate::context::FlowCtx;
28use crate::value::FlowValue;
29
30#[derive(Debug, Clone)]
31pub struct Qwen3DecodeLayerSpec {
32    pub num_heads: usize,
33    pub num_kv_heads: usize,
34    pub head_dim: usize,
35    pub kv_group_size: usize,
36    pub eps: f32,
37    pub use_custom_mask: bool,
38    pub hidden_shape: rlx_ir::Shape,
39    pub batch: usize,
40    pub qk_norm: bool,
41    pub attention_bias: bool,
42}
43
44#[derive(Debug, Clone)]
45pub struct Qwen3DecodeLayerStage {
46    pub layer_prefix: String,
47    pub spec: Qwen3DecodeLayerSpec,
48    pub layer_idx: usize,
49    pub kv_out: Arc<Mutex<Vec<rlx_ir::HirNodeId>>>,
50    pub qk_out: Option<Arc<Mutex<Vec<rlx_ir::HirNodeId>>>>,
51}
52
53impl Qwen3DecodeLayerStage {
54    pub fn layer(
55        layer_idx: usize,
56        spec: Qwen3DecodeLayerSpec,
57        kv_out: Arc<Mutex<Vec<rlx_ir::HirNodeId>>>,
58    ) -> Self {
59        Self {
60            layer_prefix: format!("model.layers.{layer_idx}"),
61            spec,
62            layer_idx,
63            kv_out,
64            qk_out: None,
65        }
66    }
67
68    pub fn layer_with_qk(
69        layer_idx: usize,
70        spec: Qwen3DecodeLayerSpec,
71        kv_out: Arc<Mutex<Vec<rlx_ir::HirNodeId>>>,
72        qk_out: Arc<Mutex<Vec<rlx_ir::HirNodeId>>>,
73    ) -> Self {
74        Self {
75            layer_prefix: format!("model.layers.{layer_idx}"),
76            spec,
77            layer_idx,
78            kv_out,
79            qk_out: Some(qk_out),
80        }
81    }
82}
83
84impl BlockStage for Qwen3DecodeLayerStage {
85    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
86        let decode = ctx
87            .state
88            .decode
89            .clone()
90            .ok_or_else(|| anyhow::anyhow!("Qwen3DecodeLayer requires BindDecodeInputs"))?;
91        let zero_beta_h = ctx
92            .state
93            .zero_beta
94            .ok_or_else(|| anyhow::anyhow!("Qwen3DecodeLayer requires ZeroBeta"))?;
95        let zero_beta_dh = ctx
96            .state
97            .named
98            .get("zero_beta.head")
99            .copied()
100            .ok_or_else(|| anyhow::anyhow!("Qwen3DecodeLayer requires zero_beta.head"))?;
101
102        let lp = &self.layer_prefix;
103        let spec = &self.spec;
104        let nh = spec.num_heads;
105        let nkv = spec.num_kv_heads;
106        let dh = spec.head_dim;
107        let batch = spec.batch;
108
109        let in_ln_g = ctx.load_param(&format!("{lp}.input_layernorm.weight"), false)?;
110        let q_w = ctx.load_param(&format!("{lp}.self_attn.q_proj.weight"), true)?;
111        let k_w = ctx.load_param(&format!("{lp}.self_attn.k_proj.weight"), true)?;
112        let v_w = ctx.load_param(&format!("{lp}.self_attn.v_proj.weight"), true)?;
113        let o_w = ctx.load_param(&format!("{lp}.self_attn.o_proj.weight"), true)?;
114        let post_ln_g = ctx.load_param(&format!("{lp}.post_attention_layernorm.weight"), false)?;
115        let gate_w = ctx.load_param(&format!("{lp}.mlp.gate_proj.weight"), true)?;
116        let up_w = ctx.load_param(&format!("{lp}.mlp.up_proj.weight"), true)?;
117        let down_w = ctx.load_param(&format!("{lp}.mlp.down_proj.weight"), true)?;
118        let (q_bias, k_bias, v_bias) = if spec.attention_bias {
119            (
120                Some(ctx.load_param(&format!("{lp}.self_attn.q_proj.bias"), false)?),
121                Some(ctx.load_param(&format!("{lp}.self_attn.k_proj.bias"), false)?),
122                Some(ctx.load_param(&format!("{lp}.self_attn.v_proj.bias"), false)?),
123            )
124        } else {
125            (None, None, None)
126        };
127        let (q_norm_g, k_norm_g) = if spec.qk_norm {
128            (
129                Some(ctx.load_param(&format!("{lp}.self_attn.q_norm.weight"), false)?),
130                Some(ctx.load_param(&format!("{lp}.self_attn.k_norm.weight"), false)?),
131            )
132        } else {
133            (None, None)
134        };
135
136        let past_k = decode.past_k[self.layer_idx];
137        let past_v = decode.past_v[self.layer_idx];
138
139        let mut gb = HirMut::new(ctx.hir());
140        let skip = input.id;
141        let normed_in = gb.rms_norm(skip, in_ln_g, zero_beta_h, spec.eps);
142        let mut q = gb.mm(normed_in, q_w);
143        let mut k = gb.mm(normed_in, k_w);
144        let mut v = gb.mm(normed_in, v_w);
145
146        if let (Some(qb), Some(kb), Some(vb)) = (q_bias, k_bias, v_bias) {
147            q = gb.add(q, qb);
148            k = gb.add(k, kb);
149            v = gb.add(v, vb);
150        }
151
152        let (q_rope_in, k_rope_in) = if let (Some(qng), Some(kng)) = (q_norm_g, k_norm_g) {
153            let q_normed = per_head_rms(&mut gb, q, qng, zero_beta_dh, batch, 1, nh, dh, spec.eps);
154            let k_normed = per_head_rms(&mut gb, k, kng, zero_beta_dh, batch, 1, nkv, dh, spec.eps);
155            (q_normed, k_normed)
156        } else {
157            (q, k)
158        };
159
160        let q_rope = gb.rope(q_rope_in, decode.cos, decode.sin, dh);
161        let k_rope = gb.rope(k_rope_in, decode.cos, decode.sin, dh);
162
163        let new_k = gb.concat_(vec![past_k, k_rope], 1);
164        let new_v = gb.concat_(vec![past_v, v], 1);
165        self.kv_out.lock().expect("kv out").push(new_k);
166        self.kv_out.lock().expect("kv out").push(new_v);
167
168        let k_rep = repeat_kv(&mut gb, new_k, nkv, dh, spec.kv_group_size);
169        let v_rep = repeat_kv(&mut gb, new_v, nkv, dh, spec.kv_group_size);
170        if let Some(ref sink) = self.qk_out {
171            sink.lock().expect("qwen3 decode qk out").push(q_rope);
172            sink.lock().expect("qwen3 decode qk out").push(k_rep);
173        }
174
175        let attn_shape = shape::attention_shape(gb.shape(q_rope));
176        let attn = if spec.use_custom_mask {
177            let mask = decode
178                .mask
179                .ok_or_else(|| anyhow::anyhow!("custom mask requested but not bound"))?;
180            gb.attention(q_rope, k_rep, v_rep, mask, nh, dh, attn_shape)
181        } else {
182            gb.attention_kind(q_rope, k_rep, v_rep, nh, dh, MaskKind::Causal, attn_shape)
183        };
184
185        let attn_out = gb.mm(attn, o_w);
186        let post_attn = gb.add(skip, attn_out);
187        let normed_post = gb.rms_norm(post_attn, post_ln_g, zero_beta_h, spec.eps);
188        let gate = gb.mm(normed_post, gate_w);
189        let up = gb.mm(normed_post, up_w);
190        let gate_act = gb.silu(gate);
191        let swiglu = gb.mul(gate_act, up);
192        let ffn_out = gb.mm(swiglu, down_w);
193        let h_id = gb.add(post_attn, ffn_out);
194
195        Ok(Some(ctx.wrap(h_id, spec.hidden_shape.clone())))
196    }
197}
rlx_flow/blocks/qwen3_decode_layer.rs

rlx_flow/blocks/
qwen3_decode_layer.rs
