metrale_gpu_runtime/buffers/
sizes_q2.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Sizes of the two keep-packed Q2_0 prefill scratch buffers, each 0 unless its lever is `1`.
4//!
5//! Owner: gpu-runtime.
6//! Invariants:
7//! - `q2_scratch_sizes` returns 0 for the dequant scratch unless
8//!   `METRALE_GGUF_NATIVE_Q2=1`, and 0 for the q8_1 scratch unless
9//!   `METRALE_GGUF_NATIVE_Q2_MMQ=1`.
10
11use metrale_config::ModelConfig;
12
13/// 2026-09-25: Bytes of the keep-packed Q2_0 dequant scratch: `hidden_size`
14/// times the widest of intermediate, QKVZ, Mamba-2 in_proj, the q projection
15/// and 2 · kv_heads · head_dim, in BF16. This covers a weight `[N, K]` with
16/// `hidden_size` as one dimension and the other at most that width; the
17/// dequant callers debug-assert `N * K * 2` against the allocated size. It
18/// does not depend on the batch.
19pub fn q2_dequant_scratch_bytes(config: &ModelConfig) -> usize {
20    let bf16 = 2;
21    let hd = config.head_dim;
22    let q_proj_mul = if config.attn_gated { 2 } else { 1 };
23    let max_n = config
24        .intermediate_size
25        .max(config.ssm_qkvz_size())
26        .max(config.mamba2_in_proj_size())
27        .max(config.num_attention_heads * q_proj_mul * hd)
28        .max(2 * config.num_key_value_heads * hd);
29    max_n * config.hidden_size * bf16
30}
31
32/// 2026-09-25: `(q2_dequant_scratch, q2_act_q8)` bytes for the arena; `m` is
33/// `max_batch_tokens`, `h` `hidden_size`, `hd` `head_dim`. The q8_1 scratch is
34/// `m * kpad * 4 + 1 MiB` for K = max(h, intermediate, q_heads * head_dim)
35/// rounded up to 256 as `kpad`, the formula of `q8_1_scratch_bytes` in
36/// metrale-model-layers.
37pub fn q2_scratch_sizes(config: &ModelConfig, m: usize, h: usize, hd: usize) -> (usize, usize) {
38    let dequant_enabled = std::env::var("METRALE_GGUF_NATIVE_Q2").ok().as_deref() == Some("1");
39    let mmq_enabled = std::env::var("METRALE_GGUF_NATIVE_Q2_MMQ").ok().as_deref() == Some("1");
40    q2_scratch_sizes_for(config, m, h, hd, dequant_enabled, mmq_enabled)
41}
42
43pub(super) fn q2_scratch_sizes_for(
44    config: &ModelConfig,
45    m: usize,
46    h: usize,
47    hd: usize,
48    dequant_enabled: bool,
49    mmq_enabled: bool,
50) -> (usize, usize) {
51    let q2_dequant_scratch = if dequant_enabled {
52        q2_dequant_scratch_bytes(config)
53    } else {
54        0
55    };
56
57    let q2_act_q8 = if mmq_enabled {
58        let kmax = h
59            .max(config.intermediate_size)
60            .max(config.num_attention_heads * hd);
61        let kpad = kmax.div_ceil(256) * 256;
62        m * kpad * 4 + (1 << 20)
63    } else {
64        0
65    };
66
67    (q2_dequant_scratch, q2_act_q8)
68}
metrale_gpu_runtime/buffers/sizes_q2.rs

metrale_gpu_runtime/buffers/
sizes_q2.rs
