metrale_model_layers/layer/
transformer_layer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The `TransformerLayer` trait, the per-layer hooks the model drives (grouped
4//! into supertraits in the child modules), and the layout constants of the batched-verify
5//! WY pointer tables.
6//!
7//! Owner: model-layers.
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_cache::kv_cache::PagedKvCache;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13
14use super::{ForwardContext, GdnPrefillBuffers, LayerState};
15
16mod aux_state;
17mod capabilities;
18mod default_loops;
19mod graph_hooks;
20mod split_prefill;
21mod weight_setup;
22mod write_on_accept;
23
24pub use aux_state::LayerAuxState;
25pub use capabilities::LayerCapabilities;
26pub use graph_hooks::LayerGraphHooks;
27pub use split_prefill::LayerSplitPrefill;
28pub use weight_setup::LayerWeightSetup;
29pub use write_on_accept::{GdnCarryBinding, LayerWriteOnAccept};
30
31/// 2026-09-25: Batched-verify WY pointer-table layout. `TransformerModel::upload_verify_wy_tables`
32/// (model-engine `verify_e2.rs`) writes it; the qwen3_ssm `GdnStates::Multi` arm and the
33/// write-on-accept fold read it.
34///
35/// Per GDN layer, `VERIFY_WY_TABLES_PER_LAYER` tables of `VERIFY_WY_TABLE_SEQS` u64
36/// entries lie back to back: table 0 holds each sequence's `h_state`, table `t + 1` its
37/// intermediate `t`. A k-row verify writes at most the first k tables and leaves every
38/// other entry zero. The strides do not depend on k.
39///
40/// This constant is the most sequences one batched verify holds: it sizes the tables and
41/// the verify hidden and catch-up stashes (`impl_a1/spec_buffers.rs`), the GDN carry and
42/// write-on-accept tables, and `can_batch_verify_dispatch` (`verify_e.rs`) admits at most
43/// this many. 2026-09-29: 128, the serve's widest batch (`DECODE_META_MAX_ROWS`), so one
44/// verify forward covers every speculating sequence: a verify split in chunks reads the MoE
45/// expert weights once per chunk. The tables cost 3,840 B per sequence; the SSM verify pools
46/// are sized by the dispatch cap, not by this (`ssm_reserve::mtp_state_slots`).
47pub const VERIFY_WY_TABLE_SEQS: usize = 128;
48/// 2026-09-25: Catch-up stash rows per sequence (`ModelLevers::mtp_kv_exact`). The batched
49/// verify step stashes at most this many accepted drafts per sequence
50/// (`verify_k4_batch_step.rs`).
51pub const MTP_CATCHUP_MAX: usize = 4;
52/// 2026-09-25: Tables per GDN layer: `h_state` plus up to 15 intermediates.
53/// `upload_verify_wy_tables` stages a verify only when k is in
54/// `2..=VERIFY_WY_TABLES_PER_LAYER`, and returns NULL otherwise.
55pub const VERIFY_WY_TABLES_PER_LAYER: usize = 16;
56/// 2026-09-25: Bytes between consecutive tables within a layer slice.
57pub const VERIFY_WY_TABLE_STRIDE_BYTES: usize = VERIFY_WY_TABLE_SEQS * 8;
58/// 2026-09-25: Bytes between consecutive GDN layers' table slices.
59pub const VERIFY_WY_LAYER_STRIDE_BYTES: usize =
60    VERIFY_WY_TABLES_PER_LAYER * VERIFY_WY_TABLE_STRIDE_BYTES;
61
62pub trait TransformerLayer:
63    Send
64    + Sync
65    + LayerCapabilities
66    + LayerWeightSetup
67    + LayerWriteOnAccept
68    + LayerGraphHooks
69    + LayerAuxState
70    + LayerSplitPrefill
71    + crate::circuit_exec::CircuitBindings
72{
73    /// 2026-09-25: Decode one token through this layer, updating `hidden` in place.
74    ///
75    /// * `hidden` - `[1, hidden_size]` BF16, read and written
76    /// * `residual` - `[1, hidden_size]` BF16 residual stream
77    /// * `state` - this layer's state for the sequence
78    /// * `seq_len` - the token's position
79    fn decode(
80        &self,
81        hidden: DevicePtr,
82        residual: DevicePtr,
83        state: &mut dyn LayerState,
84        kv_cache: &mut PagedKvCache,
85        seq_len: usize,
86        block_table: &mut Vec<u32>,
87        // 2026-09-25: `--high-speed-swap` disk block ids, one per logical block the sequence
88        // has had; `block_table` holds the newest of those blocks. Empty without the feature.
89        disk_block_ids: &mut Vec<u32>,
90        // 2026-09-25: `--high-speed-swap`: how many `disk_block_ids` entries each attention
91        // layer has offloaded, indexed by the layer's `attn_layer_idx`.
92        disk_last_offloaded_per_layer: &mut Vec<u32>,
93        ctx: &ForwardContext,
94        stream: u64,
95    ) -> Result<()>;
96
97    /// 2026-09-25: Prefill `num_tokens` tokens starting at position `seq_len_start`.
98    /// `hidden` and `residual` are `[num_tokens, hidden_size]` BF16. `kv_write_start` is
99    /// how many of these tokens already have their KV in the cache; the Qwen3 attention
100    /// cache-skip prefill writes KV only for the rest (`prefill/cache_skip.rs`). Default:
101    /// `decode` once per token (`default_loops::prefill_default`), which ignores
102    /// `kv_write_start`.
103    fn prefill(
104        &self,
105        hidden: DevicePtr,
106        residual: DevicePtr,
107        num_tokens: usize,
108        state: &mut dyn LayerState,
109        kv_cache: &mut PagedKvCache,
110        seq_len_start: usize,
111        block_table: &mut Vec<u32>,
112        disk_block_ids: &mut Vec<u32>,
113        disk_last_offloaded_per_layer: &mut Vec<u32>,
114        _kv_write_start: usize,
115        ctx: &ForwardContext,
116        stream: u64,
117    ) -> Result<()> {
118        default_loops::prefill_default(
119            self,
120            hidden,
121            residual,
122            num_tokens,
123            state,
124            kv_cache,
125            seq_len_start,
126            block_table,
127            disk_block_ids,
128            disk_last_offloaded_per_layer,
129            ctx,
130            stream,
131        )
132    }
133
134    /// 2026-09-25: Split SSM prefill, phase 1: norm, projections, gates, conv1d and L2
135    /// norm for `num_tokens` tokens, then the GDN inputs (packed QKV, gate/beta, Z) are
136    /// copied into `gdn_bufs` at `token_offset`. The recurrence runs later, in
137    /// `prefill_gdn_full`. Default: the whole `prefill`.
138    #[allow(clippy::too_many_arguments)]
139    fn prefill_phase1(
140        &self,
141        hidden: DevicePtr,
142        residual: DevicePtr,
143        num_tokens: usize,
144        state: &mut dyn LayerState,
145        kv_cache: &mut PagedKvCache,
146        seq_len_start: usize,
147        block_table: &mut Vec<u32>,
148        disk_block_ids: &mut Vec<u32>,
149        disk_last_offloaded_per_layer: &mut Vec<u32>,
150        kv_write_start: usize,
151        gdn_bufs: &GdnPrefillBuffers,
152        token_offset: usize,
153        ctx: &ForwardContext,
154        stream: u64,
155    ) -> Result<()> {
156        let _ = (gdn_bufs, token_offset);
157        self.prefill(
158            hidden,
159            residual,
160            num_tokens,
161            state,
162            kv_cache,
163            seq_len_start,
164            block_table,
165            disk_block_ids,
166            disk_last_offloaded_per_layer,
167            kv_write_start,
168            ctx,
169            stream,
170        )
171    }
172
173    /// 2026-09-25: Decode `num_tokens` consecutive tokens of one sequence, starting at
174    /// position `seq_len`. `hidden` and `residual` are `[num_tokens, hidden_size]` BF16.
175    /// Default: `decode` once per token (`default_loops::decode_batched_default`).
176    #[allow(clippy::too_many_arguments)]
177    fn decode_batched(
178        &self,
179        hidden: DevicePtr,
180        residual: DevicePtr,
181        num_tokens: usize,
182        state: &mut dyn LayerState,
183        kv_cache: &mut PagedKvCache,
184        seq_len: usize,
185        block_table: &mut Vec<u32>,
186        disk_block_ids: &mut Vec<u32>,
187        disk_last_offloaded_per_layer: &mut Vec<u32>,
188        ctx: &ForwardContext,
189        stream: u64,
190    ) -> Result<()> {
191        default_loops::decode_batched_default(
192            self,
193            hidden,
194            residual,
195            num_tokens,
196            state,
197            kv_cache,
198            seq_len,
199            block_table,
200            disk_block_ids,
201            disk_last_offloaded_per_layer,
202            ctx,
203            stream,
204        )
205    }
206
207    /// 2026-09-25: Decode one token for each of `num_seqs` rows, each with its own entry in
208    /// `seq_lens` and `block_tables`: one row per sequence in a batched decode, one per
209    /// verify row in a batched verify (`verify_e.rs`). `hidden` and `residual` are
210    /// `[num_seqs, hidden_size]` BF16. Default: `decode` once per row
211    /// (`default_loops::decode_multi_seq_default`).
212    #[allow(clippy::too_many_arguments)]
213    fn decode_multi_seq<'a, 'b: 'a>(
214        &self,
215        hidden: DevicePtr,
216        residual: DevicePtr,
217        num_seqs: usize,
218        states: &'a mut [&'b mut (dyn LayerState + 'static)],
219        kv_cache: &mut PagedKvCache,
220        seq_lens: &[usize],
221        block_tables: &[Vec<u32>],
222        ctx: &ForwardContext,
223        stream: u64,
224    ) -> Result<()> {
225        default_loops::decode_multi_seq_default(
226            self,
227            hidden,
228            residual,
229            num_seqs,
230            states,
231            kv_cache,
232            seq_lens,
233            block_tables,
234            ctx,
235            stream,
236        )
237    }
238
239    /// 2026-09-25: Batched verify of `ks[i]` rows for each of `n_seqs` sequences in one
240    /// sweep, the rows sequence-major and contiguous in `hidden`/`residual`. The verify
241    /// (`verify_e.rs`) runs attention layers through `decode_multi_seq` and every other
242    /// layer through this. `wy_tables` is this layer's slice of the staged WY pointer
243    /// tables (layout at [`VERIFY_WY_TABLE_SEQS`]), or NULL when none were staged or the
244    /// layer is not linear attention. The default returns an error.
245    #[allow(clippy::too_many_arguments)]
246    fn decode_verify_multi<'a, 'b: 'a>(
247        &self,
248        _hidden: DevicePtr,
249        _residual: DevicePtr,
250        _n_seqs: usize,
251        _ks: &[usize],
252        _states: &'a mut [&'b mut (dyn LayerState + 'static)],
253        _kv_cache: &mut PagedKvCache,
254        _wy_tables: DevicePtr,
255        _ctx: &ForwardContext,
256        _stream: u64,
257    ) -> Result<()> {
258        anyhow::bail!("decode_verify_multi: unsupported for this layer type")
259    }
260
261    /// 2026-09-25: Allocate this layer's per-sequence state. Sequence setup (model-engine
262    /// `trait_impl/meta.rs`) calls it for every layer except a linear-attention layer that
263    /// uses the SSM pool, which gets pool addresses instead.
264    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>>;
265
266    /// 2026-09-25: Free the device memory this layer allocated for one sequence, in
267    /// `alloc_state` or attached to the state later. `LayerState` holds bare `DevicePtr`s,
268    /// so dropping it frees only the host struct.
269    ///
270    /// The teardown chokepoint `free_sequence_dispatch` calls it for every layer, logs an
271    /// error and continues. An implementation must be idempotent, null what it frees,
272    /// never free pool addresses, and decide by the state's type rather than rely on a
273    /// filter at the call site. The default does nothing, which suits state that owns no
274    /// device memory and pool-backed SSM state, released through the pool slot.
275    fn release_state(&self, _state: &mut dyn LayerState, _gpu: &dyn GpuBackend) -> Result<()> {
276        Ok(())
277    }
278}
279
280#[cfg(test)]
281mod verify_wy_table_tests {
282    use super::{
283        VERIFY_WY_LAYER_STRIDE_BYTES, VERIFY_WY_TABLE_SEQS, VERIFY_WY_TABLE_STRIDE_BYTES,
284        VERIFY_WY_TABLES_PER_LAYER,
285    };
286
287    /// 2026-09-25: `upload_verify_wy_tables` stages only k in
288    /// `2..=VERIFY_WY_TABLES_PER_LAYER`; this pins k = 16 as admitted.
289    #[test]
290    fn tables_per_layer_fund_the_widest_verify_arm() {
291        assert!(
292            (2..=VERIFY_WY_TABLES_PER_LAYER).contains(&16usize),
293            "VERIFY_WY_TABLES_PER_LAYER ({VERIFY_WY_TABLES_PER_LAYER}) must admit k=16"
294        );
295    }
296
297    /// 2026-09-25: The table-form wyN launch (`trait_decode_batched_conv_gdn_multi.rs`) and
298    /// the write-on-accept fold pass `VERIFY_WY_TABLE_SEQS` as the entry count between
299    /// slabs, which reaches slab t only when slabs are back to back at that many 8-byte
300    /// entries.
301    #[test]
302    fn slab_stride_matches_the_kernel_contract() {
303        assert_eq!(VERIFY_WY_TABLE_STRIDE_BYTES, VERIFY_WY_TABLE_SEQS * 8);
304        assert_eq!(
305            VERIFY_WY_LAYER_STRIDE_BYTES,
306            VERIFY_WY_TABLES_PER_LAYER * VERIFY_WY_TABLE_STRIDE_BYTES
307        );
308    }
309
310    /// 2026-09-25: A batch of 16 sequences fits one slab.
311    #[test]
312    fn seqs_capacity_covers_c16() {
313        const { assert!(VERIFY_WY_TABLE_SEQS >= 16) };
314    }
315}
metrale_model_layers/layer/transformer_layer.rs

metrale_model_layers/layer/
transformer_layer.rs
