metrale_model_engine/model/trait_impl/
impl_circuit.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: `impl ModelCircuit for TransformerModel`: build the circuit executor from the live
4//! layers and buffers, swap it in, and run its programs in place of the layer loops: decode
5//! (`decode_forward_body`) and multi-sequence decode (`decode_batch_compute_main_with`). Also
6//! the one statement of the single-sequence decode metadata layout, which
7//! `decode_dispatch_with` uploads and the executor reads.
8//!
9//! Owner: model-engine (decode).
10//! Invariants:
11//! - A switch drops every captured decode and verify graph before the previous executor's
12//!   workspace is freed, so no graph replays freed memory.
13//! - Model features the circuit does not model refuse the build.
14
15use anyhow::{Context, Result, bail, ensure};
16use metrale_gpu_runtime::gpu::DevicePtr;
17use metrale_model_layers::circuit_exec::{
18    BoundWeight, CircuitExec, DraftFixed, Fixed, Fusions, HeadBinding, MixerFacts, TargetModules,
19    policy,
20};
21use metrale_model_layers::layer::AttnMetadataDev;
22
23use super::super::types::TransformerModel;
24use crate::traits::{ForwardDisclosure, ForwardSelect, ModelCircuit};
25
26/// 2026-09-28: Offset of the single-sequence decode metadata in the scratch buffer.
27const DECODE_META_OFFSET: usize = 32768;
28
29impl TransformerModel {
30    /// 2026-09-28: The single-sequence decode metadata at its fixed upload address.
31    pub(super) fn decode_meta(&self, max_blocks: u32, seq_slot: DevicePtr) -> AttnMetadataDev {
32        let base = self.buffers.scratch().offset(DECODE_META_OFFSET);
33        AttnMetadataDev {
34            positions: base,
35            positions_h: base,
36            positions_w: base,
37            slot: base.offset(8),
38            seq_len: base.offset(16),
39            block_table: base.offset(256),
40            max_blocks_per_seq: max_blocks,
41            num_seqs: 1,
42            seq_slot,
43            moe_row_adapter: DevicePtr::NULL,
44        }
45    }
46
47    /// 2026-09-28: Model features the circuit does not model.
48    fn circuit_unmodelled(&self) -> Vec<String> {
49        let kv_swap = self.kv_cache.lock().config().cache_blocks_per_seq.is_some();
50        [
51            (
52                self.config.tp_world_size > 1 || self.comm.is_some(),
53                "tensor or expert parallelism",
54            ),
55            (
56                !self.dflash_capture_layers.is_empty(),
57                "DFlash hidden capture",
58            ),
59            (self.lora.is_some(), "LoRA adapters"),
60            (
61                self.config.mamba_num_heads > 0,
62                "Mamba-2 state normalisation",
63            ),
64            (self.config.kv_lora_rank > 0, "latent attention"),
65            (kv_swap, "high-speed swap"),
66            (self.profile, "profiling"),
67            (
68                metrale_model_layers::ships_vanilla_norm_weights(&self.config),
69                "vanilla RMSNorm weights (the rules launch the 1 + w kernels)",
70            ),
71        ]
72        .into_iter()
73        .filter(|(present, _)| *present)
74        .map(|(_, what)| what.to_string())
75        .collect()
76    }
77
78    /// 2026-09-28: The head block's weights, and its unmodelled features.
79    fn circuit_head(&self) -> (HeadBinding, &'static str) {
80        let (dtype, quantized) = if self.lm_head_q6k.is_some() {
81            ("q6k", true)
82        } else if self.lm_head_fp8.is_some() {
83            ("fp8", true)
84        } else if self.lm_head_nvfp4.is_some() {
85            ("nvfp4", true)
86        } else {
87            ("bf16", false)
88        };
89        let (batchm_max_rows, m16_tc) = self.bf16_head_route();
90        let unmodelled = [
91            (quantized, "a quantized lm_head"),
92            (m16_tc, "the tensor-core BF16 head (lm_head_m16_tc)"),
93            (self.use_fp32_logits, "FP32 logits"),
94            (
95                self.logit_softcap_kernel.0 != 0 || self.logit_softcap_fp32_kernel.0 != 0,
96                "logit softcapping",
97            ),
98            (self.overlays.is_some(), "token overlays"),
99            (
100                self.config.final_norm_identity,
101                "a checkpoint without a final norm",
102            ),
103        ]
104        .into_iter()
105        .filter(|(present, _)| *present)
106        .map(|(_, what)| what.to_string())
107        .collect();
108        let head = HeadBinding {
109            final_norm: self.final_norm,
110            lm_head: BoundWeight::Dense(self.lm_head_weight),
111            unmodelled,
112            batchm_max_rows,
113        };
114        (head, dtype)
115    }
116
117    /// 2026-09-28: Build the executor for `instance`.
118    fn build_circuit(
119        &self,
120        instance: &metrale_circuit::Instance,
121        fusions: Fusions,
122        modules: &TargetModules,
123        config_json: &str,
124    ) -> Result<CircuitExec> {
125        let unmodelled = self.circuit_unmodelled();
126        if !unmodelled.is_empty() {
127            bail!(
128                "the circuit does not model this model: {}",
129                unmodelled.join(", ")
130            );
131        }
132        let stream = self.gpu.default_stream();
133        for l in &self.layers {
134            l.circuit_prepare(self.gpu.as_ref(), &self.config, &self.levers, stream)?;
135        }
136        // 2026-09-30: The batched MTP verify runs the carried-state GDN verify, whose buffers
137        // the layers bind once (`gdn_carry.rs`); binding them now puts them in the layers'
138        // circuit facts. Without them the executor compiles no batched verify and each
139        // sequence verifies alone.
140        let verify_batch_rows = if self.proposer.is_some() && self.gdn_carry_bind_now()? {
141            Some(
142                (4 * metrale_model_layers::speculative::mtp_max_seqs())
143                    .min(super::verify_e2::VERIFY_ROW_CAP) as u64,
144            )
145        } else {
146            None
147        };
148        let layers: Vec<_> = self
149            .layers
150            .iter()
151            .map(|l| l.circuit_layer(&self.config, &self.levers))
152            .collect();
153        let mut kv_dtypes = layers.iter().flatten().filter_map(|l| match l.mixer {
154            MixerFacts::Attention(a) => Some(a.kv_dtype),
155            MixerFacts::Gdn(_) => None,
156        });
157        let kv = kv_dtypes
158            .next()
159            .context("the model has no attention layer to read a KV dtype from")?;
160        ensure!(
161            kv_dtypes.all(|d| d == kv),
162            "attention layers use different KV-cache dtypes; the circuit states one"
163        );
164        let (head, lm_head_dtype) = self.circuit_head();
165        let draft = self
166            .proposer
167            .as_ref()
168            .and_then(|p| p.circuit_draft(&self.config, &self.levers));
169        let fixed = {
170            let cache = self.kv_cache.lock();
171            let n = cache.num_layers();
172            Fixed {
173                hidden: self.buffers.hidden_states(),
174                residual: self.buffers.residual(),
175                logits: self.buffers.logits(),
176                meta: self.decode_meta(0, DevicePtr::NULL),
177                batch_meta: self.batch_meta_at(
178                    self.batch_meta_base(),
179                    0,
180                    0,
181                    DevicePtr::NULL,
182                    DevicePtr::NULL,
183                ),
184                ffn_act_q8: self.buffers.ffn_act_q8(),
185                tokens: self.buffers.scratch(),
186                draft: draft.as_ref().map(|d| DraftFixed {
187                    embed: self.buffers.ssm_qkvz(),
188                    meta: metrale_model_layers::layers::mtp_meta::mtp_attn_meta_dev(
189                        self.buffers
190                            .scratch()
191                            .offset(metrale_model_layers::layers::mtp_meta::MTP_META_OFFSET),
192                        0,
193                    ),
194                    k_pool: d.k_pool,
195                    v_pool: d.v_pool,
196                    block_size: d.block_size,
197                    cache_stride: d.cache_stride,
198                    vocab: d.vocab,
199                    rows: d.rows.clone(),
200                }),
201                verify_meta: self.verify_meta(0, 0, DevicePtr::NULL),
202                verify_batch_meta: self.verify_batch_meta_at(),
203                verify_wy_tables: self.verify_wy_tables,
204                verify_batch_tokens: super::verify_e::mapped_argmax_host_dev(self.gpu.as_ref())
205                    .map_or(self.buffers.scratch(), |(_, d)| d),
206                k_pools: (0..n).map(|i| cache.k_pool_ptr(i)).collect(),
207                v_pools: (0..n).map(|i| cache.v_pool_ptr(i)).collect(),
208                block_size: u32::try_from(cache.block_size())?,
209                cache_stride: cache.cache_stride() as u64,
210            }
211        };
212        CircuitExec::build(metrale_model_layers::circuit_exec::Boot {
213            gpu: self.gpu.as_ref(),
214            config: &self.config,
215            config_json,
216            levers: &self.levers,
217            instance,
218            policy: policy::live_policy(&self.levers, policy::kv_dtype_name(kv)?, lm_head_dtype),
219            layers,
220            head,
221            fixed,
222            fusions,
223            modules,
224            multi_seq_rows: self.circuit_widths(),
225            // 2026-09-29: The MTP verify widths the scheduler runs one sequence at
226            // (`serial_verify_plan`), when a drafter is loaded.
227            draft: draft.map(|d| d.layer),
228            verify_rows: if self.proposer.is_some() {
229                vec![2, 3, 4]
230            } else {
231                Vec::new()
232            },
233            verify_batch_rows,
234            // 2026-09-30: The batched propose's widths, up to the sequences MTP runs at once.
235            draft_rows: self
236                .proposer
237                .as_ref()
238                .map(|p| p.propose_batch_max(&self.buffers, &self.config))
239                .filter(|&w| w >= 2)
240                .map(|w| w.min(metrale_model_layers::speculative::mtp_max_seqs()) as u64),
241        })
242    }
243
244    /// 2026-09-28: The padded widths the executor compiles: every rung of the decode ladder up
245    /// to the one the serve's widest batch pads to.
246    fn circuit_widths(&self) -> Vec<u64> {
247        let widest = crate::traits::padded_batch_n(self.levers.max_decode_seqs as usize);
248        crate::traits::DECODE_BATCH_LADDER
249            .iter()
250            .filter(|&&r| r <= widest)
251            .map(|&r| r as u64)
252            .collect()
253    }
254}
255
256impl ModelCircuit for TransformerModel {
257    fn set_forward(&self, sel: &ForwardSelect) -> Result<()> {
258        let next = match sel {
259            ForwardSelect::Legacy => None,
260            ForwardSelect::Circuit {
261                instance,
262                fusions,
263                modules,
264                config_json,
265            } => Some(self.build_circuit(instance, *fusions, modules, config_json)?),
266        };
267        self.destroy_lora_decode_graphs();
268        // 2026-09-29: The draft head drops the previous executor's program before its
269        // workspace is freed, and takes the next one's after the swap.
270        let runner = next.as_ref().and_then(CircuitExec::draft_runner);
271        if let Some(p) = self.proposer.as_ref() {
272            p.set_circuit_draft(None);
273        }
274        let prev = std::mem::replace(&mut *self.circuit.write(), next);
275        if let Some(prev) = prev {
276            prev.free(self.gpu.as_ref())?;
277        }
278        if let (Some(p), Some(r)) = (self.proposer.as_ref(), runner) {
279            p.set_circuit_draft(Some(r));
280        }
281        Ok(())
282    }
283
284    fn forward_disclosure(&self) -> ForwardDisclosure {
285        match &*self.circuit.read() {
286            None => ForwardDisclosure::legacy(),
287            Some(e) => ForwardDisclosure {
288                forward: match e.fusions {
289                    Fusions::All => "circuit",
290                    Fusions::ReferenceOnly => "circuit-reference",
291                },
292                plan_digest: Some(e.plans_digest()),
293                launches_per_step: Some(e.decode.launches.len()),
294            },
295        }
296    }
297}
metrale_model_engine/model/trait_impl/impl_circuit.rs

metrale_model_engine/model/trait_impl/
impl_circuit.rs
