metrale_model_layers/layers/qwen3_attention/
circuit.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: `CircuitBindings` for `Qwen3AttentionLayer`: the full-attention layer's weights by
4//! slot, its heads, RoPE, KV-cache layer and decode route, and the features the circuit does not
5//! model.
6//!
7//! Owner: model-layers (qwen3 attention).
8//! Invariants:
9//! - Every arm of `decode_inner` / `attention_forward` other than the gated, per-head-norm,
10//!   interleaved-MRoPE one the rules encode is reported in `unmodelled`.
11//! - `paged_decode_plain_rows` is the layer's own routing answer (`bf16_decode_is_plain`) at
12//!   each row count, never a restatement of it.
13//! - The transposed twins the multi-sequence tile GEMMs read are bound when the loader built
14//!   them; a fused `[q | k | v]` twin that the multi-sequence path would take is unmodelled.
15
16use std::collections::BTreeMap;
17
18use metrale_circuit::LinearRole;
19
20use super::Qwen3AttentionLayer;
21use crate::circuit_exec::{
22    AttnFacts, BoundWeight, CircuitBindings, CircuitLayer, MixerFacts, RopeFacts, WeightSlot,
23};
24use crate::weight_map::{DenseWeight, QuantWeight};
25
26/// 2026-09-28: A Q/K/V projection as `attention_forward` reads it: the quantized weight when the
27/// loader installed one, else the dense `attn` weight.
28fn proj(
29    w: Option<&QuantWeight>,
30    dense: DenseWeight,
31    what: &str,
32    unmodelled: &mut Vec<String>,
33) -> BoundWeight {
34    match w {
35        Some(QuantWeight::Nvfp4(q)) => BoundWeight::Nvfp4(*q),
36        Some(QuantWeight::Dense(d)) => BoundWeight::Dense(*d),
37        Some(QuantWeight::Fp8(_)) => {
38            unmodelled.push(format!("an FP8 {what} projection"));
39            BoundWeight::Dense(dense)
40        }
41        Some(QuantWeight::PackedQ2(_)) => {
42            unmodelled.push(format!("a packed-Q2 {what} projection"));
43            BoundWeight::Dense(dense)
44        }
45        None => BoundWeight::Dense(dense),
46    }
47}
48
49/// 2026-09-30: The declared W8A8 Q, K, V (the stacked input's segments) and O over the slots
50/// the NVFP4 copies took: every decode arm tries W8A8 first (`w8a8_decode_arm.rs`), at every
51/// row count the circuit plans. A stack other than Q|K|V, or kernels that did not resolve, is
52/// unmodelled.
53fn bind_w8a8(
54    w: &crate::layers::W8a8Mixer,
55    weights: &mut BTreeMap<WeightSlot, BoundWeight>,
56    unmodelled: &mut Vec<String>,
57) {
58    if w.input.segments() != 3 || !w.ctx.kernels.resolved(w.input.scale()) {
59        unmodelled.push("a W8A8 attention input that is not a resolved Q|K|V stack".to_string());
60        return;
61    }
62    let roles = [LinearRole::Q, LinearRole::K, LinearRole::V];
63    for (i, role) in roles.into_iter().enumerate() {
64        match w.input.segment(i) {
65            Ok(seg) => {
66                weights.insert(
67                    WeightSlot::Linear(role),
68                    BoundWeight::W8a8(seg, w.ctx.kernels),
69                );
70            }
71            Err(e) => unmodelled.push(format!("W8A8 {role:?} segment: {e:#}")),
72        }
73    }
74    weights.insert(
75        WeightSlot::Linear(LinearRole::O),
76        BoundWeight::W8a8(w.output, w.ctx.kernels),
77    );
78}
79
80impl CircuitBindings for Qwen3AttentionLayer {
81    fn circuit_prepare(
82        &self,
83        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
84        config: &metrale_config::ModelConfig,
85        levers: &crate::layers::ops::ModelLevers,
86        stream: u64,
87    ) -> anyhow::Result<()> {
88        self.ffn.circuit_prepare(gpu, config, levers, stream)
89    }
90
91    fn circuit_layer(
92        &self,
93        config: &metrale_config::ModelConfig,
94        levers: &crate::layers::ops::ModelLevers,
95    ) -> Option<CircuitLayer> {
96        let mut unmodelled = Vec::new();
97        let arms = [
98            (self.hc.is_some(), "hyper-connections"),
99            (self.lora.is_some(), "an attention LoRA adapter"),
100            (self.mla.is_some(), "multi-head latent attention"),
101            (self.qsa.is_some(), "a QSA indexer"),
102            (self.v_norm_weight.is_some(), "a V norm"),
103            (
104                self.head_gate_weight.is_some(),
105                "a per-head gate projection",
106            ),
107            (
108                self.post_attn_out_norm.is_some(),
109                "a post-attention output norm",
110            ),
111            (self.post_ffn_out_norm.is_some(), "a post-FFN output norm"),
112            (self.layer_scalar.is_some(), "a layer scalar"),
113            (self.moe_ffn.is_some(), "a second (MoE) FFN"),
114            (
115                self.shortcut_carry_in.is_some() || self.shortcut_carry_out.is_some(),
116                "a shortcut carry",
117            ),
118            (
119                self.attn.q_norm_full.is_some() || self.attn.k_norm_full.is_some(),
120                "a full-width Q/K norm",
121            ),
122            (!self.yarn_inv_freq.is_null(), "YaRN RoPE"),
123            (
124                self.rope_proportional && self.rope_proportional_k.0 != 0,
125                "proportional RoPE",
126            ),
127            (self.o_dense_bf16.is_some(), "a BF16 output projection"),
128            (
129                self.qkv_nvfp4_t.is_some()
130                    && super::trait_impl::multi_seq::qkv::fused_qkv_enabled(),
131                "a fused [q | k | v] transposed twin",
132            ),
133        ];
134        for (present, what) in arms {
135            if present {
136                unmodelled.push(what.to_string());
137            }
138        }
139        let dense = BoundWeight::Dense;
140        let mut weights = BTreeMap::from([
141            (WeightSlot::InputNorm, dense(self.input_norm)),
142            (WeightSlot::PostNorm, dense(self.post_attn_norm)),
143            (WeightSlot::QNorm, dense(self.attn.q_norm)),
144            (WeightSlot::KNorm, dense(self.attn.k_norm)),
145        ]);
146        let q = proj(
147            self.q_weight.as_ref(),
148            self.attn.q_proj,
149            "q",
150            &mut unmodelled,
151        );
152        let k = proj(
153            self.k_weight.as_ref(),
154            self.attn.k_proj,
155            "k",
156            &mut unmodelled,
157        );
158        let v = proj(
159            self.v_weight.as_ref(),
160            self.attn.v_proj,
161            "v",
162            &mut unmodelled,
163        );
164        weights.insert(WeightSlot::Linear(LinearRole::Q), q);
165        weights.insert(WeightSlot::Linear(LinearRole::K), k);
166        weights.insert(WeightSlot::Linear(LinearRole::V), v);
167        if matches!(
168            self.o_weight,
169            Some(QuantWeight::Fp8(_)) | Some(QuantWeight::PackedQ2(_))
170        ) {
171            unmodelled.push("an FP8 or packed-Q2 output projection".to_string());
172        }
173        weights.insert(
174            WeightSlot::Linear(LinearRole::O),
175            BoundWeight::Nvfp4(self.attn.o_proj),
176        );
177        if let Some(w) = self.w8a8 {
178            bind_w8a8(&w, &mut weights, &mut unmodelled);
179        }
180        for (role, twin) in [
181            (LinearRole::Q, self.q_nvfp4_t),
182            (LinearRole::K, self.k_nvfp4_t),
183            (LinearRole::V, self.v_nvfp4_t),
184            (LinearRole::O, self.o_nvfp4_t),
185        ] {
186            if let Some(t) = twin {
187                weights.insert(WeightSlot::Transposed(role), BoundWeight::Nvfp4(t));
188            }
189        }
190        self.ffn.circuit_bind(levers, &mut weights, &mut unmodelled);
191        let nq = self
192            .num_q_heads_override
193            .unwrap_or(config.num_attention_heads) as u32;
194        let nkv = self
195            .num_kv_heads_override
196            .unwrap_or(config.num_key_value_heads) as u32;
197        let hd = self.head_dim_override.unwrap_or(config.head_dim) as u32;
198        let facts = AttnFacts {
199            attn_layer_idx: self.attn_layer_idx,
200            kv_dtype: self.kv_dtype,
201            num_q_heads: nq,
202            num_kv_heads: nkv,
203            head_dim: hd,
204            gated: self.gated,
205            rope: RopeFacts {
206                mrope_interleaved: self.mrope_interleaved && self.rope_mrope_interleaved_k.0 != 0,
207                theta: self.rope_theta_override.unwrap_or(config.rope_theta as f32),
208                rotary_dim: self
209                    .rotary_dim_override
210                    .unwrap_or(config.rotary_dim() as u32),
211            },
212            sliding_window: self.sliding_window.unwrap_or(0),
213            softmax_scale: self.effective_attn_scale(hd),
214            paged_decode_plain_rows: (1..=128u32)
215                .filter(|&r| self.bf16_decode_is_plain(nq, nkv, hd, r, levers.max_decode_seqs))
216                .fold(0u128, |m, r| m | 1 << (r - 1)),
217        };
218        Some(CircuitLayer {
219            mixer: MixerFacts::Attention(facts),
220            weights,
221            unmodelled,
222        })
223    }
224}
metrale_model_layers/layers/qwen3_attention/circuit.rs

metrale_model_layers/layers/qwen3_attention/
circuit.rs
