[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.

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="<f2")
+    offset = 80 if kind == 10 else size - 2 if kind in (11, 14, 34, 35) else 0
+    if kind == 39:
+        raw[..., :1] = 127
+    elif kind == 40:
+        raw[..., :4] = 0x38
+    elif kind == 29:
+        words = raw[..., 48:].copy().view("<u2")
+        words &= 0x0FFF
+        words |= ((d.view("<u2")[0] >> (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)

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<uint16_t,uint2_t>'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)
 

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.
 

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