From 5d765cc78bafc06a9bbcce17a72654fa72cca64f Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Mon, 17 Aug 2026 23:04:11 +0200 Subject: [PATCH 1/7] fix(core): give packed FP4 planes nibble width and the packed flag PerGroupCodebook accepts a plane only when it is nibble-packed at nibble width with a Value role. Every chain starts at Source, which derived ElementWidth from a dtype table where packed FP4 fell into the Byte catch-all, and set is_nibble_packed to a hardcoded false. Nothing downstream writes either field: terminals rewrite only the role. So two of the codec's three preconditions were unsatisfiable and it could not be selected by any chain, on any tensor. The chain builder documented as "enables PerGroupCodebook" did not enable it. Maps Float4E2M1FNx2 to nibble width, derives is_nibble_packed from the dtype, and replaces the bytes-per-element helper with one that accumulates in bits. The old helper computed bits_per_element / 8, which floors a 4-bit nibble to zero bytes per element and would have produced zero-length descriptors rather than visibly wrong ones. Element width and the packed flag stay separate fields. Width is how wide an element is; the flag is that elements share a byte. Codecs reading packed nibbles need both, and transforms that cannot handle sharing (BurrowsWheeler, IndexPack) already reject on the flag alone. Byte and word dtypes keep their previous lengths, covered by a new regression test alongside the two FP4 cases. --- crates/ptwm-core/src/transforms/source.rs | 75 +++++++++++++++++++++-- 1 file changed, 69 insertions(+), 6 deletions(-) diff --git a/crates/ptwm-core/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index 3863843..e7e46a8 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -10,6 +10,9 @@ use crate::types::role::Role; /// public registry); see `_CHAIN_DTYPE` in /// `python/weights/preprocessing/_chains.py` and the parallel match in /// `compressor::source_descriptor_for`. +/// Chain-internal code for `Float4E2M1FNx2`: two fp4 values share a byte. +const DTYPE_FP4_E2M1FN_X2: u16 = 0x001F; + fn element_width_for(dtype_code: u16) -> ElementWidth { match dtype_code { // FP16 / BF16 / int16 / uint16 @@ -18,14 +21,31 @@ 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. +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. +/// +/// Accumulates in bits so sub-byte widths survive. Dividing +/// bits-per-element by 8 first floors a nibble to zero, which silently +/// yields a zero-length descriptor rather than a wrong-but-visible one. +fn storage_bytes(dtype_code: u16, n_elements: u64) -> u64 { + let bits = n_elements * element_width_for(dtype_code).bits_per_element() as u64; + bits.div_ceil(8) } /// Mandatory entry node of every PPG chain. Takes no inputs (the runtime @@ -66,7 +86,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 +99,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 +310,47 @@ 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); + // 8192 fp4 values occupy 4096 bytes, not 8192. + assert_eq!(d.length_bytes, 64 * 128 / 2); + } + + #[test] + fn odd_element_count_at_nibble_width_rounds_up_to_whole_bytes() { + let src = Source { + shape: vec![7], + dtype_code: 0x001F, + }; + let d = &src.propagate_descriptors(&[]).unwrap()[0]; + // 7 nibbles need 4 bytes; the trailing nibble still costs one. + assert_eq!(d.length_bytes, 4); + } + + #[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}"); + } + } } From b6705c9ed6146cf185c238c3fff9c7e2d66abcad Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Mon, 17 Aug 2026 23:57:22 +0200 Subject: [PATCH 2/7] fix(core): route the compressor's source descriptor through one dtype map source_descriptor_for carried its own copy of the dtype-code to element-width match, and that copy is the one the compression path actually uses. Fixing the transform alone changed nothing observable: the duplicate still placed packed FP4 in the byte catch-all and still hardcoded is_nibble_packed to false, so PerGroupCodebook went on accepting no planes. Both fields now come from transforms::source, which owns the mapping. The duplicate's own docstring claimed to follow that mapping while having drifted from it, which is how the gap survived review. --- crates/ptwm-core/src/compressor.rs | 30 ++++++++++------------- crates/ptwm-core/src/transforms/source.rs | 4 +-- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 4af1ee0..063336c 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![], } } diff --git a/crates/ptwm-core/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index e7e46a8..e405944 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -13,7 +13,7 @@ use crate::types::role::Role; /// Chain-internal code for `Float4E2M1FNx2`: two fp4 values share a byte. const DTYPE_FP4_E2M1FN_X2: u16 = 0x001F; -fn element_width_for(dtype_code: u16) -> ElementWidth { +pub fn element_width_for(dtype_code: u16) -> ElementWidth { match dtype_code { // FP16 / BF16 / int16 / uint16 0x0002 | 0x000F | 0x0007 | 0x0008 => ElementWidth::Word2, @@ -34,7 +34,7 @@ fn element_width_for(dtype_code: u16) -> ElementWidth { /// 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. -fn is_nibble_packed_dtype(dtype_code: u16) -> bool { +pub fn is_nibble_packed_dtype(dtype_code: u16) -> bool { dtype_code == DTYPE_FP4_E2M1FN_X2 } From 2a68958bba34529544edcd3992d2efb608af0774 Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Tue, 18 Aug 2026 05:12:14 +0200 Subject: [PATCH 3/7] fix(codecs): stop PerGroupCodebook dropping the high nibble of every byte The codec models a nibble alphabet and its internals work on one nibble per byte, but `accepts` requires `is_nibble_packed`, so the planes it is handed carry two values per byte. Encode read `plane[i] & 0x0F` per byte and decode emitted one byte per nibble with the high half zeroed, so half the data was never coded and never reconstructed. Round-tripping a packed plane returned the low nibbles intact and zeros elsewhere: 89% of bytes wrong, length preserved, caught by both the plane CRC and the payload hash. Encode now expands the packed plane before modelling, decode re-packs before returning, and the group arithmetic counts nibbles rather than bytes on both sides (`decoded_len` is a byte count; groups are 32 nibbles). The existing roundtrip tests build one nibble per byte, so their high halves are all zero and they pass whether or not those halves survive. That is why this shipped. Adds a fixture whose every byte carries data in both halves, plus a pack/unpack identity check. The codec has been unreachable in production since `accepts` could never be satisfied, so no container written by a release carries a plane encoded this way. --- .../src/codecs/per_group_codebook.rs | 121 ++++++++++++++++-- 1 file changed, 109 insertions(+), 12 deletions(-) diff --git a/crates/ptwm-core/src/codecs/per_group_codebook.rs b/crates/ptwm-core/src/codecs/per_group_codebook.rs index 8630a66..d23b6a2 100644 --- a/crates/ptwm-core/src/codecs/per_group_codebook.rs +++ b/crates/ptwm-core/src/codecs/per_group_codebook.rs @@ -26,6 +26,33 @@ pub fn histograms(nibbles: &[u8]) -> Vec<[u32; ALPHABET]> { out } +/// Expand a nibble-packed plane to one nibble per byte, low nibble first. +/// +/// The codec's internals model a nibble alphabet and work on this expanded +/// form. The planes the dispatcher hands it are packed two values per byte, +/// which is exactly what `is_nibble_packed` on the descriptor asserts, and +/// what `accepts` requires. Converting at the `PlaneCodec` boundary keeps +/// the two representations from being confused: reading a packed plane as +/// though it were already expanded silently drops every high nibble. +fn unpack_nibbles(packed: &[u8]) -> Vec { + let mut out = Vec::with_capacity(packed.len() * 2); + for &b in packed { + out.push(b & 0x0F); + out.push((b >> 4) & 0x0F); + } + out +} + +/// Inverse of [`unpack_nibbles`]. Requires an even nibble count, which +/// `GROUP_SIZE` alignment guarantees. +fn pack_nibbles(nibbles: &[u8]) -> Vec { + let mut out = Vec::with_capacity(nibbles.len() / 2); + for pair in nibbles.chunks_exact(2) { + out.push((pair[0] & 0x0F) | ((pair[1] & 0x0F) << 4)); + } + out +} + #[cfg(test)] mod tests { use super::*; @@ -425,17 +452,22 @@ impl PlaneCodec for PerGroupCodebook { shared_state: Option<&[u8]>, _layout: &PlaneLayout, ) -> Result { - if !plane.len().is_multiple_of(GROUP_SIZE) { + // `plane` is nibble-packed: two values per byte. Expand before + // modelling, and count groups in nibbles rather than bytes. + let nibbles = unpack_nibbles(plane); + if !nibbles.len().is_multiple_of(GROUP_SIZE) { return Err(PtwmCoreError::CodecDecode { codec: "PerGroupCodebook", msg: format!( - "encode: plane length {} not a multiple of GROUP_SIZE {GROUP_SIZE}", + "encode: plane holds {} nibbles ({} bytes), not a multiple of \ + GROUP_SIZE {GROUP_SIZE}", + nibbles.len(), plane.len(), ), }); } - let hists = histograms(plane); + let hists = histograms(&nibbles); // Either consume provided shared state or fit a per-tensor codebook. let (state, inline_state_bytes): (StateV0, Vec) = match shared_state { @@ -478,8 +510,8 @@ impl PlaneCodec for PerGroupCodebook { // Range-coded per-nibble payload let mut enc = RangeEncoder::new(); - for (i, &nib_byte) in plane.iter().enumerate() { - let sym = (nib_byte & 0x0F) as usize; + for (i, &nib) in nibbles.iter().enumerate() { + let sym = nib as usize; let g = i / GROUP_SIZE; let k = assignments[g] as usize; let counts = &state.codebook[k]; @@ -521,21 +553,27 @@ impl PlaneCodec for PerGroupCodebook { }); } let n_assignments = u32::from_le_bytes(payload[..4].try_into().unwrap()) as usize; - // `decoded_len` comes from the trusted spec/plane-record path; a + // `decoded_len` comes from the trusted spec/plane-record path and is + // a byte count; groups are counted in nibbles, two per byte. A // mismatch here is corruption, not a soft failure. - if !decoded_len.is_multiple_of(GROUP_SIZE) { + let decoded_nibbles = decoded_len * 2; + if !decoded_nibbles.is_multiple_of(GROUP_SIZE) { return Err(PtwmCoreError::CodecDecode { codec: "PerGroupCodebook", - msg: format!("decoded_len {decoded_len} not divisible by GROUP_SIZE {GROUP_SIZE}"), + msg: format!( + "decoded_len {decoded_len} bytes = {decoded_nibbles} nibbles, \ + not divisible by GROUP_SIZE {GROUP_SIZE}" + ), }); } - let expected_assignments = decoded_len / GROUP_SIZE; + let expected_assignments = decoded_nibbles / GROUP_SIZE; if n_assignments != expected_assignments { return Err(PtwmCoreError::CodecDecode { codec: "PerGroupCodebook", msg: format!( - "payload claims {n_assignments} groups (={} bytes), \ - plane expects {expected_assignments} (={decoded_len} bytes)", + "payload claims {n_assignments} groups (={} nibbles), \ + plane expects {expected_assignments} (={decoded_nibbles} nibbles, \ + {decoded_len} bytes)", n_assignments * GROUP_SIZE ), }); @@ -605,7 +643,11 @@ impl PlaneCodec for PerGroupCodebook { })?; out.push(sym as u8); } - Ok(out) + // Re-pack to the plane's on-the-wire form. Returning the expanded + // nibbles here is what made this codec lossy: it produced one byte + // per nibble with the high half zeroed, which round-tripped only + // for input whose high nibbles were already zero. + Ok(pack_nibbles(&out)) } } @@ -664,6 +706,61 @@ mod codec_tests { assert_eq!(dec, b); } + /// A genuinely packed plane: both nibbles of every byte carry data. + /// + /// The other roundtrip tests here build one nibble per byte, so their + /// high halves are all zero and they pass whether or not the codec + /// preserves them. That is why this codec shipped dropping every high + /// nibble without any test noticing. + fn build_packed_plane(n_bytes: usize, seed: u64) -> Vec { + let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + let mut out = Vec::with_capacity(n_bytes); + for _ in 0..n_bytes { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + // Skew toward a few values so the codebook has structure to find, + // but keep both halves populated. + let lo = (s % 16) as u8; + let hi = ((s >> 8) % 16) as u8; + out.push(lo | (hi << 4)); + } + out + } + + #[test] + fn codec_roundtrip_preserves_both_nibbles_of_every_byte() { + let c = PerGroupCodebook; + // GROUP_SIZE nibbles per group, two nibbles per byte. + let plane = build_packed_plane(64 * GROUP_SIZE / 2, 7); + assert!( + plane.iter().any(|b| b & 0xF0 != 0), + "fixture must exercise high nibbles" + ); + let enc = c.encode(&plane, None, &PlaneLayout::Flat).unwrap(); + let dec = c + .decode( + enc.state_format_version, + &enc.state_bytes, + &enc.payload, + &PlaneLayout::Flat, + plane.len(), + ) + .unwrap(); + assert_eq!( + dec.len(), + plane.len(), + "decoded length must match the plane" + ); + assert_eq!(dec, plane, "codec must be lossless over both nibbles"); + } + + #[test] + fn nibble_pack_unpack_is_an_identity() { + let plane = build_packed_plane(256, 3); + assert_eq!(pack_nibbles(&unpack_nibbles(&plane)), plane); + } + #[test] fn decoder_rejects_bad_state_version() { let c = PerGroupCodebook; From 815a52196d02d9111e95c8e2682d0c85621e70bb Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Tue, 18 Aug 2026 10:21:58 +0200 Subject: [PATCH 4/7] fix(core): keep packed-FP4 length in storage units, fit codebooks on real nibbles Two defects in the FP4 work, both from the same confusion in opposite directions. `storage_bytes` derived the byte count from element width, halving it for packed FP4. One `Float4E2M1FNx2` element *is* one byte carrying two fp4 values, which is what `Dtype::element_size` reports and what a packed tensor's shape counts. The compress path takes the raw byte count directly while the decompress path rebuilds it from the chain's Source node, so the two disagreed and a packed plane decoded at half length, failing its CRC. Adds a test asserting the two paths agree for FP4, byte and word dtypes, which is the invariant that was missing. The shared-state codebook fitter still called `histograms` on planes that are now genuinely packed: it fitted on low nibbles alone and cut group boundaries at 32 bytes where the encoder's group is 32 nibbles. Round trips stayed lossless because `encode` recomputes its own histograms, so this cost ratio rather than correctness. It now fits on expanded nibbles, with a test that two planes differing only in their high nibbles must produce different codebooks. Also restores the doc comment that a const insertion had reattached to the wrong item, and drops two stale cross-references. --- .../src/codecs/per_group_codebook.rs | 2 +- crates/ptwm-core/src/compressor.rs | 45 +++++++++++++++++++ .../ptwm-core/src/fit/per_group_codebook.rs | 25 ++++++++++- crates/ptwm-core/src/transforms/source.rs | 39 +++++++++------- python/ptwm/preprocessing/_chains.py | 2 +- 5 files changed, 95 insertions(+), 18 deletions(-) diff --git a/crates/ptwm-core/src/codecs/per_group_codebook.rs b/crates/ptwm-core/src/codecs/per_group_codebook.rs index d23b6a2..6ce09bf 100644 --- a/crates/ptwm-core/src/codecs/per_group_codebook.rs +++ b/crates/ptwm-core/src/codecs/per_group_codebook.rs @@ -34,7 +34,7 @@ pub fn histograms(nibbles: &[u8]) -> Vec<[u32; ALPHABET]> { /// what `accepts` requires. Converting at the `PlaneCodec` boundary keeps /// the two representations from being confused: reading a packed plane as /// though it were already expanded silently drops every high nibble. -fn unpack_nibbles(packed: &[u8]) -> Vec { +pub(crate) fn unpack_nibbles(packed: &[u8]) -> Vec { let mut out = Vec::with_capacity(packed.len() * 2); for &b in packed { out.push(b & 0x0F); diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 063336c..32e43a6 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -1683,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/fit/per_group_codebook.rs b/crates/ptwm-core/src/fit/per_group_codebook.rs index 3fd9a13..11e4534 100644 --- a/crates/ptwm-core/src/fit/per_group_codebook.rs +++ b/crates/ptwm-core/src/fit/per_group_codebook.rs @@ -36,7 +36,12 @@ pub fn fit(planes_value: &[&[u8]]) -> Result, PtwmCoreE } let mut all_hists: Vec<[u32; pgc::ALPHABET]> = Vec::new(); for plane in planes_value { - all_hists.extend(pgc::histograms(plane)); + // The planes are nibble-packed, two values per byte, and + // `histograms` models one nibble per byte. Expand first, exactly as + // the codec's own `encode` does: reading a packed plane directly + // would fit the codebook on the low nibbles alone and cut every + // group boundary at 32 bytes instead of 32 nibbles. + all_hists.extend(pgc::histograms(&pgc::unpack_nibbles(plane))); } let fit_hists = stride_sample(&all_hists, FIT_HIST_SAMPLE_CAP); let cb = pgc::fit_codebook_multi_seed( @@ -84,6 +89,24 @@ mod tests { assert_eq!(e.name, "per_group_codebook"); } + #[test] + fn fit_reads_the_high_nibble_of_every_byte() { + // Two planes with identical low nibbles and different high nibbles. + // Fitting on the packed bytes as though they were one nibble per + // byte would ignore the high halves and return the same codebook + // for both. + let low: Vec = (0..1024u32).map(|i| (i % 7) as u8).collect(); + let a: Vec = low.iter().map(|&b| b | (0x1 << 4)).collect(); + let b: Vec = low + .iter() + .enumerate() + .map(|(i, &v)| v | (((i % 13) as u8) << 4)) + .collect(); + let sa = fit(&[&a]).unwrap().unwrap().state_bytes; + let sb = fit(&[&b]).unwrap().unwrap().state_bytes; + assert_ne!(sa, sb, "codebook must depend on the high nibbles too"); + } + #[test] fn empty_inputs_produce_no_entry() { let entry = fit(&[]).unwrap(); diff --git a/crates/ptwm-core/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index e405944..b67a5b7 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -5,14 +5,16 @@ 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`. -/// Chain-internal code for `Float4E2M1FNx2`: two fp4 values share a byte. -const DTYPE_FP4_E2M1FN_X2: u16 = 0x001F; - +/// `python/weights/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 @@ -40,12 +42,17 @@ pub fn is_nibble_packed_dtype(dtype_code: u16) -> bool { /// Storage bytes for `n_elements` of a chain-internal dtype. /// -/// Accumulates in bits so sub-byte widths survive. Dividing -/// bits-per-element by 8 first floors a nibble to zero, which silently -/// yields a zero-length descriptor rather than a wrong-but-visible one. +/// 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 { - let bits = n_elements * element_width_for(dtype_code).bits_per_element() as u64; - bits.div_ceil(8) + 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 @@ -324,19 +331,21 @@ mod tests { let d = &src.propagate_descriptors(&[]).unwrap()[0]; assert_eq!(d.element_width, ElementWidth::Nibble); assert!(d.is_nibble_packed); - // 8192 fp4 values occupy 4096 bytes, not 8192. - assert_eq!(d.length_bytes, 64 * 128 / 2); + // 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 odd_element_count_at_nibble_width_rounds_up_to_whole_bytes() { + 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]; - // 7 nibbles need 4 bytes; the trailing nibble still costs one. - assert_eq!(d.length_bytes, 4); + assert_eq!(d.length_bytes, 7); + assert!(d.is_nibble_packed); } #[test] 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) } From 3e70bc65c34171181573db84eeaa56cd36fb75f5 Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Tue, 18 Aug 2026 14:42:16 +0200 Subject: [PATCH 5/7] fix(codecs): assert pack_nibbles gets an even count instead of truncating Unreachable from `decode`, whose nibble count is always a multiple of GROUP_SIZE, so this is not a live defect. Guarded anyway because of how it would fail rather than whether it can: `chunks_exact` drops a trailing nibble silently, the plane returns one byte short, and the container reports a CRC mismatch pointing nowhere near nibble handling. Getting from that symptom to the cause is what the packed-versus-expanded confusion in this file already cost once. A debug assertion keeps release behavior unchanged and puts the message at the site a future caller would actually hit. The should_panic test makes the contract executable rather than a comment. --- .../src/codecs/per_group_codebook.rs | 27 +++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/crates/ptwm-core/src/codecs/per_group_codebook.rs b/crates/ptwm-core/src/codecs/per_group_codebook.rs index 6ce09bf..19745c6 100644 --- a/crates/ptwm-core/src/codecs/per_group_codebook.rs +++ b/crates/ptwm-core/src/codecs/per_group_codebook.rs @@ -43,9 +43,23 @@ pub(crate) fn unpack_nibbles(packed: &[u8]) -> Vec { out } -/// Inverse of [`unpack_nibbles`]. Requires an even nibble count, which -/// `GROUP_SIZE` alignment guarantees. +/// Inverse of [`unpack_nibbles`]. Requires an even nibble count. +/// +/// The only caller decodes exactly `n_assignments * GROUP_SIZE` nibbles and +/// `GROUP_SIZE` is even, so an odd count cannot arise today. Asserted +/// anyway, because the way it would fail is expensive to diagnose: +/// `chunks_exact` drops a trailing nibble without complaint, the plane +/// returns one byte short, and the container rejects it as a CRC mismatch +/// with nothing pointing back here. Working out that a checksum failure +/// meant mishandled nibbles is what the packed-versus-expanded confusion in +/// this file already cost once. fn pack_nibbles(nibbles: &[u8]) -> Vec { + debug_assert!( + nibbles.len().is_multiple_of(2), + "pack_nibbles: {} nibbles is odd; the trailing nibble would be dropped \ + and surface later as a container integrity failure", + nibbles.len() + ); let mut out = Vec::with_capacity(nibbles.len() / 2); for pair in nibbles.chunks_exact(2) { out.push((pair[0] & 0x0F) | ((pair[1] & 0x0F) << 4)); @@ -761,6 +775,15 @@ mod codec_tests { assert_eq!(pack_nibbles(&unpack_nibbles(&plane)), plane); } + #[test] + #[should_panic(expected = "is odd")] + fn packing_an_odd_nibble_count_is_caught_rather_than_truncated() { + // Unreachable from `decode`, whose nibble count is always a multiple + // of GROUP_SIZE. Pinned so a future caller learns it here rather than + // from a container integrity failure three layers away. + pack_nibbles(&[1, 2, 3]); + } + #[test] fn decoder_rejects_bad_state_version() { let c = PerGroupCodebook; From 18cc5b99f88113216af84e137cdd01fa80325299 Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Tue, 18 Aug 2026 15:40:51 +0200 Subject: [PATCH 6/7] fix(codecs): make the pack_nibbles guard hold in release builds `debug_assert!` compiles out under `[profile.release]`, which carries no `debug-assertions` override, so the guard was absent exactly where a silent truncation would go unnoticed. Its doc comment claimed to prevent that. The `#[should_panic]` test also failed under `cargo test --release`, since nothing panicked. `pack_nibbles` runs once per plane decode rather than per byte, so an unconditional assertion costs nothing measurable. Also corrects a doc path that named a pre-rename package directory. --- crates/ptwm-core/src/codecs/per_group_codebook.rs | 2 +- crates/ptwm-core/src/transforms/source.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/ptwm-core/src/codecs/per_group_codebook.rs b/crates/ptwm-core/src/codecs/per_group_codebook.rs index 19745c6..fb6619c 100644 --- a/crates/ptwm-core/src/codecs/per_group_codebook.rs +++ b/crates/ptwm-core/src/codecs/per_group_codebook.rs @@ -54,7 +54,7 @@ pub(crate) fn unpack_nibbles(packed: &[u8]) -> Vec { /// meant mishandled nibbles is what the packed-versus-expanded confusion in /// this file already cost once. fn pack_nibbles(nibbles: &[u8]) -> Vec { - debug_assert!( + assert!( nibbles.len().is_multiple_of(2), "pack_nibbles: {} nibbles is odd; the trailing nibble would be dropped \ and surface later as a container integrity failure", diff --git a/crates/ptwm-core/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index b67a5b7..2abde55 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -12,7 +12,7 @@ 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`. This is the one owner of the +/// `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 { From d416782dac078f20c612bbda3a6b0f13c910ddfe Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Tue, 18 Aug 2026 17:58:43 +0200 Subject: [PATCH 7/7] fix(codecs): keep PerGroupCodebook on expanded planes, one convention throughout Reverses the packing change earlier in this branch. That change assumed the codec consumes nibble-packed planes; it consumes expanded ones, and three pieces of the design say so. `GROUP_SIZE` is 32 and documented as the MXFP4 block size in nibbles; `MxFp4Deinterleave` emits exactly 32 single-nibble slots per block; and `histograms` documents one nibble per byte. Under the packed reading a group spans half a block, which has no meaning for a format whose scale is per block. So the codec was right and its `accepts` was wrong: it demanded `is_nibble_packed`, which no correct producer sets for it. The flag now means one thing everywhere, "two values share a byte", and: - `accepts` takes expanded planes and declines packed ones, so a plane it would misread by dropping every high nibble is refused rather than silently halved. - `accepts` also declines lengths that are not whole groups. Dispatch drops a candidate chain when a codec errors and packed FP4 has one chain, so a misaligned tensor previously failed outright; declining leaves an `encode` error meaning a violated invariant. - `MxFp4Deinterleave` reports its value plane as not packed, which is what it emits. - `expected_byte_len` stops halving, so `length_bytes` is a storage count for every producer instead of a value count for one of them. Adds the coverage whose absence let the descriptor mismatch persist: an end-to-end test driving the deinterleave op into the codec and asserting a bit-exact round trip, plus rejection tests for the packed and misaligned cases. The existing capability tests asserted the old flag; they now assert the new one, and the role and width cases were switched to expanded descriptors so they still isolate what they name rather than passing because the plane is packed. No wire-format change, so containers written through the deinterleave op keep decoding. --- .../src/codecs/per_group_codebook.rs | 303 ++++++++++-------- crates/ptwm-core/src/dispatch.rs | 6 +- .../ptwm-core/src/fit/per_group_codebook.rs | 25 +- .../src/transforms/mxfp4_deinterleave.rs | 13 +- crates/ptwm-core/src/transforms/op.rs | 31 +- 5 files changed, 196 insertions(+), 182 deletions(-) diff --git a/crates/ptwm-core/src/codecs/per_group_codebook.rs b/crates/ptwm-core/src/codecs/per_group_codebook.rs index fb6619c..704bc87 100644 --- a/crates/ptwm-core/src/codecs/per_group_codebook.rs +++ b/crates/ptwm-core/src/codecs/per_group_codebook.rs @@ -26,47 +26,6 @@ pub fn histograms(nibbles: &[u8]) -> Vec<[u32; ALPHABET]> { out } -/// Expand a nibble-packed plane to one nibble per byte, low nibble first. -/// -/// The codec's internals model a nibble alphabet and work on this expanded -/// form. The planes the dispatcher hands it are packed two values per byte, -/// which is exactly what `is_nibble_packed` on the descriptor asserts, and -/// what `accepts` requires. Converting at the `PlaneCodec` boundary keeps -/// the two representations from being confused: reading a packed plane as -/// though it were already expanded silently drops every high nibble. -pub(crate) fn unpack_nibbles(packed: &[u8]) -> Vec { - let mut out = Vec::with_capacity(packed.len() * 2); - for &b in packed { - out.push(b & 0x0F); - out.push((b >> 4) & 0x0F); - } - out -} - -/// Inverse of [`unpack_nibbles`]. Requires an even nibble count. -/// -/// The only caller decodes exactly `n_assignments * GROUP_SIZE` nibbles and -/// `GROUP_SIZE` is even, so an odd count cannot arise today. Asserted -/// anyway, because the way it would fail is expensive to diagnose: -/// `chunks_exact` drops a trailing nibble without complaint, the plane -/// returns one byte short, and the container rejects it as a CRC mismatch -/// with nothing pointing back here. Working out that a checksum failure -/// meant mishandled nibbles is what the packed-versus-expanded confusion in -/// this file already cost once. -fn pack_nibbles(nibbles: &[u8]) -> Vec { - assert!( - nibbles.len().is_multiple_of(2), - "pack_nibbles: {} nibbles is odd; the trailing nibble would be dropped \ - and surface later as a container integrity failure", - nibbles.len() - ); - let mut out = Vec::with_capacity(nibbles.len() / 2); - for pair in nibbles.chunks_exact(2) { - out.push((pair[0] & 0x0F) | ((pair[1] & 0x0F) << 4)); - } - out -} - #[cfg(test)] mod tests { use super::*; @@ -446,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 { @@ -466,22 +443,17 @@ impl PlaneCodec for PerGroupCodebook { shared_state: Option<&[u8]>, _layout: &PlaneLayout, ) -> Result { - // `plane` is nibble-packed: two values per byte. Expand before - // modelling, and count groups in nibbles rather than bytes. - let nibbles = unpack_nibbles(plane); - if !nibbles.len().is_multiple_of(GROUP_SIZE) { + if !plane.len().is_multiple_of(GROUP_SIZE) { return Err(PtwmCoreError::CodecDecode { codec: "PerGroupCodebook", msg: format!( - "encode: plane holds {} nibbles ({} bytes), not a multiple of \ - GROUP_SIZE {GROUP_SIZE}", - nibbles.len(), + "encode: plane length {} not a multiple of GROUP_SIZE {GROUP_SIZE}", plane.len(), ), }); } - let hists = histograms(&nibbles); + let hists = histograms(plane); // Either consume provided shared state or fit a per-tensor codebook. let (state, inline_state_bytes): (StateV0, Vec) = match shared_state { @@ -524,8 +496,8 @@ impl PlaneCodec for PerGroupCodebook { // Range-coded per-nibble payload let mut enc = RangeEncoder::new(); - for (i, &nib) in nibbles.iter().enumerate() { - let sym = nib as usize; + for (i, &nib_byte) in plane.iter().enumerate() { + let sym = (nib_byte & 0x0F) as usize; let g = i / GROUP_SIZE; let k = assignments[g] as usize; let counts = &state.codebook[k]; @@ -567,27 +539,21 @@ impl PlaneCodec for PerGroupCodebook { }); } let n_assignments = u32::from_le_bytes(payload[..4].try_into().unwrap()) as usize; - // `decoded_len` comes from the trusted spec/plane-record path and is - // a byte count; groups are counted in nibbles, two per byte. A + // `decoded_len` comes from the trusted spec/plane-record path; a // mismatch here is corruption, not a soft failure. - let decoded_nibbles = decoded_len * 2; - if !decoded_nibbles.is_multiple_of(GROUP_SIZE) { + if !decoded_len.is_multiple_of(GROUP_SIZE) { return Err(PtwmCoreError::CodecDecode { codec: "PerGroupCodebook", - msg: format!( - "decoded_len {decoded_len} bytes = {decoded_nibbles} nibbles, \ - not divisible by GROUP_SIZE {GROUP_SIZE}" - ), + msg: format!("decoded_len {decoded_len} not divisible by GROUP_SIZE {GROUP_SIZE}"), }); } - let expected_assignments = decoded_nibbles / GROUP_SIZE; + let expected_assignments = decoded_len / GROUP_SIZE; if n_assignments != expected_assignments { return Err(PtwmCoreError::CodecDecode { codec: "PerGroupCodebook", msg: format!( - "payload claims {n_assignments} groups (={} nibbles), \ - plane expects {expected_assignments} (={decoded_nibbles} nibbles, \ - {decoded_len} bytes)", + "payload claims {n_assignments} groups (={} bytes), \ + plane expects {expected_assignments} (={decoded_len} bytes)", n_assignments * GROUP_SIZE ), }); @@ -657,11 +623,7 @@ impl PlaneCodec for PerGroupCodebook { })?; out.push(sym as u8); } - // Re-pack to the plane's on-the-wire form. Returning the expanded - // nibbles here is what made this codec lossy: it produced one byte - // per nibble with the high half zeroed, which round-tripped only - // for input whose high nibbles were already zero. - Ok(pack_nibbles(&out)) + Ok(out) } } @@ -720,70 +682,6 @@ mod codec_tests { assert_eq!(dec, b); } - /// A genuinely packed plane: both nibbles of every byte carry data. - /// - /// The other roundtrip tests here build one nibble per byte, so their - /// high halves are all zero and they pass whether or not the codec - /// preserves them. That is why this codec shipped dropping every high - /// nibble without any test noticing. - fn build_packed_plane(n_bytes: usize, seed: u64) -> Vec { - let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; - let mut out = Vec::with_capacity(n_bytes); - for _ in 0..n_bytes { - s ^= s << 13; - s ^= s >> 7; - s ^= s << 17; - // Skew toward a few values so the codebook has structure to find, - // but keep both halves populated. - let lo = (s % 16) as u8; - let hi = ((s >> 8) % 16) as u8; - out.push(lo | (hi << 4)); - } - out - } - - #[test] - fn codec_roundtrip_preserves_both_nibbles_of_every_byte() { - let c = PerGroupCodebook; - // GROUP_SIZE nibbles per group, two nibbles per byte. - let plane = build_packed_plane(64 * GROUP_SIZE / 2, 7); - assert!( - plane.iter().any(|b| b & 0xF0 != 0), - "fixture must exercise high nibbles" - ); - let enc = c.encode(&plane, None, &PlaneLayout::Flat).unwrap(); - let dec = c - .decode( - enc.state_format_version, - &enc.state_bytes, - &enc.payload, - &PlaneLayout::Flat, - plane.len(), - ) - .unwrap(); - assert_eq!( - dec.len(), - plane.len(), - "decoded length must match the plane" - ); - assert_eq!(dec, plane, "codec must be lossless over both nibbles"); - } - - #[test] - fn nibble_pack_unpack_is_an_identity() { - let plane = build_packed_plane(256, 3); - assert_eq!(pack_nibbles(&unpack_nibbles(&plane)), plane); - } - - #[test] - #[should_panic(expected = "is odd")] - fn packing_an_odd_nibble_count_is_caught_rather_than_truncated() { - // Unreachable from `decode`, whose nibble count is always a multiple - // of GROUP_SIZE. Pinned so a future caller learns it here rather than - // from a container integrity failure three layers away. - pack_nibbles(&[1, 2, 3]); - } - #[test] fn decoder_rejects_bad_state_version() { let c = PerGroupCodebook; @@ -913,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)); } @@ -946,7 +844,7 @@ mod capability_tests { format: ScaleFormat::E4M3, }, ElementWidth::Nibble, - true, + false, ); assert!(!c.accepts(&d)); } @@ -959,7 +857,7 @@ mod capability_tests { format: ValueFormat::Fp4E2m1, }, ElementWidth::Byte, - true, + false, ); assert!(!c.accepts(&d)); } @@ -972,7 +870,7 @@ mod capability_tests { format: ValueFormat::Fp4E2m1, }, ElementWidth::Nibble, - true, + false, ); assert_eq!(c.priority_for(&d), 10); } @@ -985,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/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/fit/per_group_codebook.rs b/crates/ptwm-core/src/fit/per_group_codebook.rs index 11e4534..3fd9a13 100644 --- a/crates/ptwm-core/src/fit/per_group_codebook.rs +++ b/crates/ptwm-core/src/fit/per_group_codebook.rs @@ -36,12 +36,7 @@ pub fn fit(planes_value: &[&[u8]]) -> Result, PtwmCoreE } let mut all_hists: Vec<[u32; pgc::ALPHABET]> = Vec::new(); for plane in planes_value { - // The planes are nibble-packed, two values per byte, and - // `histograms` models one nibble per byte. Expand first, exactly as - // the codec's own `encode` does: reading a packed plane directly - // would fit the codebook on the low nibbles alone and cut every - // group boundary at 32 bytes instead of 32 nibbles. - all_hists.extend(pgc::histograms(&pgc::unpack_nibbles(plane))); + all_hists.extend(pgc::histograms(plane)); } let fit_hists = stride_sample(&all_hists, FIT_HIST_SAMPLE_CAP); let cb = pgc::fit_codebook_multi_seed( @@ -89,24 +84,6 @@ mod tests { assert_eq!(e.name, "per_group_codebook"); } - #[test] - fn fit_reads_the_high_nibble_of_every_byte() { - // Two planes with identical low nibbles and different high nibbles. - // Fitting on the packed bytes as though they were one nibble per - // byte would ignore the high halves and return the same codebook - // for both. - let low: Vec = (0..1024u32).map(|i| (i % 7) as u8).collect(); - let a: Vec = low.iter().map(|&b| b | (0x1 << 4)).collect(); - let b: Vec = low - .iter() - .enumerate() - .map(|(i, &v)| v | (((i % 13) as u8) << 4)) - .collect(); - let sa = fit(&[&a]).unwrap().unwrap().state_bytes; - let sb = fit(&[&b]).unwrap().unwrap().state_bytes; - assert_ne!(sa, sb, "codebook must depend on the high nibbles too"); - } - #[test] fn empty_inputs_produce_no_entry() { let entry = fit(&[]).unwrap(); 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()); } }