metrale_model_engine/model/
lm_head_q6k.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: The output head served from raw Q6_K blocks (DeepSeek-V4.1 GGUF).
5//!
6//! The blocks stay as the store holds them. For each group of up to `MAX_M` hidden
7//! rows, `kquant_q8_1_rows_bf16` quantises the rows to `block_q8_1`, then
8//! `kquant_mmvq_q6_k_w` computes their logits over the whole vocab. The factory
9//! installs this head only when the store holds `lm_head.weight` as
10//! [`WeightDtype::Q6K`].
11//!
12//! Owner: model-engine.
13//! Invariants: none beyond the types.
14//!
15//! [`WeightDtype::Q6K`]: metrale_model_weights::weights::WeightDtype::Q6K
16
17use anyhow::Result;
18use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
19
20use super::types::TransformerModel;
21use metrale_model_layers::layers::ops;
22
23/// 2026-09-25: Rows of hidden state one `kquant_mmvq_*_w` launch takes (`KQ_MAX_M` in kquant_moe.cu).
24const MAX_M: u32 = 8;
25
26/// 2026-09-25: The resident Q6_K head and the two kernels that serve it.
27pub(crate) struct LmHeadQ6k {
28    /// 2026-09-25: `[vocab][hidden / 256]` raw 210-byte super-blocks.
29    blocks: DevicePtr,
30    mmvq_w: KernelHandle,
31    q8_rows: KernelHandle,
32    /// 2026-09-25: `block_q8_1 [MAX_M][hidden / 32]` scratch for the quantised hidden rows.
33    q8_scratch: DevicePtr,
34}
35
36impl TransformerModel {
37    /// 2026-09-25: Install the raw Q6_K head. `blocks` is the store's
38    /// `lm_head.weight` pointer, `[vocab, hidden]` in Q6_K super-blocks. Fails
39    /// when `hidden_size` is not a multiple of 256.
40    pub fn set_lm_head_q6k(&mut self, blocks: DevicePtr) -> Result<()> {
41        let h = self.config.hidden_size as u32;
42        let v = self.config.vocab_size as u32;
43        anyhow::ensure!(
44            h.is_multiple_of(256),
45            "Q6_K lm_head: hidden {h} is not a multiple of 256"
46        );
47        let mmvq_w = self.gpu.kernel(ops::KQUANT_MODULE, "kquant_mmvq_q6_k_w")?;
48        let q8_rows = self
49            .gpu
50            .kernel(ops::KQUANT_MODULE, "kquant_q8_1_rows_bf16")?;
51        let q8_scratch = self.gpu.alloc(ops::kquant_q8_1_rows_bytes(MAX_M, h))?;
52        tracing::info!(
53            "LM head served from raw Q6_K blocks (K-quant GEMV, vocab={v}, {} MB resident)",
54            ops::kquant_weight_bytes(v, h, ops::Q6K_BLOCK_BYTES) >> 20
55        );
56        self.lm_head_q6k = Some(LmHeadQ6k {
57            blocks,
58            mmvq_w,
59            q8_rows,
60            q8_scratch,
61        });
62        Ok(())
63    }
64
65    /// 2026-09-25: Logits for `m` bf16 hidden rows (`[m, hidden]`, contiguous)
66    /// into `logits` (`[m, vocab]` bf16) on the Q6_K head. `Ok(false)` when no
67    /// Q6_K head is installed, so the caller falls through to its own paths.
68    pub(super) fn lm_head_q6k_run(
69        &self,
70        hidden: DevicePtr,
71        m: u32,
72        logits: DevicePtr,
73        stream: u64,
74    ) -> Result<bool> {
75        let Some(ref q6) = self.lm_head_q6k else {
76            return Ok(false);
77        };
78        let h = self.config.hidden_size as u32;
79        let v = self.config.vocab_size as u32;
80        let gpu = self.gpu.as_ref();
81        let mut done = 0u32;
82        while done < m {
83            let mm = (m - done).min(MAX_M);
84            let x = hidden.offset(done as usize * h as usize * 2);
85            let out = logits.offset(done as usize * v as usize * 2);
86            ops::kquant_q8_1_rows(gpu, q6.q8_rows, x, q6.q8_scratch, mm, h, stream)?;
87            ops::kquant_mmvq_w(
88                gpu,
89                q6.mmvq_w,
90                q6.blocks,
91                q6.q8_scratch,
92                out,
93                v,
94                h,
95                mm,
96                stream,
97            )?;
98            done += mm;
99        }
100        Ok(true)
101    }
102}
metrale_model_engine/model/lm_head_q6k.rs

metrale_model_engine/model/
lm_head_q6k.rs
