metrale_model_layers/layers/
activation_quantization.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The `--activation-quantization` value in force for this process
4//! (`metrale_config::ActivationQuantization`), and the one question decode dispatch asks of it:
5//! does this projection family run one fixed activation format at every row count, and which.
6//!
7//! Owner: model-layers (quantization dispatch).
8//! Invariants:
9//! - The first publication or read wins (`OnceLock`). The serve publishes before the model is
10//!   built. A process that never publishes (a test, an example) runs `adaptive`, today's
11//!   routing, so kernel tests keep exercising the kernels they name.
12//! - `adaptive` (a family whose ladder is all `adaptive`) takes no new branch anywhere: every
13//!   dispatch site runs its pre-flag code for it.
14
15use std::sync::OnceLock;
16
17use metrale_config::{ActQuantFormat, ActivationQuantization, ProjFamily};
18
19static VALUE: OnceLock<ActivationQuantization> = OnceLock::new();
20
21/// 2026-09-30: Publish `--activation-quantization`. Returns the value in force; a caller that
22/// gets a different one should warn.
23pub fn set_activation_quantization_from_cli(
24    v: ActivationQuantization,
25) -> &'static ActivationQuantization {
26    let _ = VALUE.set(v);
27    VALUE.get().expect("just set")
28}
29
30/// 2026-09-30: The value in force: `adaptive` unless the serve published another.
31pub fn activation_quantization() -> &'static ActivationQuantization {
32    VALUE.get_or_init(ActivationQuantization::adaptive)
33}
34
35/// 2026-09-30: The format `family` runs at `rows` rows (`Adaptive`: the site's own routing).
36pub fn act_route(family: ProjFamily, rows: usize) -> ActQuantFormat {
37    activation_quantization().route(family, u32::try_from(rows).unwrap_or(u32::MAX))
38}
39
40/// 2026-09-30: `Some(format)` when `family` runs one fixed activation format at `rows` rows,
41/// `None` when that row count is `adaptive` (the site's own routing, unchanged).
42pub fn fixed_act(family: ProjFamily, rows: usize) -> Option<ActQuantFormat> {
43    match act_route(family, rows) {
44        ActQuantFormat::Adaptive => None,
45        f => Some(f),
46    }
47}
48
49/// 2026-09-30: Whether `family` runs a fixed format at some row count.
50pub fn family_fixed(family: ProjFamily) -> bool {
51    activation_quantization()
52        .ladder(family)
53        .rungs()
54        .iter()
55        .any(|r| r.format != ActQuantFormat::Adaptive)
56}
57
58/// 2026-09-30: Whether any family runs a fixed format at any row count. The serve then publishes
59/// the canonical row tiers (`row_tiers.rs`), whose single-order kernels are the fixed `bf16`
60/// path of the W8A16 projections and of an NVFP4 LM head.
61pub fn any_fixed() -> bool {
62    ProjFamily::ALL.into_iter().any(family_fixed)
63}