[HowardYangAI]

Required human contributor answers

The contributor's original claim records the scope, validation measures and collaborators below.

1. What outcome are you trying to achieve?

“INT4 KV packing/unpacking and cross-loop residual reconstruction,” adapted from the pinned FlashLoop reference for vllm-rlt.

2. What is the highest-leverage variable for that outcome?

The claim identifies “reconstruction error and storage size” as the validation measures for this codec contribution. Serving throughput is outside this module's validation scope.

3. How will this work move that variable rather than create busywork?

The claim proposes adapting FlashLoop's packing.py and cache.py and adding CPU/CUDA tests for reconstruction error and storage size. This PR implements that standalone codec and those tests.

4. Who is the customer or user, and what do they actually need?

The claim identifies “the storage and kernel contributors in the shared KV quantization PR” and commits to coordinating the format with them.

PR necessity and bugfix evidence

The contributor explicitly requested this PR to make the codec available for review. It supplies the packing/reconstruction contribution assigned in the allocation comment. This is a feature contribution, not a bugfix.

Purpose

Related to #156. Provide standalone Torch reference codecs for one request/layer's logical [token, KV head, channel] tensors, extending the INT4 contribution with the NVFP4 reference requested in the review and SM120 discussion.

• Separate INT4 and NVFP4 backends behind pack_kv / unpack_kv. The default remains INT4 G64. Packed objects expose format, axis and group_size; the axis follows kind.
• INT4 preserves FlashLoop's affine FP32 quantization, FP16 minimum/scale and nibble order. NVFP4 follows vLLM's pinned ref_nvfp4_quant / ref_nvfp4_quant_dequant at c7a7607: E2M1 nearest-even codes, E4M3 block scales and a scalar FP32 global encoding multiplier. Tiny positive scales retain a direct-division fallback when their reciprocal overflows.
• Both formats group K over token rows per head/channel and V over channels per token/head. INT4 supports configurable groups; NVFP4 uses G16. Partial groups are supported in this logical reference layout.
• Shared residual functions subtract the reconstructed predecessor in the input dtype. Compact sparse rows carry an owned int32 position map; absent rows inherit the predecessor. Int4Residual becomes KVResidual; there are no in-tree users of the old type name.
• Preserve the full FlashLoop MIT notice in NOTICE and vLLM's attribution in the NVFP4 module. No README, packaging or other documentation changes.

packed = pack_kv(values, kind="key", format="nvfp4")
previous = unpack_kv(packed)
residual = encode_residual(current, previous, kind="key", format="nvfp4")
rebuilt = decode_residual(residual, previous)

The global scale belongs to each encoded tensor and is recomputed on each encoding call. This does not define an append-time fixed-scale policy or a physical paged-cache format.

Test Plan

vllm-rlt Version / Commit: d08b565f216bd887fad4c8330f7482062e146973, based on upstream 989921ac1d3897c5c43e27ec8cddbc91fa953194.

Validated source tree: d099f3f0848757e8d49a12c25f36e614eecb57f1.

CUDA_VISIBLE_DEVICES= OMP_NUM_THREADS=1 python -m pre_commit run --all-files --show-diff-on-failure --verbose
CUDA_VISIBLE_DEVICES= OMP_NUM_THREADS=1 python -m pytest tests/test_int4_kv_quantization.py tests/test_nvfp4_kv_quantization.py -q
# On an authorized CUDA device with CUDA_VISIBLE_DEVICES supplied:
python -m pytest tests/test_int4_kv_quantization.py tests/test_nvfp4_kv_quantization.py --run-gpu -q

INT4 tests retain their independent scalar oracle. NVFP4 adds a separate scalar oracle enumerating finite E2M1/E4M3 values, including nearest-even midpoint/nextafter cases, FP8 scale rounding, signed zero, tiny values, partial/empty/noncontiguous inputs, both group axes, storage accounting, dense/nested sparse chains and mixed-format reconstruction.

GPU evaluation environment: RTX 5090 (SM120), Python 3.12.3, Torch 2.12.1+cu130, CUDA 13.0, driver 595.71.05. The local CPU regression environment is listed separately below.

Evaluation method:

1. KV reconstruction: Run real Ouro in BF16, quantize/decode its per-layer/per-loop K and V back to BF16, and compare with the originals. Measure base and residual streams separately; residuals are encoded against the reconstructed predecessor.
2. Attention output: Hold real Q/K/V fixed and compare causal attention using original versus reconstructed K/V.
3. Final prediction: In an offline numerical simulation, feed reconstructed K/V into attention and let each policy's hidden states propagate. Use identical article tokens via teacher forcing; compare logits with BF16 using KL, delta NLL and top-1 agreement.
4. Codec timing: Time pack/unpack and residual encode/decode on fixed-shape synthetic BF16 tensors after warmup. Model execution and data preparation are outside the timing boundary; these measurements do not establish serving throughput.

Test Result

• Full pre-commit: Ruff lint, Ruff format and CPU regressions passed; 1129 passed, 678 skipped with CUDA hidden. Python 3.12.3, Torch 2.13.0+cu129, pytest 9.2.0, Ruff 0.13.2.
• Targeted codec CPU tests: 218 passed, 94 skipped. RTX 5090 / SM120 tests on the same source contents: 312 passed, including all 94 CUDA cases. The INT4 test file was subsequently renamed without content changes; codec and test hashes match the GPU-validated snapshot.
• Additional G16 CPU/CUDA boundary/noncontiguous/residual checks: 144 passed. Pinned vLLM Torch reference checks: 180 passed on six distinct real Ouro QKV captures. SM120 E2M1 conversion probes: 48 raw values and 336 production payload values passed.
• Code-quality review found no additional in-scope defect. The backend split keeps data objects separate from encoding functions and retains sequential group processing. Eight historical implementations were compared separately; their legacy K layout is not mixed into the current-format results.

Real Ouro numerical comparison

Ouro-1.4B, BF16 weights/compute, four fixed loops, six distinct WikiText articles, prefixes 256 and 1024, 32 teacher-forced predictions per window: 384 predictions per policy from 12 related windows, not 384 independent samples. Each policy uses the same format for base and residuals. All 12 BF16 passthrough checks were bitwise equal.

Format	Mean KL, no BF16 window	Mean KL, first 4 / latest 16 kept BF16	Final-loop local attention relative L2, first 4 / latest 16 BF16	Complete-group bits/element
INT4 G16	0.001767	0.000588	0.6661%	6
INT4 G64	0.003246	0.000728	1.0101%	4.5
NVFP4 G16	0.014471	0.001010	0.9348%	4.5 + 4 bytes/tensor

KL is KL(P_BF16 || P_quantized). Base/residual reconstruction errors were recorded separately for K and V. INT4 G16 has lower mean KL here but larger metadata; INT4 G64 and NVFP4 G16 have similar storage. The BF16 window lowers NVFP4's mean KL about 14.3x in this sample. These six articles are a numerical screen, not a task-accuracy qualification or a universal codec ranking.

The following additional accuracy metrics were also measured. Final-loop KV reconstruction compares the fully reconstructed fourth-loop K/V with the original BF16 fourth-loop K/V, pooling squared errors and reference norms across layers and queries. These metrics use the same frozen real Q/K/V as the local attention comparison, rather than each policy's propagated hidden states.

Format	BF16 first / latest tokens	Final-loop K relative L2	Final-loop V relative L2	Mean delta NLL (quantized - BF16), nats	Top-1 agreement with BF16
INT4 G16	0 / 0	1.0142%	1.9778%	-0.004467	97.92%
INT4 G16	4 / 16	1.0007%	1.9523%	+0.000129	98.96%
INT4 G64	0 / 0	1.6427%	2.7636%	-0.003687	97.92%
INT4 G64	4 / 16	1.6239%	2.7281%	+0.001123	99.22%
NVFP4 G16	0 / 0	1.4066%	2.8799%	-0.000474	96.09%
NVFP4 G16	4 / 16	1.3875%	2.8429%	+0.000386	98.44%

Top-1 agreement measures whether the highest-probability token matches BF16; it is not ground-truth accuracy. Delta NLL uses the real next token from the article; a negative value on this small sample does not establish improved overall model quality. Lower KV relative L2 means closer cache reconstruction.

Base and residual payload errors were recorded separately. The last residual below is the third residual update, used to reconstruct loop four. Each residual is measured against current - reconstructed_previous in that policy, so policies can have different residual references. Payload relative errors are measured on quantized rows, excluding protected BF16 rows; they are distinct from the full reconstructed KV errors above.

Format	BF16 first / latest tokens	K base relative L2	K last-residual relative L2	V base relative L2	V last-residual relative L2
INT4 G16	0 / 0	3.486%	6.693%	6.805%	6.544%
INT4 G16	4 / 16	3.477%	6.691%	6.804%	6.544%
INT4 G64	0 / 0	5.106%	10.812%	10.344%	9.110%
INT4 G64	4 / 16	5.109%	10.829%	10.343%	9.110%
NVFP4 G16	0 / 0	10.644%	9.255%	9.338%	9.487%
NVFP4 G16	4 / 16	10.644%	9.250%	9.338%	9.487%

RMSE and maximum absolute error were retained as well for base/residual payloads, reconstructed K/V and local attention outputs; all three residual depths were recorded. The tables show relative L2 for comparing streams with different numerical scales. These numerical measurements do not replace downstream task-accuracy evaluation.

Unfused codec timing

BF16 synthetic inputs [2048, 16, 128]; K/V shapes follow the model. Each operation has 20 warmups and 80 balanced, order-rotated measurements per process, repeated in three sequential independent processes. Profiler disabled; CUDA event intervals include eager CPU enqueue gaps. Across T=1/128/2048, K/V, dense/sparse and four stages, 46,080 raw timings were retained.

Each cell below is Mean / P50 / P90 in milliseconds, pooled across the three processes (240 samples/cell). P90 uses linear interpolation. These P50s differ slightly from the earlier median-of-three-process-medians summary.

Format	K pack	V pack	K unpack	V unpack
INT4 G16	13.7054 / 13.5501 / 13.7507	0.9322 / 0.9326 / 0.9458	0.0904 / 0.0902 / 0.0920	0.0902 / 0.0900 / 0.0918
INT4 G64	3.5216 / 3.4796 / 3.5245	0.2941 / 0.2889 / 0.3083	0.0895 / 0.0884 / 0.0916	0.0895 / 0.0882 / 0.0919
NVFP4 G16	26.7849 / 26.4880 / 26.9839	1.8803 / 1.8769 / 1.9065	0.1445 / 0.1444 / 0.1476	0.1449 / 0.1448 / 0.1485

NVFP4 is slower in this eager reference implementation; these are not fused-kernel timings or serving latency/throughput. At T2048, K G16 executes 128 sequential groups, V G16 eight, and K G64 32. Separate ops-only/full CPU+CUDA profiles (24 captures with shapes/memory/stacks checked in full mode) corroborate the extra operators/launches; profile durations are excluded from timing results. Allocation peaks were sampled separately from latency and storage accounting.

Remaining integration work

This module is not wired into KVCacheManager, packed attention readers or serving. No fused dequantization/residual reconstruction, native FP4 MMA or end-to-end throughput claim is made. Page/group ownership, independent page release/sharing/PD/preemption, append-time scale stability, graph/async lifetimes and adaptive exits remain integration work. Current K groups follow compact row order and may cross physical pages; G16 alone does not solve page ownership. The current-format real-model comparison does not include an INT4-base/NVFP4-residual policy, though mixed-format reconstruction is unit-tested.

[Superone77]

Thanks for putting this together!

It may also be worth considering NVFP4 as an alternative to INT4, particularly on Blackwell GPUs. I previously discussed this with @0z5a , but I don’t think there is a clear winner yet.

For FlashLoop’s cross-loop residual representation, the key comparison would be speed and accuracy after fusing dequantization and residual reconstruction into the attention kernels. NVFP4’s hardware support could be useful, but the reconstruction and scaling costs may affect the tradeoff.

From algorithm perspective, NVFP4 should have better accuracy than INT4, if the speed is also better, that would be better!

[HowardYangAI]

Thanks for putting this together!

It may also be worth considering NVFP4 as an alternative to INT4, particularly on Blackwell GPUs. I previously discussed this with @0z5a , but I don’t think there is a clear winner yet.

For FlashLoop’s cross-loop residual representation, the key comparison would be speed and accuracy after fusing dequantization and residual reconstruction into the attention kernels. NVFP4’s hardware support could be useful, but the reconstruction and scaling costs may affect the tradeoff.

From algorithm perspective, NVFP4 should have better accuracy than INT4, if the speed is also better, that would be better!

Thanks for the suggestion! The INT4 codec and cross-loop residual reconstruction haven’t been integrated into the inference pipeline yet. Please keep me posted on any updates from the discussion. I’d be happy to help implement and integrate the chosen approach into vllm-rlt and evaluate its impact on inference performance, memory usage, and model accuracy.

[0z5a]

Thanks for putting this together!
It may also be worth considering NVFP4 as an alternative to INT4, particularly on Blackwell GPUs. I previously discussed this with @0z5a , but I don’t think there is a clear winner yet.
For FlashLoop’s cross-loop residual representation, the key comparison would be speed and accuracy after fusing dequantization and residual reconstruction into the attention kernels. NVFP4’s hardware support could be useful, but the reconstruction and scaling costs may affect the tradeoff.
From algorithm perspective, NVFP4 should have better accuracy than INT4, if the speed is also better, that would be better!

Thanks for the suggestion! The INT4 codec and cross-loop residual reconstruction haven’t been integrated into the inference pipeline yet. Please keep me posted on any updates from the discussion. I’d be happy to help implement and integrate the chosen approach into vllm-rlt and evaluate its impact on inference performance, memory usage, and model accuracy.

You can do it on SM120 like RTX5090 GPU, it's okay to load the model and cheap to rent. For non-blackwell GPU, the current PyTorch as a fallback is also reasonable.

[HowardYangAI]

Thanks for putting this together!
It may also be worth considering NVFP4 as an alternative to INT4, particularly on Blackwell GPUs. I previously discussed this with @0z5a , but I don’t think there is a clear winner yet.
For FlashLoop’s cross-loop residual representation, the key comparison would be speed and accuracy after fusing dequantization and residual reconstruction into the attention kernels. NVFP4’s hardware support could be useful, but the reconstruction and scaling costs may affect the tradeoff.
From algorithm perspective, NVFP4 should have better accuracy than INT4, if the speed is also better, that would be better!

Thanks for the suggestion! The INT4 codec and cross-loop residual reconstruction haven’t been integrated into the inference pipeline yet. Please keep me posted on any updates from the discussion. I’d be happy to help implement and integrate the chosen approach into vllm-rlt and evaluate its impact on inference performance, memory usage, and model accuracy.

You can do it on SM120 like RTX5090 GPU, it's okay to load the model and cheap to rent. For non-blackwell GPU, the current PyTorch as a fallback is also reasonable.

Thanks! I’ll look for an available SM120 GPU for the evaluation.

[0z5a]

Thanks, this is a clean port. I checked it against the pinned FlashLoop packing.py / cache.py (aeaadee), and these all match:

• the affine equations, and codes computed in FP32 before the FP16 metadata is rounded;
• nibble order;
• K grouped per channel over tokens, and V per token over channels;
• compact-row K groups for sparse residuals.

Encoding against the reconstructed predecessor keeps error from compounding across depths, and test_residual_chain_uses_reconstructed_predecessors pins that.

On 76c2a67:

• ruff check and ruff format --check are clean (0.13.2).
• tests/test_kv_quantization.py: 108 passed, 74 skipped (CUDA cases).
• CPU suite: 970 passed, excluding test_serving.py.

1. The FlashLoop MIT notice was dropped from NOTICE at the head

The commit the test plan names, 8f5d51b, adds the FlashLoop attribution and the full MIT permission text to NOTICE. The head 76c2a67 no longer touches NOTICE (the PR changes 2 files). The module header (kv_quantization.py:6-7) keeps only the copyright line, but MIT requires the permission notice as well. The PR description still says the notice is preserved.

Please restore that NOTICE block. While you're there, please point the test plan at the head commit: between 8f5d51b and 76c2a67 the API changed from packed.decode() to unpack_kv.

2. Before the NVFP4 hot path: make the format explicit and add an NVFP4 reference next to INT4

pack_kv derives the grouping axis from kind (kv_quantization.py:83, and again at :110), and PackedInt4KV does not record it. That works for one codec, but the next step is an NVFP4 path on SM120. A few points are worth settling first:

• Storage is a tie, so the comparison comes down to accuracy and kernel time.
  • INT4 with G=64 and an FP16 minimum and scale costs 4 + 32/64 = 4.5 bits per element.
  • NVFP4 (E2M1, one E4M3 scale per 16 elements, one FP32 scale per tensor) is also about 4.5 bits.
• Which hot path do you have in mind?
  • Dequantize NVFP4 to BF16 in registers and keep the BF16 MMA. The layout can stay as it is. The hardware gain is the FP4 to FP16 conversion.
  • Native block-scaled FP4 MMA (mma ... kind::mxf4nvf4.block_scale). As I read the PTX docs, the scale vectors run along the reduction dimension. For QKᵀ that means K is scaled per token over 16-channel blocks; for PV, V is scaled per channel over 16-token blocks. That is the transpose of the KIVI axes here, and Q and P would also have to be FP4, which is a separate accuracy question.
• NVFP4 is symmetric (no zero point).
  • That suits the zero-centred residual streams.
  • The loop-1 base K has per-channel offsets and outlier channels, which is why K is per channel here. A 16-channel NVFP4 block that contains an outlier channel loses precision on the other 15.
  • So it is worth comparing base and residual streams separately. INT4 base plus NVFP4 residuals is a legitimate option.
  • Measure attention-output error on real Ouro KV, not only reconstruction error on random tensors.
  • In my offline teacher-forced numerics on Ouro-1.4B, keeping the first 4 and the most recent 16 tokens in BF16 cut the KL of an all-4-bit policy by about 5–12x. That window probably belongs in the comparison too.

Suggestion:

• Put format / axis / group_size in the packed descriptor (upstream Packed4Bit already carries axis).
• Let encode_residual / decode_residual take the codec.
• Add an NVFP4 Torch reference with its own scalar oracle (E2M1 round-to-nearest-even, E4M3 scale rounding) next to the INT4 one.

Then the SM120 kernel has an exact oracle to test against, and the PyTorch fallback for other GPUs reads the same descriptor.

3. Group ownership vs 16-token pages

A 64-row K group spans four 16-token pages. For sparse residuals the groups follow compact row order (Int4Residual docstring, :137), so one minimum and scale can cover rows from distant positions on different pages.

The runtime would then be unable to free, offload for preemption, transfer for P/D, or prefix-share a page on its own without re-encoding its group. The pinned packing.py already has pack_k_per_channel_grouped, which avoids merging rows from different logical groups. It would be good to decide this in the shared descriptor. Note that NVFP4's 16-element blocks along channels would make K metadata page-local by construction.

The CUDA cases were not run for this update. Please run them on the SM120 box when you have it: the oracle has device-dependent division handling (scalar_reference).

[HowardYangAI]

Thanks, this is a clean port. I checked it against the pinned FlashLoop packing.py / cache.py (aeaadee), and these all match:

• the affine equations, and codes computed in FP32 before the FP16 metadata is rounded;
• nibble order;
• K grouped per channel over tokens, and V per token over channels;
• compact-row K groups for sparse residuals.

Encoding against the reconstructed predecessor keeps error from compounding across depths, and test_residual_chain_uses_reconstructed_predecessors pins that.

On 76c2a67:

• ruff check and ruff format --check are clean (0.13.2).
• tests/test_kv_quantization.py: 108 passed, 74 skipped (CUDA cases).
• CPU suite: 970 passed, excluding test_serving.py.

1. The FlashLoop MIT notice was dropped from NOTICE at the head

The commit the test plan names, 8f5d51b, adds the FlashLoop attribution and the full MIT permission text to NOTICE. The head 76c2a67 no longer touches NOTICE (the PR changes 2 files). The module header (kv_quantization.py:6-7) keeps only the copyright line, but MIT requires the permission notice as well. The PR description still says the notice is preserved.

Please restore that NOTICE block. While you're there, please point the test plan at the head commit: between 8f5d51b and 76c2a67 the API changed from packed.decode() to unpack_kv.

2. Before the NVFP4 hot path: make the format explicit and add an NVFP4 reference next to INT4

pack_kv derives the grouping axis from kind (kv_quantization.py:83, and again at :110), and PackedInt4KV does not record it. That works for one codec, but the next step is an NVFP4 path on SM120. A few points are worth settling first:

• Storage is a tie, so the comparison comes down to accuracy and kernel time.

  • INT4 with G=64 and an FP16 minimum and scale costs 4 + 32/64 = 4.5 bits per element.
  • NVFP4 (E2M1, one E4M3 scale per 16 elements, one FP32 scale per tensor) is also about 4.5 bits.

• Which hot path do you have in mind?

  • Dequantize NVFP4 to BF16 in registers and keep the BF16 MMA. The layout can stay as it is. The hardware gain is the FP4 to FP16 conversion.
  • Native block-scaled FP4 MMA (mma ... kind::mxf4nvf4.block_scale). As I read the PTX docs, the scale vectors run along the reduction dimension. For QKᵀ that means K is scaled per token over 16-channel blocks; for PV, V is scaled per channel over 16-token blocks. That is the transpose of the KIVI axes here, and Q and P would also have to be FP4, which is a separate accuracy question.

• NVFP4 is symmetric (no zero point).

  • That suits the zero-centred residual streams.
  • The loop-1 base K has per-channel offsets and outlier channels, which is why K is per channel here. A 16-channel NVFP4 block that contains an outlier channel loses precision on the other 15.
  • So it is worth comparing base and residual streams separately. INT4 base plus NVFP4 residuals is a legitimate option.
  • Measure attention-output error on real Ouro KV, not only reconstruction error on random tensors.
  • In my offline teacher-forced numerics on Ouro-1.4B, keeping the first 4 and the most recent 16 tokens in BF16 cut the KL of an all-4-bit policy by about 5–12x. That window probably belongs in the comparison too.

Suggestion:

• Put format / axis / group_size in the packed descriptor (upstream Packed4Bit already carries axis).
• Let encode_residual / decode_residual take the codec.
• Add an NVFP4 Torch reference with its own scalar oracle (E2M1 round-to-nearest-even, E4M3 scale rounding) next to the INT4 one.

Then the SM120 kernel has an exact oracle to test against, and the PyTorch fallback for other GPUs reads the same descriptor.

3. Group ownership vs 16-token pages

A 64-row K group spans four 16-token pages. For sparse residuals the groups follow compact row order (Int4Residual docstring, :137), so one minimum and scale can cover rows from distant positions on different pages.

The runtime would then be unable to free, offload for preemption, transfer for P/D, or prefix-share a page on its own without re-encoding its group. The pinned packing.py already has pack_k_per_channel_grouped, which avoids merging rows from different logical groups. It would be good to decide this in the shared descriptor. Note that NVFP4's 16-element blocks along channels would make K metadata page-local by construction.

The CUDA cases were not run for this update. Please run them on the SM120 box when you have it: the oracle has device-dependent division handling (scalar_reference).

Thanks for the suggestions! I completed the reference-level comparison on an RTX 5090. In this evaluation, INT4 had lower mean KL and faster eager codec timings. At comparable storage, keeping the first 4 and latest 16 tokens in BF16 gave mean KL of 0.000728 for INT4 G64 versus 0.001010 for NVFP4 G16. The BF16 window also reduced NVFP4’s mean KL by about 14.3×.
I’ve updated the PR with separate INT4/NVFP4 backends, explicit format/grouping metadata, and a vLLM-based NVFP4 reference with an independent scalar oracle. The FlashLoop MIT notice is preserved. All 312 codec tests passed, including 94 CUDA cases. The PR description now includes real Ouro KV reconstruction, attention-output and prediction comparisons, plus warmed-up codec timings.
These results cover six WikiText articles and the current unfused Torch implementation. They do not establish which format will perform better after attention fusion. Page/group ownership, append-time scale handling and production integration remain open.

[0z5a]

Reviewed the codec layer: the scalar-reference tests are strong (layout, partial groups, ties-to-even, tiny-value scale handling, residual scatter), the MIT/NOTICE attribution is right, and the codecs are a good reference base. The main question is the shared interface with the FP8 KV work, so here is a concrete proposal.

Sharing one KV-quantization interface

1. Add a KVCodec protocol and a registry in vllm_rlt/kernels/kv_quantization/__init__.py instead of the if format == ... dispatch:

   • format: str, storage_dtype: torch.dtype, group_axis: int, scale_shape(shape, group_size), pack(values, *, kind, group_size) -> PackedKV, unpack(packed) -> Tensor, storage_bytes.
   • Register "int4", "nvfp4", and later "fp8_e4m3". pack_kv/unpack_kv become registry[format] lookups.
   • FP8 can implement the same protocol with group_axis="channels" and one scale per (token, head); its paged storage is [blocks, block, heads, dim] with scales [blocks, block, heads], so the packed object should expose group_axis/scale_shape rather than making the KV manager branch on format.

2. Keep residual reconstruction above the codec. KVResidual / encode_residual / decode_residual is a cross-loop delta layer, not part of the storage format. Parameterize it by a codec (or format) and leave it out of the KVCodec protocol; FP8 has no residual requirement and forcing both into one class would couple two independent features. This also lets the residual path work with any future codec.

3. Add a paged adapter. The codecs are logical [tokens, heads, channels]; the KV manager and attention kernels are paged [blocks, block, heads, dim]. Please add either pack_paged/unpack_paged or a documented gather/scatter helper so the KV manager can use the codec without reimplementing layout logic. This is also where FP8's per-block/per-token scales need to line up with the paged allocation.

4. Separate "reference codec" from "serving codec". unpack is a good reference/debug path, but FP8's Triton kernel decodes in-register and never materializes bf16. Please state in the docstring which methods are reference-only, and keep the attention kernel format-specific; otherwise the shared interface implies a decode cost that doesn't exist in serving.

Review notes

• _validate_positions calls .item()/.any() on CUDA tensors, which forces a device sync on every encode_residual. For a reference codec that's acceptable, but please either mark it debug-only or switch to torch._assert_async/host-side validation so it doesn't land on the hot path.
• int4.pack/nvfp4.pack loop over groups in Python; fine for a reference implementation, but say so explicitly (and consider a vectorized path later).
• storage_bytes excludes allocator overhead and decode scratch; good. Please add a test that compares it to the actual tensor bytes and also reports bits/element, so the KV-cache memory claim is reproducible.
• The branch is based on a0efa27; main has moved. Please rebase before the next review round so CI runs the current suite.
• Integration coverage is the main gap: add a paged round-trip through KVCacheManager and a small attention reference (the FP8 PR has torch_paged_attention for this) so the codec is proven at the layout the engine actually uses.

Happy to align the FP8 CacheConfig.kv_cache_dtype values with the codec registry (int4, nvfp4, fp8_e4m3) once the protocol lands; that gives us one --kv-cache-dtype surface for all formats.

[0z5a]

Supersedes my earlier review. I finished both items: the paged helper is exercised on real cache layouts, and the accuracy check now uses real Ouro-1.4B KV (not synthetic random KV).

Changes in the stacked PR

HowardYangAI#1 (0z5a:codec-vectorized):

• vectorized packing (padded reshape reductions, no per-group Python): INT4 key 0.091 s -> 0.003 s, NVFP4 key 0.680 s -> 0.002 s on [32768, 8, 128];
• CODECS registry and pack_paged / unpack_paged for [pages, page_size, heads, dim];
• 313 tests pass (CPU + CUDA), payloads byte-identical.

Paged path on real KV

Ouro-1.4B, 24 layers, 16 KV heads, dim 128; captured past_key_values for 4k and 8k prompts, reshaped to [pages=256, page=16, heads, dim], packed/unpacked with pack_paged, attention compared per layer against bf16:

Context	Format	KV rel error	Attention rel error (mean / max)	Pack+unpack all layers
4k	INT4	3.23%	1.73% / 7.57%	0.48 s
4k	NVFP4	10.14%	3.36% / 6.98%	0.59 s
8k	INT4	3.20%	1.67% / 8.05%	0.36 s
8k	NVFP4	10.13%	3.27% / 6.73%	0.34 s

Real KV is much friendlier than the random Gaussian check (which gave 8.96% / 9.52% KV and 12.4-13.6% attention): INT4 is usable at ~1.7% attention error; NVFP4 sits at ~3.3%. Both are stable from 4k to 8k.

Memory (corrected for 16 KV heads, 24 layers)

2304 B per token per layer for both codecs vs 8192 B for bf16 -> 3.56x compression.

Context	Batch	bf16	INT4 / NVFP4
4k	64 / 128 / 256	48 / 96 / 192 GiB	13.5 / 27 / 54 GiB
8k	64 / 128 / 256	96 / 192 / 384 GiB	27 / 54 / 108 GiB
32k	64 / 128 / 256	384 / 768 / 1536 GiB	108 / 216 / 432 GiB

On 80 GB: bf16 fits 4k b64 (48 GiB); INT4/NVFP4 fit 4k b256 and 8k b128 (54 GiB), but not 8k b256 (108 GiB). On 141 GB, INT4 fits 8k b256 and 32k b64 (108 GiB).

Remaining gates

1. These are attention-output errors from captured KV, not generation-level accuracy. A real task eval (GSM8K/perplexity) is still needed before any default.
2. The codec is not yet wired into KVCacheManager writes; pack_paged proves the layout, but the manager still needs a store/load path and a decode kernel (or an unpack-to-bf16 fallback).
3. Scale-layout contract for attention backends: FA's public paged API has no per-token descale; upstream request Paged KV cache: expose per-page/per-token FP8 descale scales Dao-AILab/flash-attention#3145.

INT4 looks close to mergeable for a reference codec; NVFP4 needs a better scale placement or a quality gate before default use.

[0z5a]

Supersedes my earlier review.

I prototyped the two structural changes from that review in a stacked PR against this branch: HowardYangAI#1 (head 0z5a:codec-vectorized, base feat/int4-kv-codec).

What changed:

• Vectorized packing: padded reshape reductions instead of per-group Python loops. GH200, [32768, 8, 128] bf16: INT4 key pack 0.091 s -> 0.003 s, NVFP4 key pack 0.680 s -> 0.002 s; payloads are byte-identical and the full suite is 313 passed (CPU + CUDA).
• Registry (CODECS) instead of if format == ..., so FP8 can be added without touching dispatch.
• pack_paged / unpack_paged for [pages, page_size, heads, dim], which is the layout the KV manager and attention actually use.

E2E / memory evidence on H100 (from the codec at this PR head, before the stacked patch):

• 312 tests pass.
• 576 B per kind per token, 1152 B K+V per token per layer -> 3.56x smaller than bf16 (scales are 12.5% at dim=128).
• 80 GB fit: bf16 4k b128 (48 GiB) but not 4k b256; INT4/NVFP4 8k b256 (54 GiB) and 32k b64 (54 GiB), not 32k b128 (108 GiB).
• Synthetic random-KV error: KV rel 8.96% (INT4) / 9.52% (NVFP4), attention output rel 12.4-13.6%, stable from 2k to 8k context. This is the main quality gate left: it needs a real-KV/model eval before default use.

Remaining before merge:

1. Wire the paged helper into KVCacheManager/attention so the codec is exercised on the real cache, not only logical tensors.
2. Real-KV accuracy gate (synthetic error is too high to be a quality claim).
3. Define the scale-layout contract for attention backends; FA's public paged API has no per-token descale (upstream request Dao-AILab/flash-attention#3145).

Happy to iterate on the stacked PR if you'd rather keep the changes there than cherry-pick.