metrale_model_layers/circuit_exec/
bindings.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: `CircuitBindings`, the supertrait of `TransformerLayer` through which a layer hands
4//! the circuit executor its weights by role and the facts its kernels need. Layers keep owning
5//! their weights; a binding is a copy of pointers.
6//!
7//! Owner: model-layers circuit executor.
8//! Invariants:
9//! - A layer lists in `unmodelled` every feature it carries that the circuit does not model
10//!   (hyper-connections, LoRA, a projection format the rules do not cover, ...). The executor
11//!   refuses a model with any, so a plan never runs over a layer it misdescribes.
12//! - A bound weight carries its storage format; the executor checks it against the format the
13//!   plan's node was resolved to.
14
15use std::collections::{BTreeMap, BTreeSet};
16
17use anyhow::{Result, anyhow, ensure};
18use metrale_cache::kv_cache::KvCacheDtype;
19use metrale_circuit::{Circuit, LinearRole};
20
21use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
22
23use crate::layers::ops::{W8a8Kernels, W8a8Weight};
24use crate::weight_map::{DenseWeight, QuantizedWeight};
25
26/// 2026-09-28: Which weight of a layer a circuit node reads.
27#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
28pub enum WeightSlot {
29    /// 2026-09-28: The mixer's input RMSNorm.
30    InputNorm,
31    /// 2026-09-28: The FFN's input RMSNorm (`post_attention_layernorm`).
32    PostNorm,
33    /// 2026-09-28: A projection; `GateUp` is bound as [`WeightSlot::FfnGate`] and
34    /// [`WeightSlot::FfnUp`].
35    Linear(LinearRole),
36    /// 2026-09-28: The dense FFN's gate projection.
37    FfnGate,
38    /// 2026-09-28: The dense FFN's up projection.
39    FfnUp,
40    /// 2026-09-28: Attention's per-head Q RMSNorm.
41    QNorm,
42    /// 2026-09-28: Attention's per-head K RMSNorm.
43    KNorm,
44    /// 2026-09-28: GatedDeltaNet `A_log`.
45    GdnALog,
46    /// 2026-09-28: GatedDeltaNet `dt_bias`.
47    GdnDtBias,
48    /// 2026-09-28: GatedDeltaNet conv1d taps.
49    GdnConv1d,
50    /// 2026-09-28: GatedDeltaNet output norm.
51    GdnNorm,
52    /// 2026-09-28: The transposed twin of a projection, which the tile GEMMs read.
53    Transposed(LinearRole),
54    /// 2026-09-28: The dense FFN's gate, repacked for the NVFP4 MMQ GEMM.
55    FfnGateMmq,
56    /// 2026-09-28: The dense FFN's up, repacked for the NVFP4 MMQ GEMM.
57    FfnUpMmq,
58    /// 2026-09-28: The dense FFN's down, repacked for the NVFP4 MMQ GEMM.
59    FfnDownMmq,
60    /// 2026-09-29: The MTP draft head's norm of the token embedding (`pre_fc_norm_embedding`).
61    EmbedNorm,
62    /// 2026-09-29: The MTP draft head's norm of the target hidden (`pre_fc_norm_hidden`).
63    HiddenNorm,
64    /// 2026-09-29: The MTP draft head's final norm.
65    FinalNorm,
66    /// 2026-09-29: The MTP draft head's vocabulary projection.
67    LmHead,
68}
69
70/// 2026-09-28: A bound weight and its storage format.
71#[derive(Debug, Clone, Copy)]
72pub enum BoundWeight {
73    /// 2026-09-28: Unquantized (BF16 for projections; norms and GDN vectors as loaded).
74    Dense(DenseWeight),
75    /// 2026-09-28: NVFP4, group 16.
76    Nvfp4(QuantizedWeight),
77    /// 2026-09-28: NVFP4 repacked for the MMQ GEMM (`ops::nvfp4_mmq_repack`); its
78    /// `weight_scale_2` stays with the source weight.
79    Mmq(DevicePtr),
80    /// 2026-09-30: A declared W8A8 projection (`W8a8Mixer`, `W8a8Ffn`): the checkpoint's E4M3
81    /// weight with its scales, and the kernels its layer runs it with.
82    W8a8(W8a8Weight, W8a8Kernels),
83}
84
85impl BoundWeight {
86    /// 2026-09-28: The format family, spelled as the precision tables spell weights.
87    pub fn family(&self) -> &'static str {
88        match self {
89            BoundWeight::Dense(_) => "bf16",
90            BoundWeight::Nvfp4(_) | BoundWeight::Mmq(_) => "nvfp4",
91            BoundWeight::W8a8(..) => "fp8",
92        }
93    }
94}
95
96/// 2026-09-28: What a GatedDeltaNet layer's kernels need beyond its weights.
97#[derive(Debug, Clone, Copy, PartialEq)]
98pub struct GdnFacts {
99    /// 2026-09-28: The qkvz projection writes `[Q | K | V | Z]` itself (`sequential_qkvz`), so no
100    /// deinterleave follows it.
101    pub qkvz_deinterleaved: bool,
102    /// 2026-09-30: Bytes from one pool slot's h state to the next (`h_slot_stride_bytes`), the
103    /// pitch the batched arm's slot check and strided recurrence use.
104    pub h_slot_bytes: u64,
105    /// 2026-09-30: Bytes from one pool slot's conv window to the next (`conv_state_bytes`).
106    pub conv_state_bytes: u64,
107    /// 2026-09-30: The carried-state verify's buffers (`LayerWriteOnAccept::gdn_carry_bind`),
108    /// which the batched verify's GDN launches read; `None` until the model binds them.
109    pub carry: Option<crate::layer::GdnCarryBinding>,
110}
111
112/// 2026-09-28: A full-attention layer's rotary embedding.
113#[derive(Debug, Clone, Copy, PartialEq)]
114pub struct RopeFacts {
115    /// 2026-09-28: Interleaved MRoPE (`mrope_interleaved` with its kernel loaded).
116    pub mrope_interleaved: bool,
117    /// 2026-09-28: `rope_theta`, after any per-layer override.
118    pub theta: f32,
119    /// 2026-09-28: Rotated dims per head, after any per-layer override.
120    pub rotary_dim: u32,
121}
122
123/// 2026-09-28: What a full-attention layer's kernels need beyond its weights.
124#[derive(Debug, Clone, Copy, PartialEq)]
125pub struct AttnFacts {
126    /// 2026-09-28: This layer's index among the attention layers: its KV-cache layer.
127    pub attn_layer_idx: usize,
128    /// 2026-09-28: KV-cache dtype of this layer.
129    pub kv_dtype: KvCacheDtype,
130    /// 2026-09-28: Query heads.
131    pub num_q_heads: u32,
132    /// 2026-09-28: KV heads.
133    pub num_kv_heads: u32,
134    /// 2026-09-28: Head dim.
135    pub head_dim: u32,
136    /// 2026-09-28: Q carries an output gate (`[Q | Gate]` after the split).
137    pub gated: bool,
138    /// 2026-09-28: Rotary embedding.
139    pub rope: RopeFacts,
140    /// 2026-09-28: Sliding window in positions; 0 for none.
141    pub sliding_window: u32,
142    /// 2026-09-28: Softmax scale (`effective_attn_scale`).
143    pub softmax_scale: f32,
144    /// 2026-09-28: Bit `r - 1` is set when the layer's own decode routing at `r` rows picks
145    /// the plain non-split paged kernel (`paged_decode_k`): no split-K pair, no GQA-packed
146    /// kernel, no 512-wide head kernel.
147    pub paged_decode_plain_rows: u128,
148}
149
150impl AttnFacts {
151    /// 2026-09-28: The plain paged kernel serves `rows` rows (1..=128).
152    pub fn paged_decode_plain(&self, rows: u64) -> bool {
153        (1..=128).contains(&rows) && self.paged_decode_plain_rows >> (rows - 1) & 1 == 1
154    }
155}
156
157/// 2026-09-28: The mixer a layer runs.
158#[derive(Debug, Clone, Copy, PartialEq)]
159pub enum MixerFacts {
160    /// 2026-09-28: GatedDeltaNet.
161    Gdn(GdnFacts),
162    /// 2026-09-28: Full attention.
163    Attention(AttnFacts),
164}
165
166/// 2026-09-28: One layer, as the executor sees it.
167#[derive(Debug, Clone)]
168pub struct CircuitLayer {
169    /// 2026-09-28: The mixer.
170    pub mixer: MixerFacts,
171    /// 2026-09-28: Weights by slot.
172    pub weights: BTreeMap<WeightSlot, BoundWeight>,
173    /// 2026-09-28: Features present on this layer that the circuit does not model.
174    pub unmodelled: Vec<String>,
175}
176
177/// 2026-09-28: A supertrait of `TransformerLayer`; see the module header.
178pub trait CircuitBindings {
179    /// 2026-09-28: Build the lazily built weights the bound kernels read (the MMQ repacks a
180    /// first wide batch would otherwise build), so a binding can hand them out. Queued on
181    /// `stream`.
182    fn circuit_prepare(
183        &self,
184        _gpu: &dyn GpuBackend,
185        _config: &metrale_config::ModelConfig,
186        _levers: &crate::layers::ops::ModelLevers,
187        _stream: u64,
188    ) -> anyhow::Result<()> {
189        Ok(())
190    }
191
192    /// 2026-09-28: This layer for the circuit executor; `None` for a layer type it does not bind.
193    /// `config` and `levers` are the model's, as its decode reads them.
194    fn circuit_layer(
195        &self,
196        _config: &metrale_config::ModelConfig,
197        _levers: &crate::layers::ops::ModelLevers,
198    ) -> Option<CircuitLayer> {
199        None
200    }
201}
202
203/// 2026-09-29: The MTP draft head as the executor binds it: its one layer (weights by slot,
204/// the attention facts of its own KV cache) and that cache.
205#[derive(Debug, Clone)]
206pub struct DraftBinding {
207    /// 2026-09-29: Weights by slot, the draft attention's facts (`attn_layer_idx` is its
208    /// cache's layer), and what the circuit does not model.
209    pub layer: CircuitLayer,
210    /// 2026-09-29: The draft cache's K pool.
211    pub k_pool: DevicePtr,
212    /// 2026-09-29: The draft cache's V pool.
213    pub v_pool: DevicePtr,
214    /// 2026-09-29: Tokens per block of the draft cache.
215    pub block_size: u32,
216    /// 2026-09-29: `PagedKvCache::cache_stride` of the draft cache.
217    pub cache_stride: u64,
218    /// 2026-09-29: The vocabulary rows the draft lm_head scores and the argmax reads (the
219    /// first `mtp_vocab_size`, or all).
220    pub vocab: u32,
221    /// 2026-09-30: What the head's n-row draft (`forward_batch_position`) reads beyond the
222    /// layer; `None` for a head outside the batched propose's scope.
223    pub rows: Option<DraftRows>,
224}
225
226/// 2026-09-30: The MTP head's n-row draft facts: where its batched metadata and confidences
227/// live, and the LM-head launch it chooses at each width (`MtpHead::lm_head_batch_kernel`,
228/// `mtp_head::lm_head_rows_arm`), which the executor checks the plan against.
229#[derive(Debug, Clone)]
230pub struct DraftRows {
231    /// 2026-09-30: The batched propose's metadata allocation (`MtpHead::propose_meta`).
232    pub meta: DevicePtr,
233    /// 2026-09-30: Byte offset in scratch of the per-row top-1 log-probabilities
234    /// (`LP_SCRATCH_OFF`); the ids sit at scratch's start.
235    pub lp_offset: usize,
236    /// 2026-09-30: `lm_head_batch_kernel(n)` at index `n`; the zero handle where none serves.
237    pub lm_head_gemv: Vec<metrale_gpu_runtime::gpu::KernelHandle>,
238    /// 2026-09-30: The tile GEMM on the transposed LM-head twin is available to the head.
239    pub lm_head_twin: bool,
240}
241
242/// 2026-09-28: The model's own weights the head block reads, filled by the model.
243#[derive(Debug, Clone)]
244pub struct HeadBinding {
245    /// 2026-09-28: The final RMSNorm.
246    pub final_norm: DenseWeight,
247    /// 2026-09-28: The vocabulary projection.
248    pub lm_head: BoundWeight,
249    /// 2026-09-28: Head features the circuit does not model (a logit overlay, softcapping,
250    /// FP32 logits, a vocab-parallel head, ...).
251    pub unmodelled: Vec<String>,
252    /// 2026-09-28: The widest padded batch the BF16 head serves with the batched GEMV
253    /// (`dense_gemv_bf16_batchm`); 0 when that arm is off. Wider batches take the GEMM.
254    pub batchm_max_rows: u32,
255}
256
257/// 2026-09-28: Refuse a model the circuit misdescribes: an unbound layer, a layer or head
258/// feature no rule models, or a layer whose mixer is not the circuit's.
259pub fn check_bindings(
260    circuit: &Circuit,
261    layers: &[Option<CircuitLayer>],
262    head: &HeadBinding,
263) -> Result<Vec<CircuitLayer>> {
264    let mut problems = BTreeSet::new();
265    let mut out = Vec::with_capacity(layers.len());
266    for (i, l) in layers.iter().enumerate() {
267        let Some(l) = l else {
268            problems.insert(format!("layer {i} has no circuit binding"));
269            continue;
270        };
271        for u in &l.unmodelled {
272            problems.insert(format!("layer {i}: {u}"));
273        }
274        let want_gdn =
275            circuit.layer_kinds.get(i) == Some(&metrale_circuit::LayerKind::LinearAttention);
276        let is_gdn = matches!(l.mixer, MixerFacts::Gdn(_));
277        if want_gdn != is_gdn {
278            problems.insert(format!(
279                "layer {i}: the circuit and the model disagree on its kind"
280            ));
281        }
282        out.push(l.clone());
283    }
284    for u in &head.unmodelled {
285        problems.insert(format!("head: {u}"));
286    }
287    if !problems.is_empty() {
288        return Err(anyhow!(
289            "the circuit does not model this model:\n  {}",
290            problems.into_iter().collect::<Vec<_>>().join("\n  ")
291        ));
292    }
293    ensure!(
294        out.len() == circuit.layer_kinds.len(),
295        "the model has {} layers, the circuit {}",
296        out.len(),
297        circuit.layer_kinds.len()
298    );
299    Ok(out)
300}
metrale_model_layers/circuit_exec/bindings.rs

metrale_model_layers/circuit_exec/
bindings.rs
