metrale_model_layers/layers/qwen3_ssm/
circuit.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: `CircuitBindings` for `Qwen3SsmLayer`: the GatedDeltaNet layer's weights by slot,
4//! its qkvz layout, and the features the circuit does not model.
5//!
6//! Owner: model-layers (Qwen3 SSM layer).
7//! Invariants:
8//! - Every arm of `decode_inner` / `ssm_forward` other than the NVFP4 FP32-kernel one the rules
9//!   encode is reported in `unmodelled`.
10
11use std::collections::BTreeMap;
12
13use metrale_circuit::LinearRole;
14
15use super::Qwen3SsmLayer;
16use crate::circuit_exec::{
17    BoundWeight, CircuitBindings, CircuitLayer, GdnFacts, MixerFacts, WeightSlot,
18};
19
20impl CircuitBindings for Qwen3SsmLayer {
21    fn circuit_prepare(
22        &self,
23        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
24        config: &metrale_config::ModelConfig,
25        levers: &crate::layers::ops::ModelLevers,
26        stream: u64,
27    ) -> anyhow::Result<()> {
28        self.ffn.circuit_prepare(gpu, config, levers, stream)
29    }
30
31    fn circuit_layer(
32        &self,
33        _config: &metrale_config::ModelConfig,
34        levers: &crate::layers::ops::ModelLevers,
35    ) -> Option<CircuitLayer> {
36        let mut unmodelled = Vec::new();
37        let arms = [
38            (self.hc.is_some(), "hyper-connections"),
39            (self.ple.is_some(), "per-layer embeddings"),
40            (self.lora_out_proj.is_some(), "an out_proj LoRA adapter"),
41            (self.qkvz_fp8w.is_some(), "an FP8 qkvz projection"),
42            (self.qkvz_q2.is_some(), "a packed-Q2 qkvz projection"),
43            (self.out_proj_fp8w.is_some(), "an FP8 out_proj"),
44            (
45                levers.gdn_fused_conv,
46                "the fused GDN conv+norm kernel (METRALE_GDN_FUSED_CONV)",
47            ),
48            (
49                self.conv1d_l2norm_f32_k.0 == 0
50                    || self.gdn_f32_k.0 == 0
51                    || self.gated_rms_norm_f32_k.0 == 0,
52                "the BF16 GDN kernels (an FP32 GDN kernel is not loaded)",
53            ),
54        ];
55        for (present, what) in arms {
56            if present {
57                unmodelled.push(what.to_string());
58            }
59        }
60        let dense = BoundWeight::Dense;
61        let mut weights = BTreeMap::from([
62            (WeightSlot::InputNorm, dense(self.input_norm)),
63            (WeightSlot::PostNorm, dense(self.post_attn_norm)),
64            (
65                WeightSlot::Linear(LinearRole::Ba),
66                dense(self.ssm.in_proj_ba),
67            ),
68            (WeightSlot::GdnALog, dense(self.ssm.a_log)),
69            (WeightSlot::GdnDtBias, dense(self.ssm.dt_bias)),
70            (WeightSlot::GdnConv1d, dense(self.ssm.conv1d)),
71            (WeightSlot::GdnNorm, dense(self.ssm.norm)),
72        ]);
73        weights.insert(
74            WeightSlot::Linear(LinearRole::Qkvz),
75            match self.qkvz_nvfp4 {
76                Some(q) => BoundWeight::Nvfp4(q),
77                None => dense(self.ssm.in_proj_qkvz),
78            },
79        );
80        weights.insert(
81            WeightSlot::Linear(LinearRole::GdnOut),
82            match self.out_proj_dense {
83                Some(d) => dense(d),
84                None => BoundWeight::Nvfp4(self.ssm.out_proj),
85            },
86        );
87        // 2026-09-30: The declared W8A8 QKV|Z and out_proj, which every decode path tries first
88        // (`w8a8_decode.rs`); installed only on a sequential-QKVZ layer, the layout
89        // `qkvz_deinterleaved` states.
90        if let Some(w) = self.w8a8 {
91            if w.ctx.kernels.resolved(w.input.scale()) {
92                weights.insert(
93                    WeightSlot::Linear(LinearRole::Qkvz),
94                    BoundWeight::W8a8(w.input, w.ctx.kernels),
95                );
96                weights.insert(
97                    WeightSlot::Linear(LinearRole::GdnOut),
98                    BoundWeight::W8a8(w.output, w.ctx.kernels),
99                );
100            } else {
101                unmodelled.push("W8A8 GDN projections whose kernels did not resolve".to_string());
102            }
103        }
104        for (role, twin) in [
105            (LinearRole::Qkvz, self.qkvz_nvfp4_t),
106            (LinearRole::GdnOut, self.out_proj_nvfp4_t),
107        ] {
108            if let Some(t) = twin {
109                weights.insert(WeightSlot::Transposed(role), BoundWeight::Nvfp4(t));
110            }
111        }
112        self.ffn.circuit_bind(levers, &mut weights, &mut unmodelled);
113        Some(CircuitLayer {
114            mixer: MixerFacts::Gdn(GdnFacts {
115                qkvz_deinterleaved: self.sequential_qkvz,
116                h_slot_bytes: self.h_slot_stride_bytes() as u64,
117                conv_state_bytes: self.conv_state_bytes as u64,
118                carry: self.carry.binding.get().copied(),
119            }),
120            weights,
121            unmodelled,
122        })
123    }
124}
metrale_model_layers/layers/qwen3_ssm/circuit.rs

metrale_model_layers/layers/qwen3_ssm/
circuit.rs
