From acd07589fae8394b4836e8a124d7f6057aaff743 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 11 Nov 2026 14:07:37 -0800 Subject: [PATCH] Run GGUF Q5_K weights in the native ggml engine The engine took packed GGUF weights only as Q4_K or Q6_K, but K-quant mixes such as Q4_K_M store some tensors as Q5_K; Muse Glimmer's LM head is one. Export already carries any GGUF type as a gguf: PackedQuant, so only the engine needed it. A table of K-quant codecs (ggml type and block bytes) now replaces the Q4_K/Q6_K conditionals, with Q5_K added. --- .../native/runtime/engine/ggml/GgmlTypes.cpp | 49 ++++++------- .../test/runtime/ggml/test_ggml_engine.cpp | 68 +++++++++++++++---- 2 files changed, 79 insertions(+), 38 deletions(-) diff --git a/backends/native/runtime/engine/ggml/GgmlTypes.cpp b/backends/native/runtime/engine/ggml/GgmlTypes.cpp index cc6e1b78abd..188eb8edcbe 100644 --- a/backends/native/runtime/engine/ggml/GgmlTypes.cpp +++ b/backends/native/runtime/engine/ggml/GgmlTypes.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -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,11 +75,12 @@ 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) { @@ -74,12 +88,9 @@ size_t packed_nbytes(const TensorMeta& meta) { "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(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: diff --git a/backends/native/test/runtime/ggml/test_ggml_engine.cpp b/backends/native/test/runtime/ggml/test_ggml_engine.cpp index 0c8a9c1e146..aa234878f4f 100644 --- a/backends/native/test/runtime/ggml/test_ggml_engine.cpp +++ b/backends/native/test/runtime/ggml/test_ggml_engine.cpp @@ -83,13 +83,14 @@ flatbuffers::Offset create_none_arg( 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 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 @@ std::vector make_linear_program( : 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 @@ std::vector make_linear_program( builder.GetBufferPointer() + builder.GetSize()}; } -std::unique_ptr 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 kquant_block(size_t block_bytes, uint8_t high_bits = 0) { + std::vector 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 make_context_with_weight( GgmlEngineHost& host, const std::vector& program_bytes, + const std::vector& 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 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 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::unique_ptr make_context( std::move(program), std::move(package), options); } +// Half the weights are 1 and half are 2. +std::unique_ptr make_context( + GgmlEngineHost& host, + const std::vector& 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::monostate>; @@ -492,6 +510,26 @@ TEST(GgmlEngineTest, ExecutesPackedQ4KLinear) { 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 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);