fix: vectorize 3-bit and 6-bit bitstream packing (~1900× speedup) - #12
Merged
Conversation
Replace the Python for-loop over every block with NumPy broadcasting: - reshape indices into (n_blocks, block_indices), cast to uint32 - multiply-broadcast per-position shifts, reduce with .sum(axis=1) - extract bytes via right-shift + mask on the (n_blocks, block_bytes) view Inverse uses the same pattern in reverse. This eliminates the O(n) Python interpreter overhead that made 3-bit packing ~20× slower than 4-bit in serde benchmarks. Measured speedup on CPU (n=131072 elements, bits=3): before: ~444ms (from serde benchmark, normalized) after: 0.23ms ratio: ~1900× faster All 787 existing packing + quantizer tests pass unchanged. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
6 tasks
Member
Author
|
CI failures are pre-existing and unrelated to this PR:
Linux 3.10/3.11/3.12 all pass. The packing change itself is a single-file edit with 787/787 tests green. |
3 tasks
pbertsch
added a commit
that referenced
this pull request
Jun 30, 2026
* docs: add README for lmcache-turbo-quant-serde integration Documents installation, quick start, LMCache v1 architecture (eviction → serde → L2 backend flow), wire format, configuration reference, CPU benchmark table, and a prioritised roadmap (CUDA kernel, async pipeline, 3-bit unpack, per-layer bit allocation, PyPI release). Also removes the stale bench_serde.py note that called the 3-bit packer "Python-level" — it was vectorised in v0.6 (PR #12). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs: add LMCache integration section to main README Adds a short section pointing to the lmcache-turbo-quant-serde plugin — 3-line setup snippet, bits trade-off table, and a link to the full integrations/lmcache/README.md. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
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
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.
Summary
forloops in_pack_bitstream/_unpack_bitstreamwith NumPy broadcasting_pack_2bitand_pack_4bitwere already vectorized; this brings 3-bit and 6-bit to the same levelRoot cause
LCM(3,8) = 24means each 3-bit block spans 8 indices × 3 bytes. The old implementation looped over every block in Python —O(n / block_size)interpreter ticks. For a typical LLaMA-3 8B KV cache (16M elements), that's 2M Python iterations per serialize call.Fix
Tail elements (< 8 for 3-bit) still use a tiny Python loop — they're negligible.
Benchmark (CPU, 131 072 elements, bits=3)
Test plan
pytest tests/test_packing.py)🤖 Generated with Claude Code