feat(kv): rk2v4-e8 compressed-KV (E8-root, 208 B/head-token) on the paged-KV engine - #173
Conversation
…-KV base Add KvCacheStorage::Rk2v4E8 (E8-root cylinder K: I8/64 + FP16/4 scale, V: U8/128 packed i4 + FP16/4 scale = 208 B/head/token) to the closed-profile layout resolver, serve --kv-dtype rk2v4-e8, and drop the self-contained E8 lattice/codec kernels + codec oracle from the original port (branch feat/rk-compressed-kv-on-paged).
Adds the rk2v4-e8 persistent codec on top of the paged-KV geometry registered in
d465fd63: K takes the fixed normalized D256 Hadamard rotation and is E8-root
("cylinder") encoded into a quarter-width 64-byte plane (one code pair per 8
dims); V takes a per-64-group Hadamard rotation and is packed to int4 at half
width. Both planes carry a per-64-group FP16 scale (208 B/token/head).
The single-row encoder is shared between the standalone append kernel and the
fused append inside the attention kernels so both paths emit identical codes and
each of the four K and four V group scales is written exactly once. The E8
encoder is warp-collective over 8-lane subgroups, so a warp owns a complete D256
row and every lane reaches every encode call.
Routes both the direct and the batch append entry points on KvCacheStorage::Rk2v4E8.
QK keeps the G64 INT8 m16n8k32 contraction: the E8-root code pairs expand back to eight int8 codes that share the per-64-group scale, so the int32 MMA output is still rescaled by qs[row,g]*ks[key,g]. V is unpacked from int4 to full-width int8 and dequantized to the existing bf16 PV tile. Neither plane can use the int8 kernel's 16-byte cp_async (both are narrower than the int8 cache and a 16-byte copy would run off the end of the plane), so the code tiles are expanded synchronously into the staging arena. PV accumulates in the rotated-V basis. The family owns its own cross-split reduce kernel, which keeps FP32 partial numerators and re-applies the per-64-group Hadamard - its own inverse - before narrowing to bf16, so the output leaves in original V coordinates. Rk2v4E8 therefore joins the fp32 partial-accumulator list.
Mirrors the INT8 prompt kernel with the compressed planes expanded into the same full-width staging tiles, so the QK/PV schedule, shared-memory arena and warp split are unchanged. The prompt kernel writes `out` directly, so the inverse per-64-group rotation runs as a separate pass over the finished tile once the attention kernel returns. Routes both the fill+attend and the attend-only prompt entry points.
Both switches fell through to "unknown", so the capacity line reported the storage as unknown for --kv-dtype rk2v4-e8.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fe62250575
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if (value == "fp8") { return KvCacheStorage::Fp8E4M3Row256; } | ||
| if (value == "nvfp4") { return KvCacheStorage::Nvfp4Group16; } | ||
| if (value == "k8v4") { return KvCacheStorage::Fp8KeyNvfp4Value; } | ||
| if (value == "rk2v4-e8") { return KvCacheStorage::Rk2v4E8; } |
There was a problem hiding this comment.
When users select this profile outside ninfer-serve, the new mode is unusable: apps/cli/options.cpp:55-61, apps/perplexity/main.cpp:104-117, and bench/targets/qwen3_6_27b/ninfer_bench_support.cpp:51-57 still reject rk2v4-e8. This leaves the required CLI, perplexity, and measurement Engine routes unable to exercise the supposedly first-class storage mode, so their parsers, names, help text, and relevant tests need the same addition.
AGENTS.md reference: AGENTS.md:L111-L113
Useful? React with 👍 / 👎.
| add_executable(ninfer_kv_e8_verify | ||
| test_e8_codec.cu | ||
| verify_1m_retrieval.cu) |
There was a problem hiding this comment.
Whenever the production RK2V4E8 codec or attention path is wrong, this verifier can still pass because the target builds only the self-contained implementations in tools/test_kv: they use unrelated E8Packed2BitTile/E8Packed4BitTile layouts and never include or invoke src/ops/kv_cache/rk2v4e8_codec.cuh, the append kernels, or either attention kernel. Consequently the advertised 5/5 result checks neither the 208-byte production representation nor its V rotation, paging, or final output; compare the actual production route directly against an independent oracle.
AGENTS.md reference: AGENTS.md:L244-L257
Useful? React with 👍 / 👎.
| # 1. From the NInfer top level with CUDA already configured | ||
| # cmake -S . -B build -GNinja -DCMAKE_CUDA_ARCHITECTURES=120a | ||
| # cmake --build build --target ninfer_kv_e8_verify |
There was a problem hiding this comment.
Following this newly documented top-level build command cannot work: the repository's root CMakeLists.txt never calls add_subdirectory(tools/test_kv), so a normal top-level configuration has no ninfer_kv_e8_verify target or CTest entry. The same file acknowledges that omission on lines 10-12, while tools/test_kv/README.md repeats the invalid command; either register the target under the appropriate build option or document the standalone source directory instead.
AGENTS.md reference: AGENTS.md:L286-L295
Useful? React with 👍 / 👎.
Address Codex P1 on Neroued#173: the --kv-dtype rk2v4-e8 profile was only parsed in ninfer-serve. apps/cli/options.cpp, apps/perplexity/main.cpp, and bench/.../ninfer_bench_support.cpp rejected it, leaving the required CLI, perplexity, and measurement Engine routes unable to select the mode. Add the parse, the enum->name mapping, the usage text, and the error string to each, mirroring the serve parser (canonical name "rk2v4-e8"). Address P2 (register the verification target): wire tools/test_kv into the root CMakeLists under BUILD_TESTING so the documented top-level build and `ctest -R ninfer_kv_e8_verify` actually work, and correct the CMakeLists/README to state the oracle's true scope — it is an independent reimplementation of the E8 Conway-Sloane math, not a test of the production 208 B layout / V-rotation / paging (which the end-to-end boot + needle verification covers). Tests: rk2v4-e8 parse + report-name assertions added to test_cli_options and test_ninfer_bench_support. Verified: both pass, perplexity binary accepts --kv-dtype rk2v4-e8 (advances past dtype check), CLI/bench help list it, and ninfer_kv_e8_verify now builds + runs as a registered target.
|
Addressed all three in P1 — wire rk2v4-e8 through every product parser (fixed). The profile was only parsed in P2 — register the documented verification target (fixed). P2 — "tests separate formats, not the production codec" (scoping corrected, not a code change). You're right that the oracle does not exercise the production route: |
1 participant
What
Re-port of the rk* compressed-KV storage format — specifically rk2v4-e8 (E8-root Conway-Sloane lattice Key codes, 2-bit Key / 4-bit Value, 208 B per head-token) — onto the current paged-KV
devengine, as a first-classKvCacheStorageprofile selectable via--kv-dtype rk2v4-e8.This is the successor to PR #35, which was written against the pre-paged-KV
gqa_attention_*kernels and closed when the engine moved to theKVPageGeometry/causal_softmax_attentionplane model. This port re-derives the same density win on the current kernel architecture, using the per-format kernel-file pattern (mirroring the K8V4 family) rather than the oldKvCacheCodectemplate fold.Why it matters on 24 GB cards
NInfer targets a single RTX 5090. The 24 GB laptop 5090 (sm_120a, 95 W) is the class where KV-cache headroom is binding, because dense 27B models leave only ~3.4 GiB free after weights. The existing compressed-KV profile, k8v4 (402 B/head-token), is coarser than rk2v4-e8 (208 B/head-token, 1.93× denser). All figures below were measured on a 24 GB RTX 5090 Laptop, CUDA 13.1.80, sm_120a — not projections.
Boot: KV headroom at 256K (Qwen3.6-35B-A3B, 19.6 GiB weights)
--kv-capacity 262144@ 256Kautomode; explicit mode boots with 574.6 MiB freeQwen3.8-27B NVFP4 (19.0 GiB weights) — 256K does not fit either
k8v4-class KV wants ~3.64 GiB at 256K but only 3.39 GiB is free after weights, so a dense 27B NVFP4 cannot reach 256K at k8v4's density. rk2v4-e8's 208 B/head-token brings the 256K reservation down enough to boot at 237,568 tokens, and to run 160K with MTP speculative decoding (
--spec mtp --draft-tokens 3 --lm-head-draft), whose draft head needs ~1 GiB.Fidelity (needle-in-haystack, case-insensitive, 3 needles at 10/50/90% depth)
tools/test_kv, 1M tokens)How
include/ninfer/types.h: addKvCacheStorage::Rk2v4E8.src/core/paged_kv_storage.hpaged_kv_storage_layout(): Rk2v4E8 geometry — K = {I8, 64, FP16, 4}, V = {U8, 128, FP16, 4} = 208 B/head/token (head_dim 256).plan_cache()derives the 4 planes automatically.src/ops/kernel/{e8_lattice.cuh,e8_root_codec.cuh}: self-contained E8 Conway-Sloane codec cores (port from the 4090 fork; oracle-verified).src/ops/kv_cache/{rk2v4e8_codec.cuh,append/rk2v4e8_kernel.cuh,append/rk2v4e8_launch.cu}: codec + append kernel family (mirrorsk8v4).src/ops/softmax_attention/dense/causal_cache/{small_t_rk2v4e8,prompt_rk2v4e8}.{cuh,cu}: split-KV small-T and causal prompt attention families (mirror k8v4).KvCacheStorage::Rk2v4E8branches whereverFp8KeyNvfp4Value(k8v4) is routed — append, small-T (both cached and uncached), prompt,fp32_acclist, split-capacity bucket.src/serve/serve_options.cpp: parse--kv-dtype rk2v4-e8; capacity/request logs name the profile.tools/test_kv/: standalone oracle + 1M retrieval verifier (raw-nvcc buildable, no engine deps).The V inverse-rotation is required (the PV output lands in rotated-V space); gated on
cache.storage == Rk2v4E8at the output launch.Verified on Beast (RTX 5090 Laptop, 24 GB, sm_120a, CUDA 13.1.80)
origin/deva140e7ae— clean rebase, 6 commits, ~4K insertions, build 321/321 exit 0.runtime 1.20 GiB | free 1.59 GiB, needle 3/3 @100k.Test plan
cmake --build build --target ninfer_kv_e8_verify && ./build/tools/test_kv/ninfer_kv_e8_verify 1000000→ 5/5.--kv-dtype rk2v4-e8 --max-context 262144 --kv-capacity 262144; confirm thecapacitylog line names rk2v4-e8 and reports the expected pages/runtime.tools/test_kv/) → 3/3.Supersedes PR #35 (which targeted the pre-paged-KV engine).