From 74ea7cf330384bbee4a93d91c8dbdfa372cc5970 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Mon, 5 Oct 2026 17:52:28 +0800 Subject: [PATCH] [Kernel] Add source-sized IQ2_S gated-pair reader Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_gguf_native_pair.py | 27 ++++++- .../ops/gguf_native_pair_readers_sm70.cuh | 70 ++++++++++++++++++ .../ops/gguf_native_pair_sm70.cu | 19 ++++- docs/design/gguf_native_iq2_s_pair.md | 25 +++++++ .../test_gguf_iq2_s_pair_cursor.py | 67 +++++++++++++++++ .../quantization/test_gguf_iq2_s_records.py | 73 +++++++++++++++++++ .../layers/quantization/gguf_iq2_s_records.py | 49 +++++++++++++ 7 files changed, 325 insertions(+), 5 deletions(-) create mode 100644 docs/design/gguf_native_iq2_s_pair.md create mode 100644 tests/kernels/quantization/test_gguf_iq2_s_pair_cursor.py create mode 100644 tests/kernels/quantization/test_gguf_iq2_s_records.py create mode 100644 vllm/model_executor/layers/quantization/gguf_iq2_s_records.py diff --git a/benchmarks/kernels/benchmark_gguf_native_pair.py b/benchmarks/kernels/benchmark_gguf_native_pair.py index 6f4d541d8..f12e5dbf8 100644 --- a/benchmarks/kernels/benchmark_gguf_native_pair.py +++ b/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, diff --git a/csrc/sm70_turbomind/ops/gguf_native_pair_readers_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_native_pair_readers_sm70.cuh index 7375335ab..6ac3e9cc1 100644 --- a/csrc/sm70_turbomind/ops/gguf_native_pair_readers_sm70.cuh +++ b/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(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(payload); + const uint4 signs = *reinterpret_cast(payload + 512); + const uint2 aux = *reinterpret_cast(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 + __device__ static turbomind::Array 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(book + index * 8); + // Reuse the original decoder's exact final FP16 operand formation. + return Decoder::table_fragment(packed, signs, record.d, + (record.scales >> (4 * Segment)) & 15); + } +}; + template <> struct NativePairReader<21> { using Decoder = Iq3NibbleBookDecoder; diff --git a/csrc/sm70_turbomind/ops/gguf_native_pair_sm70.cu b/csrc/sm70_turbomind/ops/gguf_native_pair_sm70.cu index db6d6dfe0..6bf9a949e 100644 --- a/csrc/sm70_turbomind/ops/gguf_native_pair_sm70.cu +++ b/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 @@ -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(output, input, gate, up, n, k, stream); + else if (gate_type == 21 && up_type == 22) + launch_pair(output, input, gate, up, n, k, stream); + else if (gate_type == 22) + launch_pair(output, input, gate, up, n, k, stream); + else if (up_type == 22) + launch_pair(output, input, gate, up, n, k, stream); + else if (gate_type == 12) launch_pair(output, input, gate, up, n, k, stream); else if (up_type == 12) launch_pair(output, input, gate, up, n, k, stream); diff --git a/docs/design/gguf_native_iq2_s_pair.md b/docs/design/gguf_native_iq2_s_pair.md new file mode 100644 index 000000000..2fe3507a1 --- /dev/null +++ b/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. diff --git a/tests/kernels/quantization/test_gguf_iq2_s_pair_cursor.py b/tests/kernels/quantization/test_gguf_iq2_s_pair_cursor.py new file mode 100644 index 000000000..55819da7f --- /dev/null +++ b/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("> (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 diff --git a/tests/kernels/quantization/test_gguf_iq2_s_records.py b/tests/kernels/quantization/test_gguf_iq2_s_records.py new file mode 100644 index 000000000..410d32bcb --- /dev/null +++ b/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)) diff --git a/vllm/model_executor/layers/quantization/gguf_iq2_s_records.py b/vllm/model_executor/layers/quantization/gguf_iq2_s_records.py new file mode 100644 index 000000000..d448c4e26 --- /dev/null +++ b/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