spark_model/layers/deepseek_v41_layer/
trait_impl.rs1use 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 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 fn decode_rollback_unsupported(&self) -> bool {
93 true
94 }
95}