spark_model/layers/qwen3_ssm/
trait_layer.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `TransformerLayer` impl for [`Qwen3SsmLayer`] — the trait surface that
4//! forwards into the `trait_*` sibling modules holding the actual phases.
5//! Split out of `mod.rs` to keep it under the 500-LoC cap.
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9use spark_runtime::kv_cache::PagedKvCache;
10
11use super::Qwen3SsmLayer;
12use super::ple_seq::ple_seq_state;
13use crate::layer::{ForwardContext, GdnPrefillBuffers, LayerState, TransformerLayer};
14
15impl TransformerLayer for Qwen3SsmLayer {
16    fn gdn_woa_stash_seq_floats(&self) -> Option<usize> {
17        self.woa_stash_seq_floats_impl()
18    }
19
20    fn gdn_woa_bind(&self, flag: DevicePtr, stash: DevicePtr, seqs: usize) {
21        self.woa_bind_impl(flag, stash, seqs)
22    }
23
24    fn gdn_fold_accepted(
25        &self,
26        gpu: &dyn GpuBackend,
27        h_table: DevicePtr,
28        na_tab: DevicePtr,
29        k_rows: usize,
30        n: usize,
31        stream: u64,
32    ) -> Result<bool> {
33        self.fold_accepted_impl(gpu, h_table, na_tab, k_rows, n, stream)
34    }
35
36    /// Downcast hook so the LoRA install walk can reach this layer's MoE FFN
37    /// (Feature-1: routed-expert/router deltas exist on GDN layers too).
38    fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
39        Some(self)
40    }
41
42    /// PLE's host half (hash + NVMe fault-in + slot upload), hoisted before
43    /// graph replay/capture. No-op on the 47 layers without a PLE site.
44    fn decode_prestage(
45        &self,
46        token: u32,
47        state: &mut dyn LayerState,
48        gpu: &dyn GpuBackend,
49        stream: u64,
50    ) -> Result<()> {
51        if let Some(ple) = self.ple.as_ref() {
52            let st = ple_seq_state(ple, state, gpu)?;
53            ple.prestage(st, &[token], gpu, stream)?;
54        }
55        Ok(())
56    }
57
58    fn has_aux_state(&self) -> bool {
59        self.ple.is_some()
60    }
61
62    /// PLE's per-seq host hash on the hc multi-seq decode path is
63    /// capture-illegal (pageable reads); the single-decode path prestages
64    /// around it, the batched path does not — veto batched graphs.
65    fn decode_graph_unsupported(&self) -> bool {
66        self.ple.is_some()
67    }
68
69    fn snapshot_aux(
70        &self,
71        state: &dyn LayerState,
72        gpu: &dyn GpuBackend,
73        stream: u64,
74    ) -> Result<Option<Vec<u8>>> {
75        let Some(ple) = self.ple.as_ref() else {
76            return Ok(None);
77        };
78        let ssm = state
79            .as_any()
80            .downcast_ref::<crate::layer::SsmLayerState>()
81            .ok_or_else(|| anyhow::anyhow!("PLE host layer state is not SsmLayerState"))?;
82        match ssm.ple.as_ref() {
83            Some(st) => Ok(Some(ple.snapshot_aux(st, gpu, stream)?)),
84            // Sequence never ran this layer (snapshot before first pass):
85            // nothing to carry, and restore-side declines aux-less slots.
86            None => Ok(None),
87        }
88    }
89
90    fn restore_aux(
91        &self,
92        state: &mut dyn LayerState,
93        blob: &[u8],
94        gpu: &dyn GpuBackend,
95        stream: u64,
96    ) -> Result<()> {
97        let ple = self
98            .ple
99            .as_ref()
100            .ok_or_else(|| anyhow::anyhow!("restore_aux: no PLE on this layer"))?;
101        let st = ple_seq_state(ple, state, gpu)?;
102        ple.restore_aux(st, blob, gpu, stream)
103    }
104
105    fn decode_prestage_rearm(&self, state: &mut dyn LayerState) {
106        if let Some(ple) = self.ple.as_ref()
107            && let Some(ssm) = state
108                .as_any_mut()
109                .downcast_mut::<crate::layer::SsmLayerState>()
110            && let Some(st) = ssm.ple.as_mut()
111        {
112            ple.rearm(st);
113        }
114    }
115
116    fn decode(
117        &self,
118        hidden: DevicePtr,
119        residual: DevicePtr,
120        state: &mut dyn LayerState,
121        kv_cache: &mut PagedKvCache,
122        seq_len: usize,
123        block_table: &mut Vec<u32>,
124        disk_block_ids: &mut Vec<u32>,
125        disk_last_offloaded_per_layer: &mut Vec<u32>,
126        ctx: &ForwardContext,
127        stream: u64,
128    ) -> Result<()> {
129        if self.hc.is_some() {
130            return self.decode_inner_hc(hidden, state, ctx, stream);
131        }
132        self.decode_inner(
133            hidden,
134            residual,
135            state,
136            kv_cache,
137            seq_len,
138            block_table,
139            disk_block_ids,
140            disk_last_offloaded_per_layer,
141            ctx,
142            stream,
143        )
144    }
145
146    fn decode_batched(
147        &self,
148        hidden: DevicePtr,
149        residual: DevicePtr,
150        num_tokens: usize,
151        state: &mut dyn LayerState,
152        _kv_cache: &mut PagedKvCache,
153        _seq_len: usize,
154        _block_table: &mut Vec<u32>,
155        _disk_block_ids: &mut Vec<u32>,
156        _disk_last_offloaded_per_layer: &mut Vec<u32>,
157        ctx: &ForwardContext,
158        stream: u64,
159    ) -> Result<()> {
160        // v1 is C=1 only under an mHC highway: these paths keep their own
161        // residual bookkeeping, which the highway replaces. Refusing is the
162        // point — a batched GDN step running on an unmixed stream produces
163        // plausible, wrong activations. Atlas #753.
164        self.refuse_batched_under_hc("decode_batched")?;
165        self.decode_batched_inner(
166            hidden,
167            residual,
168            num_tokens,
169            super::trait_decode_batched::GdnStates::Single(state),
170            ctx,
171            stream,
172        )
173    }
174
175    fn decode_verify_multi<'a, 'b: 'a>(
176        &self,
177        hidden: DevicePtr,
178        residual: DevicePtr,
179        n_seqs: usize,
180        ks: &[usize],
181        states: &'a mut [&'b mut (dyn LayerState + 'static)],
182        _kv_cache: &mut PagedKvCache,
183        wy_tables: DevicePtr,
184        ctx: &ForwardContext,
185        stream: u64,
186    ) -> Result<()> {
187        self.refuse_batched_under_hc("decode_verify_multi")?;
188        anyhow::ensure!(
189            states.len() == n_seqs && ks.len() == n_seqs,
190            "decode_verify_multi: states/ks/n mismatch"
191        );
192        let num_tokens: usize = ks.iter().sum();
193        self.decode_batched_inner(
194            hidden,
195            residual,
196            num_tokens,
197            super::trait_decode_batched::GdnStates::Multi {
198                states,
199                ks,
200                wy_tables,
201            },
202            ctx,
203            stream,
204        )
205    }
206
207    fn decode_multi_seq<'a, 'b: 'a>(
208        &self,
209        hidden: DevicePtr,
210        residual: DevicePtr,
211        num_seqs: usize,
212        states: &'a mut [&'b mut (dyn LayerState + 'static)],
213        kv_cache: &mut PagedKvCache,
214        seq_lens: &[usize],
215        block_tables: &[Vec<u32>],
216        ctx: &ForwardContext,
217        stream: u64,
218    ) -> Result<()> {
219        if self.hc.is_some() {
220            // #753 item B milestone 2: the highway replaces the residual the
221            // non-hc path folds into its fused norm kernels; run the
222            // hc-bracketed variant instead of refusing.
223            return self.decode_multi_seq_inner_hc(hidden, num_seqs, states, seq_lens, ctx, stream);
224        }
225        self.decode_multi_seq_inner(
226            hidden,
227            residual,
228            num_seqs,
229            states,
230            kv_cache,
231            seq_lens,
232            block_tables,
233            ctx,
234            stream,
235        )
236    }
237
238    fn prefill(
239        &self,
240        hidden: DevicePtr,
241        residual: DevicePtr,
242        num_tokens: usize,
243        state: &mut dyn LayerState,
244        kv_cache: &mut PagedKvCache,
245        seq_len_start: usize,
246        block_table: &mut Vec<u32>,
247        disk_block_ids: &mut Vec<u32>,
248        disk_last_offloaded_per_layer: &mut Vec<u32>,
249        kv_write_start: usize,
250        ctx: &ForwardContext,
251        stream: u64,
252    ) -> Result<()> {
253        // Under an mHC highway the residual bookkeeping is completely
254        // different — the highway IS the residual — so this is a second entry
255        // path, not a flag on the first. See `trait_prefill_hc.rs`.
256        if self.hc.is_some() {
257            return self.prefill_inner_hc(hidden, num_tokens, state, seq_len_start, ctx, stream);
258        }
259        self.prefill_inner(
260            hidden,
261            residual,
262            num_tokens,
263            state,
264            kv_cache,
265            seq_len_start,
266            block_table,
267            disk_block_ids,
268            disk_last_offloaded_per_layer,
269            kv_write_start,
270            ctx,
271            stream,
272        )
273    }
274
275    fn is_ssm_layer(&self) -> bool {
276        self.is_ssm_layer_inner()
277    }
278
279    fn prefill_phase1(
280        &self,
281        hidden: DevicePtr,
282        residual: DevicePtr,
283        num_tokens: usize,
284        state: &mut dyn LayerState,
285        kv_cache: &mut PagedKvCache,
286        seq_len_start: usize,
287        block_table: &mut Vec<u32>,
288        disk_block_ids: &mut Vec<u32>,
289        disk_last_offloaded_per_layer: &mut Vec<u32>,
290        kv_write_start: usize,
291        gdn_bufs: &GdnPrefillBuffers,
292        token_offset: usize,
293        ctx: &ForwardContext,
294        stream: u64,
295    ) -> Result<()> {
296        self.prefill_phase1_inner(
297            hidden,
298            residual,
299            num_tokens,
300            state,
301            kv_cache,
302            seq_len_start,
303            block_table,
304            disk_block_ids,
305            disk_last_offloaded_per_layer,
306            kv_write_start,
307            gdn_bufs,
308            token_offset,
309            ctx,
310            stream,
311        )
312    }
313
314    fn prefill_phase1_proj_batched(
315        &self,
316        hidden_stacked: DevicePtr,
317        residual_stacked: DevicePtr,
318        total_tokens: usize,
319        gdn_bufs: &GdnPrefillBuffers,
320        ctx: &ForwardContext,
321        stream: u64,
322    ) -> Result<()> {
323        self.prefill_phase1_proj_batched_inner(
324            hidden_stacked,
325            residual_stacked,
326            total_tokens,
327            gdn_bufs,
328            ctx,
329            stream,
330        )
331    }
332
333    fn prefill_phase1_conv1d_one(
334        &self,
335        state: &mut dyn LayerState,
336        token_offset: usize,
337        len: usize,
338        gdn_bufs: &GdnPrefillBuffers,
339        ctx: &ForwardContext,
340        stream: u64,
341    ) -> Result<()> {
342        self.prefill_phase1_conv1d_one_inner(state, token_offset, len, gdn_bufs, ctx, stream)
343    }
344
345    fn prefill_phase1_l2_batched(
346        &self,
347        total_tokens: usize,
348        gdn_bufs: &GdnPrefillBuffers,
349        ctx: &ForwardContext,
350        stream: u64,
351    ) -> Result<()> {
352        self.prefill_phase1_l2_batched_inner(total_tokens, gdn_bufs, ctx, stream)
353    }
354
355    fn prefill_gdn_full(
356        &self,
357        state: &mut dyn LayerState,
358        gdn_bufs: &GdnPrefillBuffers,
359        ctx: &ForwardContext,
360        stream: u64,
361    ) -> Result<()> {
362        self.prefill_gdn_full_inner(state, gdn_bufs, ctx, stream)
363    }
364
365    fn prefill_gdn_full_batched(
366        &self,
367        h_state_ptrs: DevicePtr,
368        gdn_bufs: &GdnPrefillBuffers,
369        batch_size: u32,
370        chunk_len: u32,
371        ctx: &ForwardContext,
372        stream: u64,
373    ) -> Result<()> {
374        self.prefill_gdn_full_batched_inner(
375            h_state_ptrs,
376            gdn_bufs,
377            batch_size,
378            chunk_len,
379            ctx,
380            stream,
381        )
382    }
383
384    fn prefill_gdn_full_batched_fla_varlen(
385        &self,
386        h_state_ptrs: DevicePtr,
387        gdn_bufs: &GdnPrefillBuffers,
388        batch_size: u32,
389        cu_seqlens: DevicePtr,
390        max_num_chunks: u32,
391        total_nt: usize,
392        max_seqlen: u32,
393        ctx: &ForwardContext,
394        stream: u64,
395    ) -> Result<bool> {
396        self.prefill_gdn_full_batched_fla_varlen_inner(
397            h_state_ptrs,
398            gdn_bufs,
399            batch_size,
400            cu_seqlens,
401            max_num_chunks,
402            total_nt,
403            max_seqlen,
404            ctx,
405            stream,
406        )
407    }
408
409    fn prefill_phase3(
410        &self,
411        hidden: DevicePtr,
412        residual: DevicePtr,
413        num_tokens: usize,
414        gdn_bufs: &GdnPrefillBuffers,
415        token_offset: usize,
416        ctx: &ForwardContext,
417        stream: u64,
418    ) -> Result<()> {
419        self.prefill_phase3_inner(
420            hidden,
421            residual,
422            num_tokens,
423            gdn_bufs,
424            token_offset,
425            ctx,
426            stream,
427        )
428    }
429
430    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
431        self.alloc_state_inner(gpu)
432    }
433
434    /// Free the PLE carry this sequence lazily attached.
435    ///
436    /// Only the `ple` field — the h/conv state in `SsmLayerState` is pooled
437    /// and released by slot in `free_sequence_dispatch`, so freeing it here
438    /// would be a double free. The PLE conv buffer is the one piece that is
439    /// allocated per sequence and owned by nothing.
440    fn release_state(&self, state: &mut dyn LayerState, gpu: &dyn GpuBackend) -> Result<()> {
441        let Some(ssm) = state
442            .as_any_mut()
443            .downcast_mut::<crate::layer::SsmLayerState>()
444        else {
445            return Ok(());
446        };
447        let Some(mut st) = ssm.ple.take() else {
448            return Ok(());
449        };
450        let Some(ple) = self.ple.as_ref() else {
451            anyhow::bail!("release_state: PLE seq state present but layer has no PLE");
452        };
453        ple.release_seq_state(&mut st, gpu)
454    }
455}
spark_model/layers/qwen3_ssm/trait_layer.rs

spark_model/layers/qwen3_ssm/
trait_layer.rs
