fix: make PerGroupCodebook reachable and lossless on packed FP4 - #3
Merged
Conversation
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.
… 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.
…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.
…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.
…ting 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.
khwstolle
marked this pull request as draft
August 18, 2026 13:40
`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.
… 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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Correction to the original description
The original text claimed the codec "has been unselectable since
acceptscould never be satisfied, so no container written by a release carries a
plane encoded this way."
That is wrong.
MxFp4Deinterleavealready emits a value plane withRole::Value{Fp4E2m1},ElementWidth::Nibbleandis_nibble_packed: true(
mxfp4_deinterleave.rs:88), which satisfies all three conditions, and thechain explorer can build chains containing it. The claim came from a sweep
of
PRODUCTION_CHAINSonly; "absent from the production table" wasgeneralized to "unreachable", which does not follow.
Blockers
Existing containers stop decoding. A container written through
MxFp4Deinterleavestoresn_assignments = L/32. The new decode computesdecoded_nibbles = L*2and thereforeexpected_assignments = L/16, so itmismatches and returns
CodecDecode. Neither the codec id norSTATE_FORMAT_VERSIONwas bumped, so nothing distinguishes old payloadsfrom new.
Compression can abort where it previously succeeded.
encodenowrequires the plane byte count to be a multiple of 16. Packed FP4 has
exactly one production chain, and a codec's encode error propagates out of
the trial encode and drops the whole candidate chain, so a packed-FP4
tensor whose byte count is not a multiple of 16 fails outright rather than
falling back.
Underlying disagreement this exposes
Three producers disagree about what a packed value plane is.
MxFp4Deinterleavelabels its planeis_nibble_packed: truewhilephysically emitting one nibble per byte.
Sourcenow labels the same flagon genuinely packed bytes.
op.rs:148validateslength_bytesas a valuecount and agrees with neither.
The codec was correct for its actual caller and wrong for the packed
convention; this branch fixes it for packed input and regresses it for the
existing caller, which now gets expanded a second time. Picking one
convention across the three producers is the real fix, and it is larger
than this branch.
What still stands
single owner rather than a drifted duplicate.
and returned them with the high half zeroed.
nibbles, and an assertion that the compress and decompress descriptor
paths agree on
length_bytesfor every dtype.Verification
cargo test --workspace: 909 passed, 0 failed, in both debug and release.cargo clippy --all-targetsandcargo fmt --checkclean. None of thatexercises the two blockers, which is the point: the suite has no coverage
of
MxFp4Deinterleavefeeding this codec, nor of a packed plane whose bytecount is not a multiple of 16.