spark_model/layers/attn_v41/
init.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! [`AttnV41`] construction and teardown: kernel lookup, the RoPE tables and
5//! the per-call workspaces (`new`), and their release (`free`). Split from
6//! `attn_v41.rs` (500-LoC cap).
7
8use anyhow::{Result, ensure};
9use spark_runtime::gpu::{DevicePtr, GpuBackend};
10
11use super::{AttnV41, AttnV41Cfg, GEMM_MODULE, Kernels, MODULE, upload_f32};
12use crate::layers::deepseek_v41_ref::attn::freqs_cis;
13use crate::layers::deepseek_v41_ref::compress::yarn_freqs_cis;
14use crate::layers::ops::{KQUANT_MODULE, kquant_mmq_act_bytes, kquant_q8_1_rows_bytes};
15
16impl AttnV41 {
17    pub fn new(gpu: &dyn GpuBackend, cfg: AttnV41Cfg) -> Result<Self> {
18        ensure!(
19            cfg.rope_dim.is_multiple_of(2) && cfg.rope_dim / 2 <= 256,
20            "rope_dim {} not supported",
21            cfg.rope_dim
22        );
23        ensure!(
24            (cfg.head_dim * cfg.n_heads).is_multiple_of(cfg.groups),
25            "n_heads * hd not divisible by groups"
26        );
27        ensure!(
28            cfg.head_dim.is_multiple_of(32) && cfg.index_hd.is_multiple_of(32),
29            "head dims must be multiples of the 32-block"
30        );
31        let k = Kernels {
32            gemm: gpu.kernel(GEMM_MODULE, "dense_gemm_bf16")?,
33            gemv: gpu.kernel("gemv", "dense_gemv_bf16")?,
34            q8_rows: gpu.kernel(KQUANT_MODULE, "kquant_q8_1_rows_bf16")?,
35            mmvq_q2k_w: gpu.kernel(KQUANT_MODULE, "kquant_mmvq_q2_k_w")?,
36            mmvq_q2k_groups_w: gpu.kernel(KQUANT_MODULE, "kquant_mmvq_q2_k_groups_w")?,
37            mmvq_q2k_pair_w: gpu.kernel(KQUANT_MODULE, "kquant_mmvq_q2_k_pair_w")?,
38            quant_d2s6: gpu.kernel(KQUANT_MODULE, "atlas_q8_1_quantize_d2s6_bf16")?,
39            mmq_q2k_nc: gpu.kernel(KQUANT_MODULE, "atlas_q2_k_mmq128_nc")?,
40            mmq_q2k_wc: gpu.kernel(KQUANT_MODULE, "atlas_q2_k_mmq128_wc")?,
41            rmsnorm_bf16: gpu.kernel(MODULE, "attn_v41_rmsnorm_bf16")?,
42            rmsnorm_f32: gpu.kernel(MODULE, "attn_v41_rmsnorm_f32")?,
43            rope: gpu.kernel(MODULE, "attn_v41_rope")?,
44            act_quant: gpu.kernel(MODULE, "attn_v41_act_quant_fp8")?,
45            fp4_quant: gpu.kernel(MODULE, "attn_v41_fp4_quant")?,
46            gemm_f32: gpu.kernel(MODULE, "attn_v41_gemm_f32")?,
47            gemv_f32_staged: gpu.kernel(MODULE, "attn_v41_gemv_f32_staged")?,
48            pool: gpu.kernel(MODULE, "attn_v41_pool")?,
49            index_score: gpu.kernel(MODULE, "attn_v41_index_score")?,
50            sparse_attn: gpu.kernel(MODULE, "attn_v41_sparse_attn")?,
51            slice_cols: gpu.kernel(MODULE, "attn_v41_slice_cols")?,
52            scatter_cols: gpu.kernel(MODULE, "attn_v41_scatter_cols")?,
53            scale_bf16: gpu.kernel(MODULE, "attn_v41_scale_bf16")?,
54            ring_put: gpu.kernel(MODULE, "attn_v41_ring_put")?,
55        };
56        let plain = freqs_cis(cfg.rope_dim, cfg.max_seq, cfg.rope_theta);
57        let yarn = yarn_freqs_cis(
58            cfg.rope_dim,
59            cfg.max_seq,
60            cfg.orig_seq,
61            cfg.compress_rope_theta,
62            cfg.rope_factor,
63            cfg.beta_fast,
64            cfg.beta_slow,
65        );
66        let flat = |v: &[(f32, f32)]| -> Vec<f32> { v.iter().flat_map(|&(c, s)| [c, s]).collect() };
67        let fc_plain = upload_f32(gpu, &flat(&plain))?;
68        let fc_yarn = upload_f32(gpu, &flat(&yarn))?;
69        let m = cfg.max_tokens;
70        let (nh, hd, nhi, ihd) = (cfg.n_heads, cfg.head_dim, cfg.index_heads, cfg.index_hd);
71        let max_width = cfg.max_seq;
72        let max_topk = cfg.window + cfg.index_topk;
73        let alloc = |bytes: usize| gpu.alloc(bytes.max(16));
74        // `n_heads * head_dim`: the grouped `wo_a` decode path quantises the
75        // whole rotated attention output row once (primitives::wo_a_grouped).
76        let kmax = cfg
77            .dim
78            .max(cfg.q_rank)
79            .max(cfg.gw())
80            .max(cfg.groups * cfg.o_rank)
81            .max(nh * hd);
82        Ok(AttnV41 {
83            a_q8: alloc(
84                kquant_q8_1_rows_bytes(8, kmax as u32)
85                    .max(kquant_mmq_act_bytes(m as u32, kmax as u32)),
86            )?,
87            qr_raw: alloc(m * cfg.q_rank * 2)?,
88            qr: alloc(m * cfg.q_rank * 2)?,
89            q: alloc(m * nh * hd * 2)?,
90            kv_raw: alloc(m * hd * 2)?,
91            kv: alloc(m * hd * 2)?,
92            o: alloc(m * nh * hd * 2)?,
93            o_rot: alloc(m * nh * hd * 2)?,
94            og: alloc(m * cfg.groups * cfg.o_rank * 2)?,
95            slice_in: alloc(m * cfg.gw() * 2)?,
96            slice_out: alloc(m * cfg.o_rank * 2)?,
97            out: alloc(m * cfg.dim * 2)?,
98            pos: alloc(m * 4)?,
99            head_pos: alloc(m * nh * 4)?,
100            idx_pos: alloc(m * nhi * 4)?,
101            grp_pos: alloc(m * 4)?,
102            idx_dev: alloc(m * max_topk * 4)?,
103            idx_dev_win: alloc(m * max_topk * 4)?,
104            ckv: alloc(m * hd * 4)?,
105            cscore: alloc(m * hd * 4)?,
106            pooled: alloc(m * hd * 4)?,
107            latent_raw: alloc(m * hd * 2)?,
108            latent: alloc(m * hd * 2)?,
109            ik_raw: alloc(m * ihd * 2)?,
110            ik: alloc(m * ihd * 2)?,
111            iq: alloc(m * nhi * ihd * 2)?,
112            iw_raw: alloc(m * nhi * 2)?,
113            iw: alloc(m * nhi * 2)?,
114            score: alloc(m * max_width * 4)?,
115            decode_pos: None,
116            decode_idx: None,
117            decode_idx_win: None,
118            cfg,
119            k,
120            fc_plain,
121            fc_yarn,
122        })
123    }
124
125    /// The layer output buffer, bf16 `[max_tokens, dim]`: a pointer a captured
126    /// decode step bakes.
127    pub fn out_ptr(&self) -> DevicePtr {
128        self.out
129    }
130
131    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
132        for p in [
133            self.fc_plain,
134            self.fc_yarn,
135            self.qr_raw,
136            self.qr,
137            self.q,
138            self.kv_raw,
139            self.kv,
140            self.o,
141            self.o_rot,
142            self.og,
143            self.slice_in,
144            self.slice_out,
145            self.out,
146            self.pos,
147            self.head_pos,
148            self.idx_pos,
149            self.grp_pos,
150            self.idx_dev,
151            self.idx_dev_win,
152            self.ckv,
153            self.cscore,
154            self.pooled,
155            self.latent_raw,
156            self.latent,
157            self.ik_raw,
158            self.ik,
159            self.iq,
160            self.iw_raw,
161            self.iw,
162            self.score,
163        ] {
164            gpu.free(p)?;
165        }
166        Ok(())
167    }
168}
spark_model/layers/attn_v41/init.rs

spark_model/layers/attn_v41/
init.rs
