metrale_model_arch/attn_v41/
init.rs

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

metrale_model_arch/attn_v41/
init.rs
