proxima_primitives/
codec.rs

1//! The one source-neutral codec identity: which packed byte layout a
2//! quantized/packed buffer holds. Payload-less by design -- a pipe or fn
3//! pointer cannot serialize itself to a wire byte, so a sidecar wire format
4//! (`proxima-model-interop`'s expert sidecar today) needs a compile-time-
5//! exhaustive, serializable tag instead. [`Codec::tag`]/[`Codec::from_tag`]
6//! are that wire mapping.
7//!
8//! This enum is the single identity meant to replace the parallel copies
9//! that grew up independently: `proxima-model-interop`'s `PackedOwnedKind`
10//! (migrated onto this type), `omega`'s `PackedCodec`, and
11//! `proxima-tensor`'s `QuantizedBlock` tag (both still to migrate).
12
13use crate::block_layout::BlockLayout;
14use crate::float_layout::FloatLayout;
15
16/// Which packed byte layout a codec-tagged buffer holds.
17#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18pub enum Codec {
19    Q4K,
20    Q5K,
21    Q6K,
22    Q8_0,
23    Q3K,
24    Q4_0,
25    Float16,
26    BFloat16,
27    Q2K,
28    Q5_1,
29    Q5_0,
30    Q4_1,
31    Q8_1,
32    Q8K,
33    Iq1S,
34    Iq1M,
35    Iq2Xxs,
36    Iq2Xs,
37    Iq2S,
38    Iq3Xxs,
39    Iq3S,
40    Iq4Nl,
41    Iq4Xs,
42    Tq10,
43    Tq20,
44    Mxfp4,
45    Nvfp4,
46    Q1_0,
47    Q2_0,
48    BFloat8,
49    Bf8E5M2,
50    Bf4E2M1,
51    Fp6E3M2,
52}
53
54impl Codec {
55    /// The wire byte a sidecar tags this codec with. Stable across
56    /// releases -- an on-disk sidecar's bytes decode by this mapping.
57    #[must_use]
58    pub fn tag(self) -> u8 {
59        match self {
60            Codec::Q2K => 0,
61            Codec::Q3K => 1,
62            Codec::Q4K => 2,
63            Codec::Q5K => 3,
64            Codec::Q6K => 4,
65            Codec::Q8_0 => 5,
66            Codec::Q4_0 => 6,
67            Codec::Float16 => 7,
68            Codec::BFloat16 => 8,
69            Codec::Q5_1 => 9,
70            Codec::Q5_0 => 10,
71            Codec::Q4_1 => 11,
72            Codec::Q8_1 => 12,
73            Codec::Q8K => 13,
74            Codec::Iq1S => 14,
75            Codec::Iq1M => 15,
76            Codec::Iq2Xxs => 16,
77            Codec::Iq2Xs => 17,
78            Codec::Iq2S => 18,
79            Codec::Iq3Xxs => 19,
80            Codec::Iq3S => 20,
81            Codec::Iq4Nl => 21,
82            Codec::Iq4Xs => 22,
83            Codec::Tq10 => 23,
84            Codec::Tq20 => 24,
85            Codec::Mxfp4 => 25,
86            Codec::Nvfp4 => 26,
87            Codec::Q1_0 => 27,
88            Codec::Q2_0 => 28,
89            Codec::BFloat8 => 29,
90            Codec::Bf8E5M2 => 30,
91            Codec::Bf4E2M1 => 31,
92            Codec::Fp6E3M2 => 32,
93        }
94    }
95
96    /// The inverse of [`Self::tag`]. `None` for a byte no codec claims.
97    #[must_use]
98    pub fn from_tag(tag: u8) -> Option<Self> {
99        match tag {
100            0 => Some(Codec::Q2K),
101            1 => Some(Codec::Q3K),
102            2 => Some(Codec::Q4K),
103            3 => Some(Codec::Q5K),
104            4 => Some(Codec::Q6K),
105            5 => Some(Codec::Q8_0),
106            6 => Some(Codec::Q4_0),
107            7 => Some(Codec::Float16),
108            8 => Some(Codec::BFloat16),
109            9 => Some(Codec::Q5_1),
110            10 => Some(Codec::Q5_0),
111            11 => Some(Codec::Q4_1),
112            12 => Some(Codec::Q8_1),
113            13 => Some(Codec::Q8K),
114            14 => Some(Codec::Iq1S),
115            15 => Some(Codec::Iq1M),
116            16 => Some(Codec::Iq2Xxs),
117            17 => Some(Codec::Iq2Xs),
118            18 => Some(Codec::Iq2S),
119            19 => Some(Codec::Iq3Xxs),
120            20 => Some(Codec::Iq3S),
121            21 => Some(Codec::Iq4Nl),
122            22 => Some(Codec::Iq4Xs),
123            23 => Some(Codec::Tq10),
124            24 => Some(Codec::Tq20),
125            25 => Some(Codec::Mxfp4),
126            26 => Some(Codec::Nvfp4),
127            27 => Some(Codec::Q1_0),
128            28 => Some(Codec::Q2_0),
129            29 => Some(Codec::BFloat8),
130            30 => Some(Codec::Bf8E5M2),
131            31 => Some(Codec::Bf4E2M1),
132            32 => Some(Codec::Fp6E3M2),
133            _ => None,
134        }
135    }
136
137    /// Weights per [`Codec::Fp6E3M2`] block, one f16 absmax scale each.
138    pub const FP6_E3M2_BLOCK_ELEMENTS: usize = 32;
139
140    /// Bytes per [`Codec::Fp6E3M2`] block: the 6-bit [`FloatLayout::E3M2`] fields
141    /// packed, then the f16 scale.
142    pub const FP6_E3M2_BLOCK_BYTES: usize =
143        Self::FP6_E3M2_BLOCK_ELEMENTS * FloatLayout::E3M2.width() as usize / 8 + 2;
144
145    /// The block packing of the linear block codecs, `None` for every other codec.
146    #[must_use]
147    pub const fn block_layout(self) -> Option<BlockLayout> {
148        match self {
149            Codec::Q4_0 => Some(BlockLayout::Q4_0),
150            Codec::Q4_1 => Some(BlockLayout::Q4_1),
151            Codec::Q5_0 => Some(BlockLayout::Q5_0),
152            Codec::Q5_1 => Some(BlockLayout::Q5_1),
153            Codec::Q8_0 => Some(BlockLayout::Q8_0),
154            Codec::Q8_1 => Some(BlockLayout::Q8_1),
155            Codec::Q1_0 => Some(BlockLayout::Q1_0),
156            Codec::Q2_0 => Some(BlockLayout::Q2_0),
157            Codec::Q2K => Some(BlockLayout::Q2_K),
158            Codec::Q3K => Some(BlockLayout::Q3_K),
159            Codec::Q4K => Some(BlockLayout::Q4_K),
160            Codec::Q5K => Some(BlockLayout::Q5_K),
161            Codec::Q6K => Some(BlockLayout::Q6_K),
162            Codec::Q8K => Some(BlockLayout::Q8_K),
163            _ => None,
164        }
165    }
166
167    /// The element layout of the float codecs, `None` for every codec that is not a plain float
168    /// layout (block-quantized integer codecs, codebooks, and the formats with no decoder).
169    /// `Fp6E3M2` is the layout of its six-bit fields; the f16 scale per 32 elements is its block.
170    #[must_use]
171    pub const fn layout(self) -> Option<FloatLayout> {
172        match self {
173            Codec::Float16 => Some(FloatLayout::F16),
174            Codec::BFloat16 => Some(FloatLayout::BF16),
175            Codec::BFloat8 => Some(FloatLayout::E5M2),
176            Codec::Bf8E5M2 => Some(FloatLayout::E5M2_SATURATING),
177            Codec::Bf4E2M1 => Some(FloatLayout::E2M1),
178            Codec::Fp6E3M2 => Some(FloatLayout::E3M2),
179            Codec::Q4K
180            | Codec::Q5K
181            | Codec::Q6K
182            | Codec::Q8_0
183            | Codec::Q3K
184            | Codec::Q4_0
185            | Codec::Q2K
186            | Codec::Q5_1
187            | Codec::Q5_0
188            | Codec::Q4_1
189            | Codec::Q8_1
190            | Codec::Q8K
191            | Codec::Iq1S
192            | Codec::Iq1M
193            | Codec::Iq2Xxs
194            | Codec::Iq2Xs
195            | Codec::Iq2S
196            | Codec::Iq3Xxs
197            | Codec::Iq3S
198            | Codec::Iq4Nl
199            | Codec::Iq4Xs
200            | Codec::Tq10
201            | Codec::Tq20
202            | Codec::Mxfp4
203            | Codec::Nvfp4
204            | Codec::Q1_0
205            | Codec::Q2_0 => None,
206        }
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use super::Codec;
213
214    #[test]
215    fn card_02_bf8_codec_round_trips_all_stable_tags() {
216        let variants = [
217            (Codec::Q2K, 0),
218            (Codec::Q3K, 1),
219            (Codec::Q4K, 2),
220            (Codec::Q5K, 3),
221            (Codec::Q6K, 4),
222            (Codec::Q8_0, 5),
223            (Codec::Q4_0, 6),
224            (Codec::Float16, 7),
225            (Codec::BFloat16, 8),
226            (Codec::Q5_1, 9),
227            (Codec::Q5_0, 10),
228            (Codec::Q4_1, 11),
229            (Codec::Q8_1, 12),
230            (Codec::Q8K, 13),
231            (Codec::Iq1S, 14),
232            (Codec::Iq1M, 15),
233            (Codec::Iq2Xxs, 16),
234            (Codec::Iq2Xs, 17),
235            (Codec::Iq2S, 18),
236            (Codec::Iq3Xxs, 19),
237            (Codec::Iq3S, 20),
238            (Codec::Iq4Nl, 21),
239            (Codec::Iq4Xs, 22),
240            (Codec::Tq10, 23),
241            (Codec::Tq20, 24),
242            (Codec::Mxfp4, 25),
243            (Codec::Nvfp4, 26),
244            (Codec::Q1_0, 27),
245            (Codec::Q2_0, 28),
246            (Codec::BFloat8, 29),
247            (Codec::Bf8E5M2, 30),
248            (Codec::Bf4E2M1, 31),
249            (Codec::Fp6E3M2, 32),
250        ];
251        assert_eq!(variants.len(), 33);
252        for (codec, stable_tag) in variants {
253            assert_eq!(codec.tag(), stable_tag);
254            assert_eq!(Codec::from_tag(stable_tag), Some(codec));
255        }
256    }
257
258    #[test]
259    fn fp6_e3m2_codec_tag_32_round_trip() {
260        assert_eq!(Codec::Fp6E3M2.tag(), 32);
261        assert_eq!(Codec::from_tag(32), Some(Codec::Fp6E3M2));
262    }
263
264    #[test]
265    fn fp6_e3m2_codec_tag_33_rejected() {
266        assert_eq!(Codec::from_tag(33), None);
267    }
268
269    #[test]
270    fn fp6_e3m2_codec_block_geometry_26_bytes_32_elements() {
271        assert_eq!(Codec::FP6_E3M2_BLOCK_BYTES, 26);
272        assert_eq!(Codec::FP6_E3M2_BLOCK_ELEMENTS, 32);
273    }
274
275    #[test]
276    fn card_02_bf8_codec_rejects_unassigned_tag() {
277        assert_eq!(Codec::from_tag(33), None);
278        assert_eq!(Codec::from_tag(255), None);
279    }
280}
proxima_primitives/codec.rs

proxima_primitives/
codec.rs
