hurray_core/quantization/
mxfp.rs

1//! MXFP (OCP Microscaling) quantization descriptor (scheme tag `0x05`).
2//!
3//! A block quantization format standardised by the Open Compute Project
4//! Microscaling specification (OCP MX v1.0). Each contiguous block of elements
5//! along a chosen axis shares a single `float8_e8m0` exponent-only scale stored
6//! in a dedicated buffer. Unlike per-block affine, MXFP requires exact
7//! divisibility — partial trailing blocks are **not** permitted.
8//!
9//! See `docs/spec/quantization/mxfp.md` for the normative definition.
10
11use crate::{ElementType, Error, Result};
12
13// ── Wire layout constants ─────────────────────────────────────────────────────
14
15/// Scheme tag byte for MXFP quantization.
16pub(crate) const SCHEME_TAG: u8 = 0x05;
17
18/// Total descriptor length in bytes (header 4 + axis 4 + block_size 4 + scale_buf 4).
19pub(crate) const ENCODED_LEN: usize = 16;
20
21/// Version this implementation supports.
22pub(crate) const SUPPORTED_VERSION: u8 = 0x01;
23
24/// Minimum block size for MXFP (inclusive).
25///
26/// Values below 16 have no hardware Tensor Core support and are invalid under
27/// any OCP MX revision.
28pub const MXFP_MIN_BLOCK_SIZE: u32 = 16;
29
30/// Maximum block size for MXFP (inclusive).
31pub const MXFP_MAX_BLOCK_SIZE: u32 = 2048;
32
33/// OCP MX v1.0 canonical block size.
34pub const MXFP_CANONICAL_BLOCK_SIZE: u32 = 32;
35
36/// No flags are defined for this scheme; all 16 bits must be zero.
37const RESERVED_FLAGS_MASK: u16 = 0xFFFF;
38
39/// Wire sentinel that MUST NOT appear as `scale_buffer_index`.
40const INVALID_BUF_SENTINEL: u32 = 0xFFFF_FFFF;
41
42// Wire field offsets (relative to start of descriptor, including header).
43const OFFSET_AXIS: usize = 4;
44const OFFSET_BLOCK_SIZE: usize = 8;
45const OFFSET_SCALE_BUF: usize = 12;
46
47// ── Mxfp ──────────────────────────────────────────────────────────────────────
48
49/// Quantization parameters for MXFP (OCP Microscaling) block quantization.
50///
51/// Each block of `block_size` consecutive elements along `axis` shares a single
52/// `float8_e8m0` exponent-only scale. The scale values are stored as raw bytes
53/// in the buffer identified by `scale_buffer_index`.
54///
55/// The dequantization formula for element `q` with block index `b` and scale
56/// byte `e = scale[b]` is:
57///
58/// ```text
59/// s      = 2^(e - 127)          (float8_e8m0 shared exponent)
60/// x_real = s * value_of(q)      (float_value_of(q) or int_value_of(q))
61/// ```
62///
63/// Unlike per-block affine, MXFP requires exact divisibility: `shape[axis]`
64/// MUST be a positive multiple of `block_size`. No partial trailing blocks are
65/// permitted.
66///
67/// # Wire format
68///
69/// Total descriptor length: **16 bytes** (including the 4-byte header).
70///
71/// | Offset | Field | Type |
72/// |--------|-------|------|
73/// | 0 | `scheme_tag` | `uint8` (must be `0x05`) |
74/// | 1 | `scheme_version` | `uint8` (must be `0x01`) |
75/// | 2 | `flags` | `uint16` LE (must be `0x0000`) |
76/// | 4 | `axis` | `uint32` LE |
77/// | 8 | `block_size` | `uint32` LE, power-of-two in `[16, 2048]` |
78/// | 12 | `scale_buffer_index` | `uint32` LE |
79///
80/// # Design notes
81///
82/// `PartialEq`, `Eq`, and `Hash` are all derived because this struct contains
83/// no floating-point fields — all fields are integers.
84///
85/// `Copy` because the struct is ≤ 24 bytes with no `Drop` glue.
86///
87/// # Examples
88///
89/// ```
90/// use hurray_core::Mxfp;
91///
92/// let q = Mxfp::new(0, 32, 1).unwrap();
93/// assert_eq!(q.axis(), 0);
94/// assert_eq!(q.block_size(), 32);
95/// assert_eq!(q.scale_buffer_index(), 1);
96/// ```
97// WHY Eq + Hash: no float fields.
98// WHY Copy: ≤24 bytes, no Drop glue (design decision #7).
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
100#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
101pub struct Mxfp {
102    axis: u32,
103    block_size: u32,
104    scale_buffer_index: u32,
105}
106
107impl Mxfp {
108    /// Creates a new [`Mxfp`] descriptor.
109    ///
110    /// # Errors
111    ///
112    /// - [`Error::InvalidBlockSize`] — `block_size` is not a power of two, or
113    ///   lies outside `[`[`MXFP_MIN_BLOCK_SIZE`]`, `[`MXFP_MAX_BLOCK_SIZE`]`]` (`[16, 2048]`).
114    ///
115    /// # Examples
116    ///
117    /// ```
118    /// use hurray_core::{Mxfp, Error};
119    ///
120    /// assert!(Mxfp::new(0, 32, 1).is_ok());
121    /// assert!(Mxfp::new(0, 16, 1).is_ok());
122    /// assert!(Mxfp::new(0, 2048, 1).is_ok());
123    ///
124    /// // block_size < 16 is rejected.
125    /// assert!(Mxfp::new(0, 8, 1).is_err());
126    /// // block_size > 2048 is rejected.
127    /// assert!(Mxfp::new(0, 4096, 1).is_err());
128    /// // Non-power-of-two is rejected.
129    /// assert!(Mxfp::new(0, 48, 1).is_err());
130    /// ```
131    pub fn new(axis: u32, block_size: u32, scale_buffer_index: u32) -> Result<Self> {
132        validate_block_size(block_size)?;
133        Ok(Self {
134            axis,
135            block_size,
136            scale_buffer_index,
137        })
138    }
139
140    /// Returns the quantization axis index.
141    ///
142    /// # Examples
143    ///
144    /// ```
145    /// use hurray_core::Mxfp;
146    ///
147    /// let q = Mxfp::new(2, 32, 1).unwrap();
148    /// assert_eq!(q.axis(), 2);
149    /// ```
150    #[inline]
151    pub fn axis(&self) -> u32 {
152        self.axis
153    }
154
155    /// Returns the number of logical elements per block along `axis`.
156    ///
157    /// Always a power of two in `[`[`MXFP_MIN_BLOCK_SIZE`]`, `[`MXFP_MAX_BLOCK_SIZE`]`]`.
158    ///
159    /// # Examples
160    ///
161    /// ```
162    /// use hurray_core::Mxfp;
163    ///
164    /// let q = Mxfp::new(0, 64, 1).unwrap();
165    /// assert_eq!(q.block_size(), 64);
166    /// ```
167    #[inline]
168    pub fn block_size(&self) -> u32 {
169        self.block_size
170    }
171
172    /// Returns the buffer table index of the per-block `float8_e8m0` scale array.
173    ///
174    /// # Examples
175    ///
176    /// ```
177    /// use hurray_core::Mxfp;
178    ///
179    /// let q = Mxfp::new(0, 32, 3).unwrap();
180    /// assert_eq!(q.scale_buffer_index(), 3);
181    /// ```
182    #[inline]
183    pub fn scale_buffer_index(&self) -> u32 {
184        self.scale_buffer_index
185    }
186
187    /// Computes the number of blocks along `axis` for a given `shape_axis` size.
188    ///
189    /// MXFP requires exact divisibility: `shape_axis` MUST be a positive multiple
190    /// of `block_size`. This is stricter than per-block affine, which uses
191    /// `div_ceil` and allows partial trailing blocks.
192    ///
193    /// # Errors
194    ///
195    /// - [`Error::QuantizationShapeMismatch`] — `shape_axis` is not evenly
196    ///   divisible by `block_size`.
197    ///
198    /// # Examples
199    ///
200    /// ```
201    /// use hurray_core::Mxfp;
202    ///
203    /// let q = Mxfp::new(0, 32, 1).unwrap();
204    /// assert_eq!(q.num_blocks_per_axis(64).unwrap(), 2);
205    /// assert_eq!(q.num_blocks_per_axis(32).unwrap(), 1);
206    ///
207    /// // 65 is not a multiple of 32 — error.
208    /// assert!(q.num_blocks_per_axis(65).is_err());
209    /// // 0 is not a positive multiple — error.
210    /// assert!(q.num_blocks_per_axis(0).is_err());
211    /// ```
212    pub fn num_blocks_per_axis(&self, shape_axis: u64) -> Result<u64> {
213        // MXFP requires shape_axis > 0 AND exact divisibility by block_size.
214        // Stricter than per-block-affine which uses div_ceil — exact divisibility
215        // required by MXFP spec (mxfp.md § Validity Constraints).
216        if shape_axis == 0 || !shape_axis.is_multiple_of(self.block_size as u64) {
217            return Err(Error::QuantizationShapeMismatch {
218                axis: self.axis,
219                shape_axis,
220                block_size: self.block_size,
221                reason: "MXFP requires shape[axis] to be a positive multiple of block_size; partial trailing blocks are not permitted",
222            });
223        }
224        Ok(shape_axis / self.block_size as u64)
225    }
226
227    /// Validates that `shape_axis` satisfies the MXFP divisibility constraint.
228    ///
229    /// `shape_axis` MUST be greater than zero AND evenly divisible by `block_size`.
230    ///
231    /// # Errors
232    ///
233    /// - [`Error::QuantizationShapeMismatch`] — `shape_axis == 0` or
234    ///   `shape_axis % block_size != 0`.
235    ///
236    /// # Examples
237    ///
238    /// ```
239    /// use hurray_core::Mxfp;
240    ///
241    /// let q = Mxfp::new(0, 32, 1).unwrap();
242    /// assert!(q.validate_against_shape_axis(64).is_ok());
243    /// assert!(q.validate_against_shape_axis(32).is_ok());
244    ///
245    /// // Zero is rejected (not a positive multiple).
246    /// assert!(q.validate_against_shape_axis(0).is_err());
247    /// // 65 is not divisible by 32.
248    /// assert!(q.validate_against_shape_axis(65).is_err());
249    /// ```
250    pub fn validate_against_shape_axis(&self, shape_axis: u64) -> Result<()> {
251        self.num_blocks_per_axis(shape_axis).map(|_| ())
252    }
253
254    /// Validates the content of a scale buffer for MXFP conformance.
255    ///
256    /// Per `docs/spec/quantization/mxfp.md § Referenced Buffer`: the bit patterns
257    /// `0x00` and `0xFF` are reserved (`NaN` per OCP MX v1.0 § 5.6) and MUST NOT
258    /// appear in any scale byte.
259    ///
260    /// # Errors
261    ///
262    /// - [`Error::InvalidQuantization`] — any byte in `bytes` is `0x00` or `0xFF`.
263    ///
264    /// # Examples
265    ///
266    /// ```
267    /// use hurray_core::Mxfp;
268    ///
269    /// let q = Mxfp::new(0, 32, 1).unwrap();
270    ///
271    /// // Valid scale bytes: non-zero, non-0xFF.
272    /// assert!(q.validate_scale_bytes(&[0x01, 0x7F, 0x80, 0xFE]).is_ok());
273    ///
274    /// // 0x00 is forbidden.
275    /// assert!(q.validate_scale_bytes(&[0x01, 0x00, 0x7F]).is_err());
276    ///
277    /// // 0xFF is forbidden.
278    /// assert!(q.validate_scale_bytes(&[0x7F, 0xFF]).is_err());
279    /// ```
280    pub fn validate_scale_bytes(&self, bytes: &[u8]) -> Result<()> {
281        for (i, &b) in bytes.iter().enumerate() {
282            if b == 0x00 || b == 0xFF {
283                return Err(Error::InvalidQuantization(format!(
284                    "MXFP scale buffer byte at index {i} is 0x{b:02X}, which is a reserved \
285                     float8_e8m0 NaN pattern (OCP MX v1.0 § 5.6); values 0x00 and 0xFF are \
286                     forbidden in the scale buffer"
287                )));
288            }
289        }
290        Ok(())
291    }
292
293    /// Returns the set of storage [`ElementType`]s that are valid for this scheme.
294    ///
295    /// Per `docs/spec/quantization/mxfp.md § Valid Storage Types`.
296    ///
297    /// WHY `&'static [ElementType]`: no allocation per call; `slice::contains`
298    /// over ≤10 items beats any hash structure (design decision #5).
299    ///
300    /// # Examples
301    ///
302    /// ```
303    /// use hurray_core::{ElementType, Mxfp};
304    ///
305    /// assert!(Mxfp::valid_storage_types().contains(&ElementType::Float8E4M3));
306    /// assert!(Mxfp::valid_storage_types().contains(&ElementType::Int8));
307    /// assert!(Mxfp::valid_storage_types().contains(&ElementType::Int4));
308    /// assert!(!Mxfp::valid_storage_types().contains(&ElementType::Float32));
309    /// ```
310    pub fn valid_storage_types() -> &'static [ElementType] {
311        &[
312            ElementType::Float8E4M3, // MXFP8
313            ElementType::Float8E5M2, // MXFP8
314            ElementType::Float4E2M1, // MXFP4
315            ElementType::Float6E2M3, // MXFP6
316            ElementType::Float6E3M2, // MXFP6
317            ElementType::Int8,       // MXINT8
318            ElementType::Int4,       // MXINT4 / MXFP4 integer surrogate
319        ]
320    }
321
322    // ── Crate-internal encode/decode ──────────────────────────────────────────
323
324    /// Decodes the scheme-specific payload from `bytes`.
325    ///
326    /// `bytes` is the full descriptor slice (including the 4-byte header that the
327    /// caller has already validated). `flags` are the header flags (must be zero
328    /// for this scheme).
329    pub(crate) fn decode_payload(flags: u16, bytes: &[u8]) -> Result<Self> {
330        if bytes.len() < ENCODED_LEN {
331            return Err(Error::QuantizationDescriptorTooShort {
332                found: bytes.len(),
333                needed: ENCODED_LEN,
334            });
335        }
336        // No flags defined for MXFP.
337        if flags & RESERVED_FLAGS_MASK != 0 {
338            return Err(Error::ReservedQuantizationFlagsBits {
339                flags,
340                mask: RESERVED_FLAGS_MASK,
341            });
342        }
343
344        let axis = u32::from_le_bytes(
345            bytes[OFFSET_AXIS..OFFSET_AXIS + 4]
346                .try_into()
347                .map_err(|_| Error::InvalidQuantization("axis slice error".into()))?,
348        );
349        let block_size = u32::from_le_bytes(
350            bytes[OFFSET_BLOCK_SIZE..OFFSET_BLOCK_SIZE + 4]
351                .try_into()
352                .map_err(|_| Error::InvalidQuantization("block_size slice error".into()))?,
353        );
354        let scale_buf = u32::from_le_bytes(
355            bytes[OFFSET_SCALE_BUF..OFFSET_SCALE_BUF + 4]
356                .try_into()
357                .map_err(|_| Error::InvalidQuantization("scale_buffer_index slice error".into()))?,
358        );
359
360        // block_size must be a power of two in [16, 2048].
361        validate_block_size(block_size)?;
362
363        // scale_buffer_index must not be the sentinel value 0xFFFFFFFF.
364        if scale_buf == INVALID_BUF_SENTINEL {
365            return Err(Error::InvalidQuantization(
366                "MXFP scale_buffer_index must not be 0xFFFFFFFF".into(),
367            ));
368        }
369
370        Ok(Self {
371            axis,
372            block_size,
373            scale_buffer_index: scale_buf,
374        })
375    }
376
377    /// Encodes the scheme-specific payload into `out`.
378    ///
379    /// `out` must be at least [`ENCODED_LEN`] bytes. The caller writes the 4-byte
380    /// header via [`super::QuantizationHeader::write`]; this method writes
381    /// bytes 4–15.
382    pub(crate) fn encode_payload(&self, out: &mut [u8]) {
383        out[OFFSET_AXIS..OFFSET_AXIS + 4].copy_from_slice(&self.axis.to_le_bytes());
384        out[OFFSET_BLOCK_SIZE..OFFSET_BLOCK_SIZE + 4]
385            .copy_from_slice(&self.block_size.to_le_bytes());
386        out[OFFSET_SCALE_BUF..OFFSET_SCALE_BUF + 4]
387            .copy_from_slice(&self.scale_buffer_index.to_le_bytes());
388    }
389}
390
391// ── Helpers ───────────────────────────────────────────────────────────────────
392
393fn validate_block_size(block_size: u32) -> Result<()> {
394    if !block_size.is_power_of_two()
395        || !(MXFP_MIN_BLOCK_SIZE..=MXFP_MAX_BLOCK_SIZE).contains(&block_size)
396    {
397        return Err(Error::InvalidBlockSize {
398            scheme_tag: SCHEME_TAG,
399            block_size,
400            min: MXFP_MIN_BLOCK_SIZE,
401            max: MXFP_MAX_BLOCK_SIZE,
402        });
403    }
404    Ok(())
405}
406
407// ── Tests ─────────────────────────────────────────────────────────────────────
408
409#[cfg(test)]
410mod tests {
411    use super::*;
412    use crate::{ElementType, Error};
413
414    // ── Mxfp::new ─────────────────────────────────────────────────────────────
415
416    #[test]
417    fn new_block_size_below_min_is_err() {
418        // MXFP_MIN_BLOCK_SIZE = 16; block_size = 8 is below.
419        assert!(matches!(
420            Mxfp::new(0, 8, 1),
421            Err(Error::InvalidBlockSize { .. })
422        ));
423    }
424
425    #[test]
426    fn new_block_size_above_max_is_err() {
427        // MXFP_MAX_BLOCK_SIZE = 2048; block_size = 4096 is above.
428        assert!(matches!(
429            Mxfp::new(0, 4096, 1),
430            Err(Error::InvalidBlockSize { .. })
431        ));
432    }
433
434    #[test]
435    fn new_block_size_not_power_of_two_is_err() {
436        assert!(matches!(
437            Mxfp::new(0, 48, 1),
438            Err(Error::InvalidBlockSize { .. })
439        ));
440    }
441
442    #[test]
443    fn new_min_block_size_is_ok() {
444        assert!(Mxfp::new(0, MXFP_MIN_BLOCK_SIZE, 1).is_ok());
445    }
446
447    #[test]
448    fn new_canonical_block_size_is_ok() {
449        assert!(Mxfp::new(0, MXFP_CANONICAL_BLOCK_SIZE, 1).is_ok());
450    }
451
452    #[test]
453    fn new_max_block_size_is_ok() {
454        assert!(Mxfp::new(0, MXFP_MAX_BLOCK_SIZE, 1).is_ok());
455    }
456
457    // ── num_blocks_per_axis ───────────────────────────────────────────────────
458
459    #[test]
460    fn num_blocks_per_axis_zero_shape_is_err() {
461        // MXFP requires a positive multiple — shape_axis=0 is always rejected.
462        let q = Mxfp::new(0, 32, 1).unwrap();
463        assert!(matches!(
464            q.num_blocks_per_axis(0),
465            Err(Error::QuantizationShapeMismatch { .. })
466        ));
467    }
468
469    #[test]
470    fn num_blocks_per_axis_exact_divisibility_is_ok() {
471        let q = Mxfp::new(0, 32, 1).unwrap();
472        assert_eq!(q.num_blocks_per_axis(64).unwrap(), 2);
473    }
474
475    #[test]
476    fn num_blocks_per_axis_not_divisible_is_err() {
477        // 65 is not a multiple of 32 — MXFP forbids partial trailing blocks.
478        let q = Mxfp::new(0, 32, 1).unwrap();
479        assert!(matches!(
480            q.num_blocks_per_axis(65),
481            Err(Error::QuantizationShapeMismatch { .. })
482        ));
483    }
484
485    // ── validate_scale_bytes ──────────────────────────────────────────────────
486
487    #[test]
488    fn validate_scale_bytes_empty_slice_is_ok() {
489        let q = Mxfp::new(0, 32, 1).unwrap();
490        assert!(q.validate_scale_bytes(&[]).is_ok());
491    }
492
493    #[test]
494    fn validate_scale_bytes_valid_bytes_is_ok() {
495        let q = Mxfp::new(0, 32, 1).unwrap();
496        assert!(q.validate_scale_bytes(&[0x01, 0x7F, 0x80, 0xFE]).is_ok());
497    }
498
499    #[test]
500    fn validate_scale_bytes_zero_byte_is_err() {
501        let q = Mxfp::new(0, 32, 1).unwrap();
502        assert!(matches!(
503            q.validate_scale_bytes(&[0x01, 0x00, 0x7F]),
504            Err(Error::InvalidQuantization(_))
505        ));
506    }
507
508    #[test]
509    fn validate_scale_bytes_0xff_is_err() {
510        let q = Mxfp::new(0, 32, 1).unwrap();
511        assert!(matches!(
512            q.validate_scale_bytes(&[0x7F, 0xFF]),
513            Err(Error::InvalidQuantization(_))
514        ));
515    }
516
517    #[test]
518    fn validate_scale_bytes_mixed_valid_and_invalid_is_err() {
519        let q = Mxfp::new(0, 32, 1).unwrap();
520        // First byte is valid; second byte 0x00 is invalid.
521        assert!(matches!(
522            q.validate_scale_bytes(&[0x7F, 0x00]),
523            Err(Error::InvalidQuantization(_))
524        ));
525    }
526
527    // ── valid_storage_types ───────────────────────────────────────────────────
528
529    #[test]
530    fn valid_storage_types_contains_expected_float_and_int_types() {
531        let types = Mxfp::valid_storage_types();
532        assert!(types.contains(&ElementType::Float8E4M3));
533        assert!(types.contains(&ElementType::Float8E5M2));
534        assert!(types.contains(&ElementType::Float4E2M1));
535        assert!(types.contains(&ElementType::Float6E2M3));
536        assert!(types.contains(&ElementType::Float6E3M2));
537        assert!(types.contains(&ElementType::Int8));
538        assert!(types.contains(&ElementType::Int4));
539    }
540
541    #[test]
542    fn valid_storage_types_does_not_contain_float32() {
543        assert!(!Mxfp::valid_storage_types().contains(&ElementType::Float32));
544    }
545
546    // ── Round-trips ───────────────────────────────────────────────────────────
547
548    fn encode_decode(q: &Mxfp) -> Mxfp {
549        let mut buf = vec![0u8; ENCODED_LEN];
550        buf[0] = SCHEME_TAG;
551        buf[1] = SUPPORTED_VERSION;
552        buf[2] = 0;
553        buf[3] = 0;
554        q.encode_payload(&mut buf);
555        Mxfp::decode_payload(0, &buf).unwrap()
556    }
557
558    #[test]
559    fn round_trip_canonical_block_size() {
560        let original = Mxfp::new(0, MXFP_CANONICAL_BLOCK_SIZE, 3).unwrap();
561        let decoded = encode_decode(&original);
562        assert_eq!(decoded, original);
563    }
564
565    #[test]
566    fn round_trip_min_block_size() {
567        let original = Mxfp::new(1, MXFP_MIN_BLOCK_SIZE, 2).unwrap();
568        let decoded = encode_decode(&original);
569        assert_eq!(decoded, original);
570    }
571
572    #[test]
573    fn round_trip_max_block_size() {
574        let original = Mxfp::new(0, MXFP_MAX_BLOCK_SIZE, 1).unwrap();
575        let decoded = encode_decode(&original);
576        assert_eq!(decoded, original);
577    }
578}
hurray_core/quantization/mxfp.rs

hurray_core/quantization/
mxfp.rs
