Skip to content

[KV] Add INT4/NVFP4 reference codecs and cross-loop residual reconstruction - #166

Open
HowardYangAI wants to merge 2 commits into
InferMatrixLab:mainfrom
HowardYangAI:feat/int4-kv-codec
Open

HowardYangAI wants to merge 2 commits into
InferMatrixLab:mainfrom
HowardYangAI:feat/int4-kv-codec

Conversation

@HowardYangAI

@HowardYangAI HowardYangAI commented Nov 3, 2026 •

Copy link
Copy Markdown

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.

@HowardYangAI
HowardYangAI force-pushed the feat/int4-kv-codec branch 3 times, most recently from aa4d0be to 76c2a67 Compare November 3, 2026 09:51
@HowardYangAI
HowardYangAI marked this pull request as ready for review November 3, 2026 09:55
@Superone77

Copy link
Copy Markdown

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

Copy link
Copy Markdown
Author

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

0z5a commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

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

Copy link
Copy Markdown
Author

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 0z5a left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 HowardYangAI changed the title [KV] Add INT4 codec reference and cross-loop residual reconstruction [KV] Add INT4/NVFP4 reference codecs and cross-loop residual reconstruction Nov 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants