Skip to content

Invalidate CPU quantization LUTs for complete codebook changes - #2141

Draft
Afloat16 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
Afloat16:fix/cpu-codebook-cache-invalidation-20261113
Draft

Afloat16 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
Afloat16:fix/cpu-codebook-cache-invalidation-20261113

Conversation

@Afloat16

Copy link
Copy Markdown

Problem

The CPU quantization LUT cache checks the codebook pointer and only four entries: 0, 1, 127 and 255. A caller can update other entries of a valid, sorted, contiguous codebook in place while all four sampled values stay unchanged. The cache then returns the old LUT, so the encoded weights no longer agree with the new codebook saved in QuantState.

Changes

Keep a complete 256-value codebook snapshot in each cache slot. Reuse a cached LUT only when the pointer and every value match; copy the new snapshot after rebuilding the LUT.

The four-slot replacement policy, existing locking and LUT quantization algorithm remain unchanged. The snapshots add approximately 4 KB of static storage and compare 256 float values on a pointer match.

Validation

Native CPU source builds and regression tests, source 5fa7276f1f2590b7e27419c02ddd05f25b5b1612:

  • Original C++ source: 15 failed, 6 passed.
  • Patched C++ source: all 21 new cases passed.
  • Existing CPU blockwise and noncontiguous-input coverage: 33 passed, 9 existing CPU skips.
  • Repository pre-commit hooks, including clang-format, and git diff --check passed.

The runner separately rebuilt the original and patched CPU shared libraries with the repository's CMake build, verified different library hashes, and used fresh Python processes for each revision. Environment: Ubuntu 24.04, Python 3.12, PyTorch 2.14.0+cpu.

Tests cover FP16/BF16/FP32, two block sizes, changes in either codebook half that preserve the old fingerprint, an independent nearest-code reference, unchanged maps, cache-slot replacement, and actual saved quantized weights used by a linear layer.

This fixes stale-cache correctness for codebook updates. It does not change LUT approximation, cache concurrency guarantees or GPU kernels, and makes no claim of improved model quality or general speed.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant