From fbe7761c02d8d1ffd76d7f897ace0f2468581df2 Mon Sep 17 00:00:00 2001 From: howardforpr Date: Tue, 3 Nov 2026 20:26:39 +0800 Subject: [PATCH 1/2] feat: add INT4 KV codec and cross-loop residual reconstruction --- NOTICE | 27 ++ tests/test_kv_quantization.py | 403 ++++++++++++++++++++++++++++ vllm_rlt/kernels/kv_quantization.py | 208 ++++++++++++++ 3 files changed, 638 insertions(+) create mode 100644 tests/test_kv_quantization.py create mode 100644 vllm_rlt/kernels/kv_quantization.py diff --git a/NOTICE b/NOTICE index 4e818fb8..8b7cefe2 100644 --- a/NOTICE +++ b/NOTICE @@ -12,3 +12,30 @@ All rights reserved. vllm-rlt modifies the execution structure, configuration handling, model loading, and attention integration for standalone loop-level scheduling and depth-aware paged KV caching. It does not include an upstream CDB runtime implementation. + +The affine INT4 reference codec in vllm_rlt/kernels/kv_quantization.py adapts +FlashLoop's packing.py and cross-loop residual arithmetic in cache.py, revision +aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6, to vllm-rlt's logical KV tensor layout. +https://github.com/Superone77/FlashLoop + +MIT License + +Copyright (c) 2026 FlashLoop contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/tests/test_kv_quantization.py b/tests/test_kv_quantization.py new file mode 100644 index 00000000..c35d85bc --- /dev/null +++ b/tests/test_kv_quantization.py @@ -0,0 +1,403 @@ +import numpy as np +import pytest +import torch + +from vllm_rlt.kernels.kv_quantization import ( + decode_residual, + encode_residual, + pack_kv, + unpack_kv, +) + +_DEVICES = ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)] + + +def scalar_reference(values, kind, group_size): + """Scalar affine equations, independent of the codec's tensor layout operations.""" + source = values.float().cpu().numpy() + tokens, heads, dimensions = source.shape + codes = np.zeros(source.shape, dtype=np.uint8) + decoded = np.zeros(source.shape, dtype=np.float32) + extent = tokens if kind == "key" else dimensions + for head in range(heads): + for fixed in range(dimensions if kind == "key" else tokens): + for start in range(0, extent, group_size): + indices = [ + (index, head, fixed) if kind == "key" else (fixed, head, index) + for index in range(start, min(start + group_size, extent)) + ] + low = min(source[index] for index in indices) + high = max(source[index] for index in indices) + # CPU/CUDA can round division by 15 differently at code boundaries. + step = ( + np.float32((torch.tensor(high - low, device=values.device) / 15).item()) + if high != low + else np.float32(1) + ) + stored_low, stored_step = np.float32(np.float16(low)), np.float32(np.float16(step)) + for index in indices: + code = min(15, max(0, round(float(np.float32((source[index] - low) / step))))) + codes[index] = code + decoded[index] = np.float32(np.float32(code) * stored_step + stored_low) + payload = np.zeros((tokens, heads, (dimensions + 1) // 2), dtype=np.uint8) + for token in range(tokens): + for head in range(heads): + for dimension in range(dimensions): + payload[token, head, dimension // 2] |= int(codes[token, head, dimension]) << ( + 4 * (dimension % 2) + ) + return torch.from_numpy(payload), torch.from_numpy(decoded).to(values.dtype) + + +def test_known_nibble_order_and_ties_to_even(): + values = torch.tensor([0, 0.5, 1.5, 2.5, 14.5, 15]).reshape(1, 1, 6) + packed = pack_kv(values, kind="value", group_size=6) + assert packed.payload.tolist() == [[[0x00, 0x22, 0xFE]]] + torch.testing.assert_close( + unpack_kv(packed), torch.tensor([0, 0, 2, 2, 14, 15.0]).reshape(1, 1, 6) + ) + odd = pack_kv(torch.tensor([[[0.0, 1.0, 15.0]]]), kind="value", group_size=3) + assert odd.payload.tolist() == [[[0x10, 0x0F]]] + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16, torch.float64]) +@pytest.mark.parametrize("group_size", [1, 3, 64]) +def test_groups_partial_groups_and_noncontiguous_inputs(kind, dtype, group_size, device): + values = torch.randn(5, 2, 14, generator=torch.Generator().manual_seed(714)).to( + device=device, dtype=dtype + )[..., ::2] + assert not values.is_contiguous() + original = values.clone() + expected_payload, expected = scalar_reference(values, kind, group_size) + packed = pack_kv(values, kind=kind, group_size=group_size) + assert packed.original_shape == (5, 2, 7) + assert packed.original_dtype == dtype + assert packed.payload.dtype == torch.uint8 + assert packed.minimum.dtype == packed.scale.dtype == torch.float16 + groups = ( + (5 + group_size - 1) // group_size if kind == "key" else (7 + group_size - 1) // group_size + ) + assert ( + packed.minimum.shape + == packed.scale.shape + == ((groups, 2, 7) if kind == "key" else (5, 2, groups)) + ) + torch.testing.assert_close(packed.payload.cpu(), expected_payload, rtol=0, atol=0) + torch.testing.assert_close(unpack_kv(packed).cpu(), expected, rtol=0, atol=0) + torch.testing.assert_close(values, original, rtol=0, atol=0) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("extent", [63, 64, 65, 128, 129]) +def test_default_group_boundaries(kind, dtype, extent, device): + tokens, channels = (extent, 7) if kind == "key" else (5, extent) + values = torch.randn(tokens, 1, channels * 2, generator=torch.Generator().manual_seed(517)).to( + dtype + )[..., ::2] + values = values.to(device) + expected_payload, expected = scalar_reference(values, kind, 64) + packed = pack_kv(values, kind=kind) + torch.testing.assert_close(packed.payload.cpu(), expected_payload, rtol=0, atol=0) + torch.testing.assert_close(unpack_kv(packed).cpu(), expected, rtol=0, atol=0) + metadata_elements = ( + ((tokens + 63) // 64) * channels if kind == "key" else tokens * ((channels + 63) // 64) + ) + assert packed.storage_bytes == tokens * ((channels + 1) // 2) + 4 * metadata_elements + if channels % 2: + assert torch.count_nonzero(packed.payload[..., -1] >> 4) == 0 + + +def test_codes_use_parameters_before_fp16_metadata_rounding(): + values = torch.tensor([[[0.10001, 0.10005, 0.10009, 0.10021]]]) + packed = pack_kv(values, kind="value", group_size=4) + assert packed.payload.tolist() == [[[0x30, 0xF6]]] + assert packed.minimum.item() == torch.tensor(0.10001).half().item() + _, expected = scalar_reference(values, "value", 4) + torch.testing.assert_close(unpack_kv(packed), expected, rtol=0, atol=0) + assert unpack_kv(packed)[0, 0, 0] != values[0, 0, 0] + + +@pytest.mark.parametrize("kind", ["key", "value"]) +def test_constant_groups_and_scale_underflow(kind): + packed = pack_kv(torch.full((3, 2, 5), 1.0001), kind=kind, group_size=2) + assert torch.count_nonzero(packed.payload) == 0 + assert torch.all(packed.scale == 1) + torch.testing.assert_close(unpack_kv(packed), torch.ones(3, 2, 5), rtol=0, atol=0) + values = torch.tensor([0.0, 1e-8]).reshape((2, 1, 1) if kind == "key" else (1, 1, 2)) + packed = pack_kv(values, kind=kind, group_size=2) + assert torch.count_nonzero(packed.payload) > 0 + assert torch.count_nonzero(packed.scale) == 0 + torch.testing.assert_close(unpack_kv(packed), torch.zeros_like(values), rtol=0, atol=0) + + +@pytest.mark.parametrize("kind,metadata_elements", [("key", 2 * 2 * 7), ("value", 5 * 2 * 3)]) +def test_storage_bytes_include_metadata_and_channel_padding(kind, metadata_elements): + packed = pack_kv(torch.ones(5, 2, 7), kind=kind, group_size=3) + assert packed.storage_bytes == 5 * 2 * 4 + 4 * metadata_elements + assert packed.storage_bytes == sum( + tensor.numel() * tensor.element_size() + for tensor in (packed.payload, packed.minimum, packed.scale) + ) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("shape", [(0, 2, 7), (5, 0, 7)]) +def test_empty_token_or_head_dimension(kind, device, shape): + values = torch.empty(shape, dtype=torch.bfloat16, device=device) + packed = pack_kv(values, kind=kind) + assert packed.payload.shape == (*shape[:2], 4) + assert unpack_kv(packed).shape == values.shape + assert unpack_kv(packed).dtype == values.dtype + assert packed.storage_bytes == 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize("kind", ["key", "value"]) +def test_oversized_group_does_not_expand_decode_scratch(kind): + values = torch.randn(5, 2, 7, device="cuda", dtype=torch.bfloat16) + peaks = [] + for group_size in (7, 65536): + packed = pack_kv(values, kind=kind, group_size=group_size) + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + decoded = unpack_kv(packed) + torch.cuda.synchronize() + peaks.append(torch.cuda.max_memory_allocated() - allocated) + assert decoded.shape == values.shape + assert peaks[1] <= peaks[0] + + +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16, torch.float64]) +def test_residual_chain_uses_reconstructed_predecessors(kind, dtype): + generator = torch.Generator().manual_seed(429) + raw_previous = torch.randn(5, 2, 7, generator=generator).to(dtype) + reconstructed = unpack_kv(pack_kv(raw_previous, kind=kind, group_size=3)) + _, expected = scalar_reference(raw_previous, kind, 3) + torch.testing.assert_close(reconstructed, expected, rtol=0, atol=0) + distinguishes_raw_predecessor = False + for _ in range(3): + current = (raw_previous + torch.randn(5, 2, 7, generator=generator).to(dtype) * 0.3).to( + dtype + ) + saved_current, saved_previous = current.clone(), reconstructed.clone() + residual = encode_residual(current, reconstructed, kind=kind, group_size=3) + expected_payload, decoded_delta = scalar_reference((current - expected).to(dtype), kind, 3) + expected = (expected + decoded_delta).to(dtype) + _, wrong_delta = scalar_reference((current - raw_previous).to(dtype), kind, 3) + wrong = (reconstructed + wrong_delta).to(dtype) + distinguishes_raw_predecessor |= not torch.equal(expected, wrong) + torch.testing.assert_close(residual.data.payload, expected_payload, rtol=0, atol=0) + previous_input = reconstructed + reconstructed = decode_residual(residual, previous_input) + torch.testing.assert_close(reconstructed, expected, rtol=0, atol=0) + torch.testing.assert_close(current, saved_current, rtol=0, atol=0) + torch.testing.assert_close(previous_input, saved_previous, rtol=0, atol=0) + assert residual.positions is None + assert residual.num_tokens == 5 + assert residual.storage_bytes == residual.data.storage_bytes + raw_previous = current + assert distinguishes_raw_predecessor + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_quantized_base_nested_sparse_residual_chain(kind, dtype, device): + generator = torch.Generator().manual_seed(937) + raw_previous = torch.randn(129, 1, 65, generator=generator).to(dtype) + reconstructed = unpack_kv(pack_kv(raw_previous.to(device), kind=kind)) + _, expected = scalar_reference(raw_previous.to(device), kind, 64) + torch.testing.assert_close(reconstructed.cpu(), expected, rtol=0, atol=0) + active = torch.arange(0, 129, 2, dtype=torch.int32) # 65 compact K rows cross a group boundary. + for positions in (active, active[::3]): + current = ( + raw_previous[positions] + + torch.randn(len(positions), 1, 65, generator=generator).to(dtype) * 0.3 + ) + expected_payload, decoded_delta = scalar_reference( + (current - expected[positions]).to(device), kind, 64 + ) + saved_previous = reconstructed.clone() + residual = encode_residual( + current.to(device), + reconstructed, + kind=kind, + positions=positions.to(device), + ) + expected[positions] += decoded_delta + raw_previous[positions] = current + previous_input = reconstructed + reconstructed = decode_residual(residual, previous_input) + torch.testing.assert_close(previous_input, saved_previous, rtol=0, atol=0) + torch.testing.assert_close(residual.data.payload.cpu(), expected_payload, rtol=0, atol=0) + torch.testing.assert_close(reconstructed.cpu(), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "index_dtype", [torch.int8, torch.int16, torch.int32, torch.int64, torch.float32] +) +@pytest.mark.parametrize("contiguous", [False, True]) +def test_sparse_rows_use_compact_key_groups_and_clone_position_map(index_dtype, contiguous): + previous = torch.full((8, 1, 2), 10.0) + positions = torch.tensor([1, 0, 6, 0, 7, 0], dtype=index_dtype)[::2] + if contiguous: + positions = positions.contiguous() + current = torch.tensor([10.0, 25.0, 13.0]).reshape(3, 1, 1).expand(3, 1, 2) + residual = encode_residual(current, previous, kind="key", positions=positions, group_size=2) + assert residual.data.payload.tolist() == [[[0x00]], [[0xFF]], [[0x00]]] + assert residual.num_tokens == 8 + assert residual.positions.dtype == torch.int32 + assert residual.positions.device == current.device + assert residual.positions.is_contiguous() + assert residual.storage_bytes == residual.data.storage_bytes + 3 * 4 + assert residual.positions.untyped_storage().nbytes() == 3 * 4 + positions.fill_(0) + expected = previous.clone() + expected[[1, 6, 7]] = current + actual = decode_residual(residual, previous) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(previous, torch.full_like(previous, 10), rtol=0, atol=0) + + +def test_sparse_absence_differs_from_present_zero_codes(): + previous = torch.arange(24.0).reshape(6, 1, 4) + positions = torch.tensor([1, 4]) + current = previous[positions] + 2 + residual = encode_residual(current, previous, kind="value", positions=positions, group_size=4) + assert torch.count_nonzero(residual.data.payload) == 0 + expected = previous.clone() + expected[positions] += 2 + torch.testing.assert_close(decode_residual(residual, previous), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_sparse_reconstruction_preserves_subnormal_values(device, index_dtype): + tiny = torch.finfo(torch.float32).tiny / 2 + previous = torch.tensor([tiny, -tiny, 1, -1, -0.0, 0.0, tiny, -tiny], device=device).reshape( + 4, 1, 2 + ) + saved_previous = previous.clone() + positions = torch.tensor([0, 2], device=device, dtype=index_dtype) + residual = encode_residual(previous[positions], previous, kind="value", positions=positions) + expected = previous.clone() + expected[2].zero_() + assert torch.equal( + decode_residual(residual, previous).view(torch.int32), expected.view(torch.int32) + ) + assert torch.equal(previous.view(torch.int32), saved_previous.view(torch.int32)) + + +@pytest.mark.parametrize("kind", ["key", "value"]) +def test_empty_residual_inherits_all_previous_rows(kind): + previous = torch.arange(24.0).reshape(6, 1, 4) + residual = encode_residual( + previous[:0], previous, kind=kind, positions=torch.empty(0, dtype=torch.long) + ) + torch.testing.assert_close(decode_residual(residual, previous), previous, rtol=0, atol=0) + assert residual.num_tokens == 6 + assert residual.storage_bytes == 0 + + +@pytest.mark.parametrize( + "values", + [ + torch.ones(2, 3), + torch.ones(2, 1, 0), + torch.ones(2, 1, 3, dtype=torch.int32), + ], +) +def test_invalid_values_shape_and_dtype(values): + with pytest.raises((ValueError, TypeError)): + pack_kv(values, kind="value") + + +@pytest.mark.parametrize( + "options", + [ + {"group_size": 0}, + {"group_size": -2}, + {"kind": "other"}, + ], +) +def test_invalid_codec_options(options): + with pytest.raises((ValueError, TypeError)): + pack_kv(torch.ones(2, 1, 4), **{"kind": "value", **options}) + + +@pytest.mark.parametrize("group_size", [True, 2.5]) +def test_group_size_conversion_matches_flashloop(group_size): + values = torch.tensor([[[0.0, 0.5, 1.0, 15.0]]]) + packed = pack_kv(values, kind="value", group_size=group_size) + expected_payload, expected = scalar_reference(values, "value", int(group_size)) + assert packed.group_size == int(group_size) + torch.testing.assert_close(packed.payload, expected_payload, rtol=0, atol=0) + torch.testing.assert_close(unpack_kv(packed), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "positions", + [ + torch.tensor([[1, 2]]), + torch.tensor([-1, 2]), + torch.tensor([1, 6]), + torch.tensor([2, 1]), + torch.tensor([2, 2]), + torch.tensor([1]), + ], +) +def test_invalid_sparse_positions(positions): + with pytest.raises((ValueError, TypeError)): + encode_residual(torch.ones(2, 1, 4), torch.ones(6, 1, 4), kind="key", positions=positions) + + +def test_residual_shape_and_dtype_mismatches(): + previous = torch.ones(3, 1, 4) + for current in (torch.ones(2, 1, 4), torch.ones(3, 2, 4), previous.half()): + with pytest.raises((ValueError, TypeError)): + encode_residual(current, previous, kind="value") + residual = encode_residual(previous + 1, previous, kind="value") + for wrong in (previous[:2], previous.half(), torch.tensor(1.0)): + with pytest.raises((ValueError, TypeError)): + encode_residual(previous, wrong, kind="value") + with pytest.raises((ValueError, TypeError)): + encode_residual(previous, wrong, kind="value", positions=torch.arange(3)) + with pytest.raises((ValueError, TypeError)): + decode_residual(residual, wrong) + + +@pytest.mark.gpu +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_cuda_codec_and_sparse_residual_match_cpu_equations(kind, dtype): + previous = torch.randn(7, 2, 5, generator=torch.Generator().manual_seed(12)).to(dtype) + positions = torch.tensor([0, 3, 6]) + current = previous[positions] + 0.2 + expected_payload, decoded_delta = scalar_reference( + (current - previous[positions]).cuda(), kind, 2 + ) + expected = previous.clone() + expected[positions] += decoded_delta + residual = encode_residual( + current.cuda(), previous.cuda(), kind=kind, positions=positions.cuda(), group_size=2 + ) + assert residual.data.payload.is_cuda and residual.positions.is_cuda + torch.testing.assert_close(residual.data.payload.cpu(), expected_payload, rtol=0, atol=0) + torch.testing.assert_close( + decode_residual(residual, previous.cuda()).cpu(), expected, rtol=0, atol=0 + ) + transferred = encode_residual( + current.cuda(), previous.cuda(), kind=kind, positions=positions, group_size=2 + ) + assert transferred.positions.is_cuda and transferred.positions.dtype == torch.int32 + torch.testing.assert_close(transferred.data.payload, residual.data.payload, rtol=0, atol=0) + torch.testing.assert_close( + decode_residual(transferred, previous.cuda()).cpu(), expected, rtol=0, atol=0 + ) diff --git a/vllm_rlt/kernels/kv_quantization.py b/vllm_rlt/kernels/kv_quantization.py new file mode 100644 index 00000000..4ec8e1c1 --- /dev/null +++ b/vllm_rlt/kernels/kv_quantization.py @@ -0,0 +1,208 @@ +"""Affine INT4 KV packing and cross-loop residual reconstruction. + +Inputs are logical ``[token, KV head, channel]`` tensors for one request/layer. +K is grouped across stored token rows per channel; V across channels per token. + +Adapted from FlashLoop packing.py and cache.py at aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6. +Copyright (c) 2026 FlashLoop contributors. MIT License. +""" + +from dataclasses import dataclass +from typing import Literal + +import torch + +KVKind = Literal["key", "value"] +_DTYPES = {torch.float16, torch.bfloat16, torch.float32, torch.float64} + + +def _validate_inputs(values: torch.Tensor, *, kind: KVKind, group_size: int) -> None: + if group_size <= 0: + raise ValueError("group_size must be positive") + if kind not in {"key", "value"}: + raise ValueError("kind must be 'key' or 'value'") + if values.ndim != 3 or values.shape[2] <= 0: + raise ValueError("KV values must have shape [tokens, heads, positive channels]") + if values.dtype not in _DTYPES: + raise ValueError("KV values must use a floating-point dtype") + + +def _codes_for_group( + values: torch.Tensor, *, reduce_dim: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + work = values.float() + minimum = work.amin(dim=reduce_dim, keepdim=True) + span = work.amax(dim=reduce_dim, keepdim=True) - minimum + scale = torch.where(span > 0, span / 15.0, torch.ones_like(span)) + codes = (work - minimum).div_(scale).round_().clamp_(0, 15).to(torch.uint8) + return codes, minimum.to(torch.float16), scale.to(torch.float16) + + +def _pack_channel_nibbles(codes: torch.Tensor) -> torch.Tensor: + payload = codes[..., 0::2].clone(memory_format=torch.contiguous_format) + payload[..., : codes.shape[-1] // 2] |= codes[..., 1::2] << 4 + return payload + + +def _unpack_channel_nibbles(payload: torch.Tensor, channels: int) -> torch.Tensor: + codes = torch.empty((*payload.shape[:-1], channels), dtype=torch.float32, device=payload.device) + codes[..., 0::2] = payload & 0x0F + codes[..., 1::2] = (payload >> 4)[..., : channels // 2] + return codes + + +@dataclass(frozen=True) +class PackedInt4KV: + """Encoded tensor and FP16 affine parameters; treat its tensors as immutable. + + ``payload`` is uint8 ``[T, H, ceil(D/2)]``. For K, parameters are + ``[ceil(T/G), H, D]``; for V, ``[T, H, ceil(D/G)]``. An odd last channel + leaves the high nibble zero. Constant groups use code zero and scale one. + """ + + payload: torch.Tensor + minimum: torch.Tensor + scale: torch.Tensor + original_shape: tuple[int, int, int] + original_dtype: torch.dtype + kind: KVKind + group_size: int + + @property + def storage_bytes(self) -> int: + """Payload and parameter bytes, excluding allocator overhead and decode scratch.""" + return sum(t.numel() * t.element_size() for t in (self.payload, self.minimum, self.scale)) + + +@torch.no_grad() +def pack_kv(values: torch.Tensor, *, kind: KVKind, group_size: int = 64) -> PackedInt4KV: + """Quantize logical KV in FP32, storing uint8 payload and FP16 parameters.""" + group_size = int(group_size) + _validate_inputs(values, kind=kind, group_size=group_size) + shape = tuple(values.shape) + axis = 0 if kind == "key" else 2 + codes = torch.empty_like(values, dtype=torch.uint8) + minima, scales = [], [] + for start in range(0, shape[axis], group_size): + length = min(group_size, shape[axis] - start) + group_codes, minimum, scale = _codes_for_group( + values.narrow(axis, start, length), reduce_dim=axis + ) + codes.narrow(axis, start, length).copy_(group_codes) + minima.append(minimum) + scales.append(scale) + metadata_shape = list(shape) + metadata_shape[axis] = (shape[axis] + group_size - 1) // group_size + minimum = ( + torch.cat(minima, dim=axis) + if minima + else torch.empty(metadata_shape, device=values.device, dtype=torch.float16) + ) + scale = torch.cat(scales, dim=axis) if scales else torch.empty_like(minimum) + payload = _pack_channel_nibbles(codes) + return PackedInt4KV(payload, minimum, scale, shape, values.dtype, kind, group_size) + + +@torch.no_grad() +def unpack_kv(packed: PackedInt4KV) -> torch.Tensor: + """Reconstruct with stored parameters in FP32, then cast to compute dtype.""" + codes = _unpack_channel_nibbles(packed.payload, packed.original_shape[2]) + axis = 0 if packed.kind == "key" else 2 + length = packed.original_shape[axis] + group_size = min(packed.group_size, length) + minimum = packed.minimum.repeat_interleave(group_size, axis).narrow(axis, 0, length) + scale = packed.scale.repeat_interleave(group_size, axis).narrow(axis, 0, length) + return codes.mul_(scale).add_(minimum).to(packed.original_dtype) + + +def _validate_positions(positions: torch.Tensor, rows: int, tokens: int) -> None: + if positions.ndim != 1: + raise ValueError("positions must be one-dimensional") + if positions.numel() != rows: + raise ValueError("positions must identify every compact residual row") + if ( + rows + and ( + (positions[0] < 0) | (positions[-1] >= tokens) | (positions[1:] <= positions[:-1]).any() + ).item() + ): + raise ValueError("positions must be strictly increasing and within the logical token range") + + +@dataclass(frozen=True) +class Int4Residual: + """One depth's compact residual and its logical scatter map. + + ``positions=None`` denotes a dense update; an empty map denotes inheritance + at every position. Positions affect scatter only: K groups follow compact + row order, as in the pinned upstream cache. + """ + + data: PackedInt4KV + positions: torch.Tensor | None + num_tokens: int + + @property + def storage_bytes(self) -> int: + mapping_bytes = ( + 0 if self.positions is None else self.positions.numel() * self.positions.element_size() + ) + return self.data.storage_bytes + mapping_bytes + + +@torch.no_grad() +def encode_residual( + current: torch.Tensor, + reconstructed_previous: torch.Tensor, + *, + kind: KVKind, + positions: torch.Tensor | None = None, + group_size: int = 64, +) -> Int4Residual: + """Encode ``current - reconstructed_previous`` in the input dtype. + + Optional positions map compact rows into the full predecessor and are copied + as contiguous int32 on the input device, matching FlashLoop's index format. + """ + if ( + current.dtype != reconstructed_previous.dtype + or current.device != reconstructed_previous.device + ): + raise ValueError("current and predecessor must have matching dtype and device") + if positions is None: + if current.shape != reconstructed_previous.shape: + raise ValueError("dense residual requires matching current and predecessor shapes") + predecessor = reconstructed_previous + else: + if current.shape[1:] != reconstructed_previous.shape[1:]: + raise ValueError("current and predecessor must have matching heads and channels") + positions = positions.to( + device=current.device, + dtype=torch.int32, + copy=True, + memory_format=torch.contiguous_format, + ) + _validate_positions(positions, current.shape[0], reconstructed_previous.shape[0]) + predecessor = reconstructed_previous.index_select(0, positions) + data = pack_kv(current - predecessor, kind=kind, group_size=group_size) + return Int4Residual(data, positions, reconstructed_previous.shape[0]) + + +@torch.no_grad() +def decode_residual(residual: Int4Residual, previous: torch.Tensor) -> torch.Tensor: + """Return a new full tensor; absent rows retain their previous values.""" + expected_shape = (residual.num_tokens, *residual.data.original_shape[1:]) + if previous.shape != expected_shape: + raise ValueError("predecessor shape does not match residual logical extent") + if ( + previous.dtype != residual.data.original_dtype + or previous.device != residual.data.payload.device + ): + raise ValueError("predecessor dtype and device must match residual") + delta = unpack_kv(residual.data) + if residual.positions is None: + return previous + delta + result = previous.clone() + indices = residual.positions.long() + result.index_copy_(0, indices, previous.index_select(0, indices).add_(delta)) + return result From d08b565f216bd887fad4c8330f7482062e146973 Mon Sep 17 00:00:00 2001 From: Howard Yang <284181151+HowardYangAI@users.noreply.github.com> Date: Wed, 4 Nov 2026 16:07:53 +0800 Subject: [PATCH 2/2] feat: add selectable NVFP4 KV reference codec --- NOTICE | 2 +- ...zation.py => test_int4_kv_quantization.py} | 0 tests/test_nvfp4_kv_quantization.py | 403 ++++++++++++++++++ vllm_rlt/kernels/kv_quantization.py | 208 --------- vllm_rlt/kernels/kv_quantization/__init__.py | 138 ++++++ vllm_rlt/kernels/kv_quantization/_common.py | 38 ++ vllm_rlt/kernels/kv_quantization/int4.py | 93 ++++ vllm_rlt/kernels/kv_quantization/nvfp4.py | 122 ++++++ 8 files changed, 795 insertions(+), 209 deletions(-) rename tests/{test_kv_quantization.py => test_int4_kv_quantization.py} (100%) create mode 100644 tests/test_nvfp4_kv_quantization.py delete mode 100644 vllm_rlt/kernels/kv_quantization.py create mode 100644 vllm_rlt/kernels/kv_quantization/__init__.py create mode 100644 vllm_rlt/kernels/kv_quantization/_common.py create mode 100644 vllm_rlt/kernels/kv_quantization/int4.py create mode 100644 vllm_rlt/kernels/kv_quantization/nvfp4.py diff --git a/NOTICE b/NOTICE index 8b7cefe2..21e5eab1 100644 --- a/NOTICE +++ b/NOTICE @@ -13,7 +13,7 @@ vllm-rlt modifies the execution structure, configuration handling, model loading and attention integration for standalone loop-level scheduling and depth-aware paged KV caching. It does not include an upstream CDB runtime implementation. -The affine INT4 reference codec in vllm_rlt/kernels/kv_quantization.py adapts +The affine INT4 reference codec in vllm_rlt/kernels/kv_quantization/ adapts FlashLoop's packing.py and cross-loop residual arithmetic in cache.py, revision aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6, to vllm-rlt's logical KV tensor layout. https://github.com/Superone77/FlashLoop diff --git a/tests/test_kv_quantization.py b/tests/test_int4_kv_quantization.py similarity index 100% rename from tests/test_kv_quantization.py rename to tests/test_int4_kv_quantization.py diff --git a/tests/test_nvfp4_kv_quantization.py b/tests/test_nvfp4_kv_quantization.py new file mode 100644 index 00000000..5058017b --- /dev/null +++ b/tests/test_nvfp4_kv_quantization.py @@ -0,0 +1,403 @@ +import numpy as np +import pytest +import torch + +from vllm_rlt.kernels.kv_quantization import ( + KVResidual, + PackedInt4KV, + PackedNVFP4KV, + decode_residual, + encode_residual, + pack_kv, + unpack_kv, +) +from vllm_rlt.kernels.kv_quantization.int4 import PackedInt4KV as BackendPackedInt4KV + +_FP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_FP8_VALUES = tuple( + fraction / 512 if exponent == 0 else (1 + fraction / 8) * 2 ** (exponent - 7) + for exponent in range(16) + for fraction in range(8 if exponent < 15 else 7) +) + + +def _nearest_code(value, table): + return min(range(len(table)), key=lambda code: (abs(value - table[code]), code % 2)) + + +def scalar_nvfp4(values, kind): + """Enumerate finite E2M1/E4M3 values independently of Torch's casts/kernels.""" + source = values.float().cpu().numpy() + tokens, heads, channels = source.shape + maximum = np.max(np.abs(source)) if source.size else np.float32(0) + with np.errstate(over="ignore"): + global_scale = ( + np.float32(np.float32(np.float32(1) / maximum) * np.float32(2688)) + if maximum + else np.float32(1) + ) + global_scale = min(global_scale, np.float32(np.finfo(np.float32).max)) + reciprocal_global_scale = np.float32(np.float32(1) / global_scale) + payload = np.zeros((tokens, heads, (channels + 1) // 2), dtype=np.uint8) + scale_shape = ( + ((tokens + 15) // 16, heads, channels) + if kind == "key" + else (tokens, heads, (channels + 15) // 16) + ) + scales = np.zeros(scale_shape, dtype=np.uint8) + decoded = np.empty(source.shape, dtype=np.float32) + for head in range(heads): + for fixed in range(channels if kind == "key" else tokens): + for start in range(0, tokens if kind == "key" else channels, 16): + if kind == "key": + indices = [(row, head, fixed) for row in range(start, min(start + 16, tokens))] + scale_index = (start // 16, head, fixed) + else: + indices = [ + (fixed, head, col) for col in range(start, min(start + 16, channels)) + ] + scale_index = (fixed, head, start // 16) + group = np.array([source[index] for index in indices]) + block_maximum = np.max(np.abs(group)) + scale = np.float32(global_scale * np.float32(block_maximum * np.float32(1 / 6))) + scale_code = _nearest_code(min(float(scale), 448), _FP8_VALUES) + scales[scale_index] = scale_code + stored_scale = np.float32(_FP8_VALUES[scale_code]) + effective_scale = np.float32(stored_scale * reciprocal_global_scale) + decode_scale = np.float32(stored_scale / global_scale) + with np.errstate(over="ignore"): + multiplier = ( + np.float32(np.float32(1) / effective_scale) + if effective_scale + else np.float32(0) + ) + for index, value in zip(indices, group): + # Robustness extension: a positive scale can have an + # unrepresentable reciprocal while value / scale is finite. + normalized = np.float32( + value / effective_scale + if not np.isfinite(multiplier) + else value * multiplier + ) + code = _nearest_code(abs(float(normalized)), _FP4_VALUES) + negative = np.signbit(normalized) + signed_code = code | (8 if negative else 0) + token, _, channel = index + payload[token, head, channel // 2] |= signed_code << (4 * (channel % 2)) + quantized = np.float32(-_FP4_VALUES[code] if negative else _FP4_VALUES[code]) + decoded[token, head, channel] = np.float32(quantized * decode_scale) + return payload, scales, global_scale, torch.from_numpy(decoded).to(values.dtype) + + +def assert_matches_scalar(values, kind): + expected_payload, expected_scale, expected_global_scale, expected = scalar_nvfp4(values, kind) + packed = pack_kv(values, kind=kind, format="nvfp4") + assert isinstance(packed, PackedNVFP4KV) + assert packed.format == "nvfp4" and packed.axis == (0 if kind == "key" else 2) + assert packed.group_size == 16 + assert packed.original_shape == tuple(values.shape) and packed.original_dtype == values.dtype + assert packed.payload.dtype == torch.uint8 + assert packed.scale.dtype == torch.float8_e4m3fn + assert packed.global_scale.dtype == torch.float32 and packed.global_scale.ndim == 0 + torch.testing.assert_close( + packed.payload.cpu(), torch.from_numpy(expected_payload), rtol=0, atol=0 + ) + torch.testing.assert_close( + packed.scale.cpu().view(torch.uint8), torch.from_numpy(expected_scale), rtol=0, atol=0 + ) + assert packed.global_scale.item() == expected_global_scale + torch.testing.assert_close(unpack_kv(packed).cpu(), expected, rtol=0, atol=0) + tokens, heads, channels = values.shape + scale_bytes = ( + ((tokens + 15) // 16) * heads * channels + if kind == "key" + else tokens * heads * ((channels + 15) // 16) + ) + assert packed.storage_bytes == tokens * heads * ((channels + 1) // 2) + scale_bytes + 4 + return packed + + +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +@pytest.mark.parametrize( + "shape", + [ + (5, 2, 17), + (2, 1, 16), + (1, 2, 31), + (2, 1, 1), + (0, 2, 17), + (3, 0, 17), + (16, 1, 3), + (17, 2, 3), + (33, 2, 17), + ], +) +def test_nvfp4_layout_partial_empty_and_noncontiguous(shape, dtype, kind): + values = torch.randn( + (*shape[:2], shape[2] * 2), generator=torch.Generator().manual_seed(831) + ).to(dtype)[..., ::2] + original = values.clone() + packed = assert_matches_scalar(values, kind) + torch.testing.assert_close(values, original, rtol=0, atol=0) + if shape[2] % 2: + assert torch.count_nonzero(packed.payload[..., -1] >> 4) == 0 + + +@pytest.mark.parametrize("kind", ["key", "value"]) +def test_nvfp4_group_axis_and_partial_tail(kind): + values = torch.zeros(33, 1, 2) + values[:16, 0, 0] = 6 + values[16:32, 0, 0] = 12 + values[32, 0, 0] = 24 + values[:, 0, 1] = 2688 + if kind == "value": + values = values.transpose(0, 2) + packed = assert_matches_scalar(values, kind) + scales = packed.scale[:, 0, 0] if kind == "key" else packed.scale[0, 0, :] + assert packed.global_scale.item() == 1 + torch.testing.assert_close(scales.float(), torch.tensor([1.0, 2.0, 4.0]), rtol=0, atol=0) + torch.testing.assert_close(unpack_kv(packed), values, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_nvfp4_zero_groups_and_negative_zero(dtype): + values = torch.zeros(2, 1, 17, dtype=dtype) + values[..., 1::2] = -0.0 + values[..., -1] = -0.0 + packed = assert_matches_scalar(values, "key") + assert packed.global_scale.item() == 1 + assert torch.count_nonzero(packed.scale.float()) == 0 + assert packed.payload[0, 0, 0].item() == 0x80 + assert packed.payload[0, 0, -1].item() == 0x08 + assert torch.equal(torch.signbit(unpack_kv(packed)), torch.signbit(values)) + + +def test_nvfp4_tiny_inputs_keep_scales_and_reconstruction_finite(): + values = torch.tensor([1e-42, -1e-42, 0.0, -0.0]).repeat(4).reshape(1, 1, 16) + packed = pack_kv(values, kind="key", format="nvfp4") + decoded = unpack_kv(packed) + assert torch.isfinite(packed.global_scale) and torch.isfinite(packed.scale.float()).all() + assert torch.isfinite(decoded).all() and torch.count_nonzero(decoded) == 0 + assert torch.equal(torch.signbit(decoded), torch.signbit(values)) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +@pytest.mark.parametrize("anchor", [False, True]) +def test_nvfp4_tiny_values_survive_reciprocal_overflow(dtype, anchor, device): + tiny = 2**-24 if dtype == torch.float16 else 1e-38 + values = torch.tensor([tiny, -tiny, 0.0, -0.0], dtype=dtype).repeat(4).reshape(1, 1, 16) + if anchor: + anchored = torch.zeros_like(values) + anchored[..., 0] = tiny * (16 if dtype == torch.float16 else 800) + values = torch.cat((values, anchored), dim=2) + packed = pack_kv(values.to(device), kind="key", format="nvfp4") + decoded = unpack_kv(packed).cpu() + if device == "cpu": + expected_payload, expected_scale, _, expected = scalar_nvfp4(values, "key") + torch.testing.assert_close(decoded, expected, rtol=0, atol=0) + torch.testing.assert_close( + packed.payload.cpu(), torch.from_numpy(expected_payload), rtol=0, atol=0 + ) + torch.testing.assert_close( + packed.scale.cpu().view(torch.uint8), torch.from_numpy(expected_scale), rtol=0, atol=0 + ) + assert torch.isfinite(decoded).all() + nonzero = values != 0 + assert (decoded[nonzero] != 0).all() + assert ((decoded[nonzero].float() / values[nonzero].float() - 1).abs() < 0.1).all() + assert torch.equal(torch.signbit(decoded), torch.signbit(values)) + + +@pytest.mark.parametrize("kind", ["key", "value"]) +def test_nvfp4_fp4_midpoints_nextafter_and_even_codes(kind): + midpoints = np.array([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5], dtype=np.float32) + probes = np.stack( + (np.nextafter(midpoints, -np.inf), midpoints, np.nextafter(midpoints, np.inf)), axis=1 + ).reshape(-1) + values = torch.zeros(len(probes) + 1, 1, 16) + values[:-1, 0, 0] = torch.from_numpy(probes) + values[:-1, 0, 1] = -torch.from_numpy(probes) + values[:-1, 0, 2] = 6 + values[-1, 0, 0] = 2688 # Fix g=1 and each probe's effective scale=1. + if kind == "key": + values = values.transpose(0, 2) + packed = assert_matches_scalar(values, kind) + assert packed.global_scale.item() == 1 + if kind == "key": + ties = torch.stack( + [ + (packed.payload[0, 0, channel // 2] >> (4 * (channel % 2))) & 0x0F + for channel in range(1, len(probes), 3) + ] + ) + else: + ties = packed.payload[1:-1:3, 0, 0] & 0x0F + assert ties.tolist() == [0, 2, 2, 4, 4, 6, 6] + + +@pytest.mark.parametrize("source_dtype", [torch.float16, torch.bfloat16]) +def test_nvfp4_all_low_precision_magnitudes_choose_nearest_even_code(source_dtype): + bits = torch.arange(32768, dtype=torch.int32).to(torch.int16) + probes = bits.view(source_dtype).float() + probes = probes[torch.isfinite(probes) & (probes <= 6)] + values = torch.zeros(len(probes) + 1, 1, 16) + values[:-1, 0, 0] = probes + values[:-1, 0, 1] = -probes + values[:-1, 0, 2] = 6 + values[-1, 0, 0] = 2688 # Keep global and probe block scales exactly one. + expected = torch.tensor( + [_nearest_code(value, _FP4_VALUES) for value in probes.tolist()], dtype=torch.uint8 + ) + packed = pack_kv(values, kind="value", format="nvfp4") + assert packed.global_scale.item() == 1 + assert torch.equal(packed.payload[:-1, 0, 0], expected | ((expected | 8) << 4)) + + +@pytest.mark.parametrize("kind", ["key", "value"]) +def test_nvfp4_fp8_scale_rounding_and_finite_maximum(kind): + targets = np.array([0.0001, 0.0009765625, 1.0625, 1.1875, 432, 448], dtype=np.float32) + probes = np.stack( + (np.nextafter(targets, -np.inf), targets, np.nextafter(targets, np.inf)), axis=1 + ).reshape(-1) + values = torch.zeros(len(probes) + 1, 1, 16) + values[:-1, 0, 0] = torch.from_numpy(probes) * 6 + values[-1, 0, 0] = 2688 + if kind == "key": + values = values.transpose(0, 2) + packed = assert_matches_scalar(values, kind) + anchor = (0, 0, -1) if kind == "key" else (-1, 0, 0) + assert packed.scale[anchor].float().item() == 448 + assert torch.isfinite(packed.scale.float()).all() + assert 0 in packed.scale.view(torch.uint8) + + +@pytest.mark.parametrize("group_size", [1, 8, 32, 64]) +@pytest.mark.parametrize("kind", ["key", "value"]) +def test_nvfp4_requires_sixteen_element_groups(group_size, kind): + with pytest.raises(ValueError): + pack_kv(torch.ones(1, 1, 16), kind=kind, format="nvfp4", group_size=group_size) + + +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_nvfp4_dense_and_nested_sparse_residuals(kind, dtype): + generator = torch.Generator().manual_seed(712) + raw_previous = torch.randn(65, 2, 17, generator=generator).to(dtype) + reconstructed = unpack_kv(pack_kv(raw_previous, kind=kind, format="nvfp4")) + _, _, _, expected = scalar_nvfp4(raw_previous, kind) + torch.testing.assert_close(reconstructed, expected, rtol=0, atol=0) + active = torch.arange(0, 65, 3) + for positions in (None, active, active[::2], active[:0]): + raw_rows = raw_previous if positions is None else raw_previous[positions] + current = raw_rows + torch.randn(raw_rows.shape, generator=generator).to(dtype) * 0.3 + previous_rows = expected if positions is None else expected[positions] + _, _, _, expected_delta = scalar_nvfp4(current - previous_rows, kind) + saved_current, saved_previous = current.clone(), reconstructed.clone() + residual = encode_residual( + current, reconstructed, kind=kind, format="nvfp4", positions=positions + ) + assert isinstance(residual, KVResidual) and isinstance(residual.data, PackedNVFP4KV) + previous_input = reconstructed + reconstructed = decode_residual(residual, previous_input) + torch.testing.assert_close(current, saved_current, rtol=0, atol=0) + torch.testing.assert_close(previous_input, saved_previous, rtol=0, atol=0) + mapping_bytes = 0 if positions is None else 4 * positions.numel() + assert residual.storage_bytes == residual.data.storage_bytes + mapping_bytes + if positions is None: + expected = expected + expected_delta + raw_previous = current + else: + inactive = torch.ones(raw_previous.shape[0], dtype=torch.bool) + inactive[positions] = False + torch.testing.assert_close( + reconstructed[inactive], saved_previous[inactive], rtol=0, atol=0 + ) + expected[positions] = expected[positions] + expected_delta + raw_previous[positions] = current + torch.testing.assert_close(reconstructed, expected, rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.parametrize("kind", ["key", "value"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("tokens", [7, 65]) +def test_nvfp4_cuda_matches_scalar_and_sparse_reconstruction(dtype, kind, tokens): + previous = torch.randn(tokens, 2, 17, generator=torch.Generator().manual_seed(68)).to(dtype) + packed = assert_matches_scalar(previous.cuda(), kind) + previous = unpack_kv(packed).cpu() + positions = torch.arange(0, tokens, 3) + current = previous[positions] + 0.2 + _, _, _, delta = scalar_nvfp4(current - previous[positions], kind) + expected = previous.clone() + expected[positions] += delta + residual = encode_residual( + current.cuda(), previous.cuda(), kind=kind, format="nvfp4", positions=positions.cuda() + ) + torch.testing.assert_close( + decode_residual(residual, previous.cuda()).cpu(), expected, rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("kind,axis", [("key", 0), ("value", 2)]) +def test_default_int4_selection_and_exports(kind, axis): + assert PackedInt4KV is BackendPackedInt4KV + values = torch.randn(65, 1, 17, generator=torch.Generator().manual_seed(412)) + legacy = pack_kv(values, kind=kind) + explicit = pack_kv(values, kind=kind, format="int4") + assert isinstance(legacy, PackedInt4KV) + assert legacy.format == explicit.format == "int4" + assert legacy.axis == explicit.axis == axis + assert legacy.group_size == explicit.group_size == 64 + for name in ("payload", "minimum", "scale"): + torch.testing.assert_close(getattr(legacy, name), getattr(explicit, name), rtol=0, atol=0) + torch.testing.assert_close(unpack_kv(legacy), unpack_kv(explicit), rtol=0, atol=0) + residual = encode_residual(values, unpack_kv(legacy), kind=kind) + assert isinstance(residual, KVResidual) and isinstance(residual.data, PackedInt4KV) + + +def test_bfloat16_mixed_format_residual_chain(): + generator = torch.Generator().manual_seed(614) + raw_previous = torch.randn(9, 2, 17, generator=generator).bfloat16() + base = pack_kv(raw_previous, kind="key", group_size=None) + reconstructed = unpack_kv(base) + expected = reconstructed.clone() + torch.testing.assert_close(reconstructed, raw_previous, rtol=0, atol=0.25) + formats = [base.format] + for format, positions in (("nvfp4", None), ("int4", torch.tensor([0, 3, 7]))): + raw_rows = raw_previous if positions is None else raw_previous[positions] + current = raw_rows + torch.randn(raw_rows.shape, generator=generator).bfloat16() * 0.3 + previous_rows = expected if positions is None else expected[positions] + delta = current - previous_rows + expected_delta = unpack_kv(pack_kv(delta, kind="key", format=format, group_size=None)) + residual = encode_residual( + current, reconstructed, kind="key", format=format, group_size=None, positions=positions + ) + formats.append(residual.data.format) + assert residual.data.group_size == (16 if format == "nvfp4" else 64) + previous_input = reconstructed + saved_previous = previous_input.clone() + reconstructed = decode_residual(residual, previous_input) + torch.testing.assert_close(previous_input, saved_previous, rtol=0, atol=0) + if positions is None: + expected = expected + expected_delta + raw_previous = current + else: + expected[positions] = expected[positions] + expected_delta + inactive = torch.ones(9, dtype=torch.bool) + inactive[positions] = False + torch.testing.assert_close( + reconstructed[inactive], saved_previous[inactive], rtol=0, atol=0 + ) + assert reconstructed.shape == raw_previous.shape + assert reconstructed.dtype == torch.bfloat16 + torch.testing.assert_close(reconstructed, expected, rtol=0, atol=0) + active_rows = reconstructed if positions is None else reconstructed[positions] + torch.testing.assert_close(active_rows, current, rtol=0, atol=0.25) + assert formats == ["int4", "nvfp4", "int4"] + + +def test_unknown_format_is_rejected(): + with pytest.raises(ValueError, match="format"): + pack_kv(torch.ones(1, 1, 16), kind="key", format="unknown") diff --git a/vllm_rlt/kernels/kv_quantization.py b/vllm_rlt/kernels/kv_quantization.py deleted file mode 100644 index 4ec8e1c1..00000000 --- a/vllm_rlt/kernels/kv_quantization.py +++ /dev/null @@ -1,208 +0,0 @@ -"""Affine INT4 KV packing and cross-loop residual reconstruction. - -Inputs are logical ``[token, KV head, channel]`` tensors for one request/layer. -K is grouped across stored token rows per channel; V across channels per token. - -Adapted from FlashLoop packing.py and cache.py at aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6. -Copyright (c) 2026 FlashLoop contributors. MIT License. -""" - -from dataclasses import dataclass -from typing import Literal - -import torch - -KVKind = Literal["key", "value"] -_DTYPES = {torch.float16, torch.bfloat16, torch.float32, torch.float64} - - -def _validate_inputs(values: torch.Tensor, *, kind: KVKind, group_size: int) -> None: - if group_size <= 0: - raise ValueError("group_size must be positive") - if kind not in {"key", "value"}: - raise ValueError("kind must be 'key' or 'value'") - if values.ndim != 3 or values.shape[2] <= 0: - raise ValueError("KV values must have shape [tokens, heads, positive channels]") - if values.dtype not in _DTYPES: - raise ValueError("KV values must use a floating-point dtype") - - -def _codes_for_group( - values: torch.Tensor, *, reduce_dim: int -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - work = values.float() - minimum = work.amin(dim=reduce_dim, keepdim=True) - span = work.amax(dim=reduce_dim, keepdim=True) - minimum - scale = torch.where(span > 0, span / 15.0, torch.ones_like(span)) - codes = (work - minimum).div_(scale).round_().clamp_(0, 15).to(torch.uint8) - return codes, minimum.to(torch.float16), scale.to(torch.float16) - - -def _pack_channel_nibbles(codes: torch.Tensor) -> torch.Tensor: - payload = codes[..., 0::2].clone(memory_format=torch.contiguous_format) - payload[..., : codes.shape[-1] // 2] |= codes[..., 1::2] << 4 - return payload - - -def _unpack_channel_nibbles(payload: torch.Tensor, channels: int) -> torch.Tensor: - codes = torch.empty((*payload.shape[:-1], channels), dtype=torch.float32, device=payload.device) - codes[..., 0::2] = payload & 0x0F - codes[..., 1::2] = (payload >> 4)[..., : channels // 2] - return codes - - -@dataclass(frozen=True) -class PackedInt4KV: - """Encoded tensor and FP16 affine parameters; treat its tensors as immutable. - - ``payload`` is uint8 ``[T, H, ceil(D/2)]``. For K, parameters are - ``[ceil(T/G), H, D]``; for V, ``[T, H, ceil(D/G)]``. An odd last channel - leaves the high nibble zero. Constant groups use code zero and scale one. - """ - - payload: torch.Tensor - minimum: torch.Tensor - scale: torch.Tensor - original_shape: tuple[int, int, int] - original_dtype: torch.dtype - kind: KVKind - group_size: int - - @property - def storage_bytes(self) -> int: - """Payload and parameter bytes, excluding allocator overhead and decode scratch.""" - return sum(t.numel() * t.element_size() for t in (self.payload, self.minimum, self.scale)) - - -@torch.no_grad() -def pack_kv(values: torch.Tensor, *, kind: KVKind, group_size: int = 64) -> PackedInt4KV: - """Quantize logical KV in FP32, storing uint8 payload and FP16 parameters.""" - group_size = int(group_size) - _validate_inputs(values, kind=kind, group_size=group_size) - shape = tuple(values.shape) - axis = 0 if kind == "key" else 2 - codes = torch.empty_like(values, dtype=torch.uint8) - minima, scales = [], [] - for start in range(0, shape[axis], group_size): - length = min(group_size, shape[axis] - start) - group_codes, minimum, scale = _codes_for_group( - values.narrow(axis, start, length), reduce_dim=axis - ) - codes.narrow(axis, start, length).copy_(group_codes) - minima.append(minimum) - scales.append(scale) - metadata_shape = list(shape) - metadata_shape[axis] = (shape[axis] + group_size - 1) // group_size - minimum = ( - torch.cat(minima, dim=axis) - if minima - else torch.empty(metadata_shape, device=values.device, dtype=torch.float16) - ) - scale = torch.cat(scales, dim=axis) if scales else torch.empty_like(minimum) - payload = _pack_channel_nibbles(codes) - return PackedInt4KV(payload, minimum, scale, shape, values.dtype, kind, group_size) - - -@torch.no_grad() -def unpack_kv(packed: PackedInt4KV) -> torch.Tensor: - """Reconstruct with stored parameters in FP32, then cast to compute dtype.""" - codes = _unpack_channel_nibbles(packed.payload, packed.original_shape[2]) - axis = 0 if packed.kind == "key" else 2 - length = packed.original_shape[axis] - group_size = min(packed.group_size, length) - minimum = packed.minimum.repeat_interleave(group_size, axis).narrow(axis, 0, length) - scale = packed.scale.repeat_interleave(group_size, axis).narrow(axis, 0, length) - return codes.mul_(scale).add_(minimum).to(packed.original_dtype) - - -def _validate_positions(positions: torch.Tensor, rows: int, tokens: int) -> None: - if positions.ndim != 1: - raise ValueError("positions must be one-dimensional") - if positions.numel() != rows: - raise ValueError("positions must identify every compact residual row") - if ( - rows - and ( - (positions[0] < 0) | (positions[-1] >= tokens) | (positions[1:] <= positions[:-1]).any() - ).item() - ): - raise ValueError("positions must be strictly increasing and within the logical token range") - - -@dataclass(frozen=True) -class Int4Residual: - """One depth's compact residual and its logical scatter map. - - ``positions=None`` denotes a dense update; an empty map denotes inheritance - at every position. Positions affect scatter only: K groups follow compact - row order, as in the pinned upstream cache. - """ - - data: PackedInt4KV - positions: torch.Tensor | None - num_tokens: int - - @property - def storage_bytes(self) -> int: - mapping_bytes = ( - 0 if self.positions is None else self.positions.numel() * self.positions.element_size() - ) - return self.data.storage_bytes + mapping_bytes - - -@torch.no_grad() -def encode_residual( - current: torch.Tensor, - reconstructed_previous: torch.Tensor, - *, - kind: KVKind, - positions: torch.Tensor | None = None, - group_size: int = 64, -) -> Int4Residual: - """Encode ``current - reconstructed_previous`` in the input dtype. - - Optional positions map compact rows into the full predecessor and are copied - as contiguous int32 on the input device, matching FlashLoop's index format. - """ - if ( - current.dtype != reconstructed_previous.dtype - or current.device != reconstructed_previous.device - ): - raise ValueError("current and predecessor must have matching dtype and device") - if positions is None: - if current.shape != reconstructed_previous.shape: - raise ValueError("dense residual requires matching current and predecessor shapes") - predecessor = reconstructed_previous - else: - if current.shape[1:] != reconstructed_previous.shape[1:]: - raise ValueError("current and predecessor must have matching heads and channels") - positions = positions.to( - device=current.device, - dtype=torch.int32, - copy=True, - memory_format=torch.contiguous_format, - ) - _validate_positions(positions, current.shape[0], reconstructed_previous.shape[0]) - predecessor = reconstructed_previous.index_select(0, positions) - data = pack_kv(current - predecessor, kind=kind, group_size=group_size) - return Int4Residual(data, positions, reconstructed_previous.shape[0]) - - -@torch.no_grad() -def decode_residual(residual: Int4Residual, previous: torch.Tensor) -> torch.Tensor: - """Return a new full tensor; absent rows retain their previous values.""" - expected_shape = (residual.num_tokens, *residual.data.original_shape[1:]) - if previous.shape != expected_shape: - raise ValueError("predecessor shape does not match residual logical extent") - if ( - previous.dtype != residual.data.original_dtype - or previous.device != residual.data.payload.device - ): - raise ValueError("predecessor dtype and device must match residual") - delta = unpack_kv(residual.data) - if residual.positions is None: - return previous + delta - result = previous.clone() - indices = residual.positions.long() - result.index_copy_(0, indices, previous.index_select(0, indices).add_(delta)) - return result diff --git a/vllm_rlt/kernels/kv_quantization/__init__.py b/vllm_rlt/kernels/kv_quantization/__init__.py new file mode 100644 index 00000000..d3226611 --- /dev/null +++ b/vllm_rlt/kernels/kv_quantization/__init__.py @@ -0,0 +1,138 @@ +"""Logical INT4/NVFP4 KV codecs and cross-loop residual reconstruction. + +Inputs are [token, KV head, channel] tensors for one request/layer. +Residual arithmetic adapts FlashLoop cache.py at +aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6. +Copyright (c) 2026 FlashLoop contributors. MIT License; see NOTICE. +""" + +from dataclasses import dataclass +from typing import Literal + +import torch + +from . import int4, nvfp4 +from ._common import KVKind +from .int4 import PackedInt4KV +from .nvfp4 import PackedNVFP4KV + +KVFormat = Literal["int4", "nvfp4"] +PackedKV = PackedInt4KV | PackedNVFP4KV + + +def pack_kv( + values: torch.Tensor, + *, + kind: KVKind, + format: KVFormat = "int4", + group_size: int | None = None, +) -> PackedKV: + """Encode logical KV; defaults are INT4 G=64 or NVFP4 G=16.""" + if format == "int4": + return int4.pack(values, kind=kind, group_size=64 if group_size is None else group_size) + if format == "nvfp4": + return nvfp4.pack(values, kind=kind, group_size=16 if group_size is None else group_size) + raise ValueError("format must be 'int4' or 'nvfp4'") + + +def unpack_kv(packed: PackedKV) -> torch.Tensor: + """Decode using the format carried by the packed tensor.""" + if isinstance(packed, PackedInt4KV): + return int4.unpack(packed) + if isinstance(packed, PackedNVFP4KV): + return nvfp4.unpack(packed) + raise TypeError("expected PackedInt4KV or PackedNVFP4KV") + + +def _validate_positions(positions: torch.Tensor, rows: int, tokens: int) -> None: + if positions.ndim != 1: + raise ValueError("positions must be one-dimensional") + if positions.numel() != rows: + raise ValueError("positions must identify every compact residual row") + if ( + rows + and ( + (positions[0] < 0) | (positions[-1] >= tokens) | (positions[1:] <= positions[:-1]).any() + ).item() + ): + raise ValueError("positions must be strictly increasing and within the logical token range") + + +@dataclass(frozen=True) +class KVResidual: + """One depth's compact residual and its logical scatter map. + + ``positions=None`` denotes a dense update; an empty map denotes inheritance + at every position. K groups follow compact row order for both formats; + V groups remain within each compact row, along channels. + """ + + data: PackedKV + positions: torch.Tensor | None + num_tokens: int + + @property + def storage_bytes(self) -> int: + mapping_bytes = ( + 0 if self.positions is None else self.positions.numel() * self.positions.element_size() + ) + return self.data.storage_bytes + mapping_bytes + + +@torch.no_grad() +def encode_residual( + current: torch.Tensor, + reconstructed_previous: torch.Tensor, + *, + kind: KVKind, + positions: torch.Tensor | None = None, + format: KVFormat = "int4", + group_size: int | None = None, +) -> KVResidual: + """Encode ``current - reconstructed_previous`` in the input dtype. + + Optional positions map compact rows into the full predecessor and are copied + as contiguous int32 on the input device, matching FlashLoop's index format. + """ + if ( + current.dtype != reconstructed_previous.dtype + or current.device != reconstructed_previous.device + ): + raise ValueError("current and predecessor must have matching dtype and device") + if positions is None: + if current.shape != reconstructed_previous.shape: + raise ValueError("dense residual requires matching current and predecessor shapes") + predecessor = reconstructed_previous + else: + if current.shape[1:] != reconstructed_previous.shape[1:]: + raise ValueError("current and predecessor must have matching heads and channels") + positions = positions.to( + device=current.device, + dtype=torch.int32, + copy=True, + memory_format=torch.contiguous_format, + ) + _validate_positions(positions, current.shape[0], reconstructed_previous.shape[0]) + predecessor = reconstructed_previous.index_select(0, positions) + data = pack_kv(current - predecessor, kind=kind, format=format, group_size=group_size) + return KVResidual(data, positions, reconstructed_previous.shape[0]) + + +@torch.no_grad() +def decode_residual(residual: KVResidual, previous: torch.Tensor) -> torch.Tensor: + """Return a new full tensor; absent rows retain their previous values.""" + expected_shape = (residual.num_tokens, *residual.data.original_shape[1:]) + if previous.shape != expected_shape: + raise ValueError("predecessor shape does not match residual logical extent") + if ( + previous.dtype != residual.data.original_dtype + or previous.device != residual.data.payload.device + ): + raise ValueError("predecessor dtype and device must match residual") + delta = unpack_kv(residual.data) + if residual.positions is None: + return previous + delta + result = previous.clone() + indices = residual.positions.long() + result.index_copy_(0, indices, previous.index_select(0, indices).add_(delta)) + return result diff --git a/vllm_rlt/kernels/kv_quantization/_common.py b/vllm_rlt/kernels/kv_quantization/_common.py new file mode 100644 index 00000000..177e3a5b --- /dev/null +++ b/vllm_rlt/kernels/kv_quantization/_common.py @@ -0,0 +1,38 @@ +"""Shared logical KV validation and nibble packing. + +Adapted from FlashLoop packing.py at aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6. +Copyright (c) 2026 FlashLoop contributors. MIT License; see NOTICE. +""" + +from typing import Literal + +import torch + +KVKind = Literal["key", "value"] +_DTYPES = {torch.float16, torch.bfloat16, torch.float32, torch.float64} + + +def _validate_inputs(values: torch.Tensor, *, kind: KVKind, group_size: int) -> None: + if group_size <= 0: + raise ValueError("group_size must be positive") + if kind not in {"key", "value"}: + raise ValueError("kind must be 'key' or 'value'") + if values.ndim != 3 or values.shape[2] <= 0: + raise ValueError("KV values must have shape [tokens, heads, positive channels]") + if values.dtype not in _DTYPES: + raise ValueError("KV values must use a floating-point dtype") + + +def _pack_channel_nibbles(codes: torch.Tensor) -> torch.Tensor: + payload = codes[..., 0::2].clone(memory_format=torch.contiguous_format) + payload[..., : codes.shape[-1] // 2] |= codes[..., 1::2] << 4 + return payload + + +def _unpack_channel_nibbles( + payload: torch.Tensor, channels: int, *, dtype: torch.dtype = torch.float32 +) -> torch.Tensor: + codes = torch.empty((*payload.shape[:-1], channels), dtype=dtype, device=payload.device) + codes[..., 0::2] = payload & 0x0F + codes[..., 1::2] = (payload >> 4)[..., : channels // 2] + return codes diff --git a/vllm_rlt/kernels/kv_quantization/int4.py b/vllm_rlt/kernels/kv_quantization/int4.py new file mode 100644 index 00000000..b9b7a7e0 --- /dev/null +++ b/vllm_rlt/kernels/kv_quantization/int4.py @@ -0,0 +1,93 @@ +"""Affine INT4 reference codec for logical [token, KV head, channel] tensors. + +Adapted from FlashLoop packing.py at aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6. +Copyright (c) 2026 FlashLoop contributors. MIT License; see NOTICE. +""" + +from dataclasses import dataclass +from typing import ClassVar, Literal + +import torch + +from ._common import KVKind, _pack_channel_nibbles, _unpack_channel_nibbles, _validate_inputs + + +def _codes_for_group( + values: torch.Tensor, *, reduce_dim: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + work = values.float() + minimum = work.amin(dim=reduce_dim, keepdim=True) + span = work.amax(dim=reduce_dim, keepdim=True) - minimum + scale = torch.where(span > 0, span / 15.0, torch.ones_like(span)) + codes = (work - minimum).div_(scale).round_().clamp_(0, 15).to(torch.uint8) + return codes, minimum.to(torch.float16), scale.to(torch.float16) + + +@dataclass(frozen=True) +class PackedInt4KV: + """Encoded tensor and FP16 affine parameters; treat its tensors as immutable. + + ``payload`` is uint8 ``[T, H, ceil(D/2)]``. For K, parameters are + ``[ceil(T/G), H, D]``; for V, ``[T, H, ceil(D/G)]``. An odd last channel + leaves the high nibble zero. Constant groups use code zero and scale one. + """ + + format: ClassVar[Literal["int4"]] = "int4" + + payload: torch.Tensor + minimum: torch.Tensor + scale: torch.Tensor + original_shape: tuple[int, int, int] + original_dtype: torch.dtype + kind: KVKind + group_size: int + + @property + def axis(self) -> int: + return 0 if self.kind == "key" else 2 + + @property + def storage_bytes(self) -> int: + """Payload and parameter bytes, excluding allocator overhead and decode scratch.""" + return sum(t.numel() * t.element_size() for t in (self.payload, self.minimum, self.scale)) + + +@torch.no_grad() +def pack(values: torch.Tensor, *, kind: KVKind, group_size: int = 64) -> PackedInt4KV: + """Quantize logical KV in FP32, storing uint8 payload and FP16 parameters.""" + group_size = int(group_size) + _validate_inputs(values, kind=kind, group_size=group_size) + shape = tuple(values.shape) + axis = 0 if kind == "key" else 2 + codes = torch.empty_like(values, dtype=torch.uint8) + minima, scales = [], [] + for start in range(0, shape[axis], group_size): + length = min(group_size, shape[axis] - start) + group_codes, minimum, scale = _codes_for_group( + values.narrow(axis, start, length), reduce_dim=axis + ) + codes.narrow(axis, start, length).copy_(group_codes) + minima.append(minimum) + scales.append(scale) + metadata_shape = list(shape) + metadata_shape[axis] = (shape[axis] + group_size - 1) // group_size + minimum = ( + torch.cat(minima, dim=axis) + if minima + else torch.empty(metadata_shape, device=values.device, dtype=torch.float16) + ) + scale = torch.cat(scales, dim=axis) if scales else torch.empty_like(minimum) + payload = _pack_channel_nibbles(codes) + return PackedInt4KV(payload, minimum, scale, shape, values.dtype, kind, group_size) + + +@torch.no_grad() +def unpack(packed: PackedInt4KV) -> torch.Tensor: + """Reconstruct with stored parameters in FP32, then cast to compute dtype.""" + codes = _unpack_channel_nibbles(packed.payload, packed.original_shape[2]) + axis = packed.axis + length = packed.original_shape[axis] + group_size = min(packed.group_size, length) + minimum = packed.minimum.repeat_interleave(group_size, axis).narrow(axis, 0, length) + scale = packed.scale.repeat_interleave(group_size, axis).narrow(axis, 0, length) + return codes.mul_(scale).add_(minimum).to(packed.original_dtype) diff --git a/vllm_rlt/kernels/kv_quantization/nvfp4.py b/vllm_rlt/kernels/kv_quantization/nvfp4.py new file mode 100644 index 00000000..64e2c8e3 --- /dev/null +++ b/vllm_rlt/kernels/kv_quantization/nvfp4.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 reference for logical [token, KV head, channel] tensors. + +Adapted from vLLM ref_nvfp4_quant / ref_nvfp4_quant_dequant at +c7a760751977e4e3396d586956d8767a9bc44a82 in +vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py. +Unlike the hardware layout, this logical format also supports partial groups. +""" + +from dataclasses import dataclass +from typing import ClassVar, Literal + +import torch + +from ._common import KVKind, _pack_channel_nibbles, _unpack_channel_nibbles, _validate_inputs + + +@dataclass(frozen=True) +class PackedNVFP4KV: + """Uint8 payload, E4M3 block scales and an FP32 global encoding multiplier. + + K groups 16 tokens per channel; V groups 16 channels per token. + ``scale`` is [ceil(T/16), H, D] for K or [T, H, ceil(D/16)] for V. + ``global_scale`` is shared by this encoded tensor. Decode multiplies + E2M1 values by ``scale / global_scale``. Treat tensors as immutable. + """ + + format: ClassVar[Literal["nvfp4"]] = "nvfp4" + + payload: torch.Tensor + scale: torch.Tensor + global_scale: torch.Tensor + original_shape: tuple[int, int, int] + original_dtype: torch.dtype + kind: KVKind + group_size: int = 16 + + @property + def axis(self) -> int: + return 0 if self.kind == "key" else 2 + + @property + def storage_bytes(self) -> int: + return sum( + t.numel() * t.element_size() for t in (self.payload, self.scale, self.global_scale) + ) + + +@torch.no_grad() +def pack(values: torch.Tensor, *, kind: KVKind, group_size: int = 16) -> PackedNVFP4KV: + """Encode E2M1 with nearest-even rounding and FP8-rounded block scales.""" + group_size = int(group_size) + _validate_inputs(values, kind=kind, group_size=group_size) + if group_size != 16: + raise ValueError("NVFP4 requires group_size=16") + shape = tuple(values.shape) + axis = 0 if kind == "key" else 2 + global_scale = torch.ones((), device=values.device, dtype=torch.float32) + if values.numel(): + amax = values.abs().amax().float() + global_scale = torch.where(amax > 0, 2688.0 / amax, global_scale).clamp_max_( + torch.finfo(torch.float32).max + ) + inverse_global = global_scale.reciprocal() + codes = torch.empty_like(values, dtype=torch.uint8) + scale_shape = list(shape) + scale_shape[axis] = (shape[axis] + 15) // 16 + scales = torch.empty( + scale_shape, + dtype=torch.float8_e4m3fn, + device=values.device, + ) + # bucketize keeps equality in the lower bin. Move odd lower-code boundaries + # down one FP32 ULP so midpoint ties choose the even upper code instead. + boundaries = torch.tensor( + [0.25, 0.75 - 2**-24, 1.25, 1.75 - 2**-23, 2.5, 3.5 - 2**-22, 5.0], + device=values.device, + dtype=torch.float32, + ) + for start in range(0, shape[axis], 16): + length = min(16, shape[axis] - start) + work = values.narrow(axis, start, length).float() + scale = (global_scale * (work.abs().amax(dim=axis, keepdim=True) * (1.0 / 6))).clamp_( + 0, 448 + ) + scale = scale.to(torch.float8_e4m3fn) + effective_scale = scale.float() * inverse_global + multiplier = effective_scale.reciprocal() + finite_multiplier = torch.isfinite(multiplier) + normalized = work * torch.where(finite_multiplier, multiplier, 0.0) + # Preserve tiny values when only the FP32 reciprocal overflows. + normalized = torch.where( + (effective_scale > 0) & ~finite_multiplier, work / effective_scale, normalized + ).clamp_(-6, 6) + magnitude = normalized.abs().contiguous() + group_codes = torch.bucketize(magnitude, boundaries, out_int32=True).to(torch.uint8) + group_codes |= torch.signbit(normalized).to(torch.uint8) << 3 + codes.narrow(axis, start, length).copy_(group_codes) + scales.select(axis, start // 16).copy_(scale.squeeze(axis)) + return PackedNVFP4KV( + _pack_channel_nibbles(codes), scales, global_scale, shape, values.dtype, kind + ) + + +@torch.no_grad() +def unpack(packed: PackedNVFP4KV) -> torch.Tensor: + """Decode with stored block/global scales in FP32, then cast to compute dtype.""" + channels = packed.original_shape[2] + codes = _unpack_channel_nibbles(packed.payload, channels, dtype=torch.uint8) + magnitudes = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + device=packed.payload.device, + dtype=torch.float32, + ) + values = torch.index_select(magnitudes, 0, (codes & 7).flatten().to(torch.int32)).reshape( + codes.shape + ) + values.copysign_((codes << 4).view(torch.int8)) + axis = packed.axis + scale = (packed.scale.float() / packed.global_scale).repeat_interleave(16, dim=axis) + return values.mul_(scale.narrow(axis, 0, packed.original_shape[axis])).to(packed.original_dtype)