Repository navigation
[KV] Add INT4 codec reference and cross-loop residual reconstruction - #166
HowardYangAI wants to merge 1 commit into
Conversation
aa4d0be to
76c2a67
Compare
|
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
left a comment
There was a problem hiding this comment.
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 checkandruff format --checkare 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. ForQKᵀthat means K is scaled per token over 16-channel blocks; forPV, 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_sizein the packed descriptor (upstreamPacked4Bitalready carriesaxis). - Let
encode_residual/decode_residualtake 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).
6961799 to
fbe7761
Compare
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 existingNOTICE.Representation proposed for discussion with storage/kernel contributors:
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 upstream989921ac1d3897c5c43e27ec8cddbc91fa953194.Validated source tree:
0c0f1571c00cd004cf05a1cdf9db9feb080a39ba.Run with
CUDA_VISIBLE_DEVICES=andOMP_NUM_THREADS=1:Codec tests cover an independent scalar oracle, nibble order, group boundaries, FP16 parameter rounding, dense/sparse residual chains, inheritance and storage size.
Test Result
test_serving.py: 1019 passed, 658 skipped in 90.24 seconds; CUDA hidden.