metrale_model_layers/layers/qwen3_ssm/
trait_layer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `TransformerLayer` impl for [`Qwen3SsmLayer`]: the trait
4//! surface, forwarding into the `trait_*` sibling modules that hold the phases.
5//!
6//! Owner: model-layers, GDN/SSM layer (`qwen3_ssm`).
7//! Invariants:
8//! - With mHC weights attached (`hc` is `Some`), `decode_batched` and
9//!   `decode_verify_multi` return an error and launch nothing.
10
11use anyhow::Result;
12use metrale_cache::kv_cache::PagedKvCache;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
14
15use super::Qwen3SsmLayer;
16use super::ple_seq::ple_seq_state;
17use crate::layer::{ForwardContext, GdnPrefillBuffers, LayerState, TransformerLayer};
18use crate::layer::{
19    LayerAuxState, LayerCapabilities, LayerGraphHooks, LayerSplitPrefill, LayerWeightSetup,
20    LayerWriteOnAccept,
21};
22
23impl TransformerLayer for Qwen3SsmLayer {
24    fn decode(
25        &self,
26        hidden: DevicePtr,
27        residual: DevicePtr,
28        state: &mut dyn LayerState,
29        kv_cache: &mut PagedKvCache,
30        seq_len: usize,
31        block_table: &mut Vec<u32>,
32        disk_block_ids: &mut Vec<u32>,
33        disk_last_offloaded_per_layer: &mut Vec<u32>,
34        ctx: &ForwardContext,
35        stream: u64,
36    ) -> Result<()> {
37        if self.hc.is_some() {
38            return self.decode_inner_hc(hidden, state, ctx, stream);
39        }
40        self.decode_inner(
41            hidden,
42            residual,
43            state,
44            kv_cache,
45            seq_len,
46            block_table,
47            disk_block_ids,
48            disk_last_offloaded_per_layer,
49            ctx,
50            stream,
51        )
52    }
53
54    fn decode_batched(
55        &self,
56        hidden: DevicePtr,
57        residual: DevicePtr,
58        num_tokens: usize,
59        state: &mut dyn LayerState,
60        _kv_cache: &mut PagedKvCache,
61        _seq_len: usize,
62        _block_table: &mut Vec<u32>,
63        _disk_block_ids: &mut Vec<u32>,
64        _disk_last_offloaded_per_layer: &mut Vec<u32>,
65        ctx: &ForwardContext,
66        stream: u64,
67    ) -> Result<()> {
68        // 2026-09-25: `decode_batched_inner` keeps its own residual
69        // bookkeeping, which an mHC highway replaces, so it is refused there.
70        self.refuse_batched_under_hc("decode_batched")?;
71        self.decode_batched_inner(
72            hidden,
73            residual,
74            num_tokens,
75            super::trait_decode_batched::GdnStates::Single(state),
76            ctx,
77            stream,
78        )
79    }
80
81    fn decode_verify_multi<'a, 'b: 'a>(
82        &self,
83        hidden: DevicePtr,
84        residual: DevicePtr,
85        n_seqs: usize,
86        ks: &[usize],
87        states: &'a mut [&'b mut (dyn LayerState + 'static)],
88        _kv_cache: &mut PagedKvCache,
89        wy_tables: DevicePtr,
90        ctx: &ForwardContext,
91        stream: u64,
92    ) -> Result<()> {
93        self.refuse_batched_under_hc("decode_verify_multi")?;
94        anyhow::ensure!(
95            states.len() == n_seqs && ks.len() == n_seqs,
96            "decode_verify_multi: states/ks/n mismatch"
97        );
98        let num_tokens: usize = ks.iter().sum();
99        self.decode_batched_inner(
100            hidden,
101            residual,
102            num_tokens,
103            super::trait_decode_batched::GdnStates::Multi {
104                states,
105                ks,
106                wy_tables,
107            },
108            ctx,
109            stream,
110        )
111    }
112
113    fn decode_multi_seq<'a, 'b: 'a>(
114        &self,
115        hidden: DevicePtr,
116        residual: DevicePtr,
117        num_seqs: usize,
118        states: &'a mut [&'b mut (dyn LayerState + 'static)],
119        kv_cache: &mut PagedKvCache,
120        seq_lens: &[usize],
121        block_tables: &[Vec<u32>],
122        ctx: &ForwardContext,
123        stream: u64,
124    ) -> Result<()> {
125        if self.hc.is_some() {
126            // 2026-09-25: The highway replaces the residual that the non-hc
127            // path folds into its fused norms.
128            return self.decode_multi_seq_inner_hc(hidden, num_seqs, states, seq_lens, ctx, stream);
129        }
130        self.decode_multi_seq_inner(
131            hidden,
132            residual,
133            num_seqs,
134            states,
135            kv_cache,
136            seq_lens,
137            block_tables,
138            ctx,
139            stream,
140        )
141    }
142
143    fn prefill(
144        &self,
145        hidden: DevicePtr,
146        residual: DevicePtr,
147        num_tokens: usize,
148        state: &mut dyn LayerState,
149        kv_cache: &mut PagedKvCache,
150        seq_len_start: usize,
151        block_table: &mut Vec<u32>,
152        disk_block_ids: &mut Vec<u32>,
153        disk_last_offloaded_per_layer: &mut Vec<u32>,
154        kv_write_start: usize,
155        ctx: &ForwardContext,
156        stream: u64,
157    ) -> Result<()> {
158        // 2026-09-25: Under an mHC highway the highway is the residual, so the
159        // prefill takes its own entry path (`trait_prefill_hc.rs`).
160        if self.hc.is_some() {
161            return self.prefill_inner_hc(hidden, num_tokens, state, seq_len_start, ctx, stream);
162        }
163        self.prefill_inner(
164            hidden,
165            residual,
166            num_tokens,
167            state,
168            kv_cache,
169            seq_len_start,
170            block_table,
171            disk_block_ids,
172            disk_last_offloaded_per_layer,
173            kv_write_start,
174            ctx,
175            stream,
176        )
177    }
178
179    fn prefill_phase1(
180        &self,
181        hidden: DevicePtr,
182        residual: DevicePtr,
183        num_tokens: usize,
184        state: &mut dyn LayerState,
185        kv_cache: &mut PagedKvCache,
186        seq_len_start: usize,
187        block_table: &mut Vec<u32>,
188        disk_block_ids: &mut Vec<u32>,
189        disk_last_offloaded_per_layer: &mut Vec<u32>,
190        kv_write_start: usize,
191        gdn_bufs: &GdnPrefillBuffers,
192        token_offset: usize,
193        ctx: &ForwardContext,
194        stream: u64,
195    ) -> Result<()> {
196        self.prefill_phase1_inner(
197            hidden,
198            residual,
199            num_tokens,
200            state,
201            kv_cache,
202            seq_len_start,
203            block_table,
204            disk_block_ids,
205            disk_last_offloaded_per_layer,
206            kv_write_start,
207            gdn_bufs,
208            token_offset,
209            ctx,
210            stream,
211        )
212    }
213
214    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
215        self.alloc_state_inner(gpu)
216    }
217
218    /// 2026-09-25: Free the PLE carry this sequence attached on first use.
219    ///
220    /// Only the `ple` field: the h/conv state points into the SSM pool, which
221    /// `free_sequence_dispatch` releases by slot. The PLE conv buffer is
222    /// allocated per sequence (`new_seq_state`), outside that pool.
223    fn release_state(&self, state: &mut dyn LayerState, gpu: &dyn GpuBackend) -> Result<()> {
224        let Some(ssm) = state
225            .as_any_mut()
226            .downcast_mut::<crate::layer::SsmLayerState>()
227        else {
228            return Ok(());
229        };
230        let Some(mut st) = ssm.ple.take() else {
231            return Ok(());
232        };
233        let Some(ple) = self.ple.as_ref() else {
234            anyhow::bail!("release_state: PLE seq state present but layer has no PLE");
235        };
236        ple.release_seq_state(&mut st, gpu)
237    }
238}
239
240impl LayerCapabilities for Qwen3SsmLayer {
241    /// 2026-09-25: True when this layer carries PLE. The multi-seq decode
242    /// (`forward_row`) hashes ids on the host inside the forward, which a graph
243    /// capture refuses. The model ORs this into `decode_graph_veto`, which keeps
244    /// decode eager.
245    fn decode_graph_unsupported(&self) -> bool {
246        self.ple.is_some()
247    }
248
249    fn is_ssm_layer(&self) -> bool {
250        self.is_ssm_layer_inner()
251    }
252
253    /// 2026-10-01: Both prefill entry paths (`prefill_inner`, `prefill_inner_hc`) run
254    /// `prefill_block`, whose conv (`conv1d_prefill_capture`) and recurrence
255    /// (`prefill_gdn_recurrence`) split at the capture point on every build, the
256    /// recurrence's tail on the exact-replay arm.
257    fn supports_replay_tail_split(&self) -> bool {
258        true
259    }
260}
261
262impl LayerWeightSetup for Qwen3SsmLayer {
263    /// 2026-09-25: Downcast hook: the LoRA install (`impl_lora_rotate.rs`) and
264    /// the qwen4_exp loader reach this concrete layer through it.
265    fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
266        Some(self)
267    }
268}
269
270impl LayerWriteOnAccept for Qwen3SsmLayer {
271    fn gdn_verify_run_batched(
272        &self,
273        states: &[&mut (dyn crate::layer::LayerState + 'static)],
274        kk: usize,
275        gdn_wyn: bool,
276        wy_tables: metrale_gpu_runtime::gpu::DevicePtr,
277    ) -> anyhow::Result<Option<bool>> {
278        Ok(Some(matches!(
279            self.multi_run_arm(states, kk, gdn_wyn, wy_tables)?,
280            super::trait_decode_batched_conv_gdn_multi::RunArm::Batched(..)
281        )))
282    }
283
284    fn gdn_woa_stash_seq_floats(&self) -> Option<usize> {
285        self.woa_stash_seq_floats_impl()
286    }
287
288    fn gdn_woa_bind(&self, flag: DevicePtr, stash: DevicePtr, seqs: usize) {
289        self.woa_bind_impl(flag, stash, seqs)
290    }
291
292    fn gdn_carry_seq_floats(&self) -> Option<usize> {
293        self.carry.seq_floats()
294    }
295
296    fn gdn_carry_bind(&self, binding: crate::layer::GdnCarryBinding) {
297        let _ = self.carry.binding.set(binding);
298    }
299
300    fn gdn_fold_accepted(
301        &self,
302        gpu: &dyn GpuBackend,
303        h_table: DevicePtr,
304        na_tab: DevicePtr,
305        k_rows: usize,
306        n: usize,
307        stream: u64,
308    ) -> Result<bool> {
309        self.fold_accepted_impl(gpu, h_table, na_tab, k_rows, n, stream)
310    }
311}
312
313impl LayerGraphHooks for Qwen3SsmLayer {
314    /// 2026-09-25: PLE's host half (id hash, NVMe fault-in, slot upload), run
315    /// before graph replay or capture. A no-op on a layer without PLE.
316    fn decode_prestage(
317        &self,
318        token: u32,
319        state: &mut dyn LayerState,
320        gpu: &dyn GpuBackend,
321        stream: u64,
322    ) -> Result<()> {
323        if let Some(ple) = self.ple.as_ref() {
324            let st = ple_seq_state(ple, state, gpu)?;
325            ple.prestage(st, &[token], gpu, stream)?;
326        }
327        Ok(())
328    }
329
330    fn decode_prestage_rearm(&self, state: &mut dyn LayerState) {
331        if let Some(ple) = self.ple.as_ref()
332            && let Some(ssm) = state
333                .as_any_mut()
334                .downcast_mut::<crate::layer::SsmLayerState>()
335            && let Some(st) = ssm.ple.as_mut()
336        {
337            ple.rearm(st);
338        }
339    }
340}
341
342impl LayerAuxState for Qwen3SsmLayer {
343    fn has_aux_state(&self) -> bool {
344        self.ple.is_some()
345    }
346
347    fn snapshot_aux(
348        &self,
349        state: &dyn LayerState,
350        gpu: &dyn GpuBackend,
351        stream: u64,
352    ) -> Result<Option<Vec<u8>>> {
353        let Some(ple) = self.ple.as_ref() else {
354            return Ok(None);
355        };
356        let ssm = state
357            .as_any()
358            .downcast_ref::<crate::layer::SsmLayerState>()
359            .ok_or_else(|| anyhow::anyhow!("PLE host layer state is not SsmLayerState"))?;
360        match ssm.ple.as_ref() {
361            Some(st) => Ok(Some(ple.snapshot_aux(st, gpu, stream)?)),
362            // 2026-09-25: The sequence has not run this layer yet, so there is
363            // nothing to carry. A restore declines a slot without aux blobs.
364            None => Ok(None),
365        }
366    }
367
368    fn restore_aux(
369        &self,
370        state: &mut dyn LayerState,
371        blob: &[u8],
372        gpu: &dyn GpuBackend,
373        stream: u64,
374    ) -> Result<()> {
375        let ple = self
376            .ple
377            .as_ref()
378            .ok_or_else(|| anyhow::anyhow!("restore_aux: no PLE on this layer"))?;
379        let st = ple_seq_state(ple, state, gpu)?;
380        ple.restore_aux(st, blob, gpu, stream)
381    }
382}
383
384impl LayerSplitPrefill for Qwen3SsmLayer {
385    fn prefill_phase1_proj_batched(
386        &self,
387        hidden_stacked: DevicePtr,
388        residual_stacked: DevicePtr,
389        total_tokens: usize,
390        gdn_bufs: &GdnPrefillBuffers,
391        ctx: &ForwardContext,
392        stream: u64,
393    ) -> Result<()> {
394        self.prefill_phase1_proj_batched_inner(
395            hidden_stacked,
396            residual_stacked,
397            total_tokens,
398            gdn_bufs,
399            ctx,
400            stream,
401        )
402    }
403
404    fn prefill_phase1_conv1d_one(
405        &self,
406        state: &mut dyn LayerState,
407        token_offset: usize,
408        len: usize,
409        gdn_bufs: &GdnPrefillBuffers,
410        ctx: &ForwardContext,
411        stream: u64,
412    ) -> Result<()> {
413        self.prefill_phase1_conv1d_one_inner(state, token_offset, len, gdn_bufs, ctx, stream)
414    }
415
416    fn prefill_phase1_l2_batched(
417        &self,
418        total_tokens: usize,
419        gdn_bufs: &GdnPrefillBuffers,
420        ctx: &ForwardContext,
421        stream: u64,
422    ) -> Result<()> {
423        self.prefill_phase1_l2_batched_inner(total_tokens, gdn_bufs, ctx, stream)
424    }
425
426    fn prefill_gdn_full(
427        &self,
428        state: &mut dyn LayerState,
429        gdn_bufs: &GdnPrefillBuffers,
430        ctx: &ForwardContext,
431        stream: u64,
432    ) -> Result<()> {
433        self.prefill_gdn_full_inner(state, gdn_bufs, ctx, stream)
434    }
435
436    fn prefill_gdn_full_batched(
437        &self,
438        h_state_ptrs: DevicePtr,
439        gdn_bufs: &GdnPrefillBuffers,
440        batch_size: u32,
441        chunk_len: u32,
442        ctx: &ForwardContext,
443        stream: u64,
444    ) -> Result<()> {
445        self.prefill_gdn_full_batched_inner(
446            h_state_ptrs,
447            gdn_bufs,
448            batch_size,
449            chunk_len,
450            ctx,
451            stream,
452        )
453    }
454
455    fn prefill_gdn_full_batched_fla_varlen(
456        &self,
457        h_state_ptrs: DevicePtr,
458        gdn_bufs: &GdnPrefillBuffers,
459        batch_size: u32,
460        cu_seqlens: DevicePtr,
461        max_num_chunks: u32,
462        total_nt: usize,
463        max_seqlen: u32,
464        ctx: &ForwardContext,
465        stream: u64,
466    ) -> Result<bool> {
467        self.prefill_gdn_full_batched_fla_varlen_inner(
468            h_state_ptrs,
469            gdn_bufs,
470            batch_size,
471            cu_seqlens,
472            max_num_chunks,
473            total_nt,
474            max_seqlen,
475            ctx,
476            stream,
477        )
478    }
479
480    fn prefill_phase3(
481        &self,
482        hidden: DevicePtr,
483        residual: DevicePtr,
484        num_tokens: usize,
485        gdn_bufs: &GdnPrefillBuffers,
486        token_offset: usize,
487        ctx: &ForwardContext,
488        stream: u64,
489    ) -> Result<()> {
490        self.prefill_phase3_inner(
491            hidden,
492            residual,
493            num_tokens,
494            gdn_bufs,
495            token_offset,
496            ctx,
497            stream,
498        )
499    }
500}
metrale_model_layers/layers/qwen3_ssm/trait_layer.rs

metrale_model_layers/layers/qwen3_ssm/
trait_layer.rs
