metrale_model_arch/attn_v41/
forward.rs

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

metrale_model_arch/attn_v41/
forward.rs
