metrale_model_layers/layer/transformer_layer/
write_on_accept.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `LayerWriteOnAccept`, the GDN write-on-accept hooks: stash sizing, buffer
4//! binding, and the fold of the accepted rows after a batched verify; and the
5//! carried-state verify's sizing and binding.
6//!
7//! Owner: model-layers.
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12
13/// 2026-09-26: A GDN layer's device buffers for the carried-state verify, set once by the
14/// model (`model-engine gdn_carry.rs`) before any capture.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub struct GdnCarryBinding {
17    /// 2026-09-26: This layer's engaged words, one `u32` per batch position; a carry
18    /// kernel sets each position it verified to 1.
19    pub flag: DevicePtr,
20    /// 2026-09-26: This layer's stash, `slots * seq_floats` f32.
21    pub stash: DevicePtr,
22    /// 2026-09-26: This layer's pending-row counts, one `u32` per slot.
23    pub pend: DevicePtr,
24    /// 2026-09-26: Batch position to slot, shared by every layer.
25    pub slot_tab: DevicePtr,
26    pub seq_floats: usize,
27    /// 2026-09-26: This layer's conv stash, `slots * conv_seq_elems` BF16 (four input rows
28    /// per slot).
29    pub conv_stash: DevicePtr,
30    pub conv_seq_elems: usize,
31    /// 2026-09-26: This layer's conv-state pointer per batch position, staged per verify;
32    /// a declining layer folds through it.
33    pub conv_tab: DevicePtr,
34}
35
36/// 2026-09-26: A supertrait of `TransformerLayer`; see the module header.
37pub trait LayerWriteOnAccept {
38    /// 2026-09-25: Per-sequence write-on-accept stash size in f32 elements when this layer
39    /// can run write-on-accept, `None` otherwise (the default). The model sizes the stash
40    /// from the largest answer and binds each GDN layer with [`Self::gdn_woa_bind`].
41    fn gdn_woa_stash_seq_floats(&self) -> Option<usize> {
42        None
43    }
44
45    /// 2026-09-25: Bind this layer's write-on-accept flag word and stash slab (`seqs`
46    /// sequences of [`Self::gdn_woa_stash_seq_floats`] f32 each). The model calls it
47    /// once, on the first write-on-accept request and before its graph decision, and
48    /// never moves the buffers afterwards (`gdn_woa.rs`).
49    fn gdn_woa_bind(&self, _flag: DevicePtr, _stash: DevicePtr, _seqs: usize) {}
50
51    /// 2026-09-26: Per-slot carried-state stash size in f32 elements when this layer can
52    /// run the carried-state verify, `None` otherwise (the default).
53    fn gdn_carry_seq_floats(&self) -> Option<usize> {
54        None
55    }
56
57    /// 2026-09-26: Bind this layer's carried-state buffers. Called once, before any
58    /// capture; the buffers never move afterwards.
59    fn gdn_carry_bind(&self, _binding: GdnCarryBinding) {}
60
61    /// 2026-09-30: Whether a batched verify runs one run of `states` (adjacent sequences of
62    /// width `kk`, under the WY tables at `wy_tables`) through the batched conv and the
63    /// table-form WY, as this layer's own batched verify decides it; `None` for a layer
64    /// without a batched GDN verify (the default). The circuit executor picks each run's arm
65    /// with it.
66    fn gdn_verify_run_batched(
67        &self,
68        _states: &[&mut (dyn crate::layer::LayerState + 'static)],
69        _kk: usize,
70        _gdn_wyn: bool,
71        _wy_tables: DevicePtr,
72    ) -> Result<Option<bool>> {
73        Ok(None)
74    }
75
76    /// 2026-09-25: Apply the accepted rows of the last batched verify to this layer's h
77    /// states. `h_table` is the layer's WY pointer-table slice, `na_tab` a device `u32[n]`
78    /// of accepted row counts in batch order. `Ok(false)` when the layer did nothing, as
79    /// the default does.
80    fn gdn_fold_accepted(
81        &self,
82        _gpu: &dyn GpuBackend,
83        _h_table: DevicePtr,
84        _na_tab: DevicePtr,
85        _k_rows: usize,
86        _n: usize,
87        _stream: u64,
88    ) -> Result<bool> {
89        Ok(false)
90    }
91}
metrale_model_layers/layer/transformer_layer/write_on_accept.rs

metrale_model_layers/layer/transformer_layer/
write_on_accept.rs
