From 3bdbb684be07c81d4ebffbe9c5e94cee8f7d13ff Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 24 Oct 2026 00:36:05 +0800 Subject: [PATCH] [Core] Slice packed GGUF rows before transcoding on the host fallback (#1253) With device_transcode disabled, expert loading decoded each complete projection on every TP rank and sliced the canonical result afterwards. When the TP boundary keeps whole source blocks, slice the packed rows first and decode only this rank's share; when the boundary cuts an original quantization block, keep the decode-then-slice order because canonical groups are smaller than source blocks. Quantization codes, coefficients and prepared layouts are unchanged. --- .../test_gguf_tp_before_transcode.py | 125 ++++++++++++++++++ .../quantization/gguf_lattice_transcode.py | 5 +- .../layers/quantization/gguf_transcode.py | 26 ++++ .../layers/quantization/gguf_turbomind_moe.py | 14 +- 4 files changed, 165 insertions(+), 5 deletions(-) create mode 100644 tests/kernels/quantization/test_gguf_tp_before_transcode.py 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