metrale_model_layers/circuit_exec/
verify_batch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The batched MTP verify's programs (`Mode::VerifyBatch`). A batched verify's
4//! plan depends on its row table (the runs of equal `k` and their contiguity), and a serve
5//! meets many tables, so each program is compiled the first time its table is verified and kept
6//! in a bounded cache. The legacy forward's graphs are keyed by the batch's `(slot, k)` pairs,
7//! which fix the table, so a graph captured around a program replays that program's launches.
8//!
9//! Owner: model-layers circuit executor.
10//! Invariants:
11//! - Every program is laid out from the shared workspace base, and the build sized the workspace
12//!   for the widest table it admits (`max_rows`); a larger table is refused, never placed past
13//!   the workspace.
14//! - The cache holds at most [`CACHE_CAP`] programs; a miss compiles, an eviction drops the least
15//!   recently used.
16
17use std::collections::BTreeMap;
18use std::sync::Arc;
19
20use anyhow::{Context, Result, ensure};
21use metrale_circuit::planner::plan_buffers_with;
22use metrale_circuit::{AvailableKernels, Circuit, FusionPlan, Policy, RowTable, Rule, fuse_table};
23use metrale_config::ModelConfig;
24use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
25use parking_lot::Mutex;
26
27use super::bindings::{CircuitLayer, HeadBinding};
28use super::compile::{self, Fixed};
29use super::kernels::KernelTable;
30use super::program::Program;
31
32/// 2026-09-30: Programs kept per executor.
33pub const CACHE_CAP: usize = 256;
34
35/// 2026-09-30: A compiled batched-verify program and its plan.
36pub struct VerifyBatchProgram {
37    pub program: Program,
38    pub plan: FusionPlan,
39}
40
41/// 2026-09-30: What compiling a batched-verify program reads, kept from the build.
42pub struct VerifyBatch {
43    circuit: Circuit,
44    rules: Vec<Rule>,
45    available: AvailableKernels,
46    policy: Policy,
47    kernels: KernelTable,
48    fixed: Fixed,
49    layers: Vec<CircuitLayer>,
50    head: HeadBinding,
51    config: ModelConfig,
52    workspace: DevicePtr,
53    workspace_bytes: u64,
54    /// 2026-09-30: The widest table admitted, in rows.
55    pub max_rows: u64,
56    cache: Mutex<(u64, BTreeMap<RowTable, (u64, Arc<VerifyBatchProgram>)>)>,
57}
58
59/// 2026-09-30: The tables the workspace is sized over: for every row count `4..=max_rows`, one
60/// table of that many rows (runs of 2, with one run of 3 for an odd count).
61pub fn sizing_tables(max_rows: u64) -> Vec<RowTable> {
62    (4..=max_rows)
63        .map(|r| {
64            let text = if r % 2 == 0 {
65                format!("2x{}", r / 2)
66            } else {
67                // 2026-10-01: A run of one sequence is never batched, so it is marked `!`.
68                let tail = (r - 3) / 2;
69                format!("3x1! 2x{tail}{}", if tail == 1 { "!" } else { "" })
70            };
71            RowTable::parse(&text).expect("a well-formed sizing table")
72        })
73        .collect()
74}
75
76/// 2026-09-30: The arena bytes of `table`'s plan.
77pub fn arena_bytes(
78    circuit: &Circuit,
79    rules: &[Rule],
80    available: &AvailableKernels,
81    policy: &Policy,
82    table: &RowTable,
83) -> Result<u64> {
84    let plan = fuse_table(circuit, rules, available, policy, table)
85        .with_context(|| format!("fusing the batched verify of `{table}`"))?;
86    let layout = compile::layout(circuit, &plan)?;
87    Ok(plan_buffers_with(circuit, &plan, plan.rows, &layout)?.arena_bytes)
88}
89
90/// 2026-09-30: The build's state a batched-verify compile reads.
91pub struct Parts {
92    pub circuit: Circuit,
93    pub rules: Vec<Rule>,
94    pub available: AvailableKernels,
95    pub policy: Policy,
96    pub kernels: KernelTable,
97    pub fixed: Fixed,
98    pub layers: Vec<CircuitLayer>,
99    pub head: HeadBinding,
100    pub config: ModelConfig,
101}
102
103impl VerifyBatch {
104    /// 2026-09-30: The compiler over `parts`, in a workspace of `workspace_bytes` at
105    /// `workspace`, admitting tables of up to `max_rows` rows.
106    pub fn new(parts: Parts, workspace: DevicePtr, workspace_bytes: u64, max_rows: u64) -> Self {
107        let Parts {
108            circuit,
109            rules,
110            available,
111            policy,
112            kernels,
113            fixed,
114            layers,
115            head,
116            config,
117        } = parts;
118        VerifyBatch {
119            circuit,
120            rules,
121            available,
122            policy,
123            kernels,
124            fixed,
125            layers,
126            head,
127            config,
128            workspace,
129            workspace_bytes,
130            max_rows,
131            cache: Mutex::new((0, BTreeMap::new())),
132        }
133    }
134
135    /// 2026-09-30: The program for `table`, compiled on a miss.
136    pub fn program(
137        &self,
138        gpu: &dyn GpuBackend,
139        table: &RowTable,
140    ) -> Result<Arc<VerifyBatchProgram>> {
141        ensure!(
142            table.rows() <= self.max_rows,
143            "a batched verify of {} rows is past the {} this executor was sized for",
144            table.rows(),
145            self.max_rows
146        );
147        {
148            let mut c = self.cache.lock();
149            c.0 += 1;
150            let tick = c.0;
151            if let Some(e) = c.1.get_mut(table) {
152                e.0 = tick;
153                return Ok(e.1.clone());
154            }
155        }
156        let plan = fuse_table(
157            &self.circuit,
158            &self.rules,
159            &self.available,
160            &self.policy,
161            table,
162        )
163        .with_context(|| format!("fusing the batched verify of `{table}`"))?;
164        let layout = compile::layout(&self.circuit, &plan)?;
165        let buffers = plan_buffers_with(&self.circuit, &plan, plan.rows, &layout)?;
166        ensure!(
167            buffers.arena_bytes <= self.workspace_bytes,
168            "the batched verify of `{table}` needs {} workspace bytes; the build sized {}",
169            buffers.arena_bytes,
170            self.workspace_bytes
171        );
172        let inputs = compile::Inputs {
173            gpu,
174            config: &self.config,
175            kernels: &self.kernels,
176            fixed: &self.fixed,
177            layers: &self.layers,
178            head: &self.head,
179            draft: None,
180        };
181        let program = compile::compile(
182            &self.circuit,
183            &plan,
184            &layout,
185            &buffers,
186            self.workspace,
187            &inputs,
188        )
189        .with_context(|| format!("compiling the batched verify of `{table}`"))?;
190        let entry = Arc::new(VerifyBatchProgram { program, plan });
191        let mut c = self.cache.lock();
192        let tick = c.0;
193        if c.1.len() >= CACHE_CAP
194            && let Some(oldest) =
195                c.1.iter()
196                    .min_by_key(|(_, (t, _))| *t)
197                    .map(|(k, _)| k.clone())
198        {
199            c.1.remove(&oldest);
200        }
201        c.1.insert(table.clone(), (tick, entry.clone()));
202        Ok(entry)
203    }
204
205    /// 2026-09-30: The metadata the programs read (`Fixed::verify_batch_meta`).
206    pub fn fixed_meta(&self) -> &crate::layer::AttnMetadataDev {
207        &self.fixed.verify_batch_meta
208    }
209
210    /// 2026-09-30: Where the programs write each row's token (`Fixed::verify_batch_tokens`).
211    pub fn fixed_tokens(&self) -> DevicePtr {
212        self.fixed.verify_batch_tokens
213    }
214
215    /// 2026-09-30: Programs cached now.
216    pub fn cached(&self) -> usize {
217        self.cache.lock().1.len()
218    }
219}
metrale_model_layers/circuit_exec/verify_batch.rs

metrale_model_layers/circuit_exec/
verify_batch.rs
