metrale_model_arch/deepseek_v41_layer/
trait_impl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: `impl TransformerLayer for DeepSeekV41Layer`: decode and
5//! prefill both go through `step`; model-level graph capture, multi-sequence
6//! decode and decode rollback are declined.
7//!
8//! Owner: model-arch, DeepSeek-V4.1.
9//! Invariants: none beyond the types.
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13
14use super::{DeepSeekV41Layer, V41LayerState};
15use crate::attn_v41::AttnV41LayerState;
16use metrale_model_layers::layer::{ForwardContext, LayerState};
17
18impl metrale_model_layers::layer::TransformerLayer for DeepSeekV41Layer {
19    #[allow(clippy::too_many_arguments)]
20    fn decode(
21        &self,
22        hidden: DevicePtr,
23        _residual: DevicePtr,
24        state: &mut dyn LayerState,
25        _kv_cache: &mut metrale_cache::kv_cache::PagedKvCache,
26        seq_len: usize,
27        _block_table: &mut Vec<u32>,
28        _disk_block_ids: &mut Vec<u32>,
29        _disk_last_offloaded_per_layer: &mut Vec<u32>,
30        ctx: &ForwardContext,
31        stream: u64,
32    ) -> Result<()> {
33        self.step(hidden, 1, seq_len, state, ctx, stream)
34    }
35
36    #[allow(clippy::too_many_arguments)]
37    fn prefill(
38        &self,
39        hidden: DevicePtr,
40        _residual: DevicePtr,
41        num_tokens: usize,
42        state: &mut dyn LayerState,
43        _kv_cache: &mut metrale_cache::kv_cache::PagedKvCache,
44        seq_len_start: usize,
45        _block_table: &mut Vec<u32>,
46        _disk_block_ids: &mut Vec<u32>,
47        _disk_last_offloaded_per_layer: &mut Vec<u32>,
48        _kv_write_start: usize,
49        ctx: &ForwardContext,
50        stream: u64,
51    ) -> Result<()> {
52        self.step(hidden, num_tokens, seq_len_start, state, ctx, stream)
53    }
54
55    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
56        Ok(Box::new(V41LayerState {
57            attn: AttnV41LayerState::new(gpu, &self.rt.attn_cfg, self.role)?,
58            graphs: None,
59        }))
60    }
61
62    /// 2026-09-25: Destroy this sequence's captured segments first (they bake
63    /// its buffers), then free the attention buffers `alloc_state` allocated
64    /// (`AttnV41LayerState::free`). A state of another type is left alone.
65    fn release_state(&self, state: &mut dyn LayerState, gpu: &dyn GpuBackend) -> Result<()> {
66        if let Some(st) = state.as_any_mut().downcast_mut::<V41LayerState>() {
67            if let Some(g) = st.graphs.take() {
68                g.destroy(gpu)?;
69            }
70            st.attn.free(gpu)?;
71        }
72        Ok(())
73    }
74}
75
76impl metrale_model_layers::layer::LayerCapabilities for DeepSeekV41Layer {
77    fn decode_graph_unsupported(&self) -> bool {
78        true
79    }
80
81    fn decode_multi_seq_unsupported(&self) -> bool {
82        true
83    }
84
85    /// 2026-09-25: Declined: lowering `seq_len` rewinds none of this layer's
86    /// state. The window ring and the ratio > 1 compressor state keep rows of
87    /// the dropped tokens, and `SharedV41::compress_len` only grows.
88    fn decode_rollback_unsupported(&self) -> bool {
89        true
90    }
91}
92
93impl metrale_model_layers::layer::LayerWeightSetup for DeepSeekV41Layer {}
94impl metrale_model_layers::layer::LayerWriteOnAccept for DeepSeekV41Layer {}
95impl metrale_model_layers::layer::LayerGraphHooks for DeepSeekV41Layer {}
96impl metrale_model_layers::layer::LayerAuxState for DeepSeekV41Layer {}
97impl metrale_model_layers::layer::LayerSplitPrefill for DeepSeekV41Layer {}
98impl metrale_model_layers::circuit_exec::CircuitBindings for DeepSeekV41Layer {}
metrale_model_arch/deepseek_v41_layer/trait_impl.rs

metrale_model_arch/deepseek_v41_layer/
trait_impl.rs
