Compressed KV-cache stack for oxidize-cpp: HelixCache, RotorQuant, and a vendor-agnostic Vulkan backend - #33
Compressed KV-cache stack for oxidize-cpp: HelixCache, RotorQuant, and a vendor-agnostic Vulkan backend#33Jackson57279 wants to merge 3 commits into
Conversation
… Vulkan backend - HelixCache: polar 4-bit keys + Hadamard 3-bit values (7.26x vs f32), AVX-512 decode (gather-free radius, incremental RoPE query rotation, pdep 3-bit unpack) and vectorized store (polynomial log/atan2). - RotorQuant: blockwise 3D rotor decorrelation + int4 (6.4x), AVX-512 decode; default scheme via the new CompressedKvCache facade. - Vulkan compute backend (OXIDIZE_VULKAN): vendor-agnostic decode kernels (NVIDIA / all AMD / Intel Arc / iGPU), pages uploaded once, single-submit decode with batch-all-heads API. Bench (4096 tok x head_dim 128, 1 thread, Xeon Gold 5220R): rotor decode ~200us, helix 309us vs Rust TurboQuant 2411us; helix store 4.7ms. GPU batched (Radeon 680M): rotor 25us/head, helix 74us/head. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Registers the in-progress oxidize-core helix_cache port (codec/config/ encode/pack/runtime) and the KV-cache compression research report the C++ RotorQuant implementation is based on. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
There was a problem hiding this comment.
22 issues found across 26 files
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="oxidize-cpp/src/vulkan/kv_cache_vk.cpp">
<violation number="1" location="oxidize-cpp/src/vulkan/kv_cache_vk.cpp:67">
P3: The inline comment says "11 SSBOs" but the pipeline is created with 9 bindings and the descriptor set is written with 9 buffers. Stale comment may confuse future readers.</violation>
<violation number="2" location="oxidize-cpp/src/vulkan/kv_cache_vk.cpp:139">
P2: The buffer allocation uses a pre-selected `host_mem_type` without verifying that the specific buffer is compatible with it via `req.memoryTypeBits`. The Vulkan spec mandates that the chosen memory type index has its bit set in `VkMemoryRequirements::memoryTypeBits`. On implementations where not all host-visible heaps support storage buffers, `vkBindBufferMemory` will return `VK_ERROR_OUT_OF_DEVICE_MEMORY` or similar. Consider adding a check like `if (!(req.memoryTypeBits & (1u << host_mem_type)))` and falling back to scanning for a compatible type.</violation>
<violation number="3" location="oxidize-cpp/src/vulkan/kv_cache_vk.cpp:620">
P2: `helix_attention_batch` doesn't validate individual query sizes before calling `fill_helix_query`. If any `queries[h]` has fewer than `head_dim` elements, the loop in `fill_helix_query` reads out of bounds. The single-query path (`helix_attention`) already has this validation — adding a matching check here would prevent undefined behavior.</violation>
</file>
<file name="oxidize-cpp/src/rotor_quant.cpp">
<violation number="1" location="oxidize-cpp/src/rotor_quant.cpp:306">
P3: The scalar fallback path in `attention()` processes every token regardless of its softmax weight, but the AVX-512 path skips tokens with weight below 1e-12. For sparse attention distributions (common in longer sequences), the scalar path does unnecessary per-element work for tokens with negligible contribution. Adding the same early-continue check in the scalar path would be a low-risk consistency improvement.</violation>
</file>
<file name="oxidize-cpp/src/kv_compressed.cpp">
<violation number="1" location="oxidize-cpp/src/kv_compressed.cpp:62">
P2: `store_page` validates `pre_rope_keys.size()` against `tokens * head_dim_` on the RotorQuant path but leaves `values.size()` unchecked. If a caller passes a values vector with an unexpected size, the mismatch will only be caught inside the underlying `RotorQuantCache::store_page` rather than failing fast at the facade boundary with a consistent error. Add a parallel check for `values.size()` alongside the existing keys check so the facade validates all inputs uniformly.</violation>
<violation number="2" location="oxidize-cpp/src/kv_compressed.cpp:85">
P1: RotorQuant attention can read past the query buffer when `query_pre_rope.size() != head_dim_`, because RoPE preprocessing indexes by `head_dim_` without validating input length. Adding a local shape check here keeps invalid inputs on the exception path instead of undefined memory access.</violation>
</file>
<file name="oxidize-cpp/include/oxidize/kv_cache_vk.hpp">
<violation number="1" location="oxidize-cpp/include/oxidize/kv_cache_vk.hpp:41">
P2: Rotor attention can return incorrect vectors if caller passes a different `RotorQuantCache` instance than the one used for `upload`, because slot data comes from internal uploaded pages while rotate/unrotate comes from the method parameter. Binding rotor parameters to uploaded state (or removing this cache parameter from runtime attention APIs) would avoid this mismatch class.</violation>
</file>
<file name="oxidize-cpp/include/oxidize/helix_cache.hpp">
<violation number="1" location="oxidize-cpp/include/oxidize/helix_cache.hpp:51">
P2: Very long-context token indices can be corrupted on Vulkan upload because page positions are exposed as `size_t` but later truncated to `uint32_t` without validation. A checked conversion (or a fixed-width position type in the view/config) would prevent silent wraparound and wrong RoPE-relative logits on GPU.</violation>
</file>
<file name="oxidize-cpp/shaders/rotor_page_acc.comp">
<violation number="1" location="oxidize-cpp/shaders/rotor_page_acc.comp:5">
P2: Rotor page accumulation silently drops dimensions when `head_dim > 128`, producing incorrect attention outputs for larger heads. The shader ties dimension index to local invocation ID in a fixed 128-thread workgroup while host dispatch launches only one workgroup per page; consider multi-workgroup dimension tiling or explicitly rejecting `head_dim > 128` on this path.</violation>
</file>
<file name="deep-research-report.md">
<violation number="1" location="deep-research-report.md:1">
P3: This is a new 697-line standalone markdown documentation file, which goes against the project's explicit convention (CLAUDE.md, AGENTS.md) that extra markdown documentation should not be created unless explicitly asked — documented content should go into `README.md` instead. This file contains substantial design-rationale and survey content that could partially inform `README.md` or an `ARCHITECTURE.md` section. Consider if the key design decisions belong in the code's inline comments or `README.md` rather than a separate report file.</violation>
<violation number="2" location="deep-research-report.md:5">
P2: The 47 citation markers in the report use Unicode Private Use Area characters (U+E200–U+E202) that will not render in any standard viewer — GitHub, editors, and doc tools will show missing glyphs or invisible text. This makes every citation reference unreadable and undermines the document's credibility as a reference. Replace the PUA codepoints with a standard citation format — e.g. bracketed numeric keys `[19]`, anchor links, or footnote syntax — that renders portably.</violation>
</file>
<file name="oxidize-core/src/compute/helix_cache/runtime.rs">
<violation number="1" location="oxidize-core/src/compute/helix_cache/runtime.rs:131">
P1: Attention logits can be phase-shifted in the wrong direction, which distorts token weights when query/key phases are not symmetric. The RoPE relative term in `page_logits` uses `key_pos - query_pos`; using `query_pos - key_pos` aligns with pre-RoPE inputs and the standard rotated dot-product.</violation>
</file>
<file name="oxidize-core/src/compute/helix_cache/config.rs">
<violation number="1" location="oxidize-core/src/compute/helix_cache/config.rs:2">
P2: Add a `#[cfg(test)]` module at the bottom of `config.rs` per the AGENTS.md test co-location convention. The parent module tests cover config validation via `HelixCache::new`, but file-level concerns like `compression_ratio_vs_f32` edge cases, `HelixCacheStats::empty()` defaults, error-format derivation, and `HelixPageTier`/`PageKey` correctness remain untested at the unit level.</violation>
</file>
<file name="oxidize-cpp/shaders/softmax.comp">
<violation number="1" location="oxidize-cpp/shaders/softmax.comp:47">
P3: This adds an extra full-workgroup synchronization in the hot softmax path with no data dependency, which increases kernel latency. The barrier right after `inv` can be removed because later code only uses `inv` and writes `data[]`.</violation>
</file>
<file name="oxidize-core/src/compute/helix_cache/codec.rs">
<violation number="1" location="oxidize-core/src/compute/helix_cache/codec.rs:23">
P3: `recent_max_overlap` is currently write-only state, so it adds mutation/metadata without changing cache behavior. Either consume it in promotion/reporting logic or remove it to keep promotion state minimal and clear.</violation>
<violation number="2" location="oxidize-core/src/compute/helix_cache/codec.rs:131">
P1: Re-storing an existing `(layer, head, page_id)` appends a duplicate page instead of replacing it, so logits/attention can count that page twice while promotion updates only hit the first match. An upsert by `same_key(...)` in store paths would keep page identity consistent.</violation>
<violation number="3" location="oxidize-core/src/compute/helix_cache/codec.rs:182">
P3: This new `codec.rs` file has no co-located `#[cfg(test)]` module, which diverges from the compute-module test convention and makes codec-local invariants easier to regress. Adding a small bottom-of-file test module here would keep shape/packing checks close to the implementation.</violation>
</file>
<file name="oxidize-cpp/tests/helix_cache_test.cpp">
<violation number="1" location="oxidize-cpp/tests/helix_cache_test.cpp:51">
P1: The reference formula `cos(rel) + cos(0.01 * rel)` assumes both coordinate pairs contribute, but pair 1's keys have rho ≈ 0.014–0.032, all below `inactive_threshold=0.05`. Those pairs are skipped during encoding, so the actual logit will be only `cos(rel)` from pair 0. Fix: either increase the key values so rho ≥ 0.05, or correct the expected formula to just `cos(rel)`.</violation>
<violation number="2" location="oxidize-cpp/tests/helix_cache_test.cpp:51">
P2: Even if pair 1's keys were above the inactive threshold, `cos(0.01f * rel)` uses the wrong RoPE frequency. With theta=10000 and head_dim=8, the pair-1 frequency is `pow(10000, -2/8) ≈ 0.1`, not 0.01. The coefficient in the reference formula is off by 10×.</violation>
</file>
<file name="oxidize-cpp/shaders/helix_page_acc.comp">
<violation number="1" location="oxidize-cpp/shaders/helix_page_acc.comp:35">
P1: Threads that take the early return at `if (i >= pc.head_dim) return;` will never reach the `barrier()` call. In Vulkan compute shaders, all invocations in a workgroup must encounter `barrier()` uniformly — divergent control flow around a barrier is undefined behavior. Since `head_dim` can legally be 64 (it only needs to be a non-zero multiple of 8 per host validation), this will trigger UB on real hardware.
The fix is to let all threads participate up to and past the barrier, guarding only the memory accesses and final write with the `i < pc.head_dim` condition.</violation>
</file>
<file name="oxidize-cpp/src/helix_cache.cpp">
<violation number="1" location="oxidize-cpp/src/helix_cache.cpp:47">
P3: `ColdKeyTile::cos_mu` and `sin_mu` are stored during quantization but never read during decode. Both the AVX-512 path and the scalar fallback derive cos/sin from `mu_phi` directly (or use the static `kPhase` table), so these fields are dead code. Removing them (and their `assign`/store lines) will reduce struct size and eliminate reader confusion about whether they serve a purpose.</violation>
<violation number="2" location="oxidize-cpp/src/helix_cache.cpp:61">
P3: `PromotionState::access_count` is incremented on every `attention()` call for matching pages but never read by any method. It doesn't feed into `should_promote`, `stats()`, or any other query path. This is dead bookkeeping — either remove it to drop the 8 bytes/page, or add a comment if it's reserved for a future access-frequency promotion heuristic.</violation>
</file>
Tip: cubic can generate docs of your entire codebase and keep them up to date. Try it here.
Re-trigger cubic
| ->attention(layer, kv_head, query_pre_rope, query_position, | ||
| rope_theta_); | ||
| } | ||
| return rotor_->attention(layer, kv_head, rope(query_pre_rope, query_position)); |
There was a problem hiding this comment.
P1: RotorQuant attention can read past the query buffer when query_pre_rope.size() != head_dim_, because RoPE preprocessing indexes by head_dim_ without validating input length. Adding a local shape check here keeps invalid inputs on the exception path instead of undefined memory access.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-cpp/src/kv_compressed.cpp, line 85:
<comment>RotorQuant attention can read past the query buffer when `query_pre_rope.size() != head_dim_`, because RoPE preprocessing indexes by `head_dim_` without validating input length. Adding a local shape check here keeps invalid inputs on the exception path instead of undefined memory access.</comment>
<file context>
@@ -0,0 +1,93 @@
+ ->attention(layer, kv_head, query_pre_rope, query_position,
+ rope_theta_);
+ }
+ return rotor_->attention(layer, kv_head, rope(query_pre_rope, query_position));
+}
+
</file context>
| let mut logit = 0.0; | ||
| for pair in 0..pairs { | ||
| if let Some((rho, phi)) = self.page_pair(page, token, pair) { | ||
| let relative = page.positions[token] as f32 - query_position as f32; |
There was a problem hiding this comment.
P1: Attention logits can be phase-shifted in the wrong direction, which distorts token weights when query/key phases are not symmetric. The RoPE relative term in page_logits uses key_pos - query_pos; using query_pos - key_pos aligns with pre-RoPE inputs and the standard rotated dot-product.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-core/src/compute/helix_cache/runtime.rs, line 131:
<comment>Attention logits can be phase-shifted in the wrong direction, which distorts token weights when query/key phases are not symmetric. The RoPE relative term in `page_logits` uses `key_pos - query_pos`; using `query_pos - key_pos` aligns with pre-RoPE inputs and the standard rotated dot-product.</comment>
<file context>
@@ -0,0 +1,234 @@
+ let mut logit = 0.0;
+ for pair in 0..pairs {
+ if let Some((rho, phi)) = self.page_pair(page, token, pair) {
+ let relative = page.positions[token] as f32 - query_position as f32;
+ let delta = rope_frequency(pair, self.config.head_dim, rope_theta) * relative;
+ logit += q_mag[pair] * rho * (q_phi[pair] - phi + delta).cos();
</file context>
| }; | ||
| self.encode_keys(&mut page, pre_rope_keys); | ||
| self.encode_values(&mut page, values); | ||
| self.pages.push(page); |
There was a problem hiding this comment.
P1: Re-storing an existing (layer, head, page_id) appends a duplicate page instead of replacing it, so logits/attention can count that page twice while promotion updates only hit the first match. An upsert by same_key(...) in store paths would keep page identity consistent.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-core/src/compute/helix_cache/codec.rs, line 131:
<comment>Re-storing an existing `(layer, head, page_id)` appends a duplicate page instead of replacing it, so logits/attention can count that page twice while promotion updates only hit the first match. An upsert by `same_key(...)` in store paths would keep page identity consistent.</comment>
<file context>
@@ -0,0 +1,184 @@
+ };
+ self.encode_keys(&mut page, pre_rope_keys);
+ self.encode_values(&mut page, values);
+ self.pages.push(page);
+ Ok(())
+ }
</file context>
| require(logits.size() == 4, "cold logits should cover every token"); | ||
| for (size_t t = 0; t < logits.size(); ++t) { | ||
| const float rel = static_cast<float>(t) - 3.0f; | ||
| const float expected = std::cos(rel) + std::cos(0.01f * rel); |
There was a problem hiding this comment.
P1: The reference formula cos(rel) + cos(0.01 * rel) assumes both coordinate pairs contribute, but pair 1's keys have rho ≈ 0.014–0.032, all below inactive_threshold=0.05. Those pairs are skipped during encoding, so the actual logit will be only cos(rel) from pair 0. Fix: either increase the key values so rho ≥ 0.05, or correct the expected formula to just cos(rel).
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-cpp/tests/helix_cache_test.cpp, line 51:
<comment>The reference formula `cos(rel) + cos(0.01 * rel)` assumes both coordinate pairs contribute, but pair 1's keys have rho ≈ 0.014–0.032, all below `inactive_threshold=0.05`. Those pairs are skipped during encoding, so the actual logit will be only `cos(rel)` from pair 0. Fix: either increase the key values so rho ≥ 0.05, or correct the expected formula to just `cos(rel)`.</comment>
<file context>
@@ -0,0 +1,112 @@
+ require(logits.size() == 4, "cold logits should cover every token");
+ for (size_t t = 0; t < logits.size(); ++t) {
+ const float rel = static_cast<float>(t) - 3.0f;
+ const float expected = std::cos(rel) + std::cos(0.01f * rel);
+ require_close(logits[t], expected, 0.001f,
+ "quantized polar logits should preserve active pairs");
</file context>
| const float expected = std::cos(rel) + std::cos(0.01f * rel); | |
| const float expected = std::cos(rel); |
| void main() { | ||
| const uint page = gl_WorkGroupID.x; | ||
| const uint i = gl_LocalInvocationID.x; | ||
| if (i >= pc.head_dim) { |
There was a problem hiding this comment.
P1: Threads that take the early return at if (i >= pc.head_dim) return; will never reach the barrier() call. In Vulkan compute shaders, all invocations in a workgroup must encounter barrier() uniformly — divergent control flow around a barrier is undefined behavior. Since head_dim can legally be 64 (it only needs to be a non-zero multiple of 8 per host validation), this will trigger UB on real hardware.
The fix is to let all threads participate up to and past the barrier, guarding only the memory accesses and final write with the i < pc.head_dim condition.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-cpp/shaders/helix_page_acc.comp, line 35:
<comment>Threads that take the early return at `if (i >= pc.head_dim) return;` will never reach the `barrier()` call. In Vulkan compute shaders, all invocations in a workgroup must encounter `barrier()` uniformly — divergent control flow around a barrier is undefined behavior. Since `head_dim` can legally be 64 (it only needs to be a non-zero multiple of 8 per host validation), this will trigger UB on real hardware.
The fix is to let all threads participate up to and past the barrier, guarding only the memory accesses and final write with the `i < pc.head_dim` condition.</comment>
<file context>
@@ -0,0 +1,61 @@
+void main() {
+ const uint page = gl_WorkGroupID.x;
+ const uint i = gl_LocalInvocationID.x;
+ if (i >= pc.head_dim) {
+ return;
+ }
</file context>
| } | ||
| barrier(); | ||
| } | ||
| const float inv = 1.0 / sh[0]; |
There was a problem hiding this comment.
P3: This adds an extra full-workgroup synchronization in the hot softmax path with no data dependency, which increases kernel latency. The barrier right after inv can be removed because later code only uses inv and writes data[].
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-cpp/shaders/softmax.comp, line 47:
<comment>This adds an extra full-workgroup synchronization in the hot softmax path with no data dependency, which increases kernel latency. The barrier right after `inv` can be removed because later code only uses `inv` and writes `data[]`.</comment>
<file context>
@@ -0,0 +1,53 @@
+ }
+ barrier();
+ }
+ const float inv = 1.0 / sh[0];
+ barrier();
+
</file context>
| #[derive(Debug, Clone, Default, PartialEq)] | ||
| pub(super) struct PromotionState { | ||
| pub(super) uncertainty_counter: u32, | ||
| pub(super) recent_max_overlap: f32, |
There was a problem hiding this comment.
P3: recent_max_overlap is currently write-only state, so it adds mutation/metadata without changing cache behavior. Either consume it in promotion/reporting logic or remove it to keep promotion state minimal and clear.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-core/src/compute/helix_cache/codec.rs, line 23:
<comment>`recent_max_overlap` is currently write-only state, so it adds mutation/metadata without changing cache behavior. Either consume it in promotion/reporting logic or remove it to keep promotion state minimal and clear.</comment>
<file context>
@@ -0,0 +1,184 @@
+#[derive(Debug, Clone, Default, PartialEq)]
+pub(super) struct PromotionState {
+ pub(super) uncertainty_counter: u32,
+ pub(super) recent_max_overlap: f32,
+ pub(super) access_count: u64,
+}
</file context>
| @@ -0,0 +1,184 @@ | |||
| use super::config::{HelixCacheConfig, HelixCacheError, HelixPageTier, PageKey}; | |||
There was a problem hiding this comment.
P3: This new codec.rs file has no co-located #[cfg(test)] module, which diverges from the compute-module test convention and makes codec-local invariants easier to regress. Adding a small bottom-of-file test module here would keep shape/packing checks close to the implementation.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-core/src/compute/helix_cache/codec.rs, line 182:
<comment>This new `codec.rs` file has no co-located `#[cfg(test)]` module, which diverges from the compute-module test convention and makes codec-local invariants easier to regress. Adding a small bottom-of-file test module here would keep shape/packing checks close to the implementation.</comment>
<file context>
@@ -0,0 +1,184 @@
+ }
+}
+
+pub(super) fn same_key(key: PageKey, layer: usize, head: usize, page: usize) -> bool {
+ key.layer == layer && key.head == head && key.page == page
+}
</file context>
| struct PromotionState { | ||
| uint32_t uncertainty_counter = 0; | ||
| float recent_max_overlap = 0.0f; | ||
| uint64_t access_count = 0; |
There was a problem hiding this comment.
P3: PromotionState::access_count is incremented on every attention() call for matching pages but never read by any method. It doesn't feed into should_promote, stats(), or any other query path. This is dead bookkeeping — either remove it to drop the 8 bytes/page, or add a comment if it's reserved for a future access-frequency promotion heuristic.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-cpp/src/helix_cache.cpp, line 61:
<comment>`PromotionState::access_count` is incremented on every `attention()` call for matching pages but never read by any method. It doesn't feed into `should_promote`, `stats()`, or any other query path. This is dead bookkeeping — either remove it to drop the 8 bytes/page, or add a comment if it's reserved for a future access-frequency promotion heuristic.</comment>
<file context>
@@ -0,0 +1,911 @@
+struct PromotionState {
+ uint32_t uncertainty_counter = 0;
+ float recent_max_overlap = 0.0f;
+ uint64_t access_count = 0;
+};
+
</file context>
| std::vector<float> rho_r2; | ||
| std::vector<float> rho_r4; | ||
| std::vector<float> rho_r8; | ||
| std::vector<float> cos_mu; |
There was a problem hiding this comment.
P3: ColdKeyTile::cos_mu and sin_mu are stored during quantization but never read during decode. Both the AVX-512 path and the scalar fallback derive cos/sin from mu_phi directly (or use the static kPhase table), so these fields are dead code. Removing them (and their assign/store lines) will reduce struct size and eliminate reader confusion about whether they serve a purpose.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At oxidize-cpp/src/helix_cache.cpp, line 47:
<comment>`ColdKeyTile::cos_mu` and `sin_mu` are stored during quantization but never read during decode. Both the AVX-512 path and the scalar fallback derive cos/sin from `mu_phi` directly (or use the static `kPhase` table), so these fields are dead code. Removing them (and their `assign`/store lines) will reduce struct size and eliminate reader confusion about whether they serve a purpose.</comment>
<file context>
@@ -0,0 +1,911 @@
+ std::vector<float> rho_r2;
+ std::vector<float> rho_r4;
+ std::vector<float> rho_r8;
+ std::vector<float> cos_mu;
+ std::vector<float> sin_mu;
+};
</file context>
- Replace index loops with zip/enumerate iterators in helix_cache runtime - Bump anyhow 1.0.102 -> 1.0.103 to clear cargo-deny advisory failure Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
|
Triage note (maintainer pass): this is the strongest-standing PR of the batch — green CI, self-contained, genuinely well-tested (roundtrip + vs-f32 parity for both schemes, Vulkan test skips cleanly with no device), and the 6.4–7.3x compression with the 48x optimization series is impressive work. One real reservation before merge: by your own "Not in this PR" list,
I'd lean toward (1) so this doesn't become unused surface area. Either way, nice PR. |
What
Adds an experimental compressed KV-cache stack to
oxidize-cpp:src/helix_cache.cpp): polar-coordinate 4-bit keys + Hadamard-transform 3-bit values, 7.26x compression vs f32 KV. AVX-512 decode (gather-free radius reconstruction via masked multiplies, incremental RoPE query rotation with zero trig in the inner loop, BMI2pdep3-bit unpack) and AVX-512 store (single vectorized polar pass with polynomial log/atan2; sin/cos eliminated algebraically).src/rotor_quant.cpp): implementation of the blockwise 3D rotor decorrelation + int4 scheme fromdeep-research-report.md, 6.4x compression. The inverse rotation folds into the query/output, so decode is a plain quantized dot product. Lloyd-Max codebooks and the QJL residual stage are deliberately deferred (noted in the header).src/vulkan/kv_cache_vk.cpp,shaders/*.comp, CMake-DOXIDIZE_VULKAN=ON): custom decode kernels that run on any Vulkan 1.1 compute queue — NVIDIA, all AMD cards (not just ROCm-supported), Intel Arc, and iGPUs. Pages upload once after CPU quantization; decode is a single submit (logits → softmax → weighted values), with a batch-all-heads API that amortizes submit latency across a layer.CompressedKvCachefacade (src/kv_compressed.cpp): single entry point over both schemes with RotorQuant as the default; callers always pass pre-RoPE keys/queries and the facade handles each scheme's RoPE semantics.helix_cachemodule scaffold inoxidize-coreand checks in the research report the RotorQuant implementation follows.Benchmarks
CPU: 4096 tokens x head_dim 128, single thread, pinned to NUMA node 0 on the Xeon Gold 5220R box (
store= quantize K+V,decode= full attention pass over the cache):GPU (AMD Radeon 680M iGPU via RADV, 32 heads batched): RotorQuant 25 us/head, HelixCache 74 us/head.
For context, HelixCache's initial scalar implementation decoded at 14,737 us — the optimization series in this PR is a cumulative 48x.
Testing
helix_cache_test,rotor_quant_test: quantization roundtrip, logits/attention vs f32 reference, compression accounting.kv_cache_vk_test: GPU output vs CPU reference parity (skips cleanly when no Vulkan device is present).kv_compressed_test: facade default-scheme check; both schemes tracked against a plain-f32 RoPE-attention reference.-DOXIDIZE_VULKAN=ON);cargo check -p oxidize-corepasses with the Rust scaffold.Not in this PR
Wiring
CompressedKvCacheinto themodel_llama.cppinference path (it still uses the F32 KV cache), perplexity evaluation on real models, and device-local + staging memory for discrete GPUs (buffers are host-visible, host-cached for now).🤖 Generated with Claude Code
Summary by cubic
Adds a compressed KV-cache to
oxidize-cppwithRotorQuant(default) andHelixCache, plus a vendor-agnostic Vulkan backend. This cuts decode time and shrinks KV memory by 6.4–7.3x vs f32 with cross‑vendor GPU support; also fixes clippy warnings in the Rusthelix_cacheport and bumpsanyhowto1.0.103to clear a cargo-deny advisory.New Features
RotorQuant: blockwise 3D rotor + int4; inverse folded into query/output; AVX‑512 decode.HelixCache: polar 4-bit keys + Hadamard 3-bit values; AVX‑512 decode and vectorized store.-DOXIDIZE_VULKAN=ON.CompressedKvCachefacade: one API over both schemes;RotorQuantdefault; handles RoPE semantics for callers.helix_cacheinoxidize-coreand added research report.Performance
RotorQuant6.40x,HelixCache7.26x.RotorQuant~200 us,HelixCache309 us, Rust TurboQuant 2,411 us; store —HelixCache4.7 ms,RotorQuant6.9 ms.RotorQuant25 us/head,HelixCache74 us/head.Written for commit f3e129f. Summary will update on new commits.