spark_runtime/kv_cache/
catalog.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The exhaustive [`KvCacheDtype`] catalogue: every variant and its canonical
4//! CLI spelling, in one place the compiler defends.
5//!
6//! `--kv-cache-dtype` is a plain `String` in clap, so anything that wants to
7//! LIST the valid values — the CLI validator's error text, the dashboard's
8//! option picker — has historically had to copy them by hand, and a copied
9//! list drifts: it offers a value the server refuses, or hides one that
10//! works. Both consumers now read [`KvCacheDtype::ALL`], and the `name` match
11//! below is what keeps `ALL` honest — it has no wildcard arm, so adding a
12//! variant fails THIS build until the new name is written here, next to the
13//! two-line comment telling you to extend `ALL` with it.
14
15use super::KvCacheDtype;
16
17impl KvCacheDtype {
18    /// Every variant, in the order the enum declares them.
19    ///
20    /// Extend this together with [`KvCacheDtype::name`] below — the
21    /// non-exhaustive-match error a new variant raises there points here.
22    pub const ALL: [KvCacheDtype; 16] = [
23        KvCacheDtype::Bf16,
24        KvCacheDtype::Fp8,
25        KvCacheDtype::Nvfp4,
26        KvCacheDtype::Turbo4,
27        KvCacheDtype::Turbo3,
28        KvCacheDtype::Turbo2,
29        KvCacheDtype::Turbo8,
30        KvCacheDtype::Turbo4KTurbo3V,
31        KvCacheDtype::Turbo4KTurbo8V,
32        KvCacheDtype::Turbo3KTurbo8V,
33        KvCacheDtype::Bf16KTurbo4V,
34        KvCacheDtype::Bf16KTurbo3V,
35        KvCacheDtype::Fp8KTurbo4V,
36        KvCacheDtype::Fp8KTurbo3V,
37        KvCacheDtype::Bf16KTurbo2V,
38        KvCacheDtype::Fp8KTurbo2V,
39    ];
40
41    /// The canonical `--kv-cache-dtype` spelling. `Display` delegates here,
42    /// so the string a picker offers is byte-for-byte the string the flag
43    /// parser reads back — the round trip the tests below pin.
44    pub const fn name(self) -> &'static str {
45        match self {
46            KvCacheDtype::Bf16 => "bf16",
47            KvCacheDtype::Fp8 => "fp8",
48            KvCacheDtype::Nvfp4 => "nvfp4",
49            KvCacheDtype::Turbo4 => "turbo4",
50            KvCacheDtype::Turbo3 => "turbo3",
51            KvCacheDtype::Turbo2 => "turbo2",
52            KvCacheDtype::Turbo8 => "turbo8",
53            KvCacheDtype::Turbo4KTurbo3V => "turbo4k_turbo3v",
54            KvCacheDtype::Turbo4KTurbo8V => "turbo4k_turbo8v",
55            KvCacheDtype::Turbo3KTurbo8V => "turbo3k_turbo8v",
56            KvCacheDtype::Bf16KTurbo4V => "bf16k_turbo4v",
57            KvCacheDtype::Bf16KTurbo3V => "bf16k_turbo3v",
58            KvCacheDtype::Fp8KTurbo4V => "fp8k_turbo4v",
59            KvCacheDtype::Fp8KTurbo3V => "fp8k_turbo3v",
60            KvCacheDtype::Bf16KTurbo2V => "bf16k_turbo2v",
61            KvCacheDtype::Fp8KTurbo2V => "fp8k_turbo2v",
62        }
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::KvCacheDtype;
69
70    #[test]
71    fn every_listed_name_parses_back_to_the_variant_that_produced_it() {
72        // The catalogue's whole promise: a value copied out of `ALL` is a
73        // value `FromStr` accepts, and it means the same dtype. A name that
74        // fails either half is an option a picker would offer and the server
75        // would refuse.
76        for dtype in KvCacheDtype::ALL {
77            let parsed: KvCacheDtype = dtype
78                .name()
79                .parse()
80                .unwrap_or_else(|e| panic!("{} does not parse: {e:#}", dtype.name()));
81            assert_eq!(
82                parsed,
83                dtype,
84                "{} parses to a different dtype",
85                dtype.name()
86            );
87            assert_eq!(
88                dtype.to_string(),
89                dtype.name(),
90                "Display and the catalogue disagree"
91            );
92        }
93    }
94
95    #[test]
96    fn the_catalogue_has_no_duplicates() {
97        // A duplicated entry is a picker row that looks like a choice and is
98        // not one; with sixteen hand-ordered entries it is an easy slip.
99        for (i, a) in KvCacheDtype::ALL.iter().enumerate() {
100            for b in &KvCacheDtype::ALL[i + 1..] {
101                assert_ne!(a, b, "{} is listed twice", a.name());
102            }
103        }
104    }
105}
spark_runtime/kv_cache/catalog.rs

spark_runtime/kv_cache/
catalog.rs
