hurray_core/quantization/
nf4.rs

1//! NF4 (NormalFloat4) quantization descriptor (scheme tag `0x04`).
2//!
3//! A non-linear 4-bit quantization scheme introduced by the QLoRA paper. Each
4//! storage code in `[0, 15]` decodes to one of 16 fixed real-valued levels. Each
5//! block along `axis` carries a single `absmax` scale stored in a separate buffer.
6//!
7//! See `docs/spec/quantization/nf4.md` for the normative definition.
8
9use crate::{ElementType, Error, Result};
10
11// ── Wire layout constants ─────────────────────────────────────────────────────
12
13/// Scheme tag byte for NF4 quantization.
14pub(crate) const SCHEME_TAG: u8 = 0x04;
15
16/// Total descriptor length in bytes (header 4 + axis 4 + block_size 4 + scale_buf 4).
17pub(crate) const ENCODED_LEN: usize = 16;
18
19/// Version this implementation supports.
20pub(crate) const SUPPORTED_VERSION: u8 = 0x01;
21
22/// Minimum block size for NF4 (must be a power of two ≥ 8).
23///
24/// Below 8 elements per block the 16-point NF4 information content provides no
25/// statistical benefit over a plain low-bit linear quantization.
26pub const NF4_MIN_BLOCK_SIZE: u32 = 8;
27
28/// No flags are defined for this scheme; all 16 bits must be zero.
29const RESERVED_FLAGS_MASK: u16 = 0xFFFF;
30
31/// Wire sentinel that must NOT appear in `scale_buffer_index` (symmetric sentinel).
32const INVALID_BUF_SENTINEL: u32 = 0xFFFF_FFFF;
33
34// Wire field offsets.
35const OFFSET_AXIS: usize = 4;
36const OFFSET_BLOCK_SIZE: usize = 8;
37const OFFSET_SCALE_BUF: usize = 12;
38
39// ── NF4 lookup table ──────────────────────────────────────────────────────────
40
41/// The 16 NF4 quantization levels, indexed by the unsigned 4-bit storage code.
42///
43/// These values are fixed by the specification and MUST NOT be altered. They are
44/// the `float32` decimal expansions of the exact levels from the QLoRA reference
45/// implementation (`docs/spec/quantization/nf4.md § Lookup Table`).
46///
47/// The extra precision in the source literals is intentional: these are the
48/// exact decimal representations of the IEEE 754 binary32 values; truncating
49/// them would alter the bit pattern.
50///
51/// # Examples
52///
53/// ```
54/// use hurray_core::NF4_LUT;
55///
56/// assert_eq!(NF4_LUT[0], -1.0f32);
57/// assert_eq!(NF4_LUT[7], 0.0f32);
58/// assert_eq!(NF4_LUT[15], 1.0f32);
59/// assert_eq!(NF4_LUT.len(), 16);
60/// ```
61#[allow(clippy::excessive_precision)]
62pub const NF4_LUT: [f32; 16] = [
63    -1.0,
64    -0.6961928009986877,
65    -0.5250730514526367,
66    -0.39491748809814453,
67    -0.28444138169288635,
68    -0.18477343022823334,
69    -0.09105003625154495,
70    0.0,
71    0.07958029955625534,
72    0.16093020141124725,
73    0.24611230194568634,
74    0.33791524171829224,
75    0.44070982933044434,
76    0.5626170039176941,
77    0.7229568362236023,
78    1.0,
79];
80
81// ── Nf4 ───────────────────────────────────────────────────────────────────────
82
83/// Quantization parameters for NF4 (NormalFloat4) block quantization.
84///
85/// Each block along `axis` contains `block_size` elements, all sharing a single
86/// `absmax` scale stored in `scale_buffer_index`. The dequantization formula for
87/// element `q` (a 4-bit code in `[0, 15]`) with block index `b` is:
88///
89/// ```text
90/// x_real = scale[b] * NF4_LUT[q]
91/// ```
92///
93/// The block index computation follows the same rule as Per-Block Affine
94/// (see `docs/spec/quantization/per-block-affine.md § Block Layout`).
95///
96/// # Wire format
97///
98/// Total descriptor length: **16 bytes** (including the 4-byte header).
99///
100/// | Offset | Field | Type |
101/// |--------|-------|------|
102/// | 4 | `axis` | `uint32` LE |
103/// | 8 | `block_size` | `uint32` LE |
104/// | 12 | `scale_buffer_index` | `uint32` LE |
105///
106/// No flags are defined; the `flags` field MUST be `0x0000`.
107///
108/// # Design notes
109///
110/// `PartialEq`, `Eq`, and `Hash` are all derived because this struct contains
111/// no floating-point fields.
112///
113/// `Copy` because the struct is ≤ 24 bytes with no `Drop` glue.
114///
115/// # Examples
116///
117/// ```
118/// use hurray_core::Nf4;
119///
120/// let q = Nf4::new(0, 64, 1).unwrap();
121/// assert_eq!(q.axis(), 0);
122/// assert_eq!(q.block_size(), 64);
123/// assert_eq!(q.scale_buffer_index(), 1);
124/// ```
125// WHY Eq + Hash: no float fields.
126// WHY Copy: ≤24 bytes, no Drop glue (design decision #7).
127#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
128#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
129pub struct Nf4 {
130    axis: u32,
131    block_size: u32,
132    scale_buffer_index: u32,
133}
134
135impl Nf4 {
136    /// Creates a new [`Nf4`] descriptor.
137    ///
138    /// # Errors
139    ///
140    /// - [`Error::InvalidBlockSize`] — `block_size` is not a power of two or is
141    ///   less than [`NF4_MIN_BLOCK_SIZE`] (8).
142    ///
143    /// # Examples
144    ///
145    /// ```
146    /// use hurray_core::{Nf4, Error};
147    ///
148    /// assert!(Nf4::new(0, 64, 1).is_ok());
149    /// assert!(Nf4::new(0, 128, 1).is_ok());
150    ///
151    /// // block_size < 8 is rejected.
152    /// assert!(Nf4::new(0, 4, 1).is_err());
153    /// // Non-power-of-two is rejected.
154    /// assert!(Nf4::new(0, 48, 1).is_err());
155    /// ```
156    pub fn new(axis: u32, block_size: u32, scale_buffer_index: u32) -> Result<Self> {
157        validate_block_size(block_size)?;
158        Ok(Self {
159            axis,
160            block_size,
161            scale_buffer_index,
162        })
163    }
164
165    /// Returns the quantization axis index.
166    ///
167    /// # Examples
168    ///
169    /// ```
170    /// use hurray_core::Nf4;
171    ///
172    /// let q = Nf4::new(2, 64, 1).unwrap();
173    /// assert_eq!(q.axis(), 2);
174    /// ```
175    #[inline]
176    pub fn axis(&self) -> u32 {
177        self.axis
178    }
179
180    /// Returns the number of logical elements per block along `axis`.
181    ///
182    /// # Examples
183    ///
184    /// ```
185    /// use hurray_core::Nf4;
186    ///
187    /// let q = Nf4::new(0, 128, 1).unwrap();
188    /// assert_eq!(q.block_size(), 128);
189    /// ```
190    #[inline]
191    pub fn block_size(&self) -> u32 {
192        self.block_size
193    }
194
195    /// Returns the buffer table index of the per-block `absmax` scale array.
196    ///
197    /// # Examples
198    ///
199    /// ```
200    /// use hurray_core::Nf4;
201    ///
202    /// let q = Nf4::new(0, 64, 3).unwrap();
203    /// assert_eq!(q.scale_buffer_index(), 3);
204    /// ```
205    #[inline]
206    pub fn scale_buffer_index(&self) -> u32 {
207        self.scale_buffer_index
208    }
209
210    /// Computes the number of blocks along `axis` for a given `shape_axis` size.
211    ///
212    /// Uses `ceil(shape_axis / block_size)`.
213    ///
214    /// Returns `0` when `shape_axis == 0` per the ADR-007 empty-axis carve-out.
215    ///
216    /// # Examples
217    ///
218    /// ```
219    /// use hurray_core::Nf4;
220    ///
221    /// let q = Nf4::new(0, 64, 1).unwrap();
222    /// assert_eq!(q.num_blocks_per_axis(128), 2);
223    /// assert_eq!(q.num_blocks_per_axis(65), 2); // partial trailing block
224    /// assert_eq!(q.num_blocks_per_axis(0), 0);  // ADR-007 empty-axis carve-out
225    /// ```
226    pub fn num_blocks_per_axis(&self, shape_axis: u64) -> u64 {
227        if shape_axis == 0 {
228            // ADR-007 empty-axis carve-out: zero blocks, zero-byte scale buffer.
229            return 0;
230        }
231        shape_axis.div_ceil(self.block_size as u64)
232    }
233
234    /// Validates this descriptor against the resolved `shape_axis` size.
235    ///
236    /// Rejects if `shape_axis` is the DYNAMIC sentinel (`u64::MAX`), or if
237    /// `shape_axis > 0` and `block_size > shape_axis`.
238    ///
239    /// # Errors
240    ///
241    /// - [`Error::QuantizationShapeMismatch`] — `shape_axis` is the DYNAMIC sentinel,
242    ///   or `shape_axis > 0` and `block_size > shape_axis`.
243    ///
244    /// # Examples
245    ///
246    /// ```
247    /// use hurray_core::{Nf4, DYNAMIC};
248    ///
249    /// let q = Nf4::new(0, 64, 1).unwrap();
250    /// assert!(q.validate_against_shape_axis(128).is_ok());
251    /// assert!(q.validate_against_shape_axis(64).is_ok());
252    /// assert!(q.validate_against_shape_axis(0).is_ok()); // ADR-007: waived
253    /// assert!(q.validate_against_shape_axis(32).is_err()); // block_size > shape_axis
254    /// assert!(q.validate_against_shape_axis(DYNAMIC).is_err()); // dynamic dimension
255    /// ```
256    pub fn validate_against_shape_axis(&self, shape_axis: u64) -> Result<()> {
257        // Reject the DYNAMIC sentinel: a dynamic dimension cannot be validated.
258        if shape_axis == crate::shape::DYNAMIC {
259            return Err(Error::QuantizationShapeMismatch {
260                axis: self.axis,
261                shape_axis,
262                block_size: self.block_size,
263                reason: "shape[axis] must not be the DYNAMIC sentinel (0xFFFFFFFFFFFFFFFF)",
264            });
265        }
266        if shape_axis > 0 && self.block_size as u64 > shape_axis {
267            return Err(Error::QuantizationShapeMismatch {
268                axis: self.axis,
269                shape_axis,
270                block_size: self.block_size,
271                reason: "NF4 block_size must not exceed shape[axis] when shape[axis] > 0",
272            });
273        }
274        Ok(())
275    }
276
277    /// Returns the set of storage [`ElementType`]s that are valid for this scheme.
278    ///
279    /// Per `docs/spec/quantization/nf4.md § Valid Storage Types`: only `uint4`
280    /// is valid for NF4.
281    ///
282    /// WHY `&'static [ElementType]`: no allocation per call (design decision #5).
283    ///
284    /// # Examples
285    ///
286    /// ```
287    /// use hurray_core::{ElementType, Nf4};
288    ///
289    /// assert_eq!(Nf4::valid_storage_types(), &[ElementType::Uint4]);
290    /// assert!(!Nf4::valid_storage_types().contains(&ElementType::Int4));
291    /// ```
292    pub fn valid_storage_types() -> &'static [ElementType] {
293        &[ElementType::Uint4]
294    }
295
296    // ── Crate-internal encode/decode ──────────────────────────────────────────
297
298    /// Decodes the scheme-specific payload from `bytes`.
299    ///
300    /// `bytes` is the full descriptor slice (including the 4-byte header).
301    /// `flags` are the header flags already read by the caller.
302    pub(crate) fn decode_payload(flags: u16, bytes: &[u8]) -> Result<Self> {
303        if bytes.len() < ENCODED_LEN {
304            return Err(Error::QuantizationDescriptorTooShort {
305                found: bytes.len(),
306                needed: ENCODED_LEN,
307            });
308        }
309        // No flags defined for NF4.
310        if flags & RESERVED_FLAGS_MASK != 0 {
311            return Err(Error::ReservedQuantizationFlagsBits {
312                flags,
313                mask: RESERVED_FLAGS_MASK,
314            });
315        }
316
317        let axis = u32::from_le_bytes(
318            bytes[OFFSET_AXIS..OFFSET_AXIS + 4]
319                .try_into()
320                .map_err(|_| Error::InvalidQuantization("axis slice error".into()))?,
321        );
322        let block_size = u32::from_le_bytes(
323            bytes[OFFSET_BLOCK_SIZE..OFFSET_BLOCK_SIZE + 4]
324                .try_into()
325                .map_err(|_| Error::InvalidQuantization("block_size slice error".into()))?,
326        );
327        let scale_buf = u32::from_le_bytes(
328            bytes[OFFSET_SCALE_BUF..OFFSET_SCALE_BUF + 4]
329                .try_into()
330                .map_err(|_| Error::InvalidQuantization("scale_buffer_index slice error".into()))?,
331        );
332
333        // block_size must be a power of two and >= 8.
334        validate_block_size(block_size)?;
335
336        // scale_buffer_index must not be the sentinel value 0xFFFFFFFF.
337        if scale_buf == INVALID_BUF_SENTINEL {
338            return Err(Error::InvalidQuantization(
339                "NF4 scale_buffer_index must not be 0xFFFFFFFF".into(),
340            ));
341        }
342
343        Ok(Self {
344            axis,
345            block_size,
346            scale_buffer_index: scale_buf,
347        })
348    }
349
350    /// Encodes the scheme-specific payload into `out`.
351    ///
352    /// `out` must be at least [`ENCODED_LEN`] bytes. The caller writes the 4-byte
353    /// header; this method writes bytes 4–15.
354    pub(crate) fn encode_payload(&self, out: &mut [u8]) {
355        out[OFFSET_AXIS..OFFSET_AXIS + 4].copy_from_slice(&self.axis.to_le_bytes());
356        out[OFFSET_BLOCK_SIZE..OFFSET_BLOCK_SIZE + 4]
357            .copy_from_slice(&self.block_size.to_le_bytes());
358        out[OFFSET_SCALE_BUF..OFFSET_SCALE_BUF + 4]
359            .copy_from_slice(&self.scale_buffer_index.to_le_bytes());
360    }
361}
362
363// ── Helpers ───────────────────────────────────────────────────────────────────
364
365fn validate_block_size(block_size: u32) -> Result<()> {
366    if !block_size.is_power_of_two() || block_size < NF4_MIN_BLOCK_SIZE {
367        return Err(Error::InvalidBlockSize {
368            scheme_tag: SCHEME_TAG,
369            block_size,
370            min: NF4_MIN_BLOCK_SIZE,
371            // NF4 has no upper bound in the spec; use u32::MAX as sentinel.
372            max: u32::MAX,
373        });
374    }
375    Ok(())
376}
377
378// ── Tests ─────────────────────────────────────────────────────────────────────
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383    use crate::Error;
384
385    // ── NF4_LUT ───────────────────────────────────────────────────────────────
386
387    #[test]
388    fn nf4_lut_has_16_entries() {
389        assert_eq!(NF4_LUT.len(), 16);
390    }
391
392    #[test]
393    fn nf4_lut_all_values_are_finite() {
394        for (i, &v) in NF4_LUT.iter().enumerate() {
395            assert!(v.is_finite(), "NF4_LUT[{i}] = {v} is not finite");
396        }
397    }
398
399    #[test]
400    fn nf4_lut_values_are_in_ascending_order() {
401        // Spec requires monotonically non-decreasing values.
402        for i in 1..NF4_LUT.len() {
403            assert!(
404                NF4_LUT[i] >= NF4_LUT[i - 1],
405                "NF4_LUT is not ascending at index {i}: NF4_LUT[{}]={} > NF4_LUT[{}]={}",
406                i - 1,
407                NF4_LUT[i - 1],
408                i,
409                NF4_LUT[i]
410            );
411        }
412    }
413
414    #[test]
415    // Deliberately asserting properties of the const NF4_LUT; the values are known at
416    // compile time, which is exactly what this test pins down.
417    #[allow(clippy::assertions_on_constants)]
418    fn nf4_lut_straddles_zero_at_index_7_8() {
419        // Per QLoRA paper: index 7 is 0.0 (the zero-crossing boundary),
420        // index 8 is the first strictly positive value.
421        assert!(NF4_LUT[7] <= 0.0, "NF4_LUT[7] should be <= 0.0");
422        assert!(NF4_LUT[8] >= 0.0, "NF4_LUT[8] should be >= 0.0");
423    }
424
425    #[test]
426    fn nf4_lut_first_value_approx_neg_one() {
427        // Spec: first entry is approximately −1.0 (±0.01).
428        let diff = (NF4_LUT[0] - (-1.0f32)).abs();
429        assert!(
430            diff < 0.01,
431            "NF4_LUT[0] = {} is not within 0.01 of -1.0",
432            NF4_LUT[0]
433        );
434    }
435
436    #[test]
437    fn nf4_lut_last_value_approx_pos_one() {
438        // Spec: last entry is approximately +1.0 (±0.01).
439        let diff = (NF4_LUT[15] - 1.0f32).abs();
440        assert!(
441            diff < 0.01,
442            "NF4_LUT[15] = {} is not within 0.01 of 1.0",
443            NF4_LUT[15]
444        );
445    }
446
447    // ── Nf4::new ──────────────────────────────────────────────────────────────
448
449    #[test]
450    fn new_block_size_below_min_is_err() {
451        // NF4_MIN_BLOCK_SIZE = 8; block_size = 4 is below minimum.
452        assert!(matches!(
453            Nf4::new(0, 4, 1),
454            Err(Error::InvalidBlockSize { .. })
455        ));
456    }
457
458    #[test]
459    fn new_block_size_not_power_of_two_is_err() {
460        assert!(matches!(
461            Nf4::new(0, 12, 1),
462            Err(Error::InvalidBlockSize { .. })
463        ));
464    }
465
466    #[test]
467    fn new_block_size_1_is_err() {
468        assert!(matches!(
469            Nf4::new(0, 1, 1),
470            Err(Error::InvalidBlockSize { .. })
471        ));
472    }
473
474    #[test]
475    fn new_valid_min_block_size_is_ok() {
476        assert!(Nf4::new(0, NF4_MIN_BLOCK_SIZE, 1).is_ok());
477    }
478
479    #[test]
480    fn new_valid_block_size_16_is_ok() {
481        assert!(Nf4::new(0, 16, 1).is_ok());
482    }
483
484    #[test]
485    fn new_valid_block_size_64_is_ok() {
486        assert!(Nf4::new(0, 64, 1).is_ok());
487    }
488
489    // ── num_blocks_per_axis ───────────────────────────────────────────────────
490
491    #[test]
492    fn num_blocks_per_axis_zero_shape_returns_zero() {
493        // ADR-007 empty-axis carve-out.
494        let q = Nf4::new(0, 64, 1).unwrap();
495        assert_eq!(q.num_blocks_per_axis(0), 0);
496    }
497
498    #[test]
499    fn num_blocks_per_axis_exact_divisibility() {
500        let q = Nf4::new(0, 64, 1).unwrap();
501        assert_eq!(q.num_blocks_per_axis(128), 2);
502    }
503
504    #[test]
505    fn num_blocks_per_axis_partial_trailing_block_uses_div_ceil() {
506        let q = Nf4::new(0, 64, 1).unwrap();
507        // 65 / 64 = 1.015... → ceil = 2.
508        assert_eq!(q.num_blocks_per_axis(65), 2);
509    }
510
511    // ── Round-trips ───────────────────────────────────────────────────────────
512
513    fn encode_decode(q: &Nf4) -> Nf4 {
514        let mut buf = vec![0u8; ENCODED_LEN];
515        buf[0] = SCHEME_TAG;
516        buf[1] = SUPPORTED_VERSION;
517        buf[2] = 0;
518        buf[3] = 0;
519        q.encode_payload(&mut buf);
520        Nf4::decode_payload(0, &buf).unwrap()
521    }
522
523    #[test]
524    fn round_trip_preserves_axis_block_size_scale_buf() {
525        let original = Nf4::new(2, 64, 5).unwrap();
526        let decoded = encode_decode(&original);
527        assert_eq!(decoded.axis(), original.axis());
528        assert_eq!(decoded.block_size(), original.block_size());
529        assert_eq!(decoded.scale_buffer_index(), original.scale_buffer_index());
530    }
531
532    #[test]
533    fn round_trip_min_block_size() {
534        let original = Nf4::new(0, NF4_MIN_BLOCK_SIZE, 1).unwrap();
535        let decoded = encode_decode(&original);
536        assert_eq!(decoded, original);
537    }
538
539    // ── validate_against_shape_axis ───────────────────────────────────────────
540
541    #[test]
542    fn validate_against_shape_axis_ok_when_block_size_le_shape() {
543        let q = Nf4::new(0, 64, 1).unwrap();
544        assert!(q.validate_against_shape_axis(64).is_ok());
545        assert!(q.validate_against_shape_axis(128).is_ok());
546    }
547
548    #[test]
549    fn validate_against_shape_axis_ok_when_shape_is_zero_adr007() {
550        // ADR-007 carve-out: empty axis is accepted.
551        let q = Nf4::new(0, 64, 1).unwrap();
552        assert!(q.validate_against_shape_axis(0).is_ok());
553    }
554
555    #[test]
556    fn validate_against_shape_axis_err_when_block_size_exceeds_shape() {
557        let q = Nf4::new(0, 64, 1).unwrap();
558        assert!(matches!(
559            q.validate_against_shape_axis(32),
560            Err(Error::QuantizationShapeMismatch { .. })
561        ));
562    }
563
564    #[test]
565    fn validate_against_shape_axis_err_for_dynamic_sentinel() {
566        let q = Nf4::new(0, 64, 1).unwrap();
567        assert!(matches!(
568            q.validate_against_shape_axis(u64::MAX),
569            Err(Error::QuantizationShapeMismatch { .. })
570        ));
571    }
572
573    // ── decode_payload error paths ────────────────────────────────────────────
574
575    #[test]
576    fn decode_payload_too_short_is_err() {
577        let short = vec![0u8; ENCODED_LEN - 1];
578        assert!(matches!(
579            Nf4::decode_payload(0, &short),
580            Err(Error::QuantizationDescriptorTooShort { .. })
581        ));
582    }
583
584    #[test]
585    fn decode_payload_nonzero_flags_is_err() {
586        let q = Nf4::new(0, 64, 1).unwrap();
587        let mut buf = vec![0u8; ENCODED_LEN];
588        buf[0] = SCHEME_TAG;
589        buf[1] = SUPPORTED_VERSION;
590        q.encode_payload(&mut buf);
591        let nonzero_flags: u16 = 0x0002;
592        assert!(matches!(
593            Nf4::decode_payload(nonzero_flags, &buf),
594            Err(Error::ReservedQuantizationFlagsBits { .. })
595        ));
596    }
597
598    #[test]
599    fn decode_payload_invalid_block_size_is_err() {
600        let mut buf = vec![0u8; ENCODED_LEN];
601        buf[0] = SCHEME_TAG;
602        buf[1] = SUPPORTED_VERSION;
603        buf[2] = 0;
604        buf[3] = 0;
605        buf[4..8].copy_from_slice(&0u32.to_le_bytes()); // axis
606                                                        // block_size = 3 (not a power of two, and < NF4_MIN_BLOCK_SIZE=8)
607        buf[8..12].copy_from_slice(&3u32.to_le_bytes());
608        buf[12..16].copy_from_slice(&1u32.to_le_bytes()); // scale_buf
609        assert!(matches!(
610            Nf4::decode_payload(0, &buf),
611            Err(Error::InvalidBlockSize { .. })
612        ));
613    }
614
615    #[test]
616    fn decode_payload_scale_buf_sentinel_is_err() {
617        let mut buf = vec![0u8; ENCODED_LEN];
618        buf[0] = SCHEME_TAG;
619        buf[1] = SUPPORTED_VERSION;
620        buf[2] = 0;
621        buf[3] = 0;
622        buf[4..8].copy_from_slice(&0u32.to_le_bytes()); // axis
623        buf[8..12].copy_from_slice(&64u32.to_le_bytes()); // block_size
624                                                          // scale_buffer_index = 0xFFFFFFFF (invalid sentinel)
625        buf[12..16].copy_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
626        assert!(matches!(
627            Nf4::decode_payload(0, &buf),
628            Err(Error::InvalidQuantization(_))
629        ));
630    }
631}
hurray_core/quantization/nf4.rs

hurray_core/quantization/
nf4.rs
