diff --git a/crates/ptwm-core/src/codecs/per_group_codebook.rs b/crates/ptwm-core/src/codecs/per_group_codebook.rs index 8630a66..704bc87 100644 --- a/crates/ptwm-core/src/codecs/per_group_codebook.rs +++ b/crates/ptwm-core/src/codecs/per_group_codebook.rs @@ -405,10 +405,28 @@ impl PlaneCodec for PerGroupCodebook { CodecId::PerGroupCodebook } + /// Accept expanded FP4 value planes whose length divides into whole + /// groups. + /// + /// The codec models one nibble per byte, and `GROUP_SIZE` of 32 is one + /// MXFP4 block: the unit that shares a scale, and therefore the unit + /// worth giving its own codebook. `MxFp4Deinterleave` produces exactly + /// that, 32 single-nibble slots per block. + /// + /// Two exclusions are deliberate. A *packed* plane, two values per + /// byte, is declined rather than misread: reading one would take low + /// nibbles only and return them with the high half zeroed, silently + /// losing half the data. Such a plane must be deinterleaved first. And + /// a length that is not a whole number of groups is declined here + /// rather than failing inside `encode`, because dispatch drops the + /// whole candidate chain when a codec errors and packed FP4 has only + /// one chain to drop. Declining keeps an `encode` error meaning what it + /// should: an invariant was violated, not merely an unsuitable input. fn accepts(&self, descriptor: &PlaneDescriptor) -> bool { - descriptor.is_nibble_packed + !descriptor.is_nibble_packed && matches!(descriptor.role, Role::Value { .. }) && descriptor.element_width == ElementWidth::Nibble + && (descriptor.length_bytes as usize).is_multiple_of(GROUP_SIZE) } fn priority_for(&self, descriptor: &PlaneDescriptor) -> i8 { @@ -793,27 +811,27 @@ mod capability_tests { } #[test] - fn pgc_accepts_nibble_packed_value_nibble() { + fn pgc_accepts_expanded_value_nibble() { let c = PerGroupCodebook; let d = descriptor( Role::Value { format: ValueFormat::Fp4E2m1, }, ElementWidth::Nibble, - true, + false, ); assert!(c.accepts(&d)); } #[test] - fn pgc_rejects_not_nibble_packed() { + fn pgc_rejects_packed_plane() { let c = PerGroupCodebook; let d = descriptor( Role::Value { format: ValueFormat::Fp4E2m1, }, ElementWidth::Nibble, - false, + true, ); assert!(!c.accepts(&d)); } @@ -826,7 +844,7 @@ mod capability_tests { format: ScaleFormat::E4M3, }, ElementWidth::Nibble, - true, + false, ); assert!(!c.accepts(&d)); } @@ -839,7 +857,7 @@ mod capability_tests { format: ValueFormat::Fp4E2m1, }, ElementWidth::Byte, - true, + false, ); assert!(!c.accepts(&d)); } @@ -852,7 +870,7 @@ mod capability_tests { format: ValueFormat::Fp4E2m1, }, ElementWidth::Nibble, - true, + false, ); assert_eq!(c.priority_for(&d), 10); } @@ -865,8 +883,131 @@ mod capability_tests { format: ScaleFormat::E4M3, }, ElementWidth::Nibble, - true, + false, ); assert_eq!(c.priority_for(&d), i8::MIN); } } + +/// End-to-end coverage for the producer this codec is actually built for. +/// +/// The suite previously exercised the codec only on hand-built planes. It +/// never drove the op that supplies it in practice, which is how a +/// descriptor mismatch between the two survived: `MxFp4Deinterleave` +/// labelled its output nibble-packed while emitting one nibble per byte, +/// and the codec's own contract expects the expanded form. +#[cfg(test)] +mod deinterleave_integration_tests { + use super::*; + use crate::transforms::mxfp4_deinterleave::MxFp4Deinterleave; + use crate::transforms::op::{Op, Plane}; + use crate::types::descriptor::{Layout, PlaneDescriptor}; + use crate::types::role::Role; + + const BLOCK_BYTES: usize = 17; + + /// OCP MXFP4 wire bytes: 16 packed-nibble bytes then one E8M0 scale. + fn mxfp4_blocks(n_blocks: usize, seed: u64) -> Vec { + let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + let mut out = Vec::with_capacity(n_blocks * BLOCK_BYTES); + for b in 0..n_blocks { + for _ in 0..16 { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + // Skewed within a block, varying across blocks, so a + // per-block codebook has structure to find. + let lo = (s % 4 + (b as u64 % 3) * 4) as u8 & 0x0F; + let hi = ((s >> 8) % 4 + (b as u64 % 3) * 4) as u8 & 0x0F; + out.push(lo | (hi << 4)); + } + out.push(0x80 | (b as u8 & 0x0F)); + } + out + } + + fn raw_descriptor(len: usize) -> PlaneDescriptor { + PlaneDescriptor { + role: Role::Raw, + element_width: ElementWidth::Byte, + length_bytes: len as u64, + layout: Layout::Flat, + derives_from_tensor: None, + residual_of: None, + is_nibble_packed: false, + vendor_bytes: vec![], + } + } + + #[test] + fn codec_accepts_and_round_trips_the_deinterleaved_value_plane() { + let n_blocks = 8; + let raw = mxfp4_blocks(n_blocks, 11); + let op = MxFp4Deinterleave::new(32).unwrap(); + + let descs = op + .propagate_descriptors(&[raw_descriptor(raw.len())]) + .unwrap(); + let value_desc = descs[0].clone(); + + // One group per MXFP4 block is the whole premise of the codec. + assert_eq!(value_desc.length_bytes as usize, n_blocks * GROUP_SIZE); + + let c = PerGroupCodebook; + assert!( + c.accepts(&value_desc), + "the codec must accept the plane its own producer emits: {value_desc:?}" + ); + + let planes = op + .forward(&[Plane::new(raw.clone(), raw_descriptor(raw.len())).unwrap()]) + .unwrap(); + let value_plane = planes[0].bytes.to_vec(); + assert_eq!(value_plane.len(), n_blocks * GROUP_SIZE); + assert!( + value_plane.iter().all(|&b| b <= 0x0F), + "the value plane is one nibble per byte, so no byte exceeds 0x0F" + ); + + let enc = c.encode(&value_plane, None, &PlaneLayout::Flat).unwrap(); + let dec = c + .decode( + enc.state_format_version, + &enc.state_bytes, + &enc.payload, + &PlaneLayout::Flat, + value_plane.len(), + ) + .unwrap(); + assert_eq!( + dec, value_plane, + "codec must be lossless on its own producer" + ); + } + + #[test] + fn codec_declines_a_genuinely_packed_plane_instead_of_halving_it() { + // Two values per byte. Reading this as one-nibble-per-byte would + // drop every high nibble, which is exactly the silent loss the + // `accepts` guard exists to prevent. + let mut d = raw_descriptor(64); + d.role = Role::Value { + format: crate::types::role::ValueFormat::Fp4E2m1, + }; + d.element_width = ElementWidth::Nibble; + d.is_nibble_packed = true; + assert!(!PerGroupCodebook.accepts(&d)); + } + + #[test] + fn codec_declines_a_length_that_is_not_whole_groups() { + // Declined rather than failed in `encode`: packed FP4 has one + // candidate chain, and a codec error drops the chain entirely. + let mut d = raw_descriptor(GROUP_SIZE + 1); + d.role = Role::Value { + format: crate::types::role::ValueFormat::Fp4E2m1, + }; + d.element_width = ElementWidth::Nibble; + assert!(!PerGroupCodebook.accepts(&d)); + } +} diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 4af1ee0..32e43a6 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -146,12 +146,17 @@ fn legacy_plane_layout(layout: Layout) -> PlaneLayout { } } -/// Build a default source `PlaneDescriptor` keyed on dtype code. The -/// `element_width` follows the canonical mapping used elsewhere in the -/// crate (see `transforms::source::bytes_per_element` and the byte-/word- -/// aware transforms). `length_bytes` is the raw byte count; -/// `Layout::Rows{row_len}` applies when `shape` is supplied, using the -/// row-major last dimension. +/// Build a default source `PlaneDescriptor` keyed on dtype code. +/// +/// `element_width` and `is_nibble_packed` come from `transforms::source`, +/// which owns the dtype-code mapping. This function used to inline its own +/// copy of that match. The copy drifted: it kept packed FP4 in the byte +/// catch-all, so every plane reaching the trial encode was byte-width and +/// not nibble-packed, and `PerGroupCodebook` could accept none of them. +/// Call the mapping, do not restate it. +/// +/// `length_bytes` is the raw byte count; `Layout::Rows{row_len}` applies +/// when `shape` is supplied, using the row-major last dimension. /// /// Public so the PyO3 binding (`ptwm-py`) can build the same source /// descriptor without duplicating the dtype-code → element-width mapping. @@ -160,16 +165,7 @@ pub fn source_descriptor_for( raw_byte_count: u64, shape: Option<&[u64]>, ) -> PlaneDescriptor { - let element_width = match dtype_code { - // FP16 / BF16 / int16 / uint16 - 0x0002 | 0x000F | 0x0007 | 0x0008 => ElementWidth::Word2, - // FP32 / int32 / uint32 - 0x0003 | 0x0009 | 0x000A => ElementWidth::Word4, - // FP64 / int64 / uint64 - 0x0004 | 0x000B | 0x000C => ElementWidth::Word8, - // int8 / uint8 / FP8 variants — and the catch-all - _ => ElementWidth::Byte, - }; + let element_width = crate::transforms::source::element_width_for(dtype_code); let layout = match shape { Some(s) if !s.is_empty() => match u32::try_from(*s.last().unwrap()) { Ok(row_len) => Layout::rows(row_len).unwrap_or(Layout::Flat), @@ -191,7 +187,7 @@ pub fn source_descriptor_for( layout, derives_from_tensor: None, residual_of: None, - is_nibble_packed: false, + is_nibble_packed: crate::transforms::source::is_nibble_packed_dtype(dtype_code), vendor_bytes: vec![], } } @@ -1687,3 +1683,48 @@ mod tests { } } } + +#[cfg(test)] +mod source_descriptor_agreement_tests { + use super::*; + use crate::transforms::op::Op; + use crate::transforms::source::Source; + + /// The compress side builds the source descriptor from the raw byte + /// count; the decompress side rebuilds it from the chain's Source node + /// (shape + dtype code). The two must agree on `length_bytes` for every + /// dtype, or a plane decodes at the wrong length. + fn assert_paths_agree(dtype_code: u16, shape: &[u64], raw_len: u64) { + let from_compressor = source_descriptor_for(dtype_code, raw_len, Some(shape)); + let src = Source { + shape: shape.iter().map(|&d| d as u32).collect(), + dtype_code, + }; + let from_chain = &src.propagate_descriptors(&[]).unwrap()[0]; + assert_eq!( + from_compressor.length_bytes, from_chain.length_bytes, + "dtype 0x{dtype_code:04X}: compress path says {} bytes, decompress path says {}", + from_compressor.length_bytes, from_chain.length_bytes + ); + assert_eq!(from_compressor.element_width, from_chain.element_width); + assert_eq!( + from_compressor.is_nibble_packed, + from_chain.is_nibble_packed + ); + } + + #[test] + fn packed_fp4_source_descriptor_agrees_across_both_paths() { + // A packed-FP4 tensor's shape counts packed bytes (one byte per + // element, two fp4 values), matching `Dtype::element_size` and the + // shape a torch `float4_e2m1fn_x2` tensor reports. + assert_paths_agree(0x001F, &[128, 128], 128 * 128); + } + + #[test] + fn byte_and_word_source_descriptors_agree_across_both_paths() { + assert_paths_agree(0x0006, &[10, 10], 100); + assert_paths_agree(0x0002, &[10, 10], 200); + assert_paths_agree(0x0003, &[10, 10], 400); + } +} diff --git a/crates/ptwm-core/src/dispatch.rs b/crates/ptwm-core/src/dispatch.rs index 6a51510..81f8f0e 100644 --- a/crates/ptwm-core/src/dispatch.rs +++ b/crates/ptwm-core/src/dispatch.rs @@ -125,21 +125,21 @@ mod tests { } #[test] - fn dispatch_nibble_value_packed_returns_pgc_first() { + fn dispatch_expanded_nibble_value_returns_pgc_first() { let d = descriptor( Role::Value { format: ValueFormat::Fp4E2m1, }, ElementWidth::Nibble, Layout::Flat, - true, // is_nibble_packed + false, // expanded: one nibble per byte, which is what PGC models ); let result = dispatch(&d); assert!(!result.is_empty()); assert_eq!( result[0], CodecId::PerGroupCodebook, - "PGC must be first for nibble-packed Value + Nibble; got {:?}", + "PGC must be first for expanded Value + Nibble; got {:?}", result ); assert!(result.contains(&CodecId::Huffman)); diff --git a/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs b/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs index 612270c..0867a0c 100644 --- a/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs +++ b/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs @@ -95,7 +95,15 @@ impl Op for MxFp4Deinterleave { layout: value_layout, derives_from_tensor: inp.derives_from_tensor, residual_of: None, - is_nibble_packed: true, + // One nibble per byte: `forward` pushes each input byte's + // low and high halves as separate slots, so nothing shares + // a byte here. The flag means "two values share a byte", + // which describes this op's *input*; `ElementWidth::Nibble` + // already says these are 4-bit values. Claiming both left + // every consumer to choose between reading the plane as + // packed, which drops half the data, and reading it as + // expanded, which contradicted the flag. + is_nibble_packed: false, vendor_bytes: vec![], }, // Plane 1: SCALE — one E8M0 byte per block. @@ -300,7 +308,8 @@ mod tests { } ); assert_eq!(outs[0].element_width, ElementWidth::Nibble); - assert!(outs[0].is_nibble_packed); + // One nibble per byte: nothing shares a byte in this op's output. + assert!(!outs[0].is_nibble_packed); assert_eq!(outs[0].length_bytes, 32); // BLOCK_VALUES } diff --git a/crates/ptwm-core/src/transforms/op.rs b/crates/ptwm-core/src/transforms/op.rs index 8139848..d49dc9b 100644 --- a/crates/ptwm-core/src/transforms/op.rs +++ b/crates/ptwm-core/src/transforms/op.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use crate::error::PtwmCoreError; -use crate::types::descriptor::{ElementWidth, PlaneDescriptor}; +use crate::types::descriptor::PlaneDescriptor; #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[repr(u16)] @@ -114,9 +114,8 @@ pub struct Plane { impl Plane { /// Construct a `Plane` after validating that `bytes.len()` matches - /// `descriptor.length_bytes`. Expects one nibble-per-byte slot when - /// `descriptor.element_width == Nibble` and `is_nibble_packed` is - /// false; nibble-packed planes use half-byte storage. + /// `descriptor.length_bytes`, which is a storage byte count whether or + /// not the plane is nibble-packed. pub fn new(bytes: Vec, descriptor: PlaneDescriptor) -> Result { let expected = expected_byte_len(&descriptor); if bytes.len() as u64 != expected { @@ -145,13 +144,16 @@ impl Plane { } /// Expected byte length for a plane's storage given its descriptor. +/// +/// `length_bytes` is a storage count for every producer, so this is the +/// identity. It used to halve the value for nibble-packed planes, reading +/// the field as a count of values instead, which made it the third of three +/// disagreeing conventions: `Source` reports storage bytes, +/// `MxFp4Deinterleave` reports one slot per value, and this halved whatever +/// it was given. A plane built from `source_descriptor_for` was rejected as +/// twice its expected size. One convention, applied everywhere. fn expected_byte_len(descriptor: &PlaneDescriptor) -> u64 { - if descriptor.is_nibble_packed && descriptor.element_width == ElementWidth::Nibble { - // Two nibbles per byte; round up so an odd element count still fits. - descriptor.length_bytes.div_ceil(2) - } else { - descriptor.length_bytes - } + descriptor.length_bytes } /// Pure-function trait every op implements. @@ -180,6 +182,7 @@ pub trait Op { #[cfg(test)] mod tests { use super::*; + use crate::types::descriptor::ElementWidth; #[test] fn op_id_roundtrip() { @@ -265,7 +268,11 @@ mod tests { is_nibble_packed: true, vendor_bytes: vec![], }; - assert!(Plane::new(vec![0; 3], desc.clone()).is_ok()); - assert!(Plane::new(vec![0; 2], desc).is_err()); + // `length_bytes` is storage, packed or not, so a packed plane is + // sized like any other. This used to expect 3 bytes for the same + // descriptor, reading the field as a count of values; a plane built + // by `source_descriptor_for` was then rejected as twice its size. + assert!(Plane::new(vec![0; 5], desc.clone()).is_ok()); + assert!(Plane::new(vec![0; 3], desc).is_err()); } } diff --git a/crates/ptwm-core/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index 3863843..2abde55 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -5,12 +5,17 @@ use crate::transforms::op::{Op, OpId, Plane}; use crate::types::descriptor::{ElementWidth, Layout, PlaneDescriptor}; use crate::types::role::Role; +/// Chain-internal code for `Float4E2M1FNx2`: two fp4 values share a byte, +/// and one element of this dtype *is* that byte. +const DTYPE_FP4_E2M1FN_X2: u16 = 0x001F; + /// Element width inferred from a chain-internal dtype code. The chain's /// Source params use a different code space than `Dtype::from_code` (the /// public registry); see `_CHAIN_DTYPE` in -/// `python/weights/preprocessing/_chains.py` and the parallel match in -/// `compressor::source_descriptor_for`. -fn element_width_for(dtype_code: u16) -> ElementWidth { +/// `python/ptwm/preprocessing/_chains.py`. This is the one owner of the +/// mapping: `compressor::source_descriptor_for` calls it rather than +/// restating it. +pub fn element_width_for(dtype_code: u16) -> ElementWidth { match dtype_code { // FP16 / BF16 / int16 / uint16 0x0002 | 0x000F | 0x0007 | 0x0008 => ElementWidth::Word2, @@ -18,14 +23,36 @@ fn element_width_for(dtype_code: u16) -> ElementWidth { 0x0003 | 0x0009 | 0x000A => ElementWidth::Word4, // FP64 / int64 / uint64 0x0004 | 0x000B | 0x000C => ElementWidth::Word8, - // int8 / uint8 / FP8 / FP4 variants — and the catch-all + // Packed FP4: sub-byte, and the only sub-byte dtype so far. + DTYPE_FP4_E2M1FN_X2 => ElementWidth::Nibble, + // int8 / uint8 / FP8 — and the catch-all _ => ElementWidth::Byte, } } -/// Bytes per element for a chain-internal dtype code. -fn bytes_per_element(dtype_code: u16) -> u64 { - element_width_for(dtype_code).bits_per_element() as u64 / 8 +/// Whether a chain-internal dtype stores two values per byte. +/// +/// Kept separate from element width on purpose: nibble width says how wide +/// an element is, this says that elements share a byte. Codecs that read +/// packed nibbles (`PerGroupCodebook`) require both, and transforms that +/// cannot handle sharing (`BurrowsWheeler`, `IndexPack`) reject on this one. +pub fn is_nibble_packed_dtype(dtype_code: u16) -> bool { + dtype_code == DTYPE_FP4_E2M1FN_X2 +} + +/// Storage bytes for `n_elements` of a chain-internal dtype. +/// +/// A packed dtype's shape counts storage units, not values: one +/// `Float4E2M1FNx2` element is one byte carrying two fp4 values, which is +/// what `Dtype::element_size` reports and what a packed tensor's own shape +/// counts. Element width describes the *value* width for such a dtype, so +/// deriving the byte count from it would halve the length and leave the +/// decompressor rebuilding every packed plane at half its size. +fn storage_bytes(dtype_code: u16, n_elements: u64) -> u64 { + if is_nibble_packed_dtype(dtype_code) { + return n_elements; + } + n_elements * (element_width_for(dtype_code).bits_per_element() as u64 / 8) } /// Mandatory entry node of every PPG chain. Takes no inputs (the runtime @@ -66,7 +93,7 @@ impl Source { // product() on an empty iterator yields 1 (multiplicative // identity), so an empty shape represents a scalar (1 element). let n_elements: u64 = self.shape.iter().map(|&d| d as u64).product(); - let length_bytes = n_elements * bytes_per_element(self.dtype_code); + let length_bytes = storage_bytes(self.dtype_code, n_elements); let layout = if let Some(&last) = self.shape.last() { Layout::Rows { row_len: last } } else { @@ -79,7 +106,7 @@ impl Source { layout, derives_from_tensor: None, residual_of: None, - is_nibble_packed: false, + is_nibble_packed: is_nibble_packed_dtype(self.dtype_code), vendor_bytes: vec![], } } @@ -290,4 +317,49 @@ mod tests { // empty shape → scalar (1 element) → 4 bytes for fp32 assert_eq!(descs[0].length_bytes, 4); } + + #[test] + fn packed_fp4_source_is_nibble_width_and_nibble_packed() { + // Chain-internal 0x001F = Float4E2M1FNx2: two fp4 values per byte. + // PerGroupCodebook accepts a plane only when it is nibble-packed at + // nibble width, so a Byte-width descriptor here leaves that codec + // unreachable no matter how a tensor is classified or routed. + let src = Source { + shape: vec![64, 128], + dtype_code: 0x001F, + }; + let d = &src.propagate_descriptors(&[]).unwrap()[0]; + assert_eq!(d.element_width, ElementWidth::Nibble); + assert!(d.is_nibble_packed); + // A packed tensor's shape counts packed bytes, so 8192 elements + // occupy 8192 bytes and carry 16384 fp4 values. Halving here would + // disagree with the raw byte count the compressor is handed. + assert_eq!(d.length_bytes, 64 * 128); + } + + #[test] + fn packed_fp4_length_holds_for_an_odd_element_count() { + let src = Source { + shape: vec![7], + dtype_code: 0x001F, + }; + let d = &src.propagate_descriptors(&[]).unwrap()[0]; + assert_eq!(d.length_bytes, 7); + assert!(d.is_nibble_packed); + } + + #[test] + fn byte_and_word_widths_keep_their_previous_lengths() { + // Guards the bits-based length arithmetic against regressing the + // non-nibble dtypes it also now covers. + for (dtype_code, per_elem) in [(0x0006u16, 1u64), (0x0002, 2), (0x0003, 4)] { + let src = Source { + shape: vec![10, 10], + dtype_code, + }; + let d = &src.propagate_descriptors(&[]).unwrap()[0]; + assert_eq!(d.length_bytes, 100 * per_elem, "dtype 0x{dtype_code:04X}"); + assert!(!d.is_nibble_packed, "dtype 0x{dtype_code:04X}"); + } + } } diff --git a/python/ptwm/preprocessing/_chains.py b/python/ptwm/preprocessing/_chains.py index 4476dd8..0bdb34d 100644 --- a/python/ptwm/preprocessing/_chains.py +++ b/python/ptwm/preprocessing/_chains.py @@ -92,7 +92,7 @@ class _Op(IntEnum): 23: 0x000B, # Long (alias for Int64) 29: 0x0010, # Float8E4M3FN 30: 0x0011, # Float8E5M2 - 31: 0x001F, # Float4E2M1FNx2 (packed fp4, 1 byte/elem in Rust fallthrough) + 31: 0x001F, # Float4E2M1FNx2 (packed fp4: nibble width, 1 byte/elem) }