[Kernel] Add source-sized IQ2_S gated-pair reader

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

benchmarks/kernels/benchmark_gguf_native_pair.py

@@ -33,6 +33,7 @@ def main():
     )
     parser.add_argument("--prototype-q4-k", action="store_true")
     parser.add_argument("--prototype-iq3xxs-iq4", action="store_true")
+    parser.add_argument("--prototype-iq2-s", action="store_true")
     parser.add_argument(
         "--check-only",
         action="store_true",
@@ -40,7 +41,14 @@ def main():
     )
     args = parser.parse_args()
     assert (
-        sum((args.prototype_iq3_xxs, args.prototype_q4_k, args.prototype_iq3xxs_iq4))
+        sum(
+            (
+                args.prototype_iq3_xxs,
+                args.prototype_q4_k,
+                args.prototype_iq3xxs_iq4,
+                args.prototype_iq2_s,
+            )
+        )
         <= 1
     )
     torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
@@ -53,7 +61,9 @@ def main():
         names = [f"blk.{layer}.ffn_{role}.weight" for role in ("gate", "up")]
         types = [int(tensors[name].tensor_type) for name in names]
         allowed = (
-            ((18, 23),)
+            ((22, 21), (21, 22), (22, 18), (18, 22))
+            if args.prototype_iq2_s
+            else ((18, 23),)
             if args.prototype_iq3xxs_iq4
             else ((12, 23), (23, 12))
             if args.prototype_q4_k
@@ -70,7 +80,12 @@ def main():
         layer_module = torch.nn.Module()
         layer_module.prefix = f"model.layers.{layer}.mlp.gate_up_proj"
         layer_module.gguf_tm_projections = torch.nn.ModuleList(projections)
-        if args.prototype_iq3_xxs or args.prototype_q4_k or args.prototype_iq3xxs_iq4:
+        if (
+            args.prototype_iq3_xxs
+            or args.prototype_q4_k
+            or args.prototype_iq3xxs_iq4
+            or args.prototype_iq2_s
+        ):
             from vllm.model_executor.layers.quantization.gguf_iq3_records import (
                 signed_index_records,
             )
@@ -79,6 +94,12 @@ def main():
             )
 
             packers = {18: pack_iq3_xxs_records, 21: signed_index_records}
+            if args.prototype_iq2_s:
+                from vllm.model_executor.layers.quantization.gguf_iq2_s_records import (
+                    pack_iq2_s_records,
+                )
+
+                packers[22] = pack_iq2_s_records
             if args.prototype_q4_k or args.prototype_iq3xxs_iq4:
                 from vllm.model_executor.layers.quantization.gguf_iq4_native import (
                     pack_iq4_xs_records,

csrc/sm70_turbomind/ops/gguf_native_pair_readers_sm70.cuh

@@ -170,6 +170,76 @@ struct NativePairReader<18> {
   }
 };
 
+template <>
+struct NativePairReader<22> {
+  using Decoder = LatticeRawDecoder<22>;
+  static constexpr int kBlockBytes = 82;
+  static constexpr int kBookId = 22;
+  static constexpr int kBookBytes = Decoder::kCodebookBytes;
+  struct Record {
+    uint32_t indices[4];
+    uint32_t signs[4];
+    uint32_t high;
+    uint32_t scales;
+    float d;
+  };
+  const uint8_t* payload;
+  const uint8_t* metadata;
+  const half* original_d;
+  float cached_d;
+  int half_block;
+  bool first;
+
+  __device__ NativePairReader(const uint8_t* source, int tile, int blocks_k,
+                              int first_part, int col) {
+    const uint8_t* macro = source + int64_t{tile} * blocks_k * 32 * kBlockBytes;
+    payload = macro + first_part * 1280 + col * 16;
+    metadata = macro + first_part * 1280 + 1024 + col * 8;
+    original_d = reinterpret_cast<const half*>(macro + blocks_k * 2560 +
+                                               (first_part / 2) * 64 + col * 2);
+    cached_d = 0;
+    half_block = first_part & 1;
+    first = true;
+  }
+
+  __device__ static void initialize(uint8_t* book) {
+    Decoder::initialize(book);
+  }
+
+  __device__ Record load() {
+    if (first || half_block == 0) cached_d = __half2float(*original_d);
+    const uint4 indices = *reinterpret_cast<const uint4*>(payload);
+    const uint4 signs = *reinterpret_cast<const uint4*>(payload + 512);
+    const uint2 aux = *reinterpret_cast<const uint2*>(metadata);
+    Record record{{indices.x, indices.y, indices.z, indices.w},
+                  {signs.x, signs.y, signs.z, signs.w},
+                  aux.x,
+                  aux.y,
+                  cached_d};
+    payload += 1280;
+    metadata += 1280;
+    original_d += half_block * 32;
+    half_block ^= 1;
+    first = false;
+    return record;
+  }
+
+  template <int Segment, int Fragment>
+  __device__ static turbomind::Array<half, 8> fragment(const Record& record,
+                                                       const uint8_t* book) {
+    static_assert(Segment >= 0 && Segment < 8 && Fragment >= 0 && Fragment < 2);
+    constexpr int octet = 2 * Segment + Fragment;
+    const uint32_t low = (record.indices[octet / 4] >> (8 * (octet & 3))) & 255;
+    const uint32_t index = low | (((record.high >> (2 * octet)) & 3) << 8);
+    const uint32_t signs = (record.signs[octet / 4] >> (8 * (octet & 3))) & 255;
+    const uint64_t packed =
+        *reinterpret_cast<const uint64_t*>(book + index * 8);
+    // Reuse the original decoder's exact final FP16 operand formation.
+    return Decoder::table_fragment<half>(packed, signs, record.d,
+                                         (record.scales >> (4 * Segment)) & 15);
+  }
+};
+
 template <>
 struct NativePairReader<21> {
   using Decoder = Iq3NibbleBookDecoder;

csrc/sm70_turbomind/ops/gguf_native_pair_sm70.cu

@@ -11,6 +11,7 @@
 using R12 = vllm::sm70_gguf::NativePairReader<12>;
 using R18 = vllm::sm70_gguf::NativePairReader<18>;
 using R21 = vllm::sm70_gguf::NativePairReader<21>;
+using R22 = vllm::sm70_gguf::NativePairReader<22>;
 using R23 = vllm::sm70_gguf::NativePairReader<23>;
 namespace {
 template <class Gate, class Up>
@@ -33,6 +34,8 @@ void gguf_native_pair_sm70_out(torch::Tensor output, torch::Tensor input,
   TORCH_CHECK(
       (gate_type == 21 && (up_type == 23 || up_type == 18)) ||
           ((gate_type == 23 || gate_type == 18) && up_type == 21) ||
+          (gate_type == 22 && (up_type == 18 || up_type == 21)) ||
+          ((gate_type == 18 || gate_type == 21) && up_type == 22) ||
           (gate_type == 18 && up_type == 23) ||
           (gate_type == 12 && up_type == 23) ||
           (gate_type == 23 && up_type == 12),
@@ -55,7 +58,11 @@ void gguf_native_pair_sm70_out(torch::Tensor output, torch::Tensor input,
                   k % 1024 == 0 && n <= INT_MAX && k <= INT_MAX,
               "GGUF native pair requires M8/N32/K1024");
   const auto block_bytes = [](int64_t type) {
-    return type == 12 ? 144 : type == 18 ? 98 : type == 21 ? 110 : 136;
+    return type == 12   ? 144
+           : type == 18 ? 98
+           : type == 21 ? 110
+           : type == 22 ? 82
+                        : 136;
   };
   TORCH_CHECK(gate.numel() == n * (k / 256) * block_bytes(gate_type) &&
                   up.numel() == n * (k / 256) * block_bytes(up_type),
@@ -65,7 +72,15 @@ void gguf_native_pair_sm70_out(torch::Tensor output, torch::Tensor input,
   TORCH_CHECK(properties->major == 7 && properties->minor == 0,
               "GGUF native pair requires SM70");
   const auto stream = at::cuda::getCurrentCUDAStream();
-  if (gate_type == 12)
+  if (gate_type == 22 && up_type == 21)
+    launch_pair<R22, R21>(output, input, gate, up, n, k, stream);
+  else if (gate_type == 21 && up_type == 22)
+    launch_pair<R21, R22>(output, input, gate, up, n, k, stream);
+  else if (gate_type == 22)
+    launch_pair<R22, R18>(output, input, gate, up, n, k, stream);
+  else if (up_type == 22)
+    launch_pair<R18, R22>(output, input, gate, up, n, k, stream);
+  else if (gate_type == 12)
     launch_pair<R12, R23>(output, input, gate, up, n, k, stream);
   else if (up_type == 12)
     launch_pair<R23, R12>(output, input, gate, up, n, k, stream);

docs/design/gguf_native_iq2_s_pair.md

@@ -0,0 +1,25 @@
+# Original IQ2_S shared-activation pair reader
+
+The next two source-byte priorities pair IQ2_S with IQ3_S in layers5/9/17
+and IQ3_XXS in layers7/18/33. All six are TP4 N4352/K5120. Both gate/up
+orientations occur for each combination.
+
+Preserve each original 82-byte K256 block as two K128 packet sets:
+16-byte indices, 16-byte signs and eight bytes of high bits/scales per
+column. Original d follows once per K256 block. Packet planes interleave
+N32 columns with aligned 16/8-byte loads; no field expands or changes.
+The reader caches original d across halves, including odd split-K starts,
+and increments pointers in the K loop.
+
+Reuse `LatticeRawDecoder<22>` codebook initialization and exact final FP16
+operand formation, with the existing llama.cpp MIT attribution. No expanded
+FP16 coefficient or second decode formula is introduced. The shared-A
+skeleton, FP32 dot/reduction and gated epilogue are unchanged. Raw operators
+compile both orientations; model capabilities remain canonical pending
+actual-weight numerical and matched speed checks.
+
+Ten CPU inverse/cursor checks pass, covering N32/64/96, complete K256 blocks,
+strided inputs and all eight split-K starts at K5120. Six real IQ2_S tensors,
+each sampled at N64/K5120, preserve 629,760 source bytes and 1,966,080 official
+FP32 dequantized values bitwise. Native build, graph and speed checks are
+pending; no end-to-end result is claimed.

tests/kernels/quantization/test_gguf_iq2_s_pair_cursor.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import runpy
+from pathlib import Path
+
+import numpy as np
+
+pack_iq2_s_records = runpy.run_path(
+    str(
+        Path(__file__).parents[3]
+        / "vllm/model_executor/layers/quantization/gguf_iq2_s_records.py"
+    )
+)["pack_iq2_s_records"]
+
+
+def test_all_split_k_starts_reconstruct_original_decoder_fields():
+    rng = np.random.default_rng(1290)
+    source = rng.integers(0, 256, (64, 20, 82), dtype=np.uint8)
+    d = rng.uniform(-1, 1, (64, 20)).astype("<f2")
+    source[..., :2] = d.view(np.uint8).reshape(64, 20, 2)
+    records = pack_iq2_s_records(source.reshape(64, 1640))
+    for tile in range(2):
+        for col in range(32):
+            for first_part in range(0, 40, 5):
+                macro = tile * 20 * 32 * 82
+                cursor = macro + first_part * 1280
+                d_cursor = macro + 20 * 2560 + (first_part // 2) * 64 + col * 2
+                half_block = first_part & 1
+                cached_d = None
+                for part in range(first_part, first_part + 5):
+                    if part == first_part or half_block == 0:
+                        cached_d = records[d_cursor : d_cursor + 2].view("<f2")[0]
+                    indices = records[cursor + col * 16 : cursor + col * 16 + 16]
+                    signs = records[cursor + 512 + col * 16 : cursor + 528 + col * 16]
+                    metadata = records[
+                        cursor + 1024 + col * 8 : cursor + 1032 + col * 8
+                    ]
+                    high = int.from_bytes(metadata[:4], "little")
+                    scales = int.from_bytes(metadata[4:], "little")
+                    block = source[tile * 32 + col, part // 2]
+                    assert cached_d == block[:2].view("<f2")[0]
+                    for segment in range(8):
+                        for fragment in range(2):
+                            octet = 2 * segment + fragment
+                            base = half_block * 128 + segment * 16 + fragment * 8
+                            index = int(indices[octet]) | (
+                                ((high >> (2 * octet)) & 3) << 8
+                            )
+                            original = int(block[2 + base // 8]) | (
+                                (
+                                    (
+                                        int(block[66 + base // 32])
+                                        >> (2 * ((base // 8) % 4))
+                                    )
+                                    & 3
+                                )
+                                << 8
+                            )
+                            assert index == original
+                            assert int(signs[octet]) == int(block[34 + base // 8])
+                            assert (scales >> (4 * segment)) & 15 == (
+                                int(block[74 + base // 32]) >> (4 * ((base // 16) % 2))
+                            ) & 15
+                    cursor += 1280
+                    d_cursor += half_block * 64
+                    half_block ^= 1

tests/kernels/quantization/test_gguf_iq2_s_records.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib.util
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+# The layout needs no native module or CUDA context, even in a source checkout.
+_spec = importlib.util.spec_from_file_location(
+    "iq2_s_records",
+    Path(__file__).parents[3]
+    / "vllm/model_executor/layers/quantization/gguf_iq2_s_records.py",
+)
+assert _spec is not None and _spec.loader is not None
+_module = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_module)
+pack_iq2_s_records = _module.pack_iq2_s_records
+
+
+@pytest.mark.parametrize("n,blocks", [(32, 1), (64, 3), (96, 20)])
+def test_every_original_bit_survives_independent_reader(n, blocks):
+    rng = np.random.default_rng(1290)
+    raw = rng.integers(0, 256, (n, blocks * 82), dtype=np.uint8)
+    before = raw.copy()
+    records = pack_iq2_s_records(raw)
+    assert records.dtype == np.uint8 and records.flags.c_contiguous
+    assert records.nbytes == raw.nbytes
+    np.testing.assert_array_equal(raw, before)
+    recovered = np.empty_like(raw)
+    for tile in range(n // 32):
+        macro = tile * blocks * 32 * 82
+        for block in range(blocks):
+            for col in range(32):
+                row = recovered[tile * 32 + col, block * 82 : (block + 1) * 82]
+                d = macro + blocks * 2560 + block * 64 + col * 2
+                row[:2] = records[d : d + 2]
+                for half in range(2):
+                    cursor = macro + (block * 2 + half) * 1280
+                    index = cursor + col * 16
+                    sign = cursor + 512 + col * 16
+                    metadata = cursor + 1024 + col * 8
+                    assert index % 16 == 0 and sign % 16 == 0 and metadata % 8 == 0
+                    row[2 + half * 16 : 18 + half * 16] = records[index : index + 16]
+                    row[34 + half * 16 : 50 + half * 16] = records[sign : sign + 16]
+                    row[66 + half * 4 : 70 + half * 4] = records[
+                        metadata : metadata + 4
+                    ]
+                    row[74 + half * 4 : 78 + half * 4] = records[
+                        metadata + 4 : metadata + 8
+                    ]
+    np.testing.assert_array_equal(recovered, raw)
+
+
+def test_noncontiguous_tp_rows_keep_original_bytes():
+    backing = np.arange(64 * 3 * 82 * 2, dtype=np.uint32).astype(np.uint8)
+    raw = backing.reshape(64, 3 * 82 * 2)[:, ::2]
+    assert not raw.flags.c_contiguous
+    np.testing.assert_array_equal(
+        pack_iq2_s_records(raw), pack_iq2_s_records(raw.copy())
+    )
+
+
+@pytest.mark.parametrize("shape", [(31, 82), (32, 99), (0, 82), (32, 0)])
+def test_rejects_partial_quant_blocks_or_column_tiles(shape):
+    with pytest.raises(ValueError):
+        pack_iq2_s_records(np.zeros(shape, dtype=np.uint8))
+
+
+def test_rejects_converted_weights():
+    with pytest.raises(TypeError):
+        pack_iq2_s_records(np.zeros((32, 82), dtype=np.float16))

vllm/model_executor/layers/quantization/gguf_iq2_s_records.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Source-sized IQ2_S records for shared-activation gated-pair readers."""
+
+import numpy as np
+
+
+def pack_iq2_s_records(raw: np.ndarray) -> np.ndarray:
+    """Interleave original K128 index, sign and high-bit/scale packets.
+
+    Each column retains two 16-byte planes and one 8-byte metadata packet
+    per K128. Original d follows the payload, once per K256. Every source
+    bit is preserved; scales and codebook indices are never expanded.
+    """
+    if raw.dtype != np.uint8:
+        raise TypeError("IQ2_S records require original uint8 blocks")
+    if (
+        raw.ndim != 2
+        or raw.shape[0] == 0
+        or raw.shape[0] % 32
+        or raw.shape[1] == 0
+        or raw.shape[1] % 82
+    ):
+        raise ValueError("IQ2_S records require N32 and complete 82-byte blocks")
+    n, row_bytes = raw.shape
+    tiles, blocks = n // 32, row_bytes // 82
+    source = raw.reshape(tiles, 32, blocks, 82).transpose(0, 2, 1, 3)
+    indices = (
+        source[..., 2:34].reshape(tiles, blocks, 32, 2, 16).transpose(0, 1, 3, 2, 4)
+    )
+    signs = (
+        source[..., 34:66].reshape(tiles, blocks, 32, 2, 16).transpose(0, 1, 3, 2, 4)
+    )
+    high = source[..., 66:74].reshape(tiles, blocks, 32, 2, 4).transpose(0, 1, 3, 2, 4)
+    scales = (
+        source[..., 74:82].reshape(tiles, blocks, 32, 2, 4).transpose(0, 1, 3, 2, 4)
+    )
+    payload = np.concatenate(
+        (
+            indices.reshape(tiles, blocks, 2, 512),
+            signs.reshape(tiles, blocks, 2, 512),
+            np.concatenate((high, scales), axis=4).reshape(tiles, blocks, 2, 256),
+        ),
+        axis=3,
+    ).reshape(tiles, blocks * 2560)
+    original_d = source[..., :2].reshape(tiles, blocks * 64)
+    records = np.concatenate((payload, original_d), axis=1).reshape(-1)
+    assert records.nbytes == raw.nbytes
+    return records