metrale_model_layers/circuit_exec/
program.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: A compiled circuit program: the straight-line launches of one plan, with every
4//! host decision (kernel, grid, sizes, pointers into the model's fixed buffers and the
5//! workspace) made at compile time.
6//!
7//! Owner: model-layers circuit executor.
8//! Invariants:
9//! - A step reads only [`StepEnv`]: the stream, each row's GDN state pointers and the
10//!   block-table width uploaded for the step. Everything else a launch uses was fixed when it
11//!   was compiled, so a CUDA-graph capture of [`Program::run`] replays correctly for any step
12//!   with the same state pointers (the graph caches key on the sequence's slot).
13//! - `run` issues no synchronize and no device-to-host copy, so it may be captured.
14//! - A program holds exactly `FusionPlan::launches` kernel launches and `FusionPlan::copies`
15//!   copies.
16
17use anyhow::{Context, Result};
18use metrale_circuit::Mode;
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
20
21/// 2026-09-29: The most rollback points an MTP verify writes: one per row but the last, at
22/// K = 4.
23pub const MAX_VERIFY_STEPS: usize = 3;
24
25/// 2026-09-28: One sequence's recurrent state in one GatedDeltaNet layer.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct GdnState {
28    /// 2026-09-28: The h state (FP32).
29    pub h: DevicePtr,
30    /// 2026-09-28: The conv1d window.
31    pub conv: DevicePtr,
32    /// 2026-09-29: The h state after verify row `t` (`SsmLayerState::h_state_intermediates`);
33    /// NULL where the sequence has none. Only a verify reads them.
34    pub h_steps: [DevicePtr; MAX_VERIFY_STEPS],
35    /// 2026-09-29: The conv window after verify row `t`, indexed like `h_steps`.
36    pub conv_steps: [DevicePtr; MAX_VERIFY_STEPS],
37}
38
39/// 2026-09-28: What varies between two runs of one program.
40pub struct StepEnv<'a> {
41    /// 2026-09-28: The backend.
42    pub gpu: &'a dyn GpuBackend,
43    /// 2026-09-28: The stream the step runs (and is captured) on.
44    pub stream: u64,
45    /// 2026-09-28: Per layer, each row's GDN state (row `i` = sequence `i`, padding rows
46    /// included); empty for a layer without one.
47    pub gdn: &'a [Vec<GdnState>],
48    /// 2026-09-28: `AttnMetadataDev::max_blocks_per_seq` of this step.
49    pub max_blocks_per_seq: u32,
50}
51
52impl StepEnv<'_> {
53    /// 2026-09-28: Row `row`'s GDN state in layer `layer`; an error when there is none.
54    pub fn gdn_state(&self, layer: usize, row: usize) -> Result<GdnState> {
55        self.gdn
56            .get(layer)
57            .and_then(|l| l.get(row))
58            .copied()
59            .with_context(|| format!("layer {layer} has no GDN state for row {row}"))
60    }
61}
62
63/// 2026-09-28: The body of one launch.
64pub(crate) type RunFn = Box<dyn Fn(&StepEnv<'_>) -> Result<()> + Send + Sync>;
65
66/// 2026-09-29: What a launch puts on the stream.
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68pub enum LaunchKind {
69    /// 2026-09-29: A kernel.
70    Kernel,
71    /// 2026-09-29: A copy-engine transfer.
72    Copy,
73}
74
75/// 2026-09-28: One kernel launch or copy.
76pub struct Launch {
77    /// 2026-09-28: Index of the plan group it belongs to.
78    pub group: usize,
79    /// 2026-09-28: `module::func` it launches, or `copy` for a transfer.
80    pub kernel: String,
81    /// 2026-09-29: Kernel or copy.
82    pub kind: LaunchKind,
83    pub(crate) run: RunFn,
84}
85
86/// 2026-09-28: The compiled launches of one plan.
87pub struct Program {
88    /// 2026-09-28: The plan's mode.
89    pub mode: Mode,
90    /// 2026-09-28: The plan's padded rows.
91    pub rows: u64,
92    /// 2026-09-28: The plan's digest (`metrale_circuit::digest::plan_digest`).
93    pub plan_digest: String,
94    /// 2026-09-28: The launches, in order.
95    pub launches: Vec<Launch>,
96}
97
98impl Program {
99    /// 2026-09-28: Issue every launch on `env.stream`.
100    pub fn run(&self, env: &StepEnv<'_>) -> Result<()> {
101        for l in &self.launches {
102            (l.run)(env).with_context(|| format!("circuit group {} ({})", l.group, l.kernel))?;
103        }
104        Ok(())
105    }
106}
107
108/// 2026-09-29: A program another component runs in place of its own forward (the MTP draft
109/// head's step), without GDN state. 2026-09-30: One program per row count: one row for a
110/// single sequence's draft, n rows for the batched propose's.
111pub trait DraftRunner: Send + Sync {
112    /// 2026-09-30: Whether a program for `rows` draft rows was compiled.
113    fn serves(&self, rows: u64) -> bool;
114    /// 2026-09-29: Issue the `rows`-row program on `stream` for a step whose block table is
115    /// `max_blocks_per_seq` wide; an error when none was compiled for `rows`.
116    fn run_draft(
117        &self,
118        gpu: &dyn GpuBackend,
119        stream: u64,
120        rows: u64,
121        max_blocks_per_seq: u32,
122    ) -> Result<()>;
123}
124
125/// 2026-09-30: The draft head's compiled programs, one per row count, ascending (one row
126/// first).
127pub struct DraftPrograms {
128    pub programs: Vec<(Program, metrale_circuit::FusionPlan)>,
129}
130
131impl DraftRunner for DraftPrograms {
132    fn serves(&self, rows: u64) -> bool {
133        self.programs.iter().any(|(p, _)| p.rows == rows)
134    }
135
136    fn run_draft(
137        &self,
138        gpu: &dyn GpuBackend,
139        stream: u64,
140        rows: u64,
141        max_blocks_per_seq: u32,
142    ) -> Result<()> {
143        let (p, _) = self
144            .programs
145            .iter()
146            .find(|(p, _)| p.rows == rows)
147            .with_context(|| format!("no draft program was compiled for {rows} rows"))?;
148        p.run(&StepEnv {
149            gpu,
150            stream,
151            gdn: &[],
152            max_blocks_per_seq,
153        })
154    }
155}
metrale_model_layers/circuit_exec/program.rs

metrale_model_layers/circuit_exec/
program.rs
