From 5fa7276f1f2590b7e27419c02ddd05f25b5b1612 Mon Sep 17 00:00:00 2001 From: Afloat Date: Fri, 13 Nov 2026 19:25:55 +0800 Subject: [PATCH] Invalidate CPU quantization LUTs when any codebook entry changes --- csrc/cpu_ops.cpp | 24 ++------ tests/test_cpu_codebook_cache.py | 102 +++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 18 deletions(-) create mode 100644 tests/test_cpu_codebook_cache.py diff --git a/csrc/cpu_ops.cpp b/csrc/cpu_ops.cpp index dfb9046ac..c270e9af9 100644 --- a/csrc/cpu_ops.cpp +++ b/csrc/cpu_ops.cpp @@ -519,29 +519,18 @@ static void build_quantize_lut(const float* codebook, unsigned char* lut) { } } -// Per-thread LUT cache with multiple slots to avoid rebuilding when alternating codebooks +// LUT cache with multiple slots to avoid rebuilding when alternating codebooks struct LUTCache { unsigned char luts[kLUTCacheSlots][kLUTSize]; const float* cached_codes[kLUTCacheSlots] = {}; - // Store fingerprint to detect pointer reuse (ABA problem): - // when a tensor is freed and a new one reuses the same address, - // the pointer matches but the codebook content may differ. - float cached_fingerprints[kLUTCacheSlots][4] = {}; + // Keep every codebook value: in-place updates and reused tensor addresses + // can leave any small subset of entries unchanged. + float cached_codebooks[kLUTCacheSlots][kCodebookSize] = {}; int next_slot = 0; - static void compute_fingerprint(const float* code, float* fp) { - fp[0] = code[0]; - fp[1] = code[1]; - fp[2] = code[127]; - fp[3] = code[255]; - } - const unsigned char* get_lut(const float* code) { - float fp[4]; - compute_fingerprint(code, fp); for (int i = 0; i < kLUTCacheSlots; ++i) { - if (cached_codes[i] == code && cached_fingerprints[i][0] == fp[0] && cached_fingerprints[i][1] == fp[1] && - cached_fingerprints[i][2] == fp[2] && cached_fingerprints[i][3] == fp[3]) { + if (cached_codes[i] == code && std::equal(code, code + kCodebookSize, cached_codebooks[i])) { return luts[i]; } } @@ -550,8 +539,7 @@ struct LUTCache { next_slot = (next_slot + 1) % kLUTCacheSlots; build_quantize_lut(code, luts[slot]); cached_codes[slot] = code; - for (int j = 0; j < 4; ++j) - cached_fingerprints[slot][j] = fp[j]; + std::copy(code, code + kCodebookSize, cached_codebooks[slot]); return luts[slot]; } }; diff --git a/tests/test_cpu_codebook_cache.py b/tests/test_cpu_codebook_cache.py new file mode 100644 index 000000000..f2844e7d4 --- /dev/null +++ b/tests/test_cpu_codebook_cache.py @@ -0,0 +1,102 @@ +import io + +import pytest +import torch + +from bitsandbytes import functional as F + + +def _values(dtype): + # These binary fractions avoid the existing LUT discretization boundaries. + return torch.tensor([-1.0, -0.75, -0.5, -0.25, 0.125, 0.5, 0.75, 1.0], dtype=dtype).repeat(33)[:257] + + +def _change_interior(code, half): + start, end = (1, 127) if half == "negative" else (127, 255) + position = torch.arange(1, end - start, dtype=code.dtype) / (end - start) + code[start + 1 : end] = code[start] + (code[end] - code[start]) * position.square() + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("blocksize", [64, 256]) +@pytest.mark.parametrize("half", ["negative", "positive"]) +def test_cpu_quantization_refreshes_changed_codebook_interior(dtype, blocksize, half): + code = torch.linspace(-1.0, 1.0, 256) + values = _values(dtype) + original = code.clone() + pointer = code.data_ptr() + warmed, warm_state = F.quantize_blockwise(values, code=code, blocksize=blocksize) + _change_interior(code, half) + assert code.data_ptr() == pointer + torch.testing.assert_close(code[[0, 1, 127, 255]], original[[0, 1, 127, 255]], rtol=0, atol=0) + assert not torch.equal(code, original) + assert torch.all(code[1:] > code[:-1]) + + actual, state = F.quantize_blockwise(values, code=code, blocksize=blocksize) + fresh_code = code.clone() + expected, expected_state = F.quantize_blockwise(values, code=fresh_code, blocksize=blocksize) + + assert not torch.equal(warmed, expected) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(state.absmax, expected_state.absmax, rtol=0, atol=0) + torch.testing.assert_close(state.code, code, rtol=0, atol=0) + torch.testing.assert_close(warm_state.code, original, rtol=0, atol=0) + # An independent nearest-code oracle checks the changed values themselves. + scales = state.absmax.repeat_interleave(blocksize)[: values.numel()] + nearest = (values.float().div(scales)[:, None] - code).abs().argmin(dim=1).to(torch.uint8) + torch.testing.assert_close(actual, nearest, rtol=0, atol=0) + decoded = F.dequantize_blockwise(actual, state) + torch.testing.assert_close(decoded, (code[nearest.long()] * scales).to(dtype), rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_cpu_codebook_cache_preserves_unchanged_and_sampled_entry_updates(dtype): + code = torch.linspace(-1.0, 1.0, 256) + values = _values(dtype) + first, state = F.quantize_blockwise(values, code=code, blocksize=64) + repeated, repeated_state = F.quantize_blockwise(values, code=code, blocksize=64) + torch.testing.assert_close(repeated, first, rtol=0, atol=0) + torch.testing.assert_close(repeated_state.absmax, state.absmax, rtol=0, atol=0) + code[127] = 0.0 + actual, actual_state = F.quantize_blockwise(values, code=code, blocksize=64) + expected, expected_state = F.quantize_blockwise(values, code=code.clone(), blocksize=64) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(actual_state.code, expected_state.code, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_cpu_codebook_cache_keeps_alternating_and_evicted_maps_distinct(dtype): + values = _values(dtype) + codebooks = [torch.linspace(-1.0, 1.0, 256) for _ in range(7)] + for index, code in enumerate(codebooks): + positions = torch.linspace(0.0, 1.0, 256) + code.copy_(2.0 * positions.pow(1.0 + index / 10.0) - 1.0) + expected = [F.quantize_blockwise(values, code=code, blocksize=64)[0] for code in codebooks] + for index in (0, 1, 0, 1, 6, 5, 4, 3, 2, 1, 0): + actual, _ = F.quantize_blockwise(values, code=codebooks[index], blocksize=64) + torch.testing.assert_close(actual, expected[index], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_updated_cpu_codebook_persists_consistent_linear_weights(dtype): + code = torch.linspace(-1.0, 1.0, 256) + weights = _values(dtype)[:256].reshape(16, 16) + F.quantize_blockwise(weights, code=code, blocksize=64) + _change_interior(code, "negative") + actual, state = F.quantize_blockwise(weights, code=code, blocksize=64) + expected, expected_state = F.quantize_blockwise(weights, code=code.clone(), blocksize=64) + buffer = io.BytesIO() + torch.save({"weights": actual, "code": state.code, "absmax": state.absmax}, buffer) + buffer.seek(0) + restored = torch.load(buffer, weights_only=True) + decoded = F.dequantize_blockwise( + restored["weights"], code=restored["code"], absmax=restored["absmax"], blocksize=64 + ) + reference = F.dequantize_blockwise(expected, code=expected_state.code, absmax=expected_state.absmax, blocksize=64) + inputs = torch.linspace(-1.0, 1.0, 48).reshape(3, 16) + torch.testing.assert_close( + torch.nn.functional.linear(inputs, decoded), + torch.nn.functional.linear(inputs, reference), + rtol=0, + atol=0, + )