Repository navigation
[KV] Add INT4/NVFP4 reference codecs and cross-loop residual reconstruction - #166
HowardYangAI wants to merge 2 commits 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
ec6708f to
d08b565
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. 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.pack_kv/unpack_kv. The default remains INT4 G64. Packed objects exposeformat,axisandgroup_size; the axis followskind.ref_nvfp4_quant/ref_nvfp4_quant_dequantatc7a7607: 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.Int4ResidualbecomesKVResidual; there are no in-tree users of the old type name.NOTICEand vLLM's attribution in the NVFP4 module. No README, packaging or other documentation changes.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 upstream989921ac1d3897c5c43e27ec8cddbc91fa953194.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 -qINT4 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:
Test Result
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.
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.
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_previousin 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.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.
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.