Admit E8M0 block scales to the order-1 scale coder (includes the length-prefix fix)

.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

.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

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

crates/ptwm-core/src/codecs/arithmetic.rs

@@ -5,29 +5,43 @@
 //! planes).
 
 use crate::codec::{CodecId, Encoded, PlaneCodec};
+use crate::codecs::length_prefix;
 use crate::entropy::arithmetic::{
     Order0Adaptive, Order0Static, Order1Adaptive, decode_bytes, encode_bytes,
 };
 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<u8> {
+pub(crate) fn frame_payload(
+    decoded_len: usize,
+    payload: &[u8],
+    codec: &'static str,
+) -> Result<Vec<u8>, 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());
+    }
 }

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);

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<u8> {
+    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));
+    }
+}

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;