Skip to content

[KV] Add INT4 codec reference and cross-loop residual reconstruction - #166

Open
HowardYangAI wants to merge 1 commit into
InferMatrixLab:mainfrom
HowardYangAI:feat/int4-kv-codec
Open

HowardYangAI wants to merge 1 commit 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. Add a standalone Torch reference codec for one request/layer's logical [token, KV head, channel] tensors. Dense and compact sparse residuals subtract the reconstructed predecessor; absent positions inherit it. Preserve FlashLoop's attribution and full MIT license in the existing NOTICE.

Representation proposed for discussion with storage/kernel contributors:

  • INT4 codes, two per uint8 byte; FP16 minimum and scale.
  • Default group size 64: K across compact token rows per head/channel; V across channels per token/head.
  • Quantized base plus successive residuals, with subtraction/addition in the input compute dtype.
  • The reference supports partial groups. Runtime tail retention/rollover and the shared encoding descriptor remain to be agreed.

This module is not wired into KVCacheManager, packed attention readers or serving. Paging, async/graph lifetimes, PD, preemption and end-to-end quality/performance validation remain integration work for the shared KV feature.

Test Plan

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

Validated source tree: 0c0f1571c00cd004cf05a1cdf9db9feb080a39ba.

Run with CUDA_VISIBLE_DEVICES= and OMP_NUM_THREADS=1:

python -m pre_commit run --all-files --show-diff-on-failure --verbose
python -m pytest tests/test_kv_quantization.py -q

Codec tests cover an independent scalar oracle, nibble order, group boundaries, FP16 parameter rounding, dense/sparse residual chains, inheritance and storage size.

Test Result

  • Full pre-commit checks: passed (Ruff lint, Ruff format, CPU regression tests).
  • CPU regression suite, including test_serving.py: 1019 passed, 658 skipped in 90.24 seconds; CUDA hidden.
  • Targeted codec tests: 108 passed, 74 skipped in 1.31 seconds (CUDA cases skipped).
  • Validation used Python 3.12.3, Torch 2.13.0+cu129, pytest 9.2.0 and Ruff 0.13.2. The validated source tree matches the PR head exactly.
  • GPU tests were not run for this update.
  • The codec remains standalone. No end-to-end serving quality, throughput or peak-memory qualification is claimed.

@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).

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