diff --git a/tests/kernels/quantization/test_gguf_tp_before_transcode.py b/tests/kernels/quantization/test_gguf_tp_before_transcode.py new file mode 100644 index 000000000..53d89cb87 --- /dev/null +++ b/tests/kernels/quantization/test_gguf_tp_before_transcode.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import fields + +import numpy as np +import pytest + +from vllm.model_executor.layers.quantization.gguf_lattice_transcode import ( + LATTICE_TYPES, + transcode_lattice, +) +from vllm.model_executor.layers.quantization.gguf_lut_transcode import ( + LUT4_TYPES, + transcode_lut4, +) +from vllm.model_executor.layers.quantization.gguf_transcode import ( + tp_slice_packed, + transcode_affine, +) +from vllm.transformers_utils.gguf_tensor_reader import quant_size + +TYPES = ( + 2, + 3, + 6, + 7, + 8, + 10, + 11, + 12, + 13, + 14, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 29, + 34, + 35, + 39, + 40, + 41, + 42, +) + + +def packed_source(kind, n=128, k=1024): + block, size = quant_size(kind) + raw = np.random.default_rng(kind).integers( + 0, 256, (n, k // block, size), dtype=np.uint8 + ) + d = np.array([0.00390625], dtype="> (4 * np.arange(4, dtype=np.uint16))) & 15) << 12 + raw[..., 48:] = words.view(np.uint8) + else: + raw[..., offset : offset + 2] = d.view(np.uint8) + if kind in (3, 7, 12, 13): + raw[..., 2:4] = d.view(np.uint8) + elif kind == 10: + raw[..., 82:84] = d.view(np.uint8) + return raw.reshape(n, -1) + + +def convert(source, kind): + if kind in LATTICE_TYPES: + return transcode_lattice(source, kind) + if kind in LUT4_TYPES: + return transcode_lut4(source, kind) + return transcode_affine(source, kind) + + +@pytest.mark.parametrize("kind", TYPES) +@pytest.mark.parametrize("axis", [0, 1]) +def test_tp_before_transcode_preserves_every_canonical_field(kind, axis): + source = packed_source(kind) + canonical = convert(source, kind) + for rank in range(4): + reference = canonical.tp_slice(rank, 4, axis=axis) + raw = tp_slice_packed(source, kind, rank, 4, axis=axis) + assert raw is not None + candidate = convert(raw, kind) + for field in fields(reference): + expected, actual = ( + getattr(reference, field.name), + getattr(candidate, field.name), + ) + if isinstance(expected, np.ndarray): + assert actual.dtype == expected.dtype + np.testing.assert_array_equal(actual, expected) + else: + assert actual == expected + if kind in LATTICE_TYPES: + for actual, expected in zip( + candidate.mma884_storage(), reference.mma884_storage() + ): + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("kind", [18, 22, 23, 42]) +def test_unaligned_original_blocks_retain_canonical_tp(kind): + # IQ K=768 -> 192 and Q2_0 K=640 -> 160 cut source blocks, while + # the expanded scale groups remain legal TP4 boundaries. + source = packed_source(kind, k=640 if kind == 42 else 768) + assert tp_slice_packed(source, kind, 0, 4, axis=1) is None + candidate = convert(source, kind).tp_slice(0, 4, axis=1) + assert candidate.dequantize().shape[1] == (160 if kind == 42 else 192) + + +@pytest.mark.parametrize("rank,size,axis", [(4, 4, 0), (0, 0, 0), (0, 4, 2)]) +def test_invalid_tp_boundary_is_rejected(rank, size, axis): + with pytest.raises(ValueError, match="Invalid GGUF TP"): + tp_slice_packed(packed_source(18), 18, rank, size, axis=axis) diff --git a/vllm/model_executor/layers/quantization/gguf_lattice_transcode.py b/vllm/model_executor/layers/quantization/gguf_lattice_transcode.py index a146d219f..4d81c1d0f 100644 --- a/vllm/model_executor/layers/quantization/gguf_lattice_transcode.py +++ b/vllm/model_executor/layers/quantization/gguf_lattice_transcode.py @@ -96,7 +96,10 @@ class LatticeGGUFProjection: metadata |= np.bitwise_or.reduce(high << high_shifts, axis=-1) # Invert Converter's adjacent-half pairing so it # preserves the packet verbatim after operand packing. - shifts = np.array([0, 8, 2, 10, 4, 12, 6, 14], dtype=np.uint64) + # Every carrier is exactly 16 bits; widening the eight-way expansion + # to uint64 needlessly quadruples its temporary memory traffic. + shifts = np.array([0, 8, 2, 10, 4, 12, 6, 14], dtype=np.uint16) + packets = packets.astype(np.uint16) codes = ((packets[..., None] >> shifts) & 3).astype(np.uint8).reshape(n, k) return np.ascontiguousarray(codes), np.ascontiguousarray(metadata) diff --git a/vllm/model_executor/layers/quantization/gguf_transcode.py b/vllm/model_executor/layers/quantization/gguf_transcode.py index 0a97ad367..fb07ee05b 100644 --- a/vllm/model_executor/layers/quantization/gguf_transcode.py +++ b/vllm/model_executor/layers/quantization/gguf_transcode.py @@ -70,6 +70,32 @@ AFFINE_U2_TYPES = frozenset((10, 34, 35, 41, 42)) AFFINE_BITPLANE_TYPES = frozenset((6, 7, 11, 13, 14)) +def tp_slice_packed( + data: np.ndarray, weight_type: int, rank: int, size: int, *, axis: int +) -> np.ndarray | None: + """Slice before decoding when TP boundaries retain whole source blocks. + + A None result requests canonical slicing: its smaller scale groups can + support boundaries such as Q2_0 K=640/4 which cut original blocks. + """ + if size <= 0 or not 0 <= rank < size or axis not in (0, 1): + raise ValueError("Invalid GGUF TP rank, size or axis") + block, block_bytes = quant_size(weight_type) + if data.dtype != np.uint8 or data.ndim != 2 or data.shape[1] % block_bytes: + raise ValueError("GGUF projection needs complete packed rows [N,bytes]") + shape = (data.shape[0], data.shape[1] // block_bytes * block) + span, remainder = divmod(shape[axis], size) + if remainder: + raise ValueError("GGUF TP dimension is not divisible") + if axis == 1: + if span % block: + return None + span = span // block * block_bytes + selection = [slice(None), slice(None)] + selection[axis] = slice(rank * span, (rank + 1) * span) + return np.ascontiguousarray(data[tuple(selection)]) + + def transcode_affine(data: np.ndarray, weight_type: int) -> AffineGGUFProjection: """Normalize affine and ternary blocks without quantizing code values. diff --git a/vllm/model_executor/layers/quantization/gguf_turbomind_moe.py b/vllm/model_executor/layers/quantization/gguf_turbomind_moe.py index 1dd469761..d0abc0417 100644 --- a/vllm/model_executor/layers/quantization/gguf_turbomind_moe.py +++ b/vllm/model_executor/layers/quantization/gguf_turbomind_moe.py @@ -603,14 +603,20 @@ class GGUFExpertBank(torch.nn.Module): if self.device_transcode: self._stage(index, source, rank, size, axis) return + # Slice packed rows before decoding when the TP boundary keeps whole + # source blocks; a None result keeps the canonical decode-then-slice + # order for boundaries that cut an original block. + local_source = tp_slice_packed(source, self.source_type, rank, size, axis=axis) + transcode_source = source if local_source is None else local_source canonical: LatticeGGUFProjection | Lut4GGUFProjection | AffineGGUFProjection if self.source_type in LATTICE_TYPES: - canonical = transcode_lattice(source, self.source_type) + canonical = transcode_lattice(transcode_source, self.source_type) elif self.source_type in LUT4_TYPES: - canonical = transcode_lut4(source, self.source_type) + canonical = transcode_lut4(transcode_source, self.source_type) else: - canonical = transcode_affine(source, self.source_type) - canonical = canonical.tp_slice(rank, size, axis=axis) + canonical = transcode_affine(transcode_source, self.source_type) + if local_source is None: + canonical = canonical.tp_slice(rank, size, axis=axis) self.group = canonical.group_size self.n, self.k = ( canonical.shape