spark_model/model/
lm_head_q6k.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3//! The output head as raw Q6_K blocks, for the DeepSeek-V4.1 Flash GGUF.
4//!
5//! The Q2_K GGUF ships `output.weight` as Q6_K (0.5 GB). Expanding it to bf16
6//! at load made `dense_gemv_bf16` read 1.3 GB a token, 5.5 ms of a 65 ms
7//! decode step, the single largest kernel. Here the blocks stay as the file
8//! left them and the logits come from the vendored ggml `Q6_K x q8_1` GEMV
9//! (`kquant_mmvq_q6_k_w`, warp per vocab row), the same dot llama.cpp runs
10//! for this head: the hidden row is quantised to `block_q8_1` first, then one
11//! launch over the vocab. Installed by the factory only when the store holds
12//! `lm_head.weight` as [`WeightDtype::Q6K`]; every other head is untouched.
13//!
14//! [`WeightDtype::Q6K`]: spark_runtime::weights::WeightDtype::Q6K
15
16use anyhow::Result;
17use spark_runtime::gpu::{DevicePtr, KernelHandle};
18
19use super::types::TransformerModel;
20use crate::layers::ops;
21
22/// Rows of hidden state one `kquant_mmvq_*_w` launch takes (`KQ_MAX_M`).
23const MAX_M: u32 = 8;
24
25/// The resident Q6_K head and the two kernels that serve it.
26pub(crate) struct LmHeadQ6k {
27    /// `[vocab][hidden / 256]` raw 210-byte super-blocks.
28    blocks: DevicePtr,
29    mmvq_w: KernelHandle,
30    q8_rows: KernelHandle,
31    /// `block_q8_1 [MAX_M][hidden / 32]` scratch for the quantised hidden rows.
32    q8_scratch: DevicePtr,
33}
34
35impl TransformerModel {
36    /// Install the raw Q6_K head. `blocks` is the store's `lm_head.weight`
37    /// pointer, `[vocab, hidden]` in Q6_K super-blocks.
38    pub fn set_lm_head_q6k(&mut self, blocks: DevicePtr) -> Result<()> {
39        let h = self.config.hidden_size as u32;
40        let v = self.config.vocab_size as u32;
41        anyhow::ensure!(
42            h.is_multiple_of(256),
43            "Q6_K lm_head: hidden {h} is not a multiple of 256"
44        );
45        let mmvq_w = self.gpu.kernel(ops::KQUANT_MODULE, "kquant_mmvq_q6_k_w")?;
46        let q8_rows = self
47            .gpu
48            .kernel(ops::KQUANT_MODULE, "kquant_q8_1_rows_bf16")?;
49        let q8_scratch = self.gpu.alloc(ops::kquant_q8_1_rows_bytes(MAX_M, h))?;
50        tracing::info!(
51            "LM head served from raw Q6_K blocks (K-quant GEMV, vocab={v}, {} MB resident)",
52            ops::kquant_weight_bytes(v, h, ops::Q6K_BLOCK_BYTES) >> 20
53        );
54        self.lm_head_q6k = Some(LmHeadQ6k {
55            blocks,
56            mmvq_w,
57            q8_rows,
58            q8_scratch,
59        });
60        Ok(())
61    }
62
63    /// Logits for `m` bf16 hidden rows (`[m, hidden]`, contiguous) into
64    /// `logits` (`[m, vocab]` bf16) on the Q6_K head. `Ok(false)` when no
65    /// Q6_K head is installed, so the caller falls through to its own paths.
66    pub(super) fn lm_head_q6k_run(
67        &self,
68        hidden: DevicePtr,
69        m: u32,
70        logits: DevicePtr,
71        stream: u64,
72    ) -> Result<bool> {
73        let Some(ref q6) = self.lm_head_q6k else {
74            return Ok(false);
75        };
76        let h = self.config.hidden_size as u32;
77        let v = self.config.vocab_size as u32;
78        let gpu = self.gpu.as_ref();
79        let mut done = 0u32;
80        while done < m {
81            let mm = (m - done).min(MAX_M);
82            let x = hidden.offset(done as usize * h as usize * 2);
83            let out = logits.offset(done as usize * v as usize * 2);
84            ops::kquant_q8_1_rows(gpu, q6.q8_rows, x, q6.q8_scratch, mm, h, stream)?;
85            ops::kquant_mmvq_w(
86                gpu,
87                q6.mmvq_w,
88                q6.blocks,
89                q6.q8_scratch,
90                out,
91                v,
92                h,
93                mm,
94                stream,
95            )?;
96            done += mm;
97        }
98        Ok(true)
99    }
100}
spark_model/model/lm_head_q6k.rs

spark_model/model/
lm_head_q6k.rs
