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 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/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/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/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); + } } 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. 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/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)]}), " 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