spark_runtime/buffers/
sizes_q2.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Native keep-packed Q2_0 buffer sizing, split out of `sizes.rs` (≤500 LoC
4//! cap). Both scratch buffers are env-gated — 0 (→ NULL) unless the respective
5//! `ATLAS_GGUF_NATIVE_Q2*` flag is set, so non-Q2 models pay nothing.
6
7use atlas_core::config::ModelConfig;
8
9/// Bytes for the native keep-packed Q2_0 prefill transient-dequant scratch: the
10/// LARGEST keep-packed projection `[N, K]` expanded to BF16 (2 bytes/elem). The
11/// prefill dequant writes `N*K` BF16 elements into this buffer, which is then
12/// consumed by the same-stream GEMM and reused by the next projection.
13///
14/// Every keep-packed projection has exactly one dimension equal to
15/// `hidden_size` (FFN gate/up `[inter, h]`, FFN down `[h, inter]`, attention
16/// q/k/v/o `[·, h]` or `[h, ·]`, fused GDN `in_proj_qkvz [qkvz, h]`), so
17/// `N*K = max_other_dim * hidden` — an EXACT bound, not an over-estimate, for
18/// the covered families (the k/v and gated-q terms are safe upper bounds).
19/// Independent of batch tokens (this dequants the WEIGHT, not activations).
20pub fn q2_dequant_scratch_bytes(config: &ModelConfig) -> usize {
21    let bf16 = 2;
22    let hd = config.head_dim;
23    let q_proj_mul = if config.attn_gated { 2 } else { 1 };
24    let max_n = config
25        .intermediate_size
26        .max(config.ssm_qkvz_size())
27        .max(config.mamba2_in_proj_size())
28        .max(config.num_attention_heads * q_proj_mul * hd)
29        .max(2 * config.num_key_value_heads * hd);
30    max_n * config.hidden_size * bf16
31}
32
33/// `(q2_dequant_scratch, q2_act_q8)` sizes for the arena. `m` = max batch
34/// tokens, `h` = hidden_size, `hd` = head_dim.
35///
36/// - `q2_dequant_scratch` (Tier-1, `ATLAS_GGUF_NATIVE_Q2=1`): the widest
37///   keep-packed projection expanded to BF16 (see [`q2_dequant_scratch_bytes`]).
38/// - `q2_act_q8` (Tier-2 MMQ, `ATLAS_GGUF_NATIVE_Q2_MMQ=1`): the q8_1 activation
39///   scratch. Widest INPUT dim K — FFN gate/up (h) or down (intermediate), attn
40///   qkv (h) or o (q_heads*head_dim), GDN qkvz (h). q8_1_mmq is 4 bytes/elem
41///   over kpad (K rounded to 256), + 1MB margin — matches `q8_1_scratch_bytes`.
42pub fn q2_scratch_sizes(config: &ModelConfig, m: usize, h: usize, hd: usize) -> (usize, usize) {
43    let dequant_enabled = std::env::var("ATLAS_GGUF_NATIVE_Q2").ok().as_deref() == Some("1");
44    let mmq_enabled = std::env::var("ATLAS_GGUF_NATIVE_Q2_MMQ").ok().as_deref() == Some("1");
45    q2_scratch_sizes_for(config, m, h, hd, dequant_enabled, mmq_enabled)
46}
47
48pub(super) fn q2_scratch_sizes_for(
49    config: &ModelConfig,
50    m: usize,
51    h: usize,
52    hd: usize,
53    dequant_enabled: bool,
54    mmq_enabled: bool,
55) -> (usize, usize) {
56    let q2_dequant_scratch = if dequant_enabled {
57        q2_dequant_scratch_bytes(config)
58    } else {
59        0
60    };
61
62    let q2_act_q8 = if mmq_enabled {
63        let kmax = h
64            .max(config.intermediate_size)
65            .max(config.num_attention_heads * hd);
66        let kpad = kmax.div_ceil(256) * 256;
67        m * kpad * 4 + (1 << 20)
68    } else {
69        0
70    };
71
72    (q2_dequant_scratch, q2_act_q8)
73}
spark_runtime/buffers/sizes_q2.rs

spark_runtime/buffers/
sizes_q2.rs
