metrale_model_arch/engram_v41/
apply.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: [`EngramV41::apply`]: the `wkv` projection of the rows in the
5//! workspace (single token: off the raw Q2_K blocks when the layer carries
6//! them, else the bf16 GEMV; several tokens: the bf16 GEMM) and the gate onto
7//! the highway.
8//!
9//! Owner: model-arch (DeepSeek-V4.1 engram).
10//! Invariants:
11//! - `apply` refuses more than `max_tokens` tokens, a layer without weights,
12//!   and a Q2_K `in_features` that is not a multiple of 256, before any launch.
13
14use anyhow::{Context, Result, ensure};
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
16use metrale_gpu_runtime::kernel_args::KernelLaunch;
17
18use super::EngramV41;
19use metrale_model_layers::layers::ops;
20use metrale_model_layers::weight_map::DenseWeight;
21
22impl EngramV41 {
23    /// 2026-09-25: Project the rows already in the workspace and gate `tokens`
24    /// positions of `streams` (`[T, hc, dim]` f32) in place, for model layer
25    /// `layer`. Enqueued on `stream`.
26    pub fn apply(
27        &self,
28        gpu: &dyn GpuBackend,
29        layer: usize,
30        streams: DevicePtr,
31        tokens: usize,
32        stream: u64,
33    ) -> Result<()> {
34        ensure!(
35            tokens <= self.max_tokens,
36            "engram: {tokens} tokens exceeds the {} workspace",
37            self.max_tokens
38        );
39        let w = self
40            .layer(layer)
41            .with_context(|| format!("engram: layer {layer} has no weights"))?;
42        let wkv = DenseWeight { weight: w.wkv };
43        if tokens == 1 && !w.wkv_q2k.is_null() {
44            let (n, k) = (self.out_features() as u32, self.in_features() as u32);
45            ensure!(
46                k.is_multiple_of(256),
47                "engram: in_features {k} is not Q2_K blocks"
48            );
49            KernelLaunch::new(gpu, self.gemv_q2k_k)
50                .grid([n.div_ceil(4), 1, 1])
51                .block([256, 1, 1])
52                .arg_ptr(self.rows_of(w))
53                .arg_ptr(w.wkv_q2k)
54                .arg_ptr(self.kv)
55                .arg_u32(n)
56                .arg_u32(k)
57                .launch(stream)?;
58        } else if tokens == 1 {
59            ops::dense_gemv(
60                gpu,
61                self.gemv_k,
62                self.rows_of(w),
63                &wkv,
64                self.kv,
65                self.out_features() as u32,
66                self.in_features() as u32,
67                stream,
68            )?;
69        } else {
70            ops::dense_gemm(
71                gpu,
72                self.gemm_k,
73                self.rows_of(w),
74                &wkv,
75                self.kv,
76                tokens as u32,
77                self.out_features() as u32,
78                self.in_features() as u32,
79                stream,
80            )?;
81        }
82        KernelLaunch::new(gpu, self.gate_k)
83            .grid([tokens as u32, self.hc as u32, 1])
84            .block([256, 1, 1])
85            .arg_ptr(streams)
86            .arg_ptr(self.kv)
87            .arg_ptr(w.qk)
88            .arg_u32(self.dim as u32)
89            .arg_u32(self.hc as u32)
90            .arg_f32(self.eps)
91            .launch(stream)
92    }
93}
metrale_model_arch/engram_v41/apply.rs

metrale_model_arch/engram_v41/
apply.rs
