metrale_model_weights/weights/
kv_scale_keys.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The one resolver for a checkpoint's per-layer FP8 KV-cache scales.
4//!
5//! Exporters spell the pair three ways, relative to the attention prefix `p`
6//! (`model.layers.N.self_attn` and the like):
7//!
8//! | spelling       | K key              | V key              | written by                         |
9//! |----------------|--------------------|--------------------|------------------------------------|
10//! | `KProjOutput`  | `p.k_proj.k_scale` | `p.v_proj.v_scale` | ModelOpt, older llm-compressor     |
11//! | `AttnModule`   | `p.attn.k_scale`   | `p.attn.v_scale`   | vLLM's own `Attention` module name |
12//! | `Bare`         | `p.k_scale`        | `p.v_scale`        | llm-compressor `kv_cache_scheme`   |
13//!
14//! Every spelling carries the same quantity: the dequant multiplier
15//! `amax / 448`, which `reshape_and_cache_fp8` divides by on the write and the
16//! FP8 attention kernels multiply by on the read.
17//!
18//! Owner: model-weights (checkpoint key naming).
19//! Invariants:
20//! - [`resolve_kv_scale_keys`] is the only place a scale key is spelled. The
21//!   per-layer loader (`weight_map::load_kv_scales`) and the serve-time census
22//!   ([`kv_scale_census`]) both go through it, so the count serve logs is the
23//!   number of layers that load a checkpoint scale.
24//! - A layer with two spellings, or with K but not V (or V but not K), is an
25//!   error, never a silent 1.0.
26
27use std::collections::BTreeMap;
28
29use anyhow::{Result, bail};
30
31/// 2026-09-28: How a checkpoint names one layer's K/V scale pair.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
33pub enum KvScaleSpelling {
34    KProjOutput,
35    AttnModule,
36    Bare,
37}
38
39impl KvScaleSpelling {
40    /// 2026-09-28: Longest suffix first, so a name parses to the most specific
41    /// spelling: `p.k_proj.k_scale` is `KProjOutput` of `p`, not `Bare` of
42    /// `p.k_proj`.
43    pub const PRECEDENCE: [Self; 3] = [Self::KProjOutput, Self::AttnModule, Self::Bare];
44
45    /// 2026-09-28: `(K suffix, V suffix)`, each appended to `"{attn_prefix}."`.
46    pub fn suffixes(self) -> (&'static str, &'static str) {
47        match self {
48            Self::KProjOutput => ("k_proj.k_scale", "v_proj.v_scale"),
49            Self::AttnModule => ("attn.k_scale", "attn.v_scale"),
50            Self::Bare => ("k_scale", "v_scale"),
51        }
52    }
53
54    fn keys(self, attn_prefix: &str) -> (String, String) {
55        let (k, v) = self.suffixes();
56        (format!("{attn_prefix}.{k}"), format!("{attn_prefix}.{v}"))
57    }
58
59    /// 2026-09-28: The attention prefix `name` is a K or V scale of, in this
60    /// spelling.
61    fn prefix_of(self, name: &str) -> Option<&str> {
62        let (k, v) = self.suffixes();
63        [k, v].into_iter().find_map(|suffix| {
64            name.strip_suffix(suffix)
65                .and_then(|rest| rest.strip_suffix('.'))
66                .filter(|p| !p.is_empty())
67        })
68    }
69}
70
71/// 2026-09-28: The two store keys one attention layer's scales are read from.
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub struct KvScaleKeys {
74    pub spelling: KvScaleSpelling,
75    pub k: String,
76    pub v: String,
77}
78
79/// 2026-09-28: The scale keys of the attention layer at `attn_prefix`.
80///
81/// `Ok(None)` when no spelling has either key. An error when the layer has a
82/// K or V key in more than one spelling (which one the exporter meant is not
83/// ours to guess), or has K without V or V without K in its one spelling.
84pub fn resolve_kv_scale_keys(
85    has: impl Fn(&str) -> bool,
86    attn_prefix: &str,
87) -> Result<Option<KvScaleKeys>> {
88    let mut found: Option<KvScaleKeys> = None;
89    for spelling in KvScaleSpelling::PRECEDENCE {
90        let (k, v) = spelling.keys(attn_prefix);
91        let (has_k, has_v) = (has(&k), has(&v));
92        if !has_k && !has_v {
93            continue;
94        }
95        if has_k != has_v {
96            let (present, missing) = if has_k { (&k, &v) } else { (&v, &k) };
97            bail!(
98                "FP8 KV scales for `{attn_prefix}` are incomplete: the checkpoint has \
99                 `{present}` but not `{missing}`"
100            );
101        }
102        if let Some(first) = &found {
103            bail!(
104                "FP8 KV scales for `{attn_prefix}` are ambiguous: the checkpoint has both \
105                 `{}`/`{}` and `{k}`/`{v}`",
106                first.k,
107                first.v,
108            );
109        }
110        found = Some(KvScaleKeys { spelling, k, v });
111    }
112    Ok(found)
113}
114
115/// 2026-09-28: Every attention layer whose scales [`resolve_kv_scale_keys`]
116/// finds, keyed by attention prefix.
117#[derive(Debug, Default, Clone, PartialEq, Eq)]
118pub struct KvScaleCensus {
119    pub layers: BTreeMap<String, KvScaleKeys>,
120}
121
122impl KvScaleCensus {
123    /// 2026-09-28: Layers that load their scales from the checkpoint.
124    pub fn len(&self) -> usize {
125        self.layers.len()
126    }
127
128    pub fn is_empty(&self) -> bool {
129        self.layers.is_empty()
130    }
131
132    /// 2026-09-28: The distinct spellings in use, as their K suffixes, for the
133    /// serve log.
134    pub fn spellings(&self) -> Vec<&'static str> {
135        let mut s: Vec<KvScaleSpelling> = self.layers.values().map(|k| k.spelling).collect();
136        s.sort();
137        s.dedup();
138        s.into_iter().map(|s| s.suffixes().0).collect()
139    }
140}
141
142/// 2026-09-28: Parse every name that ends in a known scale suffix to its
143/// attention prefix (most specific spelling first), then resolve each prefix
144/// with [`resolve_kv_scale_keys`]. Any layer that resolver rejects fails the
145/// census.
146pub fn kv_scale_census<'a>(
147    names: impl IntoIterator<Item = &'a str>,
148    has: impl Fn(&str) -> bool,
149) -> Result<KvScaleCensus> {
150    let mut prefixes: Vec<&str> = names
151        .into_iter()
152        .filter_map(|n| {
153            KvScaleSpelling::PRECEDENCE
154                .into_iter()
155                .find_map(|s| s.prefix_of(n))
156        })
157        .collect();
158    prefixes.sort_unstable();
159    prefixes.dedup();
160    let mut layers = BTreeMap::new();
161    for p in prefixes {
162        if let Some(keys) = resolve_kv_scale_keys(&has, p)? {
163            layers.insert(p.to_string(), keys);
164        }
165    }
166    Ok(KvScaleCensus { layers })
167}
168
169#[cfg(test)]
170#[path = "kv_scale_keys_tests.rs"]
171mod tests;
metrale_model_weights/weights/kv_scale_keys.rs

metrale_model_weights/weights/
kv_scale_keys.rs
