rlx_flow/blocks/
qwen3_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::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::self_attn::repeat_kv;
26use crate::context::FlowCtx;
27use crate::value::FlowValue;
28
29#[derive(Debug, Clone)]
30pub struct Qwen3DecoderSpec {
31    pub num_heads: usize,
32    pub num_kv_heads: usize,
33    pub head_dim: usize,
34    pub eps: f32,
35    pub hidden_shape: rlx_ir::Shape,
36    pub batch: usize,
37    pub seq: usize,
38    /// Per-head Q/K RMSNorm before RoPE (Qwen3); Qwen2 skips.
39    pub qk_norm: bool,
40    /// Explicit Q/K/V bias vectors (Qwen2); Qwen3 typically false.
41    pub attention_bias: bool,
42    /// Attention mask — `Causal` for standard Qwen3, `SlidingWindow(w)` for
43    /// sliding-window models (Mistral / Gemma-style).
44    pub mask: MaskKind,
45}
46
47#[derive(Debug, Clone)]
48pub struct Qwen3DecoderStage {
49    pub layer_prefix: String,
50    pub spec: Qwen3DecoderSpec,
51    pub kv_sink: Option<Arc<Mutex<Vec<rlx_ir::HirNodeId>>>>,
52    /// Optional side tap: post-RoPE Q and GQA-expanded K per layer (AIF probe).
53    pub qk_sink: Option<Arc<Mutex<Vec<rlx_ir::HirNodeId>>>>,
54}
55
56impl Qwen3DecoderStage {
57    pub fn layer(layer_idx: usize, spec: Qwen3DecoderSpec) -> Self {
58        Self {
59            layer_prefix: format!("model.layers.{layer_idx}"),
60            spec,
61            kv_sink: None,
62            qk_sink: None,
63        }
64    }
65
66    pub fn layer_with_kv(
67        layer_idx: usize,
68        spec: Qwen3DecoderSpec,
69        kv_sink: Arc<Mutex<Vec<rlx_ir::HirNodeId>>>,
70    ) -> Self {
71        Self {
72            layer_prefix: format!("model.layers.{layer_idx}"),
73            spec,
74            kv_sink: Some(kv_sink),
75            qk_sink: None,
76        }
77    }
78
79    pub fn layer_with_kv_qk(
80        layer_idx: usize,
81        spec: Qwen3DecoderSpec,
82        kv_sink: Arc<Mutex<Vec<rlx_ir::HirNodeId>>>,
83        qk_sink: Arc<Mutex<Vec<rlx_ir::HirNodeId>>>,
84    ) -> Self {
85        Self {
86            layer_prefix: format!("model.layers.{layer_idx}"),
87            spec,
88            kv_sink: Some(kv_sink),
89            qk_sink: Some(qk_sink),
90        }
91    }
92}
93
94impl BlockStage for Qwen3DecoderStage {
95    fn emit(&self, ctx: &mut FlowCtx<'_>, input: FlowValue) -> Result<Option<FlowValue>> {
96        let lp = &self.layer_prefix;
97        let spec = &self.spec;
98        let nh = spec.num_heads;
99        let nkv = spec.num_kv_heads;
100        let dh = spec.head_dim;
101        let group = nh / nkv;
102
103        let zero_beta_h = ctx
104            .state
105            .zero_beta
106            .ok_or_else(|| anyhow::anyhow!("Qwen3Decoder requires ZeroBeta"))?;
107        let zero_beta_dh = ctx
108            .state
109            .named
110            .get("zero_beta.head")
111            .copied()
112            .ok_or_else(|| anyhow::anyhow!("Qwen3Decoder requires zero_beta.head"))?;
113        let cos = ctx
114            .state
115            .rope_cos
116            .ok_or_else(|| anyhow::anyhow!("Qwen3Decoder requires RopeTables"))?;
117        let sin = ctx
118            .state
119            .rope_sin
120            .ok_or_else(|| anyhow::anyhow!("Qwen3Decoder requires RopeTables"))?;
121
122        let in_ln_g = ctx.load_param(&format!("{lp}.input_layernorm.weight"), false)?;
123        let q_w = ctx.load_param(&format!("{lp}.self_attn.q_proj.weight"), true)?;
124        let k_w = ctx.load_param(&format!("{lp}.self_attn.k_proj.weight"), true)?;
125        let v_w = ctx.load_param(&format!("{lp}.self_attn.v_proj.weight"), true)?;
126        let o_w = ctx.load_param(&format!("{lp}.self_attn.o_proj.weight"), true)?;
127        let post_ln_g = ctx.load_param(&format!("{lp}.post_attention_layernorm.weight"), false)?;
128        let gate_w = ctx.load_param(&format!("{lp}.mlp.gate_proj.weight"), true)?;
129        let up_w = ctx.load_param(&format!("{lp}.mlp.up_proj.weight"), true)?;
130        let down_w = ctx.load_param(&format!("{lp}.mlp.down_proj.weight"), true)?;
131        let (q_bias, k_bias, v_bias) = if spec.attention_bias {
132            (
133                Some(ctx.load_param(&format!("{lp}.self_attn.q_proj.bias"), false)?),
134                Some(ctx.load_param(&format!("{lp}.self_attn.k_proj.bias"), false)?),
135                Some(ctx.load_param(&format!("{lp}.self_attn.v_proj.bias"), false)?),
136            )
137        } else {
138            (None, None, None)
139        };
140        let (q_norm_g, k_norm_g) = if spec.qk_norm {
141            (
142                Some(ctx.load_param(&format!("{lp}.self_attn.q_norm.weight"), false)?),
143                Some(ctx.load_param(&format!("{lp}.self_attn.k_norm.weight"), false)?),
144            )
145        } else {
146            (None, None)
147        };
148
149        let mut gb = HirMut::new(ctx.hir());
150        let skip = input.id;
151
152        let normed_in = gb.rms_norm(skip, in_ln_g, zero_beta_h, spec.eps);
153        let mut q = gb.mm(normed_in, q_w);
154        let mut k = gb.mm(normed_in, k_w);
155        let mut v = gb.mm(normed_in, v_w);
156
157        if let (Some(qb), Some(kb), Some(vb)) = (q_bias, k_bias, v_bias) {
158            q = gb.add(q, qb);
159            k = gb.add(k, kb);
160            v = gb.add(v, vb);
161        }
162
163        let (q_rope_in, k_rope_in) = if let (Some(qng), Some(kng)) = (q_norm_g, k_norm_g) {
164            let q_normed = per_head_rms(
165                &mut gb,
166                q,
167                qng,
168                zero_beta_dh,
169                spec.batch,
170                spec.seq,
171                nh,
172                dh,
173                spec.eps,
174            );
175            let k_normed = per_head_rms(
176                &mut gb,
177                k,
178                kng,
179                zero_beta_dh,
180                spec.batch,
181                spec.seq,
182                nkv,
183                dh,
184                spec.eps,
185            );
186            (q_normed, k_normed)
187        } else {
188            (q, k)
189        };
190
191        let q_rope = gb.rope(q_rope_in, cos, sin, dh);
192        let k_rope = gb.rope(k_rope_in, cos, sin, dh);
193        if let Some(ref sink) = self.kv_sink {
194            sink.lock().expect("qwen3 kv sink").push(k_rope);
195            sink.lock().expect("qwen3 kv sink").push(v);
196        }
197        let k_rep = repeat_kv(&mut gb, k_rope, nkv, dh, group);
198        let v_rep = repeat_kv(&mut gb, v, nkv, dh, group);
199        if let Some(ref sink) = self.qk_sink {
200            sink.lock().expect("qwen3 qk sink").push(q_rope);
201            sink.lock().expect("qwen3 qk sink").push(k_rep);
202        }
203
204        let attn_shape = shape::attention_shape(gb.shape(q_rope));
205        let attn = gb.attention_kind(q_rope, k_rep, v_rep, nh, dh, spec.mask, attn_shape);
206        let attn_out = gb.mm(attn, o_w);
207        let post_attn = gb.add(skip, attn_out);
208        let normed_post = gb.rms_norm(post_attn, post_ln_g, zero_beta_h, spec.eps);
209
210        let gate = gb.mm(normed_post, gate_w);
211        let up = gb.mm(normed_post, up_w);
212        let gate_act = gb.silu(gate);
213        let swiglu = gb.mul(gate_act, up);
214        let ffn_out = gb.mm(swiglu, down_w);
215        let out = gb.add(post_attn, ffn_out);
216
217        Ok(Some(ctx.wrap(out, spec.hidden_shape.clone())))
218    }
219}
220
221pub(crate) fn per_head_rms(
222    gb: &mut HirMut,
223    x: rlx_ir::HirNodeId,
224    gamma: rlx_ir::HirNodeId,
225    beta: rlx_ir::HirNodeId,
226    batch: usize,
227    seq: usize,
228    heads: usize,
229    head_dim: usize,
230    eps: f32,
231) -> rlx_ir::HirNodeId {
232    let flat = (batch * seq * heads) as i64;
233    let dh = head_dim as i64;
234    let r = gb.reshape_(x, vec![flat, dh]);
235    let n = gb.rms_norm(r, gamma, beta, eps);
236    gb.reshape_(n, vec![batch as i64, seq as i64, (heads * head_dim) as i64])
237}
rlx_flow/blocks/qwen3_decoder.rs

rlx_flow/blocks/
qwen3_decoder.rs
