[Native] Run GGUF Q5_K weights in the ggml engine

backends/native/runtime/engine/ggml/GgmlTypes.cpp

@@ -8,6 +8,7 @@
 
 #include <algorithm>
 #include <cstring>
+#include <iterator>
 #include <limits>
 #include <stdexcept>
 #include <string>
@@ -17,11 +18,23 @@
 namespace ptn::ggml_detail {
 namespace {
 
-constexpr std::string_view kQ4KCodec = "gguf:q4_k";
-constexpr std::string_view kQ6KCodec = "gguf:q6_k";
+struct KQuantCodec {
+  std::string_view name;
+  ggml_type type;
+  size_t block_bytes;
+};
+constexpr KQuantCodec kKQuantCodecs[] = {
+    {"gguf:q4_k", GGML_TYPE_Q4_K, 144},
+    {"gguf:q5_k", GGML_TYPE_Q5_K, 176},
+    {"gguf:q6_k", GGML_TYPE_Q6_K, 210},
+};
 constexpr int64_t kKQuantBlockElements = 256;
-constexpr size_t kQ4KBlockBytes = 144;
-constexpr size_t kQ6KBlockBytes = 210;
+
+const KQuantCodec* kquant_codec(std::string_view name) {
+  const auto found = std::ranges::find(kKQuantCodecs, name, &KQuantCodec::name);
+  return found == std::end(kKQuantCodecs) ? nullptr : found;
+}
+
 constexpr size_t kQ40BlockElements = 32;
 
 struct BlockQ40 {
@@ -62,24 +75,22 @@ const PackedQuant* packed_quant(const TensorMeta& meta) {
 
 size_t packed_nbytes(const TensorMeta& meta) {
   const PackedQuant* packed = packed_quant(meta);
-  if (packed == nullptr ||
-      (packed->codec != kQ4KCodec && packed->codec != kQ6KCodec)) {
+  const KQuantCodec* codec =
+      packed == nullptr ? nullptr : kquant_codec(packed->codec);
+  if (codec == nullptr) {
     throw std::runtime_error(
-        "ggml: only PackedQuant codecs 'gguf:q4_k' and 'gguf:q6_k' "
-        "are supported");
+        "ggml: only PackedQuant codecs 'gguf:q4_k', 'gguf:q5_k' and "
+        "'gguf:q6_k' are supported");
   }
   if (meta.sizes.size() != 2 || meta.sizes[0] <= 0 || meta.sizes[1] <= 0 ||
       meta.sizes[1] % kKQuantBlockElements != 0) {
     throw std::runtime_error(
         "ggml: K-quant weight must have logical shape [O, I] with I divisible "
         "by 256");
   }
-  const ggml_type type =
-      packed->codec == kQ4KCodec ? GGML_TYPE_Q4_K : GGML_TYPE_Q6_K;
-  const size_t block_bytes =
-      packed->codec == kQ4KCodec ? kQ4KBlockBytes : kQ6KBlockBytes;
-  if (ggml_type_size(type) != block_bytes ||
-      ggml_blck_size(type) != kKQuantBlockElements) {
+  const size_t block_bytes = codec->block_bytes;
+  if (ggml_type_size(codec->type) != block_bytes ||
+      ggml_blck_size(codec->type) != kKQuantBlockElements) {
     throw std::runtime_error("ggml: unexpected GGUF K-quant block layout");
   }
   const size_t rows = static_cast<size_t>(meta.sizes[0]);
@@ -187,15 +198,7 @@ static ggml_type to_ggml_type(const TensorMeta& meta) {
   }
   if (packed_quant(meta) != nullptr) {
     packed_nbytes(meta);
-    if (packed_quant(meta)->codec == kQ4KCodec) {
-      return GGML_TYPE_Q4_K;
-    }
-    if (packed_quant(meta)->codec == kQ6KCodec) {
-      return GGML_TYPE_Q6_K;
-    }
-    throw std::runtime_error(
-        "ggml: unsupported PackedQuant codec '" + packed_quant(meta)->codec +
-        "'");
+    return kquant_codec(packed_quant(meta)->codec)->type;
   }
   switch (meta.dtype) {
     case ScalarType::Float:

backends/native/test/runtime/ggml/test_ggml_engine.cpp

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -83,13 +83,14 @@
 
 constexpr size_t kQ4KBlockBytes = 144;
 
-// A [1, 256] x [256] linear with gguf:q4_k weights, or with symmetric int4
-// AffineGroup weights whose fp16 scales `weight_scale` cover groups of 32.
-// `cache` is an unused mutable buffer.
+// A [1, 256] x [256] linear with `codec` packed weights, or with symmetric
+// int4 AffineGroup weights whose fp16 scales `weight_scale` cover groups of
+// 32. `cache` is an unused mutable buffer.
 std::vector<uint8_t> make_linear_program(
     const char* target = "torch.ops.aten.linear.default",
     bool with_cache = false,
-    bool affine = false) {
+    bool affine = false,
+    const char* codec = "gguf:q4_k") {
   flatbuffers::FlatBufferBuilder builder;
   const auto quant = affine
       ? fbs::CreateQuantSpec(
@@ -101,7 +102,7 @@
       : fbs::CreateQuantSpec(
             builder,
             fbs::QuantScheme::PackedQuant,
-            fbs::CreatePackedQuantDirect(builder, "gguf:q4_k").Union());
+            fbs::CreatePackedQuantDirect(builder, codec).Union());
   const auto input_meta =
       create_tensor_meta(builder, fbs::ScalarType::FLOAT, {1, 256});
   const auto weight_meta =
@@ -234,20 +235,28 @@
       builder.GetBufferPointer() + builder.GetSize()};
 }
 
-std::unique_ptr<GgmlEngineContext> make_context(
+// One K-quant super-block of `block_bytes`: the 16-byte header block_q4_K
+// and block_q5_K share (d = 1.0, dmin = 0, every sub-block scale 1 and min 0),
+// then `high_bits` bytes of fifth bits, then nibbles 1 (low) / 2 (high).
+std::vector<uint8_t> kquant_block(size_t block_bytes, uint8_t high_bits = 0) {
+  std::vector<uint8_t> block(block_bytes, 0x21);
+  std::fill_n(block.begin(), 16, uint8_t{0});
+  block[1] = 0x3C;
+  std::fill_n(block.begin() + 4, 4, uint8_t{1});
+  std::fill_n(block.begin() + 12, 4, uint8_t{1});
+  std::fill(block.begin() + 16, block.end() - 128, high_bits);
+  return block;
+}
+
+std::unique_ptr<GgmlEngineContext> make_context_with_weight(
     GgmlEngineHost& host,
     const std::vector<uint8_t>& program_bytes,
+    const std::vector<uint8_t>& packed_weight,
     const EngineOptions& options = {}) {
-  // One block_q4_K: d = 1.0 (fp16), dmin = 0, every sub-block scale 1 and
-  // min 0, and nibbles 1 (low) / 2 (high), so half the weights are 1 and
-  // half are 2.
-  std::vector<uint8_t> packed_weight(kQ4KBlockBytes, 0x21);
-  std::fill_n(packed_weight.begin(), 16, uint8_t{0});
-  packed_weight[1] = 0x3C;
-  std::fill_n(packed_weight.begin() + 4, 4, uint8_t{1});
-  std::fill_n(packed_weight.begin() + 12, 4, uint8_t{1});
+  const std::string size = std::to_string(packed_weight.size());
   std::vector<uint8_t> safetensors = testing::make_safetensors(
-      R"({"weight":{"dtype":"U8","shape":[144],"data_offsets":[0,144]}})",
+      R"({"weight":{"dtype":"U8","shape":[)" + size +
+          R"(],"data_offsets":[0,)" + size + "]}}",
       packed_weight.size());
   std::copy(
       packed_weight.begin(),
@@ -264,6 +273,15 @@
       std::move(program), std::move(package), options);
 }
 
+// Half the weights are 1 and half are 2.
+std::unique_ptr<GgmlEngineContext> make_context(
+    GgmlEngineHost& host,
+    const std::vector<uint8_t>& program_bytes,
+    const EngineOptions& options = {}) {
+  return make_context_with_weight(
+      host, program_bytes, kquant_block(kQ4KBlockBytes), options);
+}
+
 // A tensor (by name), an int, a float, an int list, or None.
 using TestArg = std::
     variant<std::string, int64_t, double, std::vector<int64_t>, std::monostate>;
@@ -492,6 +510,26 @@
   EXPECT_NEAR(output, 128.0f * 1.0f + 128.0f * 2.0f, 1e-2f);
 }
 
+// Every fifth bit set adds 16 to each weight, so the dot product with ones is
+// 128 * (1 + 16) + 128 * (2 + 16).
+TEST(GgmlEngineTest, ExecutesPackedQ5KLinear) {
+  auto host = GgmlEngineHost::create();
+  auto context = make_context_with_weight(
+      *host,
+      make_linear_program(
+          "torch.ops.aten.linear.default", false, false, "gguf:q5_k"),
+      kquant_block(176, 0xFF));
+  auto executable = context->compile("forward");
+
+  EXPECT_EQ(context->materialized_constant_bytes(), 176);
+  const std::vector<float> input(256, 1.0f);
+  executable->set_input(0, input.data(), input.size(), ScalarType::Float);
+  executable->execute();
+  float output = 0.0f;
+  executable->get_output(0, &output, 1, ScalarType::Float);
+  EXPECT_NEAR(output, 128.0f * 17.0f + 128.0f * 18.0f, 1e-1f);
+}
+
 TEST(GgmlEngineTest, RepeatedExecutionsReadFreshInputs) {
   for (const size_t max_cached_graphs : {size_t{0}, size_t{1}}) {
     SCOPED_TRACE(max_cached_graphs);