diff --git a/docs/tbkern/phase9-four-row-vnni64.md b/docs/tbkern/phase9-four-row-vnni64.md new file mode 100644 index 000000000000..6bd9ac309c68 --- /dev/null +++ b/docs/tbkern/phase9-four-row-vnni64.md @@ -0,0 +1,13 @@ +# Phase 9: four-row VNNI64 decode tile + +Phase 9 adds a separately opt-in single-token Q2_0 decode route: + +```text +GGML_TBKERN_Q2_0_VNNI64_4R=1 +``` + +The VNNI64 route previously rebuilt the identical 64-byte Q8 activation vector for each output row. The four-row tile constructs that vector once for each 64-value subgroup, then applies it to four independent packed-weight rows. Each output row retains the original `b=0,1` and group accumulation order, per-32 Q8 scales, and precomputed Q8-sum correction. A thread's final one to three rows use the unchanged Phase 8 route. + +This is intentionally opt-in and does not change native Prism, Phase 7 VNNI, or Phase 8 VNNI64 selection. It requires AVX-512F and AVX-512VNNI. It changes no GGUF storage, repack allocation, or generic-ggml fallback behavior. + +Validation gate: exact native/Phase 8/Phase 9 logits and token parity on the real 27B Q2_0 model, then matched CPU decode and PPL measurements. Retain it only if matched decode improvement exceeds run-to-run noise. \ No newline at end of file diff --git a/docs/tbkern/phase9-prism-evidence.md b/docs/tbkern/phase9-prism-evidence.md new file mode 100644 index 000000000000..f33a40df9079 --- /dev/null +++ b/docs/tbkern/phase9-prism-evidence.md @@ -0,0 +1,63 @@ +# Phase 9: Prism 27B minimal-PR evidence + +This document records the evidence for the opt-in four-row CPU decode path. +It is a decode measurement, not a claim about prompt processing or GPU +throughput. + +## Scope + +The minimal branch is based on Prism's prism branch at 7529fdaaf. It keeps only +the CPU backend integration, the Q2_0 cache packer, the parity comparator, and +this evidence. The native Prism Q2_0 path remains the default unless the +selector is set: + +GGML_TBKERN_Q2_0_VNNI64_4R=1 + +## Configuration + +- Model: Ternary-Bonsai-27B-Q2_0.gguf, 7,165,121,600 bytes, + SHA-256 868c11714cf8fe47f5ec9eeb2be0ab1a337112886f92ee0ede6b855c4fa31757. +- Host: Google Cloud C3 highmem-8, Intel Xeon Platinum 8481C, 4 physical + cores plus SMT siblings, one NUMA node, AVX-512F/VNNI/VBMI. +- Decode command: llama-bench -m MODEL -p 0 -n 16 -t 8 -ngl 0 -b 32 -ub 32. +- Baseline: unmodified Prism prism branch, same model, host, command, and + eight ggml threads, with no TBKERN selector. + +## Correctness + +llama-debug with prompt Hello, deterministic generation, and the same runtime +settings produced bit-identical files relative to the native path: + +- logits SHA-256: + 65d581397ae42288fe1115e7fa434589700cacd98edc67a750c0fba1765062f2 +- generated-token SHA-256: + 7a7748eacf971049271242b9d921628019d6c44698574e9301da9b8c88026381 +- NMSE: 0 +- top-1 mismatches: 0 +- one-chunk PPL: 1.0645 +/- 0.03006 for both routes + +## Matched decode triplicate + +The recorded three-run matched set used the exact command above: + +| Route | Run 1 | Run 2 | Run 3 | +| --- | ---: | ---: | ---: | +| Native Prism | 1.09 tok/s | 1.09 tok/s | 1.09 tok/s | +| Phase 9 four-row VNNI64 | 1.52 tok/s | 1.53 tok/s | 1.53 tok/s | + +The Phase 9 route is approximately 40% faster for target decode on this host. +These are historical C3 measurements from the equivalent Phase 9 implementation +before this cleanup. This cleanup removes unused files and narrows the packing +header; it does not intentionally change the arithmetic or dispatch. + +## Reproduction status + +The current Windows workspace has no CMake, GCC, Clang, or llama-bench, so the +cleaned branch could not be rebuilt and rerun here. Before marking the PR ready +for review, rebuild this exact commit on the Linux C3 host and repeat native and +Phase 9 llama-bench three times, then rerun the deterministic parity check. + +The route remains opt-in because the recorded perplexity workload was slower +even though the PPL value matched. Attention, norms, RoPE, softmax, GDN +recurrence/convolution, KV cache, sampling, scheduling, and GPU execution +remain in Prism's native runtime. \ No newline at end of file diff --git a/ggml/src/ggml-cpu/CMakeLists.txt b/ggml/src/ggml-cpu/CMakeLists.txt index 8c735a045b35..7f40377c7c97 100644 --- a/ggml/src/ggml-cpu/CMakeLists.txt +++ b/ggml/src/ggml-cpu/CMakeLists.txt @@ -1,3 +1,7 @@ +# Keep this at directory scope: CMAKE_CURRENT_FUNCTION_LIST_DIR was added +# after the project's CMake 3.14 minimum and is not available there. +set(TBKERN_ROOT "${CMAKE_SOURCE_DIR}/third_party/tbkern") + function(ggml_add_cpu_backend_features cpu_name arch) # The feature detection code is compiled as a separate target so that # it can be built without the architecture flags @@ -56,7 +60,14 @@ function(ggml_add_cpu_backend_variant_impl tag_name) target_compile_features(${GGML_CPU_NAME} PRIVATE c_std_11 cxx_std_17) target_include_directories(${GGML_CPU_NAME} PRIVATE . ggml-cpu) - + # Resolve from the project root inside this function to avoid caller-scope shadowing. + set(TBKERN_ROOT "${CMAKE_SOURCE_DIR}/third_party/tbkern") + # Q2_0 cache packing is opt-in. Scalar and AVX2 routes are selected by env; + # Prism's native Q2_0 vec-dot remains the fallback. + target_sources(${GGML_CPU_NAME} PRIVATE + "${TBKERN_ROOT}/src/q2_pack.c" + ) + target_include_directories(${GGML_CPU_NAME} PRIVATE "${TBKERN_ROOT}/include") if (APPLE AND GGML_ACCELERATE) find_library(ACCELERATE_FRAMEWORK Accelerate) if (ACCELERATE_FRAMEWORK) diff --git a/ggml/src/ggml-cpu/repack.cpp b/ggml/src/ggml-cpu/repack.cpp index ef2265251af2..7d5b20724306 100644 --- a/ggml/src/ggml-cpu/repack.cpp +++ b/ggml/src/ggml-cpu/repack.cpp @@ -8,6 +8,7 @@ #include "ggml-cpu-impl.h" #include "simd-mappings.h" #include "traits.h" +#include "quants.h" #include "arch-fallback.h" @@ -15,6 +16,11 @@ #include #include #include // for GGML_ASSERT +#include +#if defined(__AVX2__) || defined(__AVX512F__) || defined(__AVX512VNNI__) || (defined(_MSC_VER) && (defined(GGML_AVX2) || defined(GGML_AVX512))) +#include +#endif +#include "tbkern/q2_pack.h" #include "repack.h" @@ -4762,7 +4768,395 @@ class tensor_traits_base : public ggml::cpu::tensor_traits { public: virtual int repack(struct ggml_tensor * t, const void * data, size_t data_size) = 0; }; +static bool tbkern_q2_0_scalar_enabled() { + const char * v = std::getenv("GGML_TBKERN_Q2_0"); + return v && (std::strcmp(v, "1") == 0 || std::strcmp(v, "true") == 0 || std::strcmp(v, "on") == 0); +} + +static bool tbkern_q2_0_avx2_enabled() { +#if defined(__AVX2__) || (defined(_MSC_VER) && defined(GGML_AVX2)) + const char * v = std::getenv("GGML_TBKERN_Q2_0_AVX2"); + return v && (std::strcmp(v, "1") == 0 || std::strcmp(v, "true") == 0 || std::strcmp(v, "on") == 0); +#else + return false; +#endif +} + +static bool tbkern_q2_0_vnni_enabled() { +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + const char * v = std::getenv("GGML_TBKERN_Q2_0_VNNI"); + return v && (std::strcmp(v, "1") == 0 || std::strcmp(v, "true") == 0 || std::strcmp(v, "on") == 0) && ggml_cpu_has_avx512_vnni(); +#else + return false; +#endif +} + +static bool tbkern_q2_0_vnni64_enabled() { +#if defined(__AVX512F__) && defined(__AVX512VNNI__) + const char * v = std::getenv("GGML_TBKERN_Q2_0_VNNI64"); + return v && (std::strcmp(v, "1") == 0 || std::strcmp(v, "true") == 0 || std::strcmp(v, "on") == 0) && ggml_cpu_has_avx512_vnni(); +#else + return false; +#endif +} + +static bool tbkern_q2_0_vnni64_4r_enabled() { +#if defined(__AVX512F__) && defined(__AVX512VNNI__) + const char * v = std::getenv("GGML_TBKERN_Q2_0_VNNI64_4R"); + return v && (std::strcmp(v, "1") == 0 || std::strcmp(v, "true") == 0 || std::strcmp(v, "on") == 0) && ggml_cpu_has_avx512_vnni(); +#else + return false; +#endif +} + +static bool tbkern_q2_0_enabled() { + return tbkern_q2_0_scalar_enabled() || tbkern_q2_0_avx2_enabled() || tbkern_q2_0_vnni_enabled() || tbkern_q2_0_vnni64_enabled() || tbkern_q2_0_vnni64_4r_enabled(); +} + +static bool tbkern_q2_0_repackable(const ggml_tensor * t) { + if (!t || t->type != GGML_TYPE_Q2_0 || !tbkern_q2_0_enabled()) { + return false; + } + // Token embeddings are consumed by GET_ROWS and must retain native Prism layout. + const char * name = t->name; + if (name && (std::strstr(name, "token_embd") || std::strstr(name, "embeddings"))) { + return false; + } + // Phase 1 calls the native Q8_0 quantizer on the source row. Restrict + // cache ownership to complete, 2D Q2_0 blocks so Kp == K and it never + // reads past the logical activation row. + return t->ne[0] > 0 && t->ne[1] > 0 && + t->ne[0] <= INT32_MAX && t->ne[1] <= INT32_MAX && + t->ne[2] == 1 && t->ne[3] == 1 && + t->ne[0] % QK2_0 == 0; +} + +// The CPU backend owns temporary activations through params->wdata. Keeping +// this format identical to Prism's native Q8_0 input preserves the fp16 scale +// and int8 quantization semantics of the fallback vec-dot path. +static inline size_t tbkern_q2_0_q8_bytes(int K) { + return (size_t) (K / QK8_0) * sizeof(block_q8_0); +} + +static inline size_t tbkern_q2_0_sum_offset(int K) { + return GGML_PAD(tbkern_q2_0_q8_bytes(K), alignof(int32_t)); +} + +static inline size_t tbkern_q2_0_scratch_bytes(int K) { + return tbkern_q2_0_sum_offset(K) + (size_t) (K / QK8_0) * sizeof(int32_t); +} + +// Native Q2_0 bytes remain at tensor->data for generic ggml fallback. The +// opt-in decode cache follows that native payload in the repack allocation. +static inline const uint8_t * tbkern_q2_0_cache_data(const ggml_tensor * t) { + return static_cast(t->data) + ggml_nbytes(t); +} + +static inline uint8_t * tbkern_q2_0_row_ptr(ggml_tensor * t, int64_t row) { + const int64_t i1 = row % t->ne[1]; + row /= t->ne[1]; + const int64_t i2 = row % t->ne[2]; + const int64_t i3 = row / t->ne[2]; + return reinterpret_cast(t->data) + i1*t->nb[1] + i2*t->nb[2] + i3*t->nb[3]; +} + +static inline const uint8_t * tbkern_q2_0_row_ptr(const ggml_tensor * t, int64_t row) { + return tbkern_q2_0_row_ptr(const_cast(t), row); +} + +// Layout B packs each 64 weights into 16 bytes. Byte b holds k=b,b+16,b+32, +// and b+48, respectively. Decode directly from the cache: no per-call row +// expansion or heap allocation is needed. +static inline int tbkern_q2_0_code_at(const uint8_t * row, int k) { + const int within_64 = k & 63; + const uint8_t packed = row[(size_t) (k / 64) * 16 + (within_64 & 15)]; + return (packed >> (2 * (within_64 / 16))) & 0x3; +} +#if defined(__AVX2__) || (defined(_MSC_VER) && defined(GGML_AVX2)) +// Compute sum((code - 1) * q8) for one QK8_0 block. The unsigned code +// values allow maddubs to perform the code*q8 products; subtracting sum(q8) +// applies Prism's signed code mapping without an unsigned/signed mismatch. +static inline int32_t tbkern_q2_0_dot_avx2(const uint8_t * weight_row, int base, const int8_t * q8) { + alignas(32) uint8_t codes[QK8_0]; + for (int j = 0; j < QK8_0; ++j) { + codes[j] = (uint8_t) tbkern_q2_0_code_at(weight_row, base + j); + } + const __m256i c = _mm256_load_si256(reinterpret_cast(codes)); + const __m256i x = _mm256_loadu_si256(reinterpret_cast(q8)); + const __m256i pairs = _mm256_maddubs_epi16(c, x); + const __m256i sums = _mm256_madd_epi16(pairs, _mm256_set1_epi16(1)); + __m128i total = _mm_add_epi32(_mm256_castsi256_si128(sums), _mm256_extracti128_si256(sums, 1)); + total = _mm_hadd_epi32(total, total); + total = _mm_hadd_epi32(total, total); + return _mm_cvtsi128_si32(total); +} +#endif + +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) +// Compute sum((code - 1) * q8) for one QK8_0 block with AVX-512 VNNI. +// dpbusd consumes unsigned code bytes and signed q8 bytes in four-byte +// groups. Subtracting a second dot with an all-ones vector applies Prism's +// signed (code - 1) mapping without changing the native quantizer. +static inline int32_t tbkern_q2_0_dot_vnni(const uint8_t * weight_row, int base, const int8_t * q8) { + // A Prism Q2_0 row stores four 2-bit planes per 64 weights. A Q8_0 + // dot covers either the first or second half of that group, so one + // unaligned 16-byte load contains all 32 codes. Expand the two planes + // in registers rather than materializing one byte at a time. + const int within_64 = base & 63; + GGML_ASSERT(within_64 == 0 || within_64 == 32); + const __m128i packed = _mm_loadu_si128(reinterpret_cast( + weight_row + (size_t) (base / 64) * 16)); + const int shift_lo = (within_64 >> 4) * 2; + const int shift_hi = shift_lo + 2; + const __m128i mask = _mm_set1_epi8(3); + // A 16-bit shift is safe here: after masking with 0x03 per byte, bits + // shifted in from the neighbouring byte are outside the two low bits. + const __m128i lo = _mm_and_si128(_mm_srl_epi16(packed, _mm_cvtsi32_si128(shift_lo)), mask); + const __m128i hi = _mm_and_si128(_mm_srl_epi16(packed, _mm_cvtsi32_si128(shift_hi)), mask); + const __m256i c = _mm256_set_m128i(hi, lo); + const __m256i x = _mm256_loadu_si256(reinterpret_cast(q8)); + const __m256i zero = _mm256_setzero_si256(); + const __m256i prod = _mm256_dpbusd_epi32(zero, c, x); + const __m256i diff = prod; + __m128i total = _mm_add_epi32(_mm256_castsi256_si128(diff), + _mm256_extracti128_si256(diff, 1)); + total = _mm_hadd_epi32(total, total); + total = _mm_hadd_epi32(total, total); + return _mm_cvtsi128_si32(total); +} +#endif + +#if defined(__AVX512F__) && defined(__AVX512VNNI__) +// The Q8 pair is shared by every output row for a single-token decode. +static inline __m512i tbkern_q2_0_q8_pair_vnni64(const int8_t * q8_lo, const int8_t * q8_hi) { + const __m256i xlo = _mm256_loadu_si256(reinterpret_cast(q8_lo)); + const __m256i xhi = _mm256_loadu_si256(reinterpret_cast(q8_hi)); + __m512i x = _mm512_castsi256_si512(xlo); + return _mm512_inserti64x4(x, xhi, 1); +} + +// Decode one packed 64-weight subgroup and dot it against a prepared Q8 pair. +static inline void tbkern_q2_0_dot2_vnni64_x(const uint8_t * weight_row, int base, __m512i x, + int32_t * part_lo, int32_t * part_hi) { + GGML_ASSERT((base & 63) == 0); + const __m128i packed = _mm_loadu_si128(reinterpret_cast( + weight_row + (size_t) (base / 64) * 16)); + const __m128i mask = _mm_set1_epi8(3); + const __m128i c0 = _mm_and_si128(packed, mask); + const __m128i c1 = _mm_and_si128(_mm_srl_epi16(packed, _mm_cvtsi32_si128(2)), mask); + const __m128i c2 = _mm_and_si128(_mm_srl_epi16(packed, _mm_cvtsi32_si128(4)), mask); + const __m128i c3 = _mm_and_si128(_mm_srl_epi16(packed, _mm_cvtsi32_si128(6)), mask); + __m512i codes = _mm512_castsi128_si512(c0); + codes = _mm512_inserti32x4(codes, c1, 1); + codes = _mm512_inserti32x4(codes, c2, 2); + codes = _mm512_inserti32x4(codes, c3, 3); + const __m512i prod = _mm512_dpbusd_epi32(_mm512_setzero_si512(), codes, x); + const __m256i lo = _mm512_castsi512_si256(prod); + const __m256i hi = _mm512_extracti64x4_epi64(prod, 1); + auto reduce = [](__m256i v) { + __m128i total = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1)); + total = _mm_hadd_epi32(total, total); + total = _mm_hadd_epi32(total, total); + return _mm_cvtsi128_si32(total); + }; + *part_lo = reduce(lo); + *part_hi = reduce(hi); +} + +// Decode one packed 64-weight subgroup once, then compute its two Q8_0 +// products together. The cache byte layout is four 16-byte 2-bit planes. +static inline void tbkern_q2_0_dot2_vnni64(const uint8_t * weight_row, int base, + const int8_t * q8_lo, const int8_t * q8_hi, + int32_t * part_lo, int32_t * part_hi) { + tbkern_q2_0_dot2_vnni64_x(weight_row, base, tbkern_q2_0_q8_pair_vnni64(q8_lo, q8_hi), part_lo, part_hi); +} +#endif + +class tbkern_q2_0_traits : public tensor_traits_base { + public: + bool work_size(int, const ggml_tensor * op, size_t & size) override { + if (!op || op->op != GGML_OP_MUL_MAT || !op->src[0] || !op->src[1] || + op->src[0]->type != GGML_TYPE_Q2_0 || op->src[1]->type != GGML_TYPE_F32 || + ggml_nrows(op->src[1]) != 1) { + return false; + } + const int64_t K = op->src[0]->ne[0]; + if (K <= 0 || K > INT32_MAX || K % QK2_0 != 0) { + return false; + } + size = GGML_PAD(tbkern_q2_0_scratch_bytes((int) K), GGML_MEM_ALIGN); + return true; + } + + bool compute_forward(ggml_compute_params * params, ggml_tensor * op) override { + if (!op || op->op != GGML_OP_MUL_MAT) { + return false; + } + const ggml_tensor * src0 = op->src[0]; + const ggml_tensor * src1 = op->src[1]; + ggml_tensor * dst = op; + if (!src0 || !src1 || src0->type != GGML_TYPE_Q2_0 || src1->type != GGML_TYPE_F32 || + ggml_nrows(src1) != 1 || src0->ne[2] != 1 || src0->ne[3] != 1) { + return false; + } + const int K = (int) src0->ne[0]; + const int M = (int) src0->ne[1]; + if (K <= 0 || M <= 0 || K % QK2_0 != 0 || src1->ne[0] != K || dst->ne[0] != M) { + return false; + } + GGML_ASSERT(src1->nb[0] == sizeof(float) && dst->nb[0] == sizeof(float)); + + const int G = QK2_0; + const int NG = K / G; + const int n_q8 = K / QK8_0; + const size_t sum_offset = tbkern_q2_0_sum_offset(K); + const size_t scratch_bytes = tbkern_q2_0_scratch_bytes(K); + GGML_ASSERT(params->wdata && params->wsize >= scratch_bytes); + const uint8_t * cache = tbkern_q2_0_cache_data(src0); + const auto * native = reinterpret_cast(src0->data); + const size_t row_stride = tbk_row_stride_bytes(K); + + auto * xq = reinterpret_cast(params->wdata); + auto * q8_sums = reinterpret_cast(static_cast(params->wdata) + sum_offset); + if (params->ith == 0) { + const float * x = reinterpret_cast(tbkern_q2_0_row_ptr(src1, 0)); + quantize_row_q8_0(x, xq, K); + for (int ib = 0; ib < n_q8; ++ib) { + int32_t sum = 0; + for (int j = 0; j < QK8_0; ++j) { + sum += (int32_t) xq[ib].qs[j]; + } + q8_sums[ib] = sum; + } + } + ggml_barrier(params->threadpool); +#if defined(__AVX2__) || (defined(_MSC_VER) && defined(GGML_AVX2)) + const bool use_avx2 = tbkern_q2_0_avx2_enabled(); +#endif +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + const bool use_vnni = tbkern_q2_0_vnni_enabled(); +#endif +#if defined(__AVX512F__) && defined(__AVX512VNNI__) + const bool use_vnni64_4r = tbkern_q2_0_vnni64_4r_enabled(); + // A four-row tile's final 1--3 rows retain the Phase 8 VNNI64 route. + const bool use_vnni64 = tbkern_q2_0_vnni64_enabled() || use_vnni64_4r; +#endif + const int row_begin = M * params->ith / params->nth; + const int row_end = M * (params->ith + 1) / params->nth; + float * y = reinterpret_cast(tbkern_q2_0_row_ptr(dst, 0)); + int tiled_row = row_begin; +#if defined(__AVX512F__) && defined(__AVX512VNNI__) + if (use_vnni64_4r) { + for (; tiled_row + 3 < row_end; tiled_row += 4) { + const uint8_t * weight_rows[4] = { + cache + (size_t) (tiled_row + 0) * row_stride, + cache + (size_t) (tiled_row + 1) * row_stride, + cache + (size_t) (tiled_row + 2) * row_stride, + cache + (size_t) (tiled_row + 3) * row_stride, + }; + float acc[4] = { 0.0f, 0.0f, 0.0f, 0.0f }; + for (int g = 0; g < NG; ++g) { + const int base = g * G; + float gacc[4] = { 0.0f, 0.0f, 0.0f, 0.0f }; + for (int b = 0; b < G / 64; ++b) { + const int base_b = base + b * 64; + const int ib = base_b / QK8_0; + const __m512i xpair = tbkern_q2_0_q8_pair_vnni64(xq[ib].qs, xq[ib + 1].qs); + for (int r = 0; r < 4; ++r) { + int32_t part_lo; + int32_t part_hi; + tbkern_q2_0_dot2_vnni64_x(weight_rows[r], base_b, xpair, &part_lo, &part_hi); + gacc[r] += GGML_CPU_FP16_TO_FP32(xq[ib].d) * (float) (part_lo - q8_sums[ib]); + gacc[r] += GGML_CPU_FP16_TO_FP32(xq[ib + 1].d) * (float) (part_hi - q8_sums[ib + 1]); + } + } + for (int r = 0; r < 4; ++r) { + acc[r] += GGML_CPU_FP16_TO_FP32(native[(size_t) (tiled_row + r) * NG + g].d) * gacc[r]; + } + } + y[tiled_row + 0] = acc[0]; + y[tiled_row + 1] = acc[1]; + y[tiled_row + 2] = acc[2]; + y[tiled_row + 3] = acc[3]; + } + } +#endif + for (int row = tiled_row; row < row_end; ++row) { + const uint8_t * weight_row = cache + (size_t) row * row_stride; + float acc = 0.0f; + for (int g = 0; g < NG; ++g) { + const int base = g * G; + float gacc = 0.0f; +#if defined(__AVX512F__) && defined(__AVX512VNNI__) + if (use_vnni64) { + // Keep Q8_0 scales and floating-point accumulation in the + // native 32-element order even though code expansion is 64-wide. + for (int b = 0; b < G / 64; ++b) { + const int base_b = base + b * 64; + const int ib = base_b / QK8_0; + int32_t part_lo; + int32_t part_hi; + tbkern_q2_0_dot2_vnni64(weight_row, base_b, xq[ib].qs, xq[ib + 1].qs, &part_lo, &part_hi); + gacc += GGML_CPU_FP16_TO_FP32(xq[ib].d) * (float) (part_lo - q8_sums[ib]); + gacc += GGML_CPU_FP16_TO_FP32(xq[ib + 1].d) * (float) (part_hi - q8_sums[ib + 1]); + } + } else +#endif + { + for (int b = 0; b < G / QK8_0; ++b) { + const int base_b = base + b * QK8_0; + const int ib = base_b / QK8_0; + int32_t part; +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + if (use_vnni) { + part = tbkern_q2_0_dot_vnni(weight_row, base_b, xq[ib].qs); + } else +#endif +#if defined(__AVX2__) || (defined(_MSC_VER) && defined(GGML_AVX2)) + if (use_avx2) { + part = tbkern_q2_0_dot_avx2(weight_row, base_b, xq[ib].qs); + } else +#endif + { + part = 0; + for (int j = 0; j < QK8_0; ++j) { + part += tbkern_q2_0_code_at(weight_row, base_b + j) * (int32_t) xq[ib].qs[j]; + } + } + gacc += GGML_CPU_FP16_TO_FP32(xq[ib].d) * (float) (part - q8_sums[ib]); + } + } + acc += GGML_CPU_FP16_TO_FP32(native[(size_t) row * NG + g].d) * gacc; + } + y[row] = acc; + } + return true; + } + + int repack(ggml_tensor * t, const void * data, size_t data_size) override { + if (!t || !data || !tbkern_q2_0_repackable(t)) { + return TBK_EINVAL; + } + const int K = (int) t->ne[0]; + const int M = (int) t->ne[1]; + const size_t native_bytes = ggml_nbytes(t); + if (data_size != native_bytes || K % QK2_0 != 0) { + return TBK_EINVAL; + } + std::memmove(t->data, data, native_bytes); + tbk_mat packed{}; + const int rc = tbk_pack_from_q2(reinterpret_cast(t->data), M, K, QK2_0, TBK_LAYOUT_CODES, &packed); + if (rc != TBK_OK) { + return rc; + } + GGML_ASSERT(packed.Kp == K && packed.row_stride_bytes == tbk_row_stride_bytes(K)); + const size_t cache_bytes = (size_t) M * packed.row_stride_bytes; + std::memcpy(static_cast(t->data) + native_bytes, packed.data, cache_bytes); + tbk_mat_free(&packed); + return TBK_OK; + } +}; template class tensor_traits : public tensor_traits_base { bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override { @@ -5134,6 +5528,7 @@ template q4_0_4x4_q8_0; static const ggml::cpu::repack::tensor_traits q4_0_4x8_q8_0; @@ -5185,7 +5580,9 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons static const ggml::cpu::repack::tensor_traits q2_K_16x1_q8_K; #endif - if (cur->type == GGML_TYPE_Q4_0) { + if (ggml::cpu::repack::tbkern_q2_0_repackable(cur)) { + return &tbkern_q2_0; + } else if (cur->type == GGML_TYPE_Q4_0) { if (ggml_cpu_has_avx2() || (ggml_cpu_has_sve() && ggml_cpu_has_matmul_int8() && ggml_cpu_get_sve_cnt() == QK8_0)) { if (cur->ne[1] % 8 == 0) { return &q4_0_8x8_q8_0; @@ -5385,6 +5782,17 @@ static const char * ggml_backend_cpu_repack_buffer_type_get_name(ggml_backend_bu GGML_UNUSED(buft); } +static size_t ggml_backend_cpu_repack_buffer_type_get_alloc_size(ggml_backend_buffer_type_t, const ggml_tensor * tensor) { + const size_t native_bytes = ggml_nbytes(tensor); + if (ggml::cpu::repack::tbkern_q2_0_repackable(tensor)) { + const size_t K = (size_t) tensor->ne[0]; + const size_t M = (size_t) tensor->ne[1]; + const size_t row_stride = tbk_row_stride_bytes((int) K); + GGML_ASSERT(M <= (SIZE_MAX - native_bytes) / row_stride); + return native_bytes + M * row_stride; + } + return native_bytes; +} static ggml_backend_buffer_t ggml_backend_cpu_repack_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); @@ -5462,7 +5870,7 @@ ggml_backend_buffer_type_t ggml_backend_cpu_repack_buffer_type(void) { /* .alloc_buffer = */ ggml_backend_cpu_repack_buffer_type_alloc_buffer, /* .get_alignment = */ ggml_backend_cpu_repack_buffer_type_get_alignment, /* .get_max_size = */ nullptr, // defaults to SIZE_MAX - /* .get_alloc_size = */ nullptr, // defaults to ggml_nbytes + /* .get_alloc_size = */ ggml_backend_cpu_repack_buffer_type_get_alloc_size, /* .is_host = */ nullptr, }, /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0), diff --git a/scripts/tbkern_compare.py b/scripts/tbkern_compare.py new file mode 100644 index 000000000000..e725d34eb830 --- /dev/null +++ b/scripts/tbkern_compare.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Compare Prism native and opt-in tbkern llama-debug artifacts. + +The comparator is deliberately read-only: it never invokes llama.cpp or +modifies either input. Logit files are little-endian float32 values as +written by examples/debug; token files are compared byte-for-byte because +the debug utility's token representation is part of the artifact contract. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import struct +import sys +from pathlib import Path +from typing import Any + + +SCHEMA = "tbkern.parity.v1" + + +def _read_f32(path: Path) -> tuple[list[float], bytes]: + raw = path.read_bytes() + if not raw: + raise ValueError(f"empty logits file: {path}") + if len(raw) % 4: + raise ValueError(f"logits file is not a multiple of 4 bytes: {path}") + return list(struct.unpack(f"<{len(raw) // 4}f", raw)), raw + + +def _sha256(raw: bytes) -> str: + return hashlib.sha256(raw).hexdigest() + + +def _file_metadata(path: Path, raw: bytes, values: list[float]) -> dict[str, Any]: + finite = [value for value in values if math.isfinite(value)] + return { + "path": str(path), + "bytes": len(raw), + "sha256": _sha256(raw), + "float_count": len(values), + "finite_count": len(finite), + "min": min(finite) if finite else None, + "max": max(finite) if finite else None, + } + + +def _argmax(values: list[float]) -> int: + if not values: + raise ValueError("logits contain no values") + return max(range(len(values)), key=values.__getitem__) + + +def _row_top1(native: list[float], tbkern: list[float], vocab_size: int | None) -> dict[str, Any]: + # llama-debug writes the final logits vector, not a sequence of rows. + # Keep the optional argument as an explicit artifact-shape assertion. + if vocab_size is not None and vocab_size != len(native): + raise ValueError("--vocab-size must equal the logits vector length") + native_ids = [_argmax(native)] + tbkern_ids = [_argmax(tbkern)] + return { + "rows": 1, + "vocab_size": len(native), + "mismatched_rows": int(native_ids[0] != tbkern_ids[0]), + "native": native_ids, + "tbkern": tbkern_ids, + } + + +def compare(args: argparse.Namespace) -> dict[str, Any]: + for name in ("nmse_tolerance", "relative_floor"): + value = getattr(args, name) + if not math.isfinite(value) or value < 0: + raise ValueError(f"--{name.replace('_', '-')} must be a finite non-negative number") + native_path = Path(args.native_logits) + tbkern_path = Path(args.tbkern_logits) + native, native_raw = _read_f32(native_path) + tbkern, tbkern_raw = _read_f32(tbkern_path) + if len(native) != len(tbkern): + raise ValueError( + f"logit lengths differ: native={len(native)} tbkern={len(tbkern)}" + ) + if any(not math.isfinite(value) for value in native + tbkern): + raise ValueError("logit artifacts contain NaN or infinity") + + deltas = [abs(a - b) for a, b in zip(native, tbkern)] + denominator = sum(value * value for value in native) + nmse = sum(delta * delta for delta in deltas) / denominator if denominator else 0.0 + nonzero = [delta / abs(value) for value, delta in zip(native, deltas) if abs(value) > args.relative_floor] + top1 = _row_top1(native, tbkern, args.vocab_size) + + token_result: dict[str, Any] | None = None + if bool(args.native_tokens) != bool(args.tbkern_tokens): + raise ValueError("provide both --native-tokens and --tbkern-tokens, or neither") + if args.native_tokens: + native_token_path = Path(args.native_tokens) + tbkern_token_path = Path(args.tbkern_tokens) + native_token_raw = native_token_path.read_bytes() + tbkern_token_raw = tbkern_token_path.read_bytes() + token_result = { + "native": { + "path": str(native_token_path), + "bytes": len(native_token_raw), + "sha256": _sha256(native_token_raw), + }, + "tbkern": { + "path": str(tbkern_token_path), + "bytes": len(tbkern_token_raw), + "sha256": _sha256(tbkern_token_raw), + }, + "equal": native_token_raw == tbkern_token_raw, + } + + # NMSE has no meaning when the reference vector is all zeros. Only an + # exactly equal vector passes that degenerate case. + nmse_pass = nmse <= args.nmse_tolerance and (denominator != 0.0 or max(deltas) == 0.0) + top1_pass = top1["mismatched_rows"] == 0 + token_pass = token_result is None or token_result["equal"] + result: dict[str, Any] = { + "schema": SCHEMA, + "gates": { + "nmse": {"value": nmse, "limit": args.nmse_tolerance, "pass": nmse_pass}, + "top1": {"mismatched_rows": top1["mismatched_rows"], "pass": top1_pass}, + "tokens": {"present": token_result is not None, "pass": token_pass}, + }, + "logits": { + "native": _file_metadata(native_path, native_raw, native), + "tbkern": _file_metadata(tbkern_path, tbkern_raw, tbkern), + "max_abs": max(deltas), + "max_rel": max(nonzero) if nonzero else 0.0, + "nmse": nmse, + "reference_sum_sq": denominator, + "top1": top1, + }, + "tokens": token_result, + "tolerance": { + "nmse": args.nmse_tolerance, + "relative_floor": args.relative_floor, + }, + "pass": nmse_pass and top1_pass and token_pass, + } + return result + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--native-logits", required=True, help="Prism fallback float32 artifact") + parser.add_argument("--tbkern-logits", required=True, help="GGML_TBKERN_Q2_0=1 float32 artifact") + parser.add_argument("--native-tokens", help="native token artifact (must be paired)") + parser.add_argument("--tbkern-tokens", help="tbkern token artifact (must be paired)") + parser.add_argument("--vocab-size", type=int, help="compare top-1 independently for each row") + parser.add_argument("--nmse-tolerance", type=float, default=1e-4) + parser.add_argument("--relative-floor", type=float, default=1e-6) + parser.add_argument("--json-out", type=Path, help="optional path for the JSON report") + args = parser.parse_args() + try: + result = compare(args) + except (OSError, ValueError, struct.error) as error: + print(f"tbkern parity: ERROR: {error}", file=sys.stderr) + return 2 + rendered = json.dumps(result, indent=2, sort_keys=True) + print(rendered) + if args.json_out: + try: + args.json_out.write_text(rendered + "\n", encoding="utf-8") + except OSError as error: + print(f"tbkern parity: ERROR: cannot write JSON report: {error}", file=sys.stderr) + return 2 + return 0 if result["pass"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/third_party/tbkern/include/tbkern/format.h b/third_party/tbkern/include/tbkern/format.h new file mode 100644 index 000000000000..8593517203ff --- /dev/null +++ b/third_party/tbkern/include/tbkern/format.h @@ -0,0 +1,155 @@ +/* tbkern/format.h — FROZEN Phase-0 contract. + * + * Low-level, dependency-free format constants and helpers shared by every + * translation unit (C kernels, packers) and mirrored by the Python reference. + * + * Ground truth: mainline llama.cpp GGML_TYPE_Q2_0 (ggml-common.h / ggml-quants.c). + * See docs/FORMAT.md for the normative spec and bit-exact diagrams. + * + * RULE (docs/CONTEXT.md §Rules): after Phase 0 this header may only be changed + * by the orchestrator. Code against it; do not edit it. + */ +#ifndef TBKERN_FORMAT_H +#define TBKERN_FORMAT_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* ------------------------------------------------------------------ * + * Block / group constants + * ------------------------------------------------------------------ */ + +/* Mainline GGUF Q2_0 native block covers 64 weights (QK2_0). PrismML's native + * variant uses group-128 blocks with identical code semantics. All tbkern code + * parameterizes the group size G at runtime; these are the only legal values. */ +#define TBK_QK2_0 64 /* mainline GGUF Q2_0 block length */ +#define TBK_G_MIN 64 +#define TBK_G_MAX 128 + +/* A GGUF Q2_0 block is: 1x fp16 scale ("d") followed by G/4 code bytes + * (2 bits per weight, 4 weights per byte). Byte count per block: */ +#define TBK_GGUF_SCALE_BYTES 2 /* sizeof(ggml_half) */ +#define TBK_GGUF_BLOCK_BYTES(G) (TBK_GGUF_SCALE_BYTES + (size_t)(G) / 4u) +/* G=64 -> 2 + 16 = 18 bytes (2.250 bits/wt) + * G=128 -> 2 + 32 = 34 bytes (2.125 bits/wt) */ + +/* Code semantics (2-bit quant code -> ternary/quaternary weight value): + * 00 = -1, 01 = 0, 10 = +1, 11 = +2. value = (code - 1) * scale. + * The unsigned datum u = code = (w + 1) in {0..3} is what the VNNI path feeds + * to vpdpbusd (see docs/FORMAT.md §VNNI). +2 (code 3) is legal Q2_0 but never + * appears in real ternary Bonsai tensors; the bitplane path rejects it. */ +#define TBK_CODE_NEG1 0u /* w = -1 */ +#define TBK_CODE_ZERO 1u /* w = 0 (also the pad code) */ +#define TBK_CODE_POS1 2u /* w = +1 */ +#define TBK_CODE_POS2 3u /* w = +2 (legal, ternary-only paths reject)*/ + +/* Padding: logical K is padded up to Kp, the next multiple of G, using pad + * code 01 (w = 0). Activations are zero-padded to Kp. See docs/FORMAT.md. */ +#define TBK_CODE_PAD TBK_CODE_ZERO + +/* Row alignment for packed kernel-native layouts (both bitplane and codes): + * every row starts on a 64-byte boundary; the packed bytes per row equal + * Kp/4 (both layouts), rounded up to this alignment for the row stride. */ +#define TBK_ROW_ALIGN 64u + +/* Number of quant groups in a padded row (one fp32 scale per group). */ +#define TBK_NUM_GROUPS(Kp, G) ((size_t)(Kp) / (size_t)(G)) + +/* Number of 64-weight subgroups in a padded row (codes-layout storage unit). */ +#define TBK_NUM_SUBGROUPS(Kp) ((size_t)(Kp) / 64u) + +/* Packed bytes actually used per row (before 64-byte stride rounding). Both + * layouts store 2 bits/weight = Kp/4 bytes. */ +#define TBK_ROW_USED_BYTES(Kp) ((size_t)(Kp) / 4u) + +/* Round x up to a multiple of a (a must be a power of two). */ +static inline size_t tbk_align_up(size_t x, size_t a) { + return (x + (a - 1u)) & ~(a - 1u); +} + +/* 64-byte-aligned row stride (bytes) for a packed layout of Kp columns. */ +static inline size_t tbk_row_stride_bytes(int Kp) { + return tbk_align_up(TBK_ROW_USED_BYTES(Kp), TBK_ROW_ALIGN); +} + +/* Pad K up to the next multiple of G. */ +static inline int tbk_pad_k(int K, int G) { + return (int)(((K + G - 1) / G) * G); +} + +/* Number of GGUF Q2_0 blocks per row for a K-wide row (ceil division). */ +static inline int tbk_gguf_blocks_per_row(int K, int G) { + return (K + G - 1) / G; +} + +/* ------------------------------------------------------------------ * + * IEEE-754 binary16 <-> binary32 (portable, ggml/numpy bit-compatible) + * + * Verified bit-exact against numpy.float16 over all 65536 half patterns + * (widening) and all finite round-trips (narrowing, round-to-nearest-even). + * Software only: usable from baseline (no-F16C) translation units. + * ------------------------------------------------------------------ */ + +static inline float tbk_fp16_to_fp32(uint16_t h) { + uint32_t sign = (uint32_t)(h & 0x8000u) << 16; + uint32_t exp = (h >> 10) & 0x1Fu; + uint32_t mant = h & 0x3FFu; + uint32_t f; + if (exp == 0u) { + if (mant == 0u) { + f = sign; /* +/- zero */ + } else { /* subnormal -> normalize */ + exp = 1u; + while ((mant & 0x400u) == 0u) { mant <<= 1; exp--; } + mant &= 0x3FFu; + f = sign | ((exp + 112u) << 23) | (mant << 13); + } + } else if (exp == 0x1Fu) { + f = sign | 0x7F800000u | (mant << 13); /* inf / nan */ + } else { + f = sign | ((exp + 112u) << 23) | (mant << 13); /* normal (112=127-15)*/ + } + union { uint32_t u; float f; } o; + o.u = f; + return o.f; +} + +static inline uint16_t tbk_fp32_to_fp16(float value) { + union { float f; uint32_t u; } in; + in.f = value; + uint32_t x = in.u; + uint32_t sign = (x >> 16) & 0x8000u; + uint32_t e = (x >> 23) & 0xFFu; + uint32_t mant = x & 0x7FFFFFu; + int32_t exp = (int32_t)e - 127 + 15; + if (e == 0xFFu) { /* inf / nan */ + return (uint16_t)(sign | 0x7C00u | (mant ? 0x200u : 0u)); + } + if (exp >= 0x1F) { /* overflow -> inf */ + return (uint16_t)(sign | 0x7C00u); + } + if (exp <= 0) { /* subnormal / underflow */ + if (exp < -10) return (uint16_t)sign; /* rounds to +/- zero */ + mant |= 0x800000u; + int shift = 14 - exp; + uint32_t half = 1u << (shift - 1); + uint32_t res = mant >> shift; + uint32_t rem = mant & ((1u << shift) - 1u); + if (rem > half || (rem == half && (res & 1u))) res++; /* RNE */ + return (uint16_t)(sign | res); + } + uint16_t h = (uint16_t)(sign | ((uint32_t)exp << 10) | (mant >> 13)); + uint32_t rem = mant & 0x1FFFu; + if (rem > 0x1000u || (rem == 0x1000u && (h & 1u))) h++; /* RNE */ + return h; +} + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* TBKERN_FORMAT_H */ diff --git a/third_party/tbkern/include/tbkern/q2_pack.h b/third_party/tbkern/include/tbkern/q2_pack.h new file mode 100644 index 000000000000..bdb881548325 --- /dev/null +++ b/third_party/tbkern/include/tbkern/q2_pack.h @@ -0,0 +1,35 @@ +/* q2_pack.h - minimal Q2_0 repack contract used by Prism's CPU backend. + * + * This is intentionally limited to the cache packer consumed by + * ggml/src/ggml-cpu/repack.cpp. It is not a general TBKERN kernel API. + */ +#ifndef TBKERN_Q2_PACK_H +#define TBKERN_Q2_PACK_H +#include +#include +#include "tbkern/format.h" +#ifdef __cplusplus +extern "C" { +#endif +#define TBK_OK 0 +#define TBK_EINVAL -1 +#define TBK_ENOMEM -2 +#define TBK_ENOTTERNARY -4 +typedef enum { TBK_LAYOUT_BITPLANE = 0, TBK_LAYOUT_CODES = 1 } tbk_layout_kind; +typedef struct { + int32_t M, K, Kp, G; + tbk_layout_kind layout; + void * data; + size_t row_stride_bytes; + float * scales; + float * bias; + int32_t ternary_only; + void * _owner; +} tbk_mat; +int tbk_pack_from_q2(const uint8_t * gguf_q2_blocks, int M, int K, int G, + tbk_layout_kind layout, tbk_mat * out); +void tbk_mat_free(tbk_mat * mat); +#ifdef __cplusplus +} +#endif +#endif diff --git a/third_party/tbkern/src/q2_pack.c b/third_party/tbkern/src/q2_pack.c new file mode 100644 index 000000000000..87c487f305a9 --- /dev/null +++ b/third_party/tbkern/src/q2_pack.c @@ -0,0 +1,155 @@ +/* q2_pack.c - GGUF Q2_0 to kernel-native cache packing. */ +#include "tbkern/q2_pack.h" + +#include +#include +#include +#if defined(_MSC_VER) +#include +#endif + +#if defined(_MSC_VER) +static void * tbk_aligned_alloc(size_t alignment, size_t size) { return _aligned_malloc(size, alignment); } +static void tbk_aligned_free(void * p) { _aligned_free(p); } +#else +static void * tbk_aligned_alloc(size_t alignment, size_t size) { return aligned_alloc(alignment, size); } +static void tbk_aligned_free(void * p) { free(p); } +#endif + +/* ------------------------------------------------------------------ * + * GGUF block reading helpers + * ------------------------------------------------------------------ */ + +/* Read 2-bit code of local weight j (0..G-1) from a block's qs bytes. */ +static inline uint8_t block_code(const uint8_t *qs, int j) { + return (uint8_t)((qs[j >> 2] >> ((j & 3) * 2)) & 0x3u); +} + +/* Expand one GGUF row (nb = Kp/G blocks) into Kp codes; pad columns [K,Kp) are + * forced to TBK_CODE_PAD (01 = w0). Also fills the row's fp32 group scales. */ +static void read_row(const uint8_t *row_blocks, int K, int G, int Kp, + uint8_t *codes /* len Kp */, float *scales /* len Kp/G */) { + const int nb = Kp / G; /* == ceil(K/G) */ + const size_t bb = TBK_GGUF_BLOCK_BYTES(G); + for (int i = 0; i < nb; i++) { + const uint8_t *blk = row_blocks + (size_t)i * bb; + uint16_t hbits = (uint16_t)(blk[0] | ((uint16_t)blk[1] << 8)); /* LE fp16 */ + scales[i] = tbk_fp16_to_fp32(hbits); + const uint8_t *qs = blk + TBK_GGUF_SCALE_BYTES; + for (int j = 0; j < G; j++) { + int k = i * G + j; + codes[k] = (k < K) ? block_code(qs, j) : (uint8_t)TBK_CODE_PAD; + } + } +} + +/* ------------------------------------------------------------------ * + * Layout packers (from a flat Kp-code array) + * ------------------------------------------------------------------ */ + +/* Bitplane (Layout A): per 16-weight chunk two LE mask words [P, N]. + * Returns TBK_ENOTTERNARY on a code-3 (+2) weight. */ +static int pack_row_bitplane(const uint8_t *codes, int Kp, uint8_t *row) { + const int nchunks = Kp / 16; + for (int c = 0; c < nchunks; c++) { + uint16_t P = 0, N = 0; + for (int i = 0; i < 16; i++) { + uint8_t code = codes[16 * c + i]; + if (code == TBK_CODE_POS1) P |= (uint16_t)(1u << i); + else if (code == TBK_CODE_NEG1) N |= (uint16_t)(1u << i); + else if (code == TBK_CODE_POS2) return TBK_ENOTTERNARY; + /* TBK_CODE_ZERO: leave both clear */ + } + uint8_t *w = row + (size_t)c * 4; + w[0] = (uint8_t)(P & 0xFF); w[1] = (uint8_t)(P >> 8); + w[2] = (uint8_t)(N & 0xFF); w[3] = (uint8_t)(N >> 8); + } + return TBK_OK; +} + +/* Interleaved codes (Layout B): per 64-weight subgroup 16 bytes, byte b packs + * codes for k = base+b, +16, +32, +48. Sets *any_plus2 if a code-3 is present. */ +static void pack_row_codes(const uint8_t *codes, int Kp, uint8_t *row, + int *any_plus2) { + const int ns = Kp / 64; + for (int s = 0; s < ns; s++) { + const int base = 64 * s; + for (int b = 0; b < 16; b++) { + uint8_t c0 = codes[base + 0 + b]; + uint8_t c1 = codes[base + 16 + b]; + uint8_t c2 = codes[base + 32 + b]; + uint8_t c3 = codes[base + 48 + b]; + if ((c0 | c1 | c2 | c3) & TBK_CODE_POS2) { + if (c0 == TBK_CODE_POS2 || c1 == TBK_CODE_POS2 || + c2 == TBK_CODE_POS2 || c3 == TBK_CODE_POS2) *any_plus2 = 1; + } + row[(size_t)s * 16 + b] = (uint8_t)(c0 | (c1 << 2) | (c2 << 4) | (c3 << 6)); + } + } +} + +/* ------------------------------------------------------------------ * + * Public packing API + * ------------------------------------------------------------------ */ +int tbk_pack_from_q2(const uint8_t *gguf_q2_blocks, int M, int K, int G, + tbk_layout_kind layout, tbk_mat *out) { + if (!gguf_q2_blocks || !out || M <= 0 || K <= 0) return TBK_EINVAL; + if (G != 64 && G != 128) return TBK_EINVAL; + if (layout != TBK_LAYOUT_BITPLANE && layout != TBK_LAYOUT_CODES) return TBK_EINVAL; + + const int Kp = tbk_pad_k(K, G); + const int nb = Kp / G; /* blocks == groups/row */ + const size_t bb = TBK_GGUF_BLOCK_BYTES(G); + const size_t stride = tbk_row_stride_bytes(Kp); + const size_t data_sz = (size_t)M * stride; + const size_t scales_sz = (size_t)M * nb * sizeof(float); + + memset(out, 0, sizeof(*out)); + + uint8_t *data = (uint8_t *)tbk_aligned_alloc(TBK_ROW_ALIGN, + tbk_align_up(data_sz, TBK_ROW_ALIGN)); + float *scales = (float *)malloc(scales_sz); + uint8_t *codes = (uint8_t *)malloc((size_t)Kp); /* scratch, one row */ + if (!data || !scales || !codes) { + tbk_aligned_free(data); free(scales); free(codes); + return TBK_ENOMEM; + } + memset(data, 0, data_sz); + + int any_plus2 = 0; + int rc = TBK_OK; + for (int r = 0; r < M; r++) { + const uint8_t *row_blocks = gguf_q2_blocks + (size_t)r * nb * bb; + uint8_t *row = data + (size_t)r * stride; + read_row(row_blocks, K, G, Kp, codes, scales + (size_t)r * nb); + if (layout == TBK_LAYOUT_BITPLANE) { + rc = pack_row_bitplane(codes, Kp, row); + if (rc != TBK_OK) break; + } else { + pack_row_codes(codes, Kp, row, &any_plus2); + } + } + free(codes); + if (rc != TBK_OK) { tbk_aligned_free(data); free(scales); return rc; } + + out->M = M; + out->K = K; + out->Kp = Kp; + out->G = G; + out->layout = layout; + out->data = data; + out->row_stride_bytes = stride; + out->scales = scales; + out->bias = NULL; + out->ternary_only = (layout == TBK_LAYOUT_BITPLANE) ? 1 : (any_plus2 ? 0 : 1); + out->_owner = data; + return TBK_OK; +} + +void tbk_mat_free(tbk_mat *mat) { + if (!mat || !mat->_owner) { if (mat) memset(mat, 0, sizeof(*mat)); return; } + tbk_aligned_free(mat->data); + free(mat->scales); + /* bias is not owned by the packer (always NULL here) */ + memset(mat, 0, sizeof(*mat)); +}