metrale_model_layers/layers/
weight_quantization.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The `--weight-quantization` tier in force for this process
4//! (`metrale_config::WeightQuantTier`): `declared` or `nvfp4`, and under `nvfp4` the
5//! `--w4a4-downcast` lever. Loaders pair it with a model's plan in a
6//! `metrale_config::WeightQuantPolicy`; decode dispatch reads it beside each weight's stamp
7//! (`QuantizedWeight::act`). It has no environment fallback.
8//!
9//! Owner: model-layers (quantization dispatch).
10//! Invariants:
11//! - The first publication or read wins (`OnceLock`). The serve publishes before the model is
12//!   built; anything that reads first (a test, an example) fixes the default, `declared`.
13
14use std::sync::OnceLock;
15
16use metrale_config::WeightQuantTier;
17
18static TIER: OnceLock<WeightQuantTier> = OnceLock::new();
19
20/// 2026-09-28: Publish `--weight-quantization` and its lever. Returns the tier in force; a
21/// caller that gets a different one should warn.
22pub fn set_weight_quantization_from_cli(tier: WeightQuantTier) -> WeightQuantTier {
23    let _ = TIER.set(tier);
24    *TIER.get().expect("just set")
25}
26
27/// 2026-09-28: The tier in force: `declared` unless the serve published another.
28pub fn weight_quantization() -> WeightQuantTier {
29    *TIER.get_or_init(WeightQuantTier::default)
30}
31
32/// 2026-09-28: Which decode families the policy may rely on. On: the dense per-channel W8A8
33/// family (`ops::w8a8_decode`). 2026-09-29: On as well, validated on Qwen3.6-35B-A3B-FP8 by the
34/// full BFCL draw and agentic-webserver (`declared` against `nvfp4`, one binary, dgx3): the MoE
35/// expert W8A8 family (`moe/fp8_grouped_tc_w8a8.rs`). The block-scaled attention/GDN W8A8
36/// adoption (`qwen35/load_layers/w8a8_adopt.rs`) is off again (2026-09-29, see below). Also on: the batched FP8 lm_head (`model/lm_head_fp8_rows.rs`, one weight pass per 128 rows
37/// at every row count), so `declared` takes a checkpoint's declared FP8 head. Each W8A8
38/// family still checks at run time that its kernels resolved and the shapes fit.
39pub fn kernel_caps() -> metrale_config::weight_quantization::KernelCaps {
40    metrale_config::weight_quantization::KernelCaps {
41        w8a8_decode: true,
42        w8a8_moe_decode: true,
43        // 2026-09-29: Off again. With the expert and the block-scaled attention/GDN W8A8 both on,
44        // Qwen3.6-35B-A3B-FP8 failed the ssm-state-poisoning-gate reference round (dgx2, PR #51):
45        // a greedy tie at turn 3 (" a" and " computed" both at logprob -0.979) went to the reply
46        // that drops Section 4's field list. With either family alone the gate passes (2/2
47        // each). The owner kept the experts; the attention/GDN projections decode W8A16 until
48        // this path is re-validated.
49        w8a8_block_scaled_decode: false,
50        fp8_lm_head_batched: true,
51    }
52}