spark_model/layers/engram_v41/
apply.rs

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

spark_model/layers/engram_v41/
apply.rs
