spark_model/layers/deepseek_v41_layer/
trait_impl.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! `impl TransformerLayer for DeepSeekV41Layer`: decode and prefill both go
5//! through `step`; multi-sequence decode and graph capture are declined.
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9
10use super::{DeepSeekV41Layer, V41LayerState};
11use crate::layer::{ForwardContext, LayerState};
12use crate::layers::attn_v41::AttnV41LayerState;
13
14impl crate::layer::TransformerLayer for DeepSeekV41Layer {
15    #[allow(clippy::too_many_arguments)]
16    fn decode(
17        &self,
18        hidden: DevicePtr,
19        _residual: DevicePtr,
20        state: &mut dyn LayerState,
21        _kv_cache: &mut spark_runtime::kv_cache::PagedKvCache,
22        seq_len: usize,
23        _block_table: &mut Vec<u32>,
24        _disk_block_ids: &mut Vec<u32>,
25        _disk_last_offloaded_per_layer: &mut Vec<u32>,
26        ctx: &ForwardContext,
27        stream: u64,
28    ) -> Result<()> {
29        self.step(hidden, 1, seq_len, state, ctx, stream)
30    }
31
32    #[allow(clippy::too_many_arguments)]
33    fn prefill(
34        &self,
35        hidden: DevicePtr,
36        _residual: DevicePtr,
37        num_tokens: usize,
38        state: &mut dyn LayerState,
39        _kv_cache: &mut spark_runtime::kv_cache::PagedKvCache,
40        seq_len_start: usize,
41        _block_table: &mut Vec<u32>,
42        _disk_block_ids: &mut Vec<u32>,
43        _disk_last_offloaded_per_layer: &mut Vec<u32>,
44        _kv_write_start: usize,
45        ctx: &ForwardContext,
46        stream: u64,
47    ) -> Result<()> {
48        self.step(hidden, num_tokens, seq_len_start, state, ctx, stream)
49    }
50
51    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
52        Ok(Box::new(V41LayerState {
53            attn: AttnV41LayerState::new(gpu, &self.rt.attn_cfg, self.role)?,
54            graphs: None,
55        }))
56    }
57
58    fn decode_graph_unsupported(&self) -> bool {
59        true
60    }
61
62    fn decode_multi_seq_unsupported(&self) -> bool {
63        true
64    }
65
66    /// The buffers go with the sequence. `alloc_state` allocates the window
67    /// ring, the compressor's partial group, the latent cache and the
68    /// index-key cache per sequence with `gpu.alloc`, so a release that did
69    /// nothing left all four on the device for the life of the process. The
70    /// latent cache alone is `max_seq / ratio` rows of 512 bf16 (4 MiB at
71    /// 4096 tokens, ratio 1), so every request left on the order of 15 to
72    /// 25 MiB across the forty layers, and it never came back.
73    /// The captured segments bake this sequence's buffers, so they go first.
74    fn release_state(&self, state: &mut dyn LayerState, gpu: &dyn GpuBackend) -> Result<()> {
75        if let Some(st) = state.as_any_mut().downcast_mut::<V41LayerState>() {
76            if let Some(g) = st.graphs.take() {
77                g.destroy(gpu)?;
78            }
79            st.attn.free(gpu)?;
80        }
81        Ok(())
82    }
83
84    /// Lowering `seq_len` rewinds none of this layer's state: the shared
85    /// compressed-row count only grows (`SharedV41::compress_len`), the
86    /// ratio-2 compressor holds a partial group of the dropped tokens, and
87    /// the engram hasher's n-gram history ends in them. A watchdog rollback
88    /// then re-steers from a corrupted context (09-16: the second MinHeap
89    /// request spliced `_wHere` at the boundary and degenerated from
90    /// there). Declined until a rewind exists; the loop watchdog still
91    /// ends the response.
92    fn decode_rollback_unsupported(&self) -> bool {
93        true
94    }
95}