spark_model/layers/attn_v41/
forward.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! [`AttnV41::forward`]: one layer's attention for `m` tokens — projections,
5//! the window ring, the compressed rows and index selection, sparse attention
6//! with the sink, and the grouped output projection. Split from
7//! `attn_v41.rs` (500-LoC cap).
8
9use anyhow::{Context, Result, ensure};
10use spark_runtime::gpu::{DevicePtr, GpuBackend};
11
12use super::{
13    AttnV41, AttnV41LayerState, AttnV41LayerWeights, AttnV41Run, SharedV41, at, upload_i32,
14};
15use crate::layers::deepseek_v41_ref::attn::window_topk_idxs;
16use crate::layers::deepseek_v41_ref::compress::LATENT_BLOCK;
17
18impl AttnV41 {
19    /// One layer's attention for `m` tokens at `start_pos`: `x` is the normed
20    /// input `[m, dim]` bf16; the output is `[m, dim]` bf16 in `run.out`.
21    pub fn forward(
22        &self,
23        gpu: &dyn GpuBackend,
24        w: &AttnV41LayerWeights,
25        st: &mut AttnV41LayerState,
26        shared: &mut SharedV41,
27        x: DevicePtr,
28        m: usize,
29        start_pos: usize,
30        stream: u64,
31    ) -> Result<AttnV41Run> {
32        let c = &self.cfg;
33        ensure!(
34            m >= 1 && m <= c.max_tokens,
35            "attn_v41: {m} tokens outside 1..={}",
36            c.max_tokens
37        );
38        ensure!(
39            start_pos + m <= c.max_seq,
40            "attn_v41: position {} beyond max_seq {}",
41            start_pos + m,
42            c.max_seq
43        );
44        let (nh, hd) = (c.n_heads, c.head_dim);
45        let yarn = w.role.ratio > 0;
46        let pos: Vec<i32> = (0..m).map(|t| (start_pos + t) as i32).collect();
47        let hpos: Vec<i32> = pos
48            .iter()
49            .flat_map(|&p| std::iter::repeat_n(p, nh))
50            .collect();
51        upload_i32(gpu, self.pos, &pos)?;
52        upload_i32(gpu, self.head_pos, &hpos)?;
53
54        self.q_kv_projections(gpu, w, x, m, yarn, stream)?;
55
56        // the window ring (stream-ordered copies, no host sync)
57        let win = c.window;
58        let row = hd * 2;
59        let (rows_a, rows_a_len) = if start_pos == 0 {
60            if m <= win {
61                gpu.copy_d2d_async(self.kv, st.window, m * row, stream)?;
62            } else {
63                let cutoff = m % win;
64                let tail = at(self.kv, (m - win) * row);
65                gpu.copy_d2d_async(
66                    tail,
67                    at(st.window, cutoff * row),
68                    (win - cutoff) * row,
69                    stream,
70                )?;
71                gpu.copy_d2d_async(
72                    at(tail, (win - cutoff) * row),
73                    st.window,
74                    cutoff * row,
75                    stream,
76                )?;
77            }
78            (self.kv, m)
79        } else {
80            let slot = start_pos % win;
81            gpu.copy_d2d_async(self.kv, at(st.window, slot * row), row, stream)?;
82            (st.window, win)
83        };
84        let (mut idx, mut topk) = window_topk_idxs(win, m, start_pos);
85
86        let (mut rows_b, mut rows_b_len) = (None, 0usize);
87        if w.role.ratio > 0 {
88            let ratio = w.role.ratio;
89            let offset = rows_a_len;
90            let compress_len = (start_pos + m) / ratio;
91            let latent_groups = if w.role.is_kv_source {
92                let g = self.compressor(gpu, w, st, x, m, start_pos, stream)?;
93                shared.compress_kv = st.compress_kv;
94                g
95            } else {
96                None
97            };
98            let (cidx, ctopk) = if !w.role.is_index_source {
99                (shared.topk_idxs.clone(), shared.topk)
100            } else if compress_len == 0 {
101                (Vec::new(), 0)
102            } else {
103                let r = self.indexer(
104                    gpu,
105                    w,
106                    st,
107                    shared,
108                    x,
109                    latent_groups,
110                    m,
111                    start_pos,
112                    offset,
113                    stream,
114                )?;
115                shared.topk_idxs = r.0.clone();
116                shared.topk = r.1;
117                r
118            };
119            if let Some(groups) = latent_groups {
120                let cache = st.compress_kv.context("kv source without a latent cache")?;
121                upload_i32(
122                    gpu,
123                    self.grp_pos,
124                    &Self::group_positions(groups, ratio, start_pos),
125                )?;
126                self.rope(
127                    gpu,
128                    self.latent,
129                    self.grp_pos,
130                    groups,
131                    hd,
132                    true,
133                    false,
134                    stream,
135                )?;
136                self.fp4_quant(gpu, self.latent, groups * hd, LATENT_BLOCK, true, stream)?;
137                gpu.copy_d2d_async(
138                    self.latent,
139                    at(cache, (start_pos / ratio) * row),
140                    groups * row,
141                    stream,
142                )?;
143                shared.compress_kv = Some(cache);
144                shared.compress_len = shared.compress_len.max(start_pos / ratio + groups);
145            }
146            rows_b = Some(
147                shared
148                    .compress_kv
149                    .context("compressed layer before any kv source published")?,
150            );
151            rows_b_len = compress_len;
152            ensure!(
153                cidx.len() == m * ctopk,
154                "index selection is {} for {m} x {ctopk}",
155                cidx.len()
156            );
157            let mut merged = Vec::with_capacity(m * (topk + ctopk));
158            for t in 0..m {
159                merged.extend_from_slice(&idx[t * topk..(t + 1) * topk]);
160                merged.extend_from_slice(&cidx[t * ctopk..(t + 1) * ctopk]);
161            }
162            idx = merged;
163            topk += ctopk;
164        }
165        ensure!(
166            topk <= win + c.index_topk && topk <= 2048,
167            "attn_v41: topk {topk} exceeds the workspace"
168        );
169        upload_i32(gpu, self.idx_dev, &idx)?;
170
171        self.attend_and_project(
172            gpu,
173            w,
174            rows_a,
175            rows_a_len,
176            rows_b,
177            self.idx_dev,
178            topk,
179            m,
180            yarn,
181            stream,
182        )?;
183        // no host sync here: everything downstream runs on the same stream,
184        // and the routing download in the MoE block drains it before any
185        // expert slot can be rewritten
186        Ok(AttnV41Run {
187            q: self.q,
188            rows_a,
189            rows_a_len,
190            rows_b,
191            rows_b_len,
192            idx,
193            topk,
194            o: self.o,
195            out: self.out,
196        })
197    }
198}
spark_model/layers/attn_v41/forward.rs

spark_model/layers/attn_v41/
forward.rs
