spark_model/layers/
deepseek_v41_layer.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! DeepSeek-V4.1 Flash: one transformer block on the generic model.
5//!
6//! The generic `TransformerModel` owns the embedding, the FP32 mHC highway
7//! (`BufferArena::hc_streams`, `[T, hc, H]`), the final norm and the LM head;
8//! each layer here does everything in between, in the reference's order
9//! (`deepseek_v41_ref::model::forward`, 39/39 against DeepSeek's golden):
10//!
11//! ```text
12//! layer 0 only:        hc_expand(embedding -> streams)
13//! engram layers only:  streams += engram(token n-grams)          (EngramV41)
14//! attention:           (pre_a, post_a, comb_a) = hc_mixes(streams, attn site)
15//!                      x = rmsnorm(collapse(streams, pre_prev))   <- DELAYED pre
16//!                      streams = hc_post(attn(x), streams, post_a, comb_a)
17//! ffn:                 (pre_f, post_f, comb_f) = hc_mixes(streams, ffn site)
18//!                      x = rmsnorm(collapse(streams, pre_a))
19//!                      streams = hc_post(moe(x), streams, post_f, comb_f)
20//!                      pre_prev = pre_f
21//! last layer only:     hidden = collapse(streams, pre_prev)      (no learned head)
22//! ```
23//!
24//! `pre_prev` starts as the one-hot on stream 0 (the reference's initial
25//! pre-mix) and travels with the sequence through [`V41Runtime`]. So do the
26//! shared attention slots, the engram hasher and the expert cache: V4.1's
27//! attention is shared ACROSS layers (four kv sources write, forty read), which
28//! no per-layer state can express, so one runtime object is shared by all
29//! forty layers through an `Arc` and guarded by mutexes. One sequence at a
30//! time; multi-sequence decode and the MODEL-level CUDA graph are declined
31//! through the layer hooks. The layer captures its own single-token step
32//! instead (`ATLAS_DS41_GRAPH=1`, see [`GraphMode`]): the host work in the
33//! middle of every layer (the routing download and the expert fetch; on the
34//! twelve kv/index source layers also the compressor group and the index
35//! top-k) splits the step into graph segments with the host spans between.
36//!
37//! Routed experts never sit in HBM: the cache is a page-locked, device-visible
38//! arena the GPU reads in place (`ExpertLru` + `ExpertSliceMap`), filled by
39//! pread from the seven shards; the engram tables are read by row (`EngramRowReader`).
40
41use std::sync::atomic::AtomicBool;
42use std::sync::{Arc, Mutex, OnceLock};
43
44use anyhow::Result;
45
46use spark_runtime::gpu::{DevicePtr, GpuBackend, GraphHandle, KernelHandle};
47use spark_runtime::weights::expert_stream::{
48    EngramRowReader, ExpertArena, ExpertLru, ExpertSliceMap,
49};
50
51use crate::layer::LayerState;
52use crate::layers::attn_v41::{
53    AttnV41, AttnV41Cfg, AttnV41LayerState, AttnV41LayerWeights, LayerRole, SharedV41,
54};
55use crate::layers::engram_v41::{EngramHashTables, EngramHasher, EngramV41};
56use crate::layers::moe_v41::{MoeV41, MoeV41Cfg, MoeV41LayerWeights, RouterWeights};
57use crate::layers::qwen3_attention::HcSiteWeights;
58use crate::weight_map::DenseWeight;
59
60mod graph;
61mod hc_launch;
62mod step;
63mod step_ffn;
64mod step_seg;
65mod step_seg_run;
66mod trait_impl;
67
68pub use step_seg::{SegState, step_graph_on};
69
70/// Everything the forty layers share for the ONE sequence in flight.
71pub struct V41Runtime {
72    pub attn_cfg: AttnV41Cfg,
73    pub moe_cfg: MoeV41Cfg,
74    pub attn: Mutex<AttnV41>,
75    pub moe: Mutex<MoeV41>,
76    pub engram: Mutex<EngramV41>,
77    pub lru: Mutex<ExpertLru>,
78    /// Kept alive for the cache's lifetime; freed with the runtime.
79    pub arena: ExpertArena,
80    pub slices: Arc<ExpertSliceMap>,
81    pub rows: EngramRowReader,
82    pub tables: Arc<EngramHashTables>,
83    pub hasher: Mutex<EngramHasher>,
84    pub shared: Mutex<SharedV41>,
85    /// The step's engram hashes `[tokens, n_engram_layers, cols]`, computed by
86    /// the first engram layer and reused by the second.
87    pub step_hashes: Mutex<Option<Vec<i64>>>,
88    /// The delayed `pre` mix `[max_tokens, hc]` f32 on the device.
89    pub pre_prev: DevicePtr,
90    /// `[max_tokens, (2 + hc) * hc]` f32 scratch for the hc mixes (dot -> finish)
91    pub mixes_s: DevicePtr,
92    pub reader_threads: usize,
93    pub n_layers: usize,
94    pub hc_mult: usize,
95    pub hidden: usize,
96    pub sinkhorn_iters: usize,
97    pub hc_eps: f32,
98    pub norm_eps: f32,
99    // per-step scratch, guarded by the same one-sequence contract
100    pub pre_a: DevicePtr,
101    pub pre_f: DevicePtr,
102    pub post_s: DevicePtr,
103    pub comb_s: DevicePtr,
104    pub attn_in: DevicePtr,
105    pub max_tokens: usize,
106    /// per-step totals across layers, printed by the last layer when ATLAS_DS41_DIAG=1
107    pub step_moe: Mutex<crate::layers::moe_v41::MoeV41Timing>,
108    pub step_attn_ms: Mutex<f64>,
109    pub step_engram_ms: Mutex<f64>,
110    pub step_start: Mutex<Option<std::time::Instant>>,
111    /// Set when a capture failed: every segment from then on runs eagerly
112    /// (graphs already captured keep replaying; they are valid).
113    pub graph_disabled: AtomicBool,
114    /// The whole-step segment graphs (`ATLAS_DS41_STEP_GRAPH=1`, `step_seg.rs`).
115    pub seg: Mutex<SegState>,
116    /// Every layer's role, filled as the layers are built (the segment owner
117    /// reads the roles of the layers its graph spans).
118    pub roles: Mutex<Vec<Option<LayerRole>>>,
119    /// `(layer, hash index)` of the engram layers.
120    pub engram_layers: Mutex<Vec<(usize, usize)>>,
121    /// `ATLAS_DS41_PREDICT_TRACE`: the last three layers' MoE inputs
122    /// (`[3, hidden]` bf16, slot `layer % 3`) and the trace file.
123    pub pred_x: DevicePtr,
124    pub pred_trace: Mutex<Option<std::io::BufWriter<std::fs::File>>>,
125}
126
127/// How the single-token step runs. `ATLAS_DS41_GRAPH_ORACLE=1` runs every
128/// layer BOTH ways and compares the highway bit for bit; `ATLAS_DS41_GRAPH=1`
129/// replays the captured segments; anything else is the eager step. Read once.
130#[derive(Clone, Copy, Debug, PartialEq, Eq)]
131pub enum GraphMode {
132    Off,
133    On,
134    Oracle,
135}
136
137pub fn graph_mode() -> GraphMode {
138    static MODE: OnceLock<GraphMode> = OnceLock::new();
139    *MODE.get_or_init(|| {
140        let is_one = |k: &str| std::env::var(k).is_ok_and(|v| v == "1");
141        if is_one("ATLAS_DS41_GRAPH_ORACLE") {
142            GraphMode::Oracle
143        } else if is_one("ATLAS_DS41_GRAPH") {
144            GraphMode::On
145        } else {
146            GraphMode::Off
147        }
148    })
149}
150
151/// The pointers a layer's captured segments bake that are not the runtime's
152/// own (those live as long as the runtime). Checked on every replay: a
153/// different set means the graphs describe other buffers and are recaptured.
154#[derive(Clone, Copy, Debug, PartialEq, Eq)]
155struct Baked {
156    hidden: DevicePtr,
157    streams: DevicePtr,
158    normed: DevicePtr,
159    window: DevicePtr,
160    /// the kv source's latent cache the captured `sparse_attn` reads
161    rows_b: Option<DevicePtr>,
162    attn_out: DevicePtr,
163    moe_out: DevicePtr,
164}
165
166/// One layer's captured single-token step, per sequence (it bakes the
167/// sequence's window ring and the kv source's cache).
168///
169/// Segment A = the attention site's mixes, collapse and norm, the attention
170/// (capturable layers only), `hc_post`, the ffn site's mixes, collapse and
171/// norm, the router GEMV. On the twelve kv/index source layers the attention
172/// runs eagerly between `a[0]` (through the norm) and `a[1]` (from `hc_post`).
173/// Host span: the routing download, the expert fetch, the plan uploads.
174/// Segment B = the expert compute, `hc_post`, the delayed-mix copy, and on
175/// the last layer the final collapse.
176struct LayerGraphs {
177    a: [Option<GraphHandle>; 2],
178    b: Option<GraphHandle>,
179    baked: Baked,
180}
181
182impl LayerGraphs {
183    fn destroy(self, gpu: &dyn GpuBackend) -> Result<()> {
184        for g in self.a.into_iter().chain([self.b]).flatten() {
185            gpu.destroy_graph(g)?;
186        }
187        Ok(())
188    }
189}
190
191// SAFETY: every raw device/host pointer here names memory the runtime owns
192// for its whole life; access is serialised by the mutexes and the
193// one-sequence contract.
194unsafe impl Send for V41Runtime {}
195unsafe impl Sync for V41Runtime {}
196
197pub struct V41LayerState {
198    pub attn: AttnV41LayerState,
199    graphs: Option<LayerGraphs>,
200}
201
202impl LayerState for V41LayerState {
203    fn as_any(&self) -> &dyn std::any::Any {
204        self
205    }
206    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
207        self
208    }
209}
210
211pub struct DeepSeekV41Layer {
212    pub idx: usize,
213    pub role: LayerRole,
214    pub rt: Arc<V41Runtime>,
215    pub attn_w: AttnV41LayerWeights,
216    pub moe_w: MoeV41LayerWeights,
217    /// The next layer's router, run on this layer's MoE input to predict
218    /// (and start reading) the experts the next layer will ask for.
219    pub next_router: Option<RouterWeights>,
220    /// `Some(index into engram_layer_ids)` on engram layers.
221    pub engram_index: Option<usize>,
222    pub hc_attn: HcSiteWeights,
223    pub hc_ffn: HcSiteWeights,
224    pub attn_norm: DenseWeight,
225    pub ffn_norm: DenseWeight,
226    pub k_hc_expand: KernelHandle,
227    pub k_hc_post: KernelHandle,
228    pub k_mixes_dot: KernelHandle,
229    pub k_mixes_finish: KernelHandle,
230    pub k_collapse: KernelHandle,
231    /// the site's finish and the block's collapse in one wide launch
232    pub k_finish_collapse: KernelHandle,
233    /// hc_v41_collapse / hc_post over ceil(H/256) blocks a token
234    pub k_collapse_wide: KernelHandle,
235    pub k_post_wide: KernelHandle,
236    pub k_rms_norm: KernelHandle,
237}
spark_model/layers/deepseek_v41_layer.rs

spark_model/layers/
deepseek_v41_layer.rs
