From 090b9f96d5a85e96457b62c2fe8d89d33004f882 Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Fri, 23 Oct 2026 09:51:36 +0200 Subject: [PATCH 1/5] fix(core): let encoders and decoders share one length-prefix limit rANS, tANS, and the arithmetic coders refused to decode any plane declaring more than 1 GiB, while their encoders wrote such planes without checking. A 2.1 GB BF16 embedding (128256 x 8192) therefore encoded and then failed to decode with "implausible orig_len". One rule in codecs::length_prefix now decides which declared lengths are accepted, and both sides apply it. Up to 1 GiB any ratio is accepted, as before. Above that, the encoded body must be at least 1/64 of the declared length, which still bounds what a short or corrupted payload can make the decoder allocate. rANS and tANS store a plane raw when the rule rejects it; the arithmetic coders, which have no raw form, return an encode error instead of an undecodable payload. The rANS and Huffman encoders also cast the plane length to u32 unchecked, so a plane over 4 GiB got a wrapped prefix. They now store such planes raw. Encoder output is unchanged for every plane up to 1 GiB. --- crates/ptwm-core/src/codecs/arithmetic.rs | 61 ++++++++++--- crates/ptwm-core/src/codecs/huffman.rs | 11 +-- crates/ptwm-core/src/codecs/length_prefix.rs | 89 +++++++++++++++++++ crates/ptwm-core/src/codecs/mod.rs | 1 + .../ptwm-core/src/codecs/order1_arithmetic.rs | 6 +- crates/ptwm-core/src/codecs/rans.rs | 51 ++++++++--- crates/ptwm-core/src/codecs/tans.rs | 39 ++++++-- 7 files changed, 223 insertions(+), 35 deletions(-) create mode 100644 crates/ptwm-core/src/codecs/length_prefix.rs diff --git a/crates/ptwm-core/src/codecs/arithmetic.rs b/crates/ptwm-core/src/codecs/arithmetic.rs index 8578369..03c2a1d 100644 --- a/crates/ptwm-core/src/codecs/arithmetic.rs +++ b/crates/ptwm-core/src/codecs/arithmetic.rs @@ -5,6 +5,7 @@ //! planes). use crate::codec::{CodecId, Encoded, PlaneCodec}; +use crate::codecs::length_prefix; use crate::entropy::arithmetic::{ Order0Adaptive, Order0Static, Order1Adaptive, decode_bytes, encode_bytes, }; @@ -12,22 +13,35 @@ use crate::error::PtwmCoreError; use crate::layout::PlaneLayout; use crate::types::descriptor::PlaneDescriptor; -/// Absolute cap on the framed decoded length, mirroring `rans.rs`. Guards a -/// hostile length prefix from driving a multi-GiB decode loop. -const MAX_DECODED_LEN: usize = 1 << 30; // 1 GiB - /// Prepend the true decoded length (u32 LE) so decode is self-delimiting. /// +/// Fails when the decoder would refuse the length (see +/// [`length_prefix::accepted`]); these coders have no raw fallback. +/// /// The container's `decode` receives a `decoded_len` derived from the /// *source-tensor* descriptor, which for byte-split planes is larger than the /// actual plane length. Self-delimiting codecs (Huffman, rANS, …) ignore that /// argument and recover the length from their own payload; the arithmetic /// coders do the same here rather than trusting `decoded_len`. -pub(crate) fn frame_payload(decoded_len: usize, payload: &[u8]) -> Vec { +pub(crate) fn frame_payload( + decoded_len: usize, + payload: &[u8], + codec: &'static str, +) -> Result, PtwmCoreError> { + if !length_prefix::accepted(decoded_len, payload.len()) { + return Err(PtwmCoreError::CodecEncode { + codec, + msg: format!( + "decoded length {decoded_len} with {} encoded bytes exceeds the length-prefix limit", + payload.len() + ), + }); + } let mut out = Vec::with_capacity(4 + payload.len()); + // `accepted` guarantees the length fits the u32 prefix. out.extend_from_slice(&(decoded_len as u32).to_le_bytes()); out.extend_from_slice(payload); - out + Ok(out) } /// Inverse of [`frame_payload`]: read the length prefix and return the rest. @@ -42,13 +56,15 @@ pub(crate) fn unframe_payload<'a>( }); } let n = u32::from_le_bytes([payload[0], payload[1], payload[2], payload[3]]) as usize; - if n > MAX_DECODED_LEN { + let body = &payload[4..]; + // Guards a hostile length prefix from driving a multi-GiB decode loop. + if !length_prefix::accepted(n, body.len()) { return Err(PtwmCoreError::CodecDecode { codec, msg: format!("implausible decoded length {n}"), }); } - Ok((n, &payload[4..])) + Ok((n, body)) } /// Order-0 arithmetic coder with a serialized normalized table. @@ -85,7 +101,11 @@ impl PlaneCodec for ArithmeticO0 { }); } let mut model = Order0Static::fit(plane); - let payload = frame_payload(plane.len(), &encode_bytes(&mut model, plane)); + let payload = frame_payload( + plane.len(), + &encode_bytes(&mut model, plane), + "arithmetic_o0", + )?; Ok(Encoded { state_bytes: model.serialize(), state_format_version: 1, @@ -156,7 +176,8 @@ impl PlaneCodec for ArithmeticO0Adaptive { let payload = frame_payload( plane.len(), &encode_bytes(&mut Order0Adaptive::new(), plane), - ); + "arithmetic_o0_adaptive", + )?; Ok(Encoded { state_bytes: Vec::new(), state_format_version: 0, @@ -229,7 +250,8 @@ impl PlaneCodec for ArithmeticO1 { let payload = frame_payload( plane.len(), &encode_bytes(&mut Order1Adaptive::new(), plane), - ); + "arithmetic_o1", + )?; Ok(Encoded { state_bytes: Vec::new(), state_format_version: 0, @@ -435,4 +457,21 @@ mod tests { "arithmetic O0 ({ac_total}) must not exceed Huffman ({huff_total}) on a structured plane", ); } + + #[test] + fn unframe_rejects_a_short_payload_declaring_a_large_plane() { + let mut payload = length_prefix::SHORT_BODY_DECLARED_LEN + .to_le_bytes() + .to_vec(); + payload.extend_from_slice(&[0u8; 64]); + assert!(unframe_payload(&payload, "arithmetic_o0").is_err()); + } + + #[test] + fn frame_and_unframe_agree_on_the_length() { + let framed = frame_payload(5, &[1, 2, 3], "arithmetic_o0").unwrap(); + let (n, body) = unframe_payload(&framed, "arithmetic_o0").unwrap(); + assert_eq!((n, body), (5, &[1u8, 2, 3][..])); + assert!(frame_payload(length_prefix::RATIO_FREE_LEN + 1, &[0], "arithmetic_o0").is_err()); + } } diff --git a/crates/ptwm-core/src/codecs/huffman.rs b/crates/ptwm-core/src/codecs/huffman.rs index 538421c..df4a943 100644 --- a/crates/ptwm-core/src/codecs/huffman.rs +++ b/crates/ptwm-core/src/codecs/huffman.rs @@ -57,10 +57,10 @@ impl PlaneCodec for Huffman { } let mut buf = vec![0u8; huffman_output_bound(plane.len())]; - match huffman::compress_outcome(&mut buf, plane)? { - CompressOutcome::Encoded(n) => { + let outcome = huffman::compress_outcome(&mut buf, plane)?; + match (outcome, u32::try_from(plane.len())) { + (CompressOutcome::Encoded(n), Ok(orig_len)) => { // TAG_HUF | orig_len (4 bytes LE) | huffman-encoded bytes - let orig_len = plane.len() as u32; let mut payload = Vec::with_capacity(1 + 4 + n); payload.push(TAG_HUF); payload.extend_from_slice(&orig_len.to_le_bytes()); @@ -71,9 +71,10 @@ impl PlaneCodec for Huffman { payload, }) } - _outcome => { + _ => { // NotBeneficial / Incompressible / DstTooSmall / - // LenOverflow → store raw + // LenOverflow, or a plane longer than the u32 length + // prefix → store raw let mut payload = Vec::with_capacity(1 + plane.len()); payload.push(TAG_RAW); payload.extend_from_slice(plane); diff --git a/crates/ptwm-core/src/codecs/length_prefix.rs b/crates/ptwm-core/src/codecs/length_prefix.rs new file mode 100644 index 0000000..f2e1f59 --- /dev/null +++ b/crates/ptwm-core/src/codecs/length_prefix.rs @@ -0,0 +1,89 @@ +//! Acceptance rule for the decoded length a plane payload declares. +//! +//! rANS, tANS, and the arithmetic coders prefix their payload with the +//! decoded length as a little-endian `u32`, and the decoder sizes its output +//! (or its decode loop) from that prefix before reading the body. A corrupted +//! or hostile prefix must not drive an arbitrarily large allocation, so the +//! decoder checks it against [`accepted`]. Encoders apply the same rule before +//! writing a prefix, so every plane they emit is one their decoder accepts. + +/// Decoded length accepted at any compression ratio. +/// +/// These coders reach very high ratios on near-degenerate distributions (an +/// all-zero 64 KiB plane encodes to about 30 bytes), so below this size the +/// rule bounds the length alone. +pub(crate) const RATIO_FREE_LEN: usize = 1 << 30; // 1 GiB + +/// Largest decoded-to-encoded length ratio accepted above [`RATIO_FREE_LEN`]. +/// +/// A payload that declares more than [`RATIO_FREE_LEN`] bytes must carry at +/// least `1 / MAX_RATIO` of that length as encoded body, which bounds the +/// allocation a short payload can request. Weight planes this large (a BF16 +/// embedding or output projection of a large-vocabulary model) compress far +/// less than this. +pub(crate) const MAX_RATIO: usize = 64; + +/// Whether a payload declaring `decoded_len` bytes from an encoded body of +/// `body_len` bytes is accepted. +/// +/// Lengths that do not fit the `u32` prefix are never accepted. +pub(crate) fn accepted(decoded_len: usize, body_len: usize) -> bool { + u32::try_from(decoded_len).is_ok() + && (decoded_len <= RATIO_FREE_LEN || decoded_len <= body_len.saturating_mul(MAX_RATIO)) +} + +/// Bytes from a skewed distribution (about 2 bits per byte), long enough to +/// exercise lengths above [`RATIO_FREE_LEN`] while staying far below +/// [`MAX_RATIO`]. +#[cfg(test)] +pub(crate) fn skewed_bytes(len: usize) -> Vec { + let mut x: u64 = 0x9E37_79B9_7F4A_7C15; + (0..len) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + x.trailing_zeros().min(15) as u8 + }) + .collect() +} + +/// A payload body shorter than `1 / MAX_RATIO` of the 3 GiB it declares. +#[cfg(test)] +pub(crate) const SHORT_BODY_DECLARED_LEN: u32 = 3 << 30; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn any_ratio_is_accepted_up_to_the_ratio_free_length() { + assert!(accepted(0, 0)); + assert!(accepted(RATIO_FREE_LEN, 1)); + } + + #[test] + fn a_large_plane_at_a_weight_like_ratio_is_accepted() { + // A BF16 embedding of a 128256-token, 8192-wide model, at the ratio + // rANS reached on it. + let decoded = 128_256 * 8192 * 2; + assert!(decoded > RATIO_FREE_LEN); + assert!(accepted(decoded, 1_611_325_931)); + } + + #[test] + fn a_short_body_cannot_declare_a_large_plane() { + let decoded = RATIO_FREE_LEN + 1; + assert!(!accepted(decoded, 16)); + assert!(!accepted(decoded, decoded / MAX_RATIO - 1)); + assert!(accepted(decoded, decoded.div_ceil(MAX_RATIO))); + } + + #[test] + #[cfg(target_pointer_width = "64")] + fn lengths_beyond_the_u32_prefix_are_rejected() { + let too_long = u32::MAX as usize + 1; + assert!(!accepted(too_long, too_long)); + assert!(accepted(u32::MAX as usize, u32::MAX as usize)); + } +} diff --git a/crates/ptwm-core/src/codecs/mod.rs b/crates/ptwm-core/src/codecs/mod.rs index d8689ce..39d3049 100644 --- a/crates/ptwm-core/src/codecs/mod.rs +++ b/crates/ptwm-core/src/codecs/mod.rs @@ -4,6 +4,7 @@ pub mod fpc; pub mod huff_llm; pub mod huffman; pub mod identity; +pub(crate) mod length_prefix; pub mod neural_predictor; pub mod order1_arithmetic; pub mod order1_scale_ac; diff --git a/crates/ptwm-core/src/codecs/order1_arithmetic.rs b/crates/ptwm-core/src/codecs/order1_arithmetic.rs index 74956fc..c1f16c9 100644 --- a/crates/ptwm-core/src/codecs/order1_arithmetic.rs +++ b/crates/ptwm-core/src/codecs/order1_arithmetic.rs @@ -72,7 +72,11 @@ impl PlaneCodec for Order1Arithmetic { }); } let mut model = Order1Static::fit(plane); - let payload = frame_payload(plane.len(), &encode_bytes(&mut model, plane)); + let payload = frame_payload( + plane.len(), + &encode_bytes(&mut model, plane), + "order1_arithmetic", + )?; Ok(Encoded { state_bytes: model.serialize(), state_format_version: STATE_FORMAT_VERSION, diff --git a/crates/ptwm-core/src/codecs/rans.rs b/crates/ptwm-core/src/codecs/rans.rs index 375b718..03205ab 100644 --- a/crates/ptwm-core/src/codecs/rans.rs +++ b/crates/ptwm-core/src/codecs/rans.rs @@ -6,6 +6,7 @@ use crate::codec::{CodecId, Encoded, PlaneCodec}; use crate::codec_tagged::{rans_decode, rans_encode}; +use crate::codecs::length_prefix; use crate::error::PtwmCoreError; use crate::layout::PlaneLayout; use crate::types::descriptor::PlaneDescriptor; @@ -63,7 +64,11 @@ impl PlaneCodec for Rans { // decode entry point needs, since PlaneCodec.decode has no // expected_len. let inner_tag = *tagged.first().unwrap_or(&0u8); - if inner_tag == 0 { + // tagged[1..] is the rANS-encoded bytes. Store the plane raw when + // rans_encode fell back, or when the decoder would refuse the + // declared length. + let encoded = &tagged[1..]; + if inner_tag == 0 || !length_prefix::accepted(plane.len(), encoded.len()) { // raw fallback: drop the inner tagged framing and re-wrap. let mut payload = Vec::with_capacity(1 + plane.len()); payload.push(TAG_RAW); @@ -74,9 +79,10 @@ impl PlaneCodec for Rans { payload, }); } - // tagged[1..] is the rANS-encoded bytes. - let encoded = &tagged[1..]; - let orig_len = plane.len() as u32; + let orig_len = u32::try_from(plane.len()).map_err(|_| PtwmCoreError::CodecEncode { + codec: "rans", + msg: "plane length exceeds u32::MAX".into(), + })?; let mut payload = Vec::with_capacity(1 + 4 + encoded.len()); payload.push(TAG_RANS); payload.extend_from_slice(&orig_len.to_le_bytes()); @@ -113,14 +119,9 @@ impl PlaneCodec for Rans { } let orig_len = u32::from_le_bytes([rest[0], rest[1], rest[2], rest[3]]) as usize; let encoded = &rest[4..]; - // Absolute cap to defend against adversarial headers that - // would drive a multi-GiB allocation. Bounded by size, not - // compression ratio: rANS reaches arbitrarily high ratios - // on near-degenerate distributions (an all-zero 64 KiB - // byte plane encodes to ~30 bytes), so a ratio cap rejects - // legitimate inputs. - const MAX_ORIG_LEN: usize = 1 << 30; // 1 GiB - if orig_len > MAX_ORIG_LEN { + // Defend against headers that would drive a multi-GiB + // allocation; the encoder applies the same rule. + if !length_prefix::accepted(orig_len, encoded.len()) { return Err(PtwmCoreError::CodecDecode { codec: "rans", msg: format!( @@ -245,4 +246,30 @@ mod tests { let result = c.encode(&[0u8; 32], Some(&[1, 2, 3]), &PlaneLayout::Flat); assert!(result.is_err()); } + + #[test] + fn rans_rejects_a_short_payload_declaring_a_large_plane() { + let mut payload = vec![TAG_RANS]; + payload.extend_from_slice(&length_prefix::SHORT_BODY_DECLARED_LEN.to_le_bytes()); + payload.extend_from_slice(&[0u8; 64]); + let err = Rans + .decode(0, &[], &payload, &PlaneLayout::Flat, 0) + .unwrap_err(); + assert!(err.to_string().contains("implausible orig_len"), "{err}"); + } + + #[test] + #[ignore = "allocates about 4 GiB; run with --ignored"] + fn rans_roundtrip_above_the_ratio_free_length() { + let data = length_prefix::skewed_bytes(length_prefix::RATIO_FREE_LEN + 4097); + let enc = Rans.encode(&data, None, &PlaneLayout::Flat).unwrap(); + assert_eq!( + enc.payload[0], TAG_RANS, + "plane must be rANS-coded, not raw" + ); + let dec = Rans + .decode(0, &enc.state_bytes, &enc.payload, &PlaneLayout::Flat, 0) + .unwrap(); + assert!(dec == data); + } } diff --git a/crates/ptwm-core/src/codecs/tans.rs b/crates/ptwm-core/src/codecs/tans.rs index ec3f22e..0302582 100644 --- a/crates/ptwm-core/src/codecs/tans.rs +++ b/crates/ptwm-core/src/codecs/tans.rs @@ -6,6 +6,7 @@ //! tANS can't shrink is stored verbatim and decode stays self-describing. use crate::codec::{CodecId, Encoded, PlaneCodec}; +use crate::codecs::length_prefix; use crate::entropy::tans; use crate::error::PtwmCoreError; use crate::layout::PlaneLayout; @@ -60,7 +61,9 @@ impl PlaneCodec for Tans { // that case store the plane verbatim under TAG_RAW. let mut scratch = vec![0u8; plane.len()]; let encoded_len = tans::compress(&mut scratch, plane)?; - if encoded_len == 0 { + // Also store the plane raw when the decoder would refuse the + // declared length. + if encoded_len == 0 || !length_prefix::accepted(plane.len(), encoded_len) { let mut payload = Vec::with_capacity(1 + plane.len()); payload.push(TAG_RAW); payload.extend_from_slice(plane); @@ -111,11 +114,9 @@ impl PlaneCodec for Tans { } let orig_len = u32::from_le_bytes([rest[0], rest[1], rest[2], rest[3]]) as usize; let encoded = &rest[4..]; - // Absolute size cap (not a ratio cap): tANS reaches very - // high ratios on near-degenerate distributions, so bound the - // allocation by an implausible-size guard rather than ratio. - const MAX_ORIG_LEN: usize = 1 << 30; // 1 GiB - if orig_len > MAX_ORIG_LEN { + // Defend against headers that would drive a multi-GiB + // allocation; the encoder applies the same rule. + if !length_prefix::accepted(orig_len, encoded.len()) { return Err(PtwmCoreError::CodecDecode { codec: "tans", msg: format!( @@ -233,4 +234,30 @@ mod tests { let r = Tans.decode(0, &[], &[], &PlaneLayout::Flat, 0); assert!(r.is_err()); } + + #[test] + fn decode_rejects_a_short_payload_declaring_a_large_plane() { + let mut payload = vec![TAG_TANS]; + payload.extend_from_slice(&length_prefix::SHORT_BODY_DECLARED_LEN.to_le_bytes()); + payload.extend_from_slice(&[0u8; 64]); + let err = Tans + .decode(0, &[], &payload, &PlaneLayout::Flat, 0) + .unwrap_err(); + assert!(err.to_string().contains("implausible orig_len"), "{err}"); + } + + #[test] + #[ignore = "allocates about 4 GiB; run with --ignored"] + fn roundtrip_above_the_ratio_free_length() { + let data = length_prefix::skewed_bytes(length_prefix::RATIO_FREE_LEN + 4097); + let enc = Tans.encode(&data, None, &PlaneLayout::Flat).unwrap(); + assert_eq!( + enc.payload[0], TAG_TANS, + "plane must be tANS-coded, not raw" + ); + let dec = Tans + .decode(0, &enc.state_bytes, &enc.payload, &PlaneLayout::Flat, 0) + .unwrap(); + assert!(dec == data); + } } From d6dec5bd16d214d5c804c01d9e338fbdb27cb768 Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Fri, 23 Oct 2026 15:13:08 +0200 Subject: [PATCH 2/5] feat(core): admit E8M0 block scales to the order-1 scale coder Order1ScaleAC accepted only FP8 scale planes, on the assumption that MXFP4's E8M0 block scales have no row structure. They do: across gpt-oss-20b's 597M E8M0 scales the entropy conditioned on the previous scale in the row is 1.20 bits, against 1.89 bits unconditioned. E8M0 Scale and GlobalScale planes now reach the coder, and the per-plane trial encode still keeps whichever codec is smallest. --- README.md | 4 +-- .../ptwm-core/src/codecs/order1_scale_ac.rs | 31 ++++++++++--------- crates/ptwm-core/src/dispatch.rs | 11 +++---- 3 files changed, 23 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 41abeef..7ad9f65 100644 --- a/README.md +++ b/README.md @@ -314,8 +314,8 @@ fields; see its docstring for the full list. | `bfloat16` / `float16` | Bit reorder + 2-plane byte split | Exponent → plane 1 | | `float8_e4m3fn` / `float8_e5m2` | Bit reorder + 2-plane nibble split | Exponent → high nibble | | `float4_e2m1fn_x2` (weight nibbles) | Passthrough, per-plane trial-encode | `PerGroupCodebook` competes against Huffman / rANS on the nibble-packed byte stream | -| MXFP4 scale (`uint8`, E8M0) | Passthrough, per-plane trial-encode | Pure 8-bit exponent with no row structure to exploit; compresses to ≈ 0.16 of original via the generic codec menu | -| NVFP4 scale (FP8-E4M3) | `Order1ScaleAC` (per-row lag-1 arithmetic coding) | Mantissa bits give exploitable row autocorrelation that E8M0 lacks | +| MXFP4 scale (`uint8`, E8M0) | `Order1ScaleAC` (per-row lag-1 arithmetic coding), per-plane trial-encode | Rows are autocorrelated: on gpt-oss-20b the entropy given the previous scale in the row is 1.20 bits, against 1.89 bits unconditioned | +| NVFP4 scale (FP8-E4M3) | `Order1ScaleAC` (per-row lag-1 arithmetic coding) | Mantissa bits give exploitable row autocorrelation | | `int8`, `uint8`, `int16`, `int32`, `int64`, `bool` | Single stream | No byte split or bit reorder yet; whatever wins per-plane trial-encode on the raw bytes | MXFP4 / NVFP4 tensors are compressed as a pair: the packed FP4 weight diff --git a/crates/ptwm-core/src/codecs/order1_scale_ac.rs b/crates/ptwm-core/src/codecs/order1_scale_ac.rs index 3d08114..c8561ae 100644 --- a/crates/ptwm-core/src/codecs/order1_scale_ac.rs +++ b/crates/ptwm-core/src/codecs/order1_scale_ac.rs @@ -453,16 +453,19 @@ impl PlaneCodec for Order1ScaleAC { fn accepts(&self, descriptor: &PlaneDescriptor) -> bool { // O1SAC requires byte-width Scale or GlobalScale planes with row layout - // AND an FP8 scale format (E4M3 or E5M2). The codec models per-row - // lag-1 PMFs over scale bytes — a structure that exists in FP8 scale - // distributions (mantissa bits give a continuous distribution with - // exploitable autocorrelation) but not in E8M0 scales, which are pure - // 8-bit exponents (no sign, no mantissa) used by MXFP4 block scales. - // Allow-listing E4M3/E5M2 keeps the codec inside its design domain; - // future scale formats need explicit vetting before being admitted. + // and a vetted scale format. The codec models per-row lag-1 PMFs over + // scale bytes. FP8 scales (E4M3, E5M2) carry that structure, and so do + // MXFP4's E8M0 block scales: across gpt-oss-20b's 597M E8M0 scales the + // entropy conditioned on the previous scale in the row is 1.20 + // bits against 1.89 bits unconditioned. Trial encoding still picks the + // smallest codec per plane. Future scale formats need the same vetting + // before being admitted. let role_ok = match descriptor.role { Role::Scale { format } | Role::GlobalScale { format } => { - matches!(format, ScaleFormat::E4M3 | ScaleFormat::E5M2) + matches!( + format, + ScaleFormat::E4M3 | ScaleFormat::E5M2 | ScaleFormat::E8M0 + ) } _ => false, }; @@ -1081,10 +1084,8 @@ mod capability_tests { } #[test] - fn o1sac_rejects_e8m0_scale() { - // E8M0 (MXFP4 block scale) is a pure 8-bit exponent — structurally - // different from FP8 scales (E4M3 / E5M2) and outside O1SAC's design - // domain. + fn o1sac_accepts_e8m0_scale() { + // E8M0 (MXFP4 block scale) rows are autocorrelated like FP8 scale rows. let c = Order1ScaleAC; let d = descriptor( Role::Scale { @@ -1094,11 +1095,11 @@ mod capability_tests { Layout::Rows { row_len: 32 }, false, ); - assert!(!c.accepts(&d)); + assert!(c.accepts(&d)); } #[test] - fn o1sac_rejects_e8m0_global_scale() { + fn o1sac_accepts_e8m0_global_scale() { let c = Order1ScaleAC; let d = descriptor( Role::GlobalScale { @@ -1108,7 +1109,7 @@ mod capability_tests { Layout::Rows { row_len: 16 }, false, ); - assert!(!c.accepts(&d)); + assert!(c.accepts(&d)); } #[test] diff --git a/crates/ptwm-core/src/dispatch.rs b/crates/ptwm-core/src/dispatch.rs index 81f8f0e..a279bdd 100644 --- a/crates/ptwm-core/src/dispatch.rs +++ b/crates/ptwm-core/src/dispatch.rs @@ -99,10 +99,9 @@ mod tests { } #[test] - fn dispatch_e8m0_scale_byte_rows_excludes_o1sac() { - // uint8 weight_scale planes (MXFP4 E8M0 block scales) must not - // surface Order1ScaleAC in the dispatch menu — the codec is FP8-only - // by design. + fn dispatch_e8m0_scale_byte_rows_includes_o1sac() { + // uint8 weight_scale planes (MXFP4 E8M0 block scales) reach + // Order1ScaleAC alongside the generic coders. let d = descriptor( Role::Scale { format: ScaleFormat::E8M0, @@ -113,8 +112,8 @@ mod tests { ); let result = dispatch(&d); assert!( - !result.contains(&CodecId::Order1ScaleAC), - "O1SAC must not appear for E8M0 scales; got {:?}", + result.contains(&CodecId::Order1ScaleAC), + "O1SAC must appear for E8M0 scales; got {:?}", result ); // Standard codecs still present. From 740da1aba22753213efe45660ba007af349a80a4 Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Sat, 24 Oct 2026 19:50:35 +0200 Subject: [PATCH 3/5] test(integrations): expect E8M0 scale planes to reach Order1ScaleAC The two E8M0 scale-block tests asserted the old FP8-only gate. With E8M0 admitted, the effective menu includes Order1ScaleAC, which wins the synthetic NVFP4 fixture plane. --- .../integrations/test_compress_safetensors.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/ptwm/integrations/test_compress_safetensors.py b/tests/ptwm/integrations/test_compress_safetensors.py index 0057241..6d07b50 100644 --- a/tests/ptwm/integrations/test_compress_safetensors.py +++ b/tests/ptwm/integrations/test_compress_safetensors.py @@ -160,12 +160,12 @@ def test_nvfp4_hf_quant_classifier_round_trip(tmp_path): def test_scale_block_tensor_uses_order1_scale_ac(tmp_path): - """uint8 weight_scale (E8M0 MXFP4 block scale) must not be encoded with O1SAC. + """uint8 weight_scale (E8M0 MXFP4 block scale) reaches Order1ScaleAC. - Order1ScaleAC is FP8-only by design (E4M3 / E5M2). E8M0 is a pure 8-bit - exponent — structurally different — and is excluded at the - ``Order1ScaleAC::accepts()`` gate. The effective menu for an E8M0 plane is - therefore {Huffman, Rans, Zstd}; trial-encode picks the smallest. + E8M0 scale rows are autocorrelated like FP8 scale rows, so + ``Order1ScaleAC::accepts()`` admits them. The effective menu for an E8M0 + plane is {Huffman, Rans, Zstd, Order1ScaleAC}; trial-encode picks the + smallest. """ # Use a larger fixture so the plane exceeds the O1SAC minimum threshold # (4 096 bytes). out_features=128, in_features=512 → n_groups=32, @@ -183,8 +183,9 @@ def test_scale_block_tensor_uses_order1_scale_ac(tmp_path): base = "model.layers.0.self_attn.q_proj" scale_codecs = _codec_ids_for_tensor(blob, f"{base}.weight_scale") assert scale_codecs, "weight_scale has no planes in container" - # O1SAC is filtered for PlaneRole::Value planes; effective menu is {Huffman, Rans, Zstd}. - _effective_menu = frozenset({_CODEC_HUFFMAN, _CODEC_RANS, _CODEC_ZSTD}) + _effective_menu = frozenset( + {_CODEC_HUFFMAN, _CODEC_RANS, _CODEC_ZSTD, _CODEC_ORDER1_SCALE_AC} + ) assert all(c in _effective_menu for c in scale_codecs), ( f"expected all planes of weight_scale to use a hint-menu codec " f"({[hex(c) for c in sorted(_effective_menu)]}), " @@ -225,9 +226,9 @@ def test_mxfp4_heuristic_scale_tensor_uses_order1_scale_ac(tmp_path): Use 128 × 512 = 65 536 bytes for the scale plane (well above the 4 096-byte O1SAC threshold). - Order1ScaleAC is FP8-only by design; E8M0 (uint8 MXFP4 block scale) is - excluded at the ``accepts()`` gate. Trial-encode picks from the effective - menu {Huffman, Rans, Zstd}. + E8M0 (uint8 MXFP4 block scale) reaches Order1ScaleAC alongside the + generic coders; trial-encode picks from the effective menu + {Huffman, Rans, Zstd, Order1ScaleAC}. """ from ptwm.classify import HeuristicClassifier @@ -250,8 +251,9 @@ def test_mxfp4_heuristic_scale_tensor_uses_order1_scale_ac(tmp_path): blob = _ptwm_blob(out_dir) scale_codecs = _codec_ids_for_tensor(blob, "block.weight_scale") assert scale_codecs, "weight_scale has no planes in container" - # O1SAC is filtered for PlaneRole::Value planes; effective menu is {Huffman, Rans, Zstd}. - _effective_menu = frozenset({_CODEC_HUFFMAN, _CODEC_RANS, _CODEC_ZSTD}) + _effective_menu = frozenset( + {_CODEC_HUFFMAN, _CODEC_RANS, _CODEC_ZSTD, _CODEC_ORDER1_SCALE_AC} + ) assert all(c in _effective_menu for c in scale_codecs), ( f"expected hint-menu codec on weight_scale " f"({[hex(c) for c in sorted(_effective_menu)]}), " From 04204817ee75f7eaea351a06c849eb12046aeb24 Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Sat, 24 Oct 2026 19:50:35 +0200 Subject: [PATCH 4/5] style: merge a nested device check and sort test imports The pinned ruff flags both (SIM102, I001). --- python/ptwm/core/_compressor.py | 23 ++++++++++++----------- tests/test_codec_selection.py | 1 - 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/python/ptwm/core/_compressor.py b/python/ptwm/core/_compressor.py index 56e98fc..51915cf 100644 --- a/python/ptwm/core/_compressor.py +++ b/python/ptwm/core/_compressor.py @@ -83,17 +83,18 @@ def compress( once; the reference's BLAKE3 digest is recorded as a Dependency so the decoder can verify that the caller supplied the matching reference. """ - if isinstance(data, torch.Tensor) and data.device.type == "cuda": - if ( - self.config.device is not None - and self.config.device != data.device.index - ): - msg = ( - f"device mismatch: config.device={self.config.device} but the " - f"input tensor is on cuda:{data.device.index}. Move the tensor " - f"or change config.device; ptwm will not relocate it implicitly." - ) - raise ValueError(msg) + if ( + isinstance(data, torch.Tensor) + and data.device.type == "cuda" + and self.config.device is not None + and self.config.device != data.device.index + ): + msg = ( + f"device mismatch: config.device={self.config.device} but the " + f"input tensor is on cuda:{data.device.index}. Move the tensor " + f"or change config.device; ptwm will not relocate it implicitly." + ) + raise ValueError(msg) if self.config.codec is not None: msg = ( diff --git a/tests/test_codec_selection.py b/tests/test_codec_selection.py index 87b4021..52659cf 100644 --- a/tests/test_codec_selection.py +++ b/tests/test_codec_selection.py @@ -1,6 +1,5 @@ import pytest import torch - from ptwm import CompressionConfig, Compressor From 43433647fbe693e456bf9e0122a67cfda1f3066e Mon Sep 17 00:00:00 2001 From: "Kurt H. W. Stolle" Date: Sat, 24 Oct 2026 19:50:35 +0200 Subject: [PATCH 5/5] ci: pin Python 3.13 for the Rust jobs The Rust test, coverage, and lint jobs build pyo3-ffi against the runner default Python, now 3.14, which PyO3 0.24 does not support. --- .github/workflows/lint-and-format.yml | 18 ++++-------------- .github/workflows/rust-tests.yml | 23 ++++++++--------------- 2 files changed, 12 insertions(+), 29 deletions(-) diff --git a/.github/workflows/lint-and-format.yml b/.github/workflows/lint-and-format.yml index 0064b0f..e6a6a15 100644 --- a/.github/workflows/lint-and-format.yml +++ b/.github/workflows/lint-and-format.yml @@ -1,62 +1,52 @@ name: Lint and format - on: pull_request: # No push trigger: the PR already gates merges into master, so re-running # lint/format on the merge commit just doubles the Actions spend. - permissions: contents: read - # Cancel superseded in-progress runs on the same PR branch (pushes to # master always run to completion). concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: ${{ github.event_name == 'pull_request' }} - jobs: rust-lint: # Direct toolchain: cargo fmt + clippy on the workspace. runs-on: ubuntu-latest timeout-minutes: 25 - steps: - name: Checkout repository uses: actions/checkout@v4 - + - name: Setup Python (PyO3 0.24 builds against at most 3.13) + uses: actions/setup-python@v5 + with: + python-version: '3.13' - name: Setup Rust (stable, rustfmt + clippy components) uses: dtolnay/rust-toolchain@stable with: components: rustfmt, clippy - - name: Cache cargo registry and target uses: Swatinem/rust-cache@v2 - - name: cargo fmt --check run: cargo fmt --check - - name: cargo clippy run: cargo clippy --workspace --all-features -- -D warnings - lint: # Project-level ruff check on tracked sources. runs-on: ubuntu-latest timeout-minutes: 15 - steps: - name: Checkout repository uses: actions/checkout@v4 - - name: Setup Python uses: actions/setup-python@v5 with: python-version: '3.13' - - name: Install uv uses: astral-sh/setup-uv@v6 with: enable-cache: true - - name: Ruff lint # `uv run` resolves the project's pinned ruff (the dev # dependency group declares ruff>=0.15.4 in pyproject.toml), so diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml index bbf6896..99d9ded 100644 --- a/.github/workflows/rust-tests.yml +++ b/.github/workflows/rust-tests.yml @@ -1,64 +1,57 @@ name: Rust tests - on: pull_request: # No push trigger: the PR already gates merges into master, so re-running # the full suite on the merge commit just doubles the Actions spend. - permissions: contents: read - # Cancel superseded in-progress runs on the same PR branch (pushes to # master always run to completion). concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: ${{ github.event_name == 'pull_request' }} - jobs: test: runs-on: ubuntu-latest timeout-minutes: 30 - steps: - name: Checkout repository uses: actions/checkout@v4 - + - name: Setup Python (PyO3 0.24 builds against at most 3.13) + uses: actions/setup-python@v5 + with: + python-version: '3.13' - name: Setup Rust (stable) uses: dtolnay/rust-toolchain@stable - - name: Cache cargo registry and target uses: Swatinem/rust-cache@v2 - - name: Run Rust workspace tests run: cargo test --workspace --all-targets - coverage: runs-on: ubuntu-latest timeout-minutes: 30 # Non-blocking: coverage reporting is informational until a baseline # threshold is set. Failures here don't block PRs. continue-on-error: true - steps: - name: Checkout repository uses: actions/checkout@v4 - + - name: Setup Python (PyO3 0.24 builds against at most 3.13) + uses: actions/setup-python@v5 + with: + python-version: '3.13' - name: Setup Rust (stable, llvm-tools required by llvm-cov) uses: dtolnay/rust-toolchain@stable with: components: llvm-tools-preview - - name: Cache cargo registry and target uses: Swatinem/rust-cache@v2 - - name: Install cargo-llvm-cov uses: taiki-e/install-action@v2 with: tool: cargo-llvm-cov - - name: Generate coverage (lcov) run: cargo llvm-cov --workspace --lcov --output-path lcov.info - - name: Upload coverage artifact if: always() uses: actions/upload-artifact@v4