metrale_model_layers/layer/transformer_layer/
split_prefill.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `LayerSplitPrefill`, the phased and batched prefill entry points: the SSM phase 1
4//! pieces, the GDN recurrence per stream, batched and varlen, phase 3, and the batched
5//! attention prefill.
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;
13
14use crate::layer::{BatchedAttnMetadata, ForwardContext, GdnPrefillBuffers, LayerState};
15
16/// 2026-09-26: A supertrait of `TransformerLayer`; see the module header.
17pub trait LayerSplitPrefill {
18    /// 2026-09-25: Phase 1 projections (norm, QKVZ, gates) over all stacked tokens of a
19    /// batched prefill at once. The caller then runs `prefill_phase1_conv1d_one` per
20    /// request and `prefill_phase1_l2_batched`. The default returns an error.
21    fn prefill_phase1_proj_batched(
22        &self,
23        hidden_stacked: DevicePtr,
24        residual_stacked: DevicePtr,
25        total_tokens: usize,
26        gdn_bufs: &GdnPrefillBuffers,
27        ctx: &ForwardContext,
28        stream: u64,
29    ) -> Result<()> {
30        let _ = (
31            hidden_stacked,
32            residual_stacked,
33            total_tokens,
34            gdn_bufs,
35            ctx,
36            stream,
37        );
38        anyhow::bail!("prefill_phase1_proj_batched: only implemented for SSM layers")
39    }
40
41    /// 2026-09-25: Conv1d over one request's tokens of a batched prefill, advancing its
42    /// conv state and writing its slice of `gdn_bufs.qkv`. The default returns an error.
43    fn prefill_phase1_conv1d_one(
44        &self,
45        state: &mut dyn LayerState,
46        token_offset: usize,
47        len: usize,
48        gdn_bufs: &GdnPrefillBuffers,
49        ctx: &ForwardContext,
50        stream: u64,
51    ) -> Result<()> {
52        let _ = (state, token_offset, len, gdn_bufs, ctx, stream);
53        anyhow::bail!("prefill_phase1_conv1d_one: only implemented for SSM layers")
54    }
55
56    /// 2026-09-25: L2 norm over the whole stacked `gdn_bufs.qkv`, after every request's
57    /// conv1d. The default returns an error.
58    fn prefill_phase1_l2_batched(
59        &self,
60        total_tokens: usize,
61        gdn_bufs: &GdnPrefillBuffers,
62        ctx: &ForwardContext,
63        stream: u64,
64    ) -> Result<()> {
65        let _ = (total_tokens, gdn_bufs, ctx, stream);
66        anyhow::bail!("prefill_phase1_l2_batched: only implemented for SSM layers")
67    }
68
69    /// 2026-09-25: Split SSM prefill, phase 2: the GDN recurrence over all
70    /// `gdn_bufs.total_len` tokens, reading packed QKV and gate/beta and writing `output`.
71    /// Default: nothing.
72    fn prefill_gdn_full(
73        &self,
74        _state: &mut dyn LayerState,
75        _gdn_bufs: &GdnPrefillBuffers,
76        _ctx: &ForwardContext,
77        _stream: u64,
78    ) -> Result<()> {
79        Ok(())
80    }
81
82    /// 2026-09-25: Prefill of one attention layer over the stacked tokens of several
83    /// streams, with per-stream metadata from `batched_meta`. The default returns an
84    /// error; the Qwen3 attention layer overrides it.
85    fn prefill_inner_batched_q12(
86        &self,
87        _hidden_stacked: DevicePtr,
88        _residual_stacked: DevicePtr,
89        _num_tokens: usize,
90        _kv_cache: &mut PagedKvCache,
91        _seq_len_start: usize,
92        _batched_meta: &BatchedAttnMetadata,
93        _ctx: &ForwardContext,
94        _stream: u64,
95    ) -> Result<()> {
96        anyhow::bail!("prefill_inner_batched_q12: not implemented for this layer type")
97    }
98
99    /// 2026-09-25: GDN recurrence for `batch_size` streams of equal length `chunk_len` in
100    /// one call. `h_state_ptrs` is a device array of one h-state pointer per stream
101    /// (`TransformerModel::stage_h_state_ptrs`), and the streams' tokens lie back to back
102    /// in `gdn_bufs`. `prefill_ssm_batched_layer` calls it only when every stream has
103    /// the same length. The default returns an error.
104    fn prefill_gdn_full_batched(
105        &self,
106        _h_state_ptrs: DevicePtr,
107        _gdn_bufs: &GdnPrefillBuffers,
108        _batch_size: u32,
109        _chunk_len: u32,
110        _ctx: &ForwardContext,
111        _stream: u64,
112    ) -> Result<()> {
113        anyhow::bail!(
114            "prefill_gdn_full_batched: layer does not implement batched GDN \
115             — caller should fall back to per-stream prefill_gdn_full"
116        )
117    }
118
119    /// 2026-09-25: GDN recurrence for streams of different lengths in one FLA call,
120    /// driven by `cu_seqlens`. `Ok(true)` if it ran; `Ok(false)` if not eligible, and the
121    /// caller then runs `prefill_gdn_full` per stream. Default `Ok(false)`.
122    #[allow(clippy::too_many_arguments)]
123    fn prefill_gdn_full_batched_fla_varlen(
124        &self,
125        _h_state_ptrs: DevicePtr,
126        _gdn_bufs: &GdnPrefillBuffers,
127        _batch_size: u32,
128        _cu_seqlens: DevicePtr,
129        _max_num_chunks: u32,
130        _total_nt: usize,
131        _max_seqlen: u32,
132        _ctx: &ForwardContext,
133        _stream: u64,
134    ) -> Result<bool> {
135        Ok(false)
136    }
137
138    /// 2026-09-25: Split SSM prefill, phase 3: for `num_tokens` tokens at `token_offset`,
139    /// gated RMSNorm of the GDN output with Z, the output projection, the residual add and
140    /// the FFN. Default: nothing.
141    #[allow(clippy::too_many_arguments)]
142    fn prefill_phase3(
143        &self,
144        _hidden: DevicePtr,
145        _residual: DevicePtr,
146        _num_tokens: usize,
147        _gdn_bufs: &GdnPrefillBuffers,
148        _token_offset: usize,
149        _ctx: &ForwardContext,
150        _stream: u64,
151    ) -> Result<()> {
152        Ok(())
153    }
154}
metrale_model_layers/layer/transformer_layer/split_prefill.rs

metrale_model_layers/layer/transformer_layer/
split_prefill.rs
