metrale_model_layers/layer/transformer_layer/
graph_hooks.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `LayerGraphHooks`, the host-side work a layer does around a CUDA-graph decode
4//! step: the per-step prestage and its re-arm, and the room check and bookkeeping of a
5//! replayed step.
6//!
7//! Owner: model-layers.
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::GpuBackend;
12
13use crate::layer::LayerState;
14
15/// 2026-09-26: A supertrait of `TransformerLayer`; see the module header.
16pub trait LayerGraphHooks {
17    /// 2026-09-25: Per-step host work of a layer that computes on the host at decode (the
18    /// PLE hash and slot upload). The model's single-token decode calls it before any
19    /// CUDA graph replay or capture, when the token id is on the host (`decode_a.rs`), so
20    /// a graph reads only device buffers refreshed here. Default: nothing.
21    fn decode_prestage(
22        &self,
23        _token: u32,
24        _state: &mut dyn LayerState,
25        _gpu: &dyn GpuBackend,
26        _stream: u64,
27    ) -> Result<()> {
28        Ok(())
29    }
30
31    /// 2026-09-25: Re-arm prestaged state that a failed CUDA-graph capture consumed, so the
32    /// same step can re-run eagerly (`decode_a.rs`). It must not recompute the prestage
33    /// work: PLE's history already advanced in `decode_prestage`.
34    fn decode_prestage_rearm(&self, _state: &mut dyn LayerState) {}
35
36    /// 2026-09-25: Update host-side per-sequence bookkeeping for a step served by a
37    /// replayed CUDA graph, in which the layer's `decode` did not run. `seq_len` is the
38    /// length before the step's `k` rows. GLM-5.3's DSA layer rewinds its indexer-cache
39    /// length to `seq_len` when it is ahead, errors when it is behind, and then advances it
40    /// by `k` (`Glm5NextDsaState::sync_to`). The model calls it after a replay in
41    /// `decode_a.rs`, `verify_b.rs`, `verify_c.rs` and `verify_c2.rs`. Default: nothing.
42    fn sync_replayed_step(
43        &self,
44        _state: &mut dyn LayerState,
45        _seq_len: usize,
46        _k: usize,
47    ) -> Result<()> {
48        Ok(())
49    }
50
51    /// 2026-09-25: Return an error for a replayed step whose writes would land past the end
52    /// of a host-tracked cache. The model calls it before launching the graph;
53    /// `sync_replayed_step` runs only after the launch. The step ends at `seq_len + k`.
54    /// Default: nothing.
55    fn check_replay_room(&self, _state: &dyn LayerState, _seq_len: usize, _k: usize) -> Result<()> {
56        Ok(())
57    }
58}
metrale_model_layers/layer/transformer_layer/graph_hooks.rs

metrale_model_layers/layer/transformer_layer/
graph_hooks.rs
