metrale_model_layers/circuit_exec/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The circuit executor: run a model's decode from its fused circuit plan instead of
4//! the hand-written layer loops. At boot it instantiates the model's circuit, fuses it under the
5//! live policy with the kernels the loaded modules contain, lays its buffers out in one
6//! workspace and compiles a straight-line [`program::Program`] for single-sequence decode and
7//! one per multi-sequence batch width; a decode step then runs the program for its width
8//! between the legacy prologue (embedding, metadata upload) and the logits.
9//!
10//! Owner: model-layers circuit executor.
11//! Invariants:
12//! - Every emitter launches through an existing `ops::*` function; the executor adds no launch
13//!   code of its own.
14//! - Anything the circuit does not model (a layer feature, a switch, a head feature, a plan
15//!   kernel no emitter launches) refuses the build; a decode never runs a plan that
16//!   misdescribes the model.
17//! - The workspace is allocated once, at build, and never moves, so captured graphs stay valid.
18//!   Every program lays its buffers out from the workspace base: programs run one at a time on
19//!   the model's stream, so they share it, and it is sized for the largest.
20
21pub mod bindings;
22pub mod compile;
23mod draft_rows;
24mod emitters;
25pub mod fixed;
26pub mod kernels;
27pub mod policy;
28pub mod program;
29pub mod routes;
30pub mod sources;
31pub mod spec_key;
32pub mod verify_batch;
33
34use anyhow::{Context, Result, bail};
35use metrale_circuit::planner::plan_buffers_with;
36use metrale_circuit::{FusionPlan, Instance, Mode, Numerics, Policy, fuse};
37use metrale_config::ModelConfig;
38use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
39
40pub use bindings::{
41    AttnFacts, BoundWeight, CircuitBindings, CircuitLayer, DraftBinding, DraftRows, GdnFacts,
42    HeadBinding, MixerFacts, RopeFacts, WeightSlot,
43};
44pub use compile::{DraftFixed, Fixed};
45pub use program::{DraftPrograms, DraftRunner, GdnState, Program, StepEnv};
46
47/// 2026-09-28: Which rules a build may select.
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub enum Fusions {
50    /// 2026-09-28: Every rule whose kernel is present.
51    All,
52    /// 2026-09-28: Only `reference` rules: the kernels of `bit_identical` rules are treated as
53    /// absent, so the plan is today's routing exactly. The parity harness compares the two.
54    ReferenceOnly,
55}
56
57/// 2026-09-28: The compiled kernel modules of the served target, `(module, PTX)`: the fuser's
58/// view of which kernels exist ([`kernels::available_in`]).
59#[derive(Clone)]
60pub struct TargetModules(pub Vec<(&'static str, &'static [u8])>);
61
62impl std::fmt::Debug for TargetModules {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        write!(f, "TargetModules({} modules)", self.0.len())
65    }
66}
67
68/// 2026-09-28: What a build reads.
69pub struct Boot<'a> {
70    pub gpu: &'a dyn GpuBackend,
71    pub config: &'a ModelConfig,
72    /// 2026-09-30: The checkpoint's `config.json`, which the circuit's config map turns into the
73    /// served shape ([`sources::checkpoint_shape`]).
74    pub config_json: &'a str,
75    /// 2026-09-28: The model's levers, as its legacy layers read them (`ForwardContext::levers`).
76    pub levers: &'a crate::layers::ops::ModelLevers,
77    pub instance: &'a Instance,
78    pub policy: Policy,
79    pub layers: Vec<Option<CircuitLayer>>,
80    pub head: HeadBinding,
81    pub fixed: Fixed,
82    pub fusions: Fusions,
83    pub modules: &'a TargetModules,
84    /// 2026-09-28: The padded multi-sequence batch widths to compile (the decode graph ladder
85    /// up to the serve's widest batch).
86    pub multi_seq_rows: Vec<u64>,
87    /// 2026-09-29: The MTP verify widths `K` to compile (none without speculative decode).
88    pub verify_rows: Vec<u64>,
89    /// 2026-09-29: The MTP draft head, when one is loaded; its buffers are `fixed.draft`.
90    pub draft: Option<CircuitLayer>,
91    /// 2026-09-30: The widest batched MTP verify, in rows (`Σ k`), when the serve verifies
92    /// several sequences in one forward; `None` otherwise.
93    pub verify_batch_rows: Option<u64>,
94    /// 2026-09-30: The widest batched propose (`propose_batch_max`), whose widths from 2 get
95    /// n-row draft programs; `None` when the head proposes one sequence at a time.
96    pub draft_rows: Option<u64>,
97}
98
99/// 2026-09-28: A built executor: the decode program and the workspace it runs in.
100pub struct CircuitExec {
101    /// 2026-09-28: Decode at one row.
102    pub decode: Program,
103    /// 2026-09-28: The plan `decode` was compiled from.
104    pub decode_plan: FusionPlan,
105    /// 2026-09-28: Multi-sequence decode, one program per padded width, ascending; each with
106    /// the plan it was compiled from.
107    pub multi_seq: Vec<(Program, FusionPlan)>,
108    /// 2026-09-29: The single-sequence MTP verify, one program per `K`, ascending.
109    pub verify: Vec<(Program, FusionPlan)>,
110    /// 2026-09-29: The MTP draft head's single-row step and (2026-09-30) its n-row ones, which
111    /// the head runs itself (`DraftRunner`).
112    pub draft: Option<std::sync::Arc<DraftPrograms>>,
113    /// 2026-09-30: The runtime routes' arms, beside the primary program of each mode and row
114    /// count they apply to ([`routes`]).
115    pub routes: Vec<routes::RoutedProgram>,
116    /// 2026-09-30: Per layer, the `(h, conv)` slot pitch of its GDN state; `None` elsewhere.
117    pub gdn_pitch: Vec<Option<(usize, usize)>>,
118    /// 2026-09-30: The batched MTP verify's programs, compiled per row table.
119    pub verify_batch: Option<verify_batch::VerifyBatch>,
120    /// 2026-09-28: SHA-256 of the FUSIONS.toml the plan was chosen from.
121    pub rules_digest: String,
122    /// 2026-09-28: Which rules the build allowed.
123    pub fusions: Fusions,
124    workspace: DevicePtr,
125    workspace_bytes: u64,
126}
127
128impl CircuitExec {
129    /// 2026-09-28: Build the executor.
130    pub fn build(b: Boot<'_>) -> Result<Self> {
131        let mut served = b.instance.clone();
132        let from_checkpoint = sources::checkpoint_shape(b.config_json, b.config.vocab_size as u64)?;
133        served.shape = sources::served_shape(&b.instance.shape, &from_checkpoint)?;
134        let src = sources::sources(&served)?;
135        let loaded = match &served.precision {
136            // 2026-09-28: The served model's own policy: its parsed quantization_config under
137            // the published tier and kernel capabilities, with the instance's engine formats.
138            metrale_circuit::PrecisionSpec::Policy { engine, .. } => {
139                let policy = metrale_config::WeightQuantPolicy::for_checkpoint(
140                    crate::layers::weight_quantization(),
141                    b.config.quantization_config.as_ref(),
142                    crate::layers::kernel_caps(),
143                );
144                metrale_circuit::load_with(
145                    &served,
146                    src,
147                    &metrale_circuit::precision_policy::PolicyPrecision::new(policy, engine),
148                )?
149            }
150            metrale_circuit::PrecisionSpec::Table(_) => metrale_circuit::load(&served, src)?,
151        };
152        let unmodelled = policy::unmodelled_switches(b.levers);
153        if !unmodelled.is_empty() {
154            bail!(
155                "the circuit does not model these switches: {}",
156                unmodelled.join(", ")
157            );
158        }
159        let layers = bindings::check_bindings(&loaded.circuit, &b.layers, &b.head)?;
160        if let Some(d) = &b.draft {
161            if !d.unmodelled.is_empty() {
162                bail!(
163                    "the circuit does not model this draft head: {}",
164                    d.unmodelled.join(", ")
165                );
166            }
167            anyhow::ensure!(b.fixed.draft.is_some(), "a draft head without its buffers");
168        }
169        let present = kernels::available_in(&loaded.rules, &b.modules.0)?;
170        let mut available = present.clone();
171        if b.fusions == Fusions::ReferenceOnly {
172            for r in &loaded.rules {
173                if matches!(r.numerics, Numerics::BitIdentical { .. }) {
174                    for k in &r.kernels {
175                        available.kernels.remove(k);
176                    }
177                }
178            }
179        }
180        let table = kernels::KernelTable::resolve(b.gpu, &present);
181        let inputs = compile::Inputs {
182            gpu: b.gpu,
183            config: b.config,
184            kernels: &table,
185            fixed: &b.fixed,
186            layers: &layers,
187            head: &b.head,
188            draft: b.draft.as_ref(),
189        };
190        let shapes = std::iter::once((Mode::Decode, 1))
191            .chain(b.multi_seq_rows.iter().map(|&r| (Mode::MultiSeq, r)))
192            .chain(b.verify_rows.iter().map(|&r| (Mode::Verify, r)))
193            .chain(b.draft.iter().map(|_| (Mode::Draft, 1)));
194        routes::check_known(&loaded.runtime)?;
195        let mut laid = Vec::new();
196        for (mode, rows) in shapes {
197            let plan = fuse(
198                &loaded.circuit,
199                &loaded.rules,
200                &available,
201                &b.policy,
202                mode,
203                rows,
204            )
205            .with_context(|| format!("fusing {mode:?} at {rows} rows"))?;
206            let set = (loaded.rules.as_slice(), loaded.runtime.as_slice());
207            let arms = metrale_circuit::runtime::route_arms(
208                &loaded.circuit,
209                set,
210                &available,
211                &b.policy,
212                &plan,
213            )
214            .with_context(|| format!("fusing the runtime routes of {mode:?} at {rows} rows"))?;
215            for (route, plan) in
216                std::iter::once((None, plan)).chain(arms.into_iter().map(|(r, p)| (Some(r), p)))
217            {
218                let layout = compile::layout(&loaded.circuit, &plan)?;
219                let buffers = plan_buffers_with(&loaded.circuit, &plan, rows, &layout)?;
220                laid.push((route, plan, layout, buffers));
221            }
222        }
223        let mut workspace_bytes = laid
224            .iter()
225            .map(|(_, _, _, buf)| buf.arena_bytes)
226            .max()
227            .unwrap_or(0)
228            .max(1);
229        // 2026-09-30: The batched verify compiles later, per table, in this same workspace; only
230        // for an instance whose batched-verify plans are checked in (`Instance::verify_batch`).
231        let verify_batch_rows = b
232            .verify_batch_rows
233            .filter(|_| !b.instance.verify_batch.is_empty());
234        if b.verify_batch_rows.is_some() && verify_batch_rows.is_none() {
235            tracing::info!(
236                "circuit: `{}` states no batched-verify plans; each sequence verifies alone",
237                b.instance.recipe
238            );
239        }
240        if let Some(max) = verify_batch_rows {
241            for t in verify_batch::sizing_tables(max) {
242                let a = verify_batch::arena_bytes(
243                    &loaded.circuit,
244                    &loaded.rules,
245                    &available,
246                    &b.policy,
247                    &t,
248                )?;
249                workspace_bytes = workspace_bytes.max(a);
250            }
251        }
252        let draft_laid = match (&b.fixed.draft, b.draft_rows) {
253            (Some(d), Some(max)) if b.draft.is_some() && d.rows.is_some() => {
254                draft_rows::lay_out(&loaded, &available, &b.policy, max)
255            }
256            _ => Vec::new(),
257        };
258        for l in &draft_laid {
259            workspace_bytes = workspace_bytes.max(l.buffers.arena_bytes);
260        }
261        let workspace = b
262            .gpu
263            .alloc(workspace_bytes as usize)
264            .context("allocating the circuit workspace")?;
265        let mut programs = Vec::with_capacity(laid.len());
266        let mut routed = Vec::new();
267        for (route, plan, layout, buffers) in laid {
268            match compile::compile(
269                &loaded.circuit,
270                &plan,
271                &layout,
272                &buffers,
273                workspace,
274                &inputs,
275            ) {
276                Ok(program) => match route {
277                    None => programs.push((program, plan)),
278                    Some(route) => routed.push(routes::RoutedProgram {
279                        route,
280                        program,
281                        plan,
282                    }),
283                },
284                Err(e) => {
285                    b.gpu.free(workspace).ok();
286                    return Err(
287                        e.context(format!("compiling {:?} at {} rows", plan.mode, plan.rows))
288                    );
289                }
290            }
291        }
292        let mut programs = programs.into_iter();
293        let (decode, plan) = programs.next().context("no decode program")?;
294        let (draft, rest): (Vec<_>, Vec<_>) = programs.partition(|(p, _)| p.mode == Mode::Draft);
295        let (multi_seq, verify): (Vec<_>, Vec<_>) = rest
296            .into_iter()
297            .partition(|(p, _)| p.mode == Mode::MultiSeq);
298        let draft = draft.into_iter().next().map(|one| {
299            let mut programs = vec![one];
300            programs.extend(draft_rows::compile_all(
301                &loaded.circuit,
302                draft_laid,
303                workspace,
304                &inputs,
305            ));
306            std::sync::Arc::new(DraftPrograms { programs })
307        });
308        tracing::info!(
309            "circuit decode: {} launches/step, plan {} ({} rules, {:?}), {} multi-seq widths, \
310             {} verify widths, {} draft widths, {} runtime-route programs, workspace {} KiB",
311            decode.launches.len(),
312            &plan.digest[..12],
313            loaded.rules.len(),
314            b.fusions,
315            multi_seq.len(),
316            verify.len(),
317            draft.as_ref().map_or(0, |d| d.programs.len()),
318            routed.len(),
319            workspace_bytes / 1024
320        );
321        let gdn_pitch = layers
322            .iter()
323            .map(|l| match l.mixer {
324                MixerFacts::Gdn(g) => Some((g.h_slot_bytes as usize, g.conv_state_bytes as usize)),
325                MixerFacts::Attention(_) => None,
326            })
327            .collect();
328        let verify_batch = verify_batch_rows.map(|max| {
329            verify_batch::VerifyBatch::new(
330                verify_batch::Parts {
331                    circuit: loaded.circuit.clone(),
332                    rules: loaded.rules.clone(),
333                    available: available.clone(),
334                    policy: b.policy.clone(),
335                    kernels: table.clone(),
336                    fixed: b.fixed.clone(),
337                    layers: layers.clone(),
338                    head: b.head.clone(),
339                    config: b.config.clone(),
340                },
341                workspace,
342                workspace_bytes,
343                max,
344            )
345        });
346        Ok(CircuitExec {
347            decode,
348            decode_plan: plan,
349            multi_seq,
350            verify,
351            draft,
352            routes: routed,
353            gdn_pitch,
354            verify_batch,
355            rules_digest: loaded.rules_digest,
356            fusions: b.fusions,
357            workspace,
358            workspace_bytes,
359        })
360    }
361
362    /// 2026-09-28: The multi-sequence program for `rows` padded rows, if one was compiled.
363    pub fn multi_seq_program(&self, rows: u64) -> Option<&Program> {
364        self.multi_seq
365            .iter()
366            .find(|(p, _)| p.rows == rows)
367            .map(|(p, _)| p)
368    }
369
370    /// 2026-09-30: The program a multi-sequence step of `rows` padded rows runs, given its GDN
371    /// states: the arm of the first runtime route whose condition holds, else the primary one.
372    pub fn multi_seq_step(&self, rows: u64, gdn: &[Vec<GdnState>]) -> Result<&Program> {
373        for r in self
374            .routes
375            .iter()
376            .filter(|r| r.program.mode == Mode::MultiSeq && r.program.rows == rows)
377        {
378            if routes::holds(&r.route, &self.gdn_pitch, gdn)? {
379                return Ok(&r.program);
380            }
381        }
382        self.multi_seq_program(rows)
383            .with_context(|| format!("no circuit program was compiled for {rows} rows"))
384    }
385
386    /// 2026-09-29: The draft head's program, if one was compiled.
387    pub fn draft_runner(&self) -> Option<std::sync::Arc<dyn DraftRunner>> {
388        self.draft
389            .as_ref()
390            .map(|d| d.clone() as std::sync::Arc<dyn DraftRunner>)
391    }
392
393    /// 2026-09-29: The verify program for `k` rows, if one was compiled.
394    pub fn verify_program(&self, k: u64) -> Option<&Program> {
395        self.verify
396            .iter()
397            .find(|(p, _)| p.rows == k)
398            .map(|(p, _)| p)
399    }
400
401    /// 2026-09-28: One digest over every compiled plan, in order (decode, then the
402    /// multi-sequence widths ascending, then the verify widths, the draft step and, 2026-09-30,
403    /// the runtime routes' arms): what a record of this forward discloses.
404    pub fn plans_digest(&self) -> String {
405        metrale_circuit::digest::plans_digest(
406            std::iter::once(self.decode_plan.digest.as_str()).chain(
407                self.multi_seq
408                    .iter()
409                    .chain(&self.verify)
410                    .map(|(_, p)| p.digest.as_str())
411                    .chain(
412                        self.draft
413                            .iter()
414                            .flat_map(|d| d.programs.iter().map(|(_, p)| p.digest.as_str())),
415                    )
416                    .chain(self.routes.iter().map(|r| r.plan.digest.as_str())),
417            ),
418        )
419    }
420
421    /// 2026-09-28: Bytes of the workspace.
422    pub fn workspace_bytes(&self) -> u64 {
423        self.workspace_bytes
424    }
425
426    /// 2026-09-28: Free the workspace. The caller must first destroy every graph that captured
427    /// a program of this executor.
428    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
429        if let Some(p) = &self.draft {
430            anyhow::ensure!(
431                std::sync::Arc::strong_count(p) == 1,
432                "the draft program is still installed; remove it before freeing the workspace"
433            );
434        }
435        gpu.free(self.workspace)
436    }
437}
438
439#[cfg(test)]
440#[path = "exec_declared_tests.rs"]
441mod exec_declared_tests;
442#[cfg(test)]
443#[path = "exec_draft_rows_tests.rs"]
444mod exec_draft_rows_tests;
445#[cfg(test)]
446#[path = "exec_draft_tests.rs"]
447mod exec_draft_tests;
448#[cfg(test)]
449#[path = "exec_fixture.rs"]
450mod exec_fixture;
451#[cfg(test)]
452#[path = "exec_fixture_run.rs"]
453mod exec_fixture_run;
454#[cfg(test)]
455#[path = "exec_multi_tests.rs"]
456mod exec_multi_tests;
457#[cfg(test)]
458#[path = "exec_tests.rs"]
459mod exec_tests;
460#[cfg(test)]
461#[path = "exec_verify_batch_tests.rs"]
462mod exec_verify_batch_tests;
463#[cfg(test)]
464#[path = "exec_verify_tests.rs"]
465mod exec_verify_tests;
metrale_model_layers/circuit_exec/mod.rs

metrale_model_layers/circuit_exec/
mod.rs
