// PQ2_0 rows are 128 wide, so a model may legitimately hold tensors whose row length is any
// multiple of 128 -- 128 and 384 as much as 256. The AVX2 kernels added for PQ2_0 consume Q8_K
// activations, which only exist in 256-element blocks, so routing every PQ2_0 matmul through
// Q8_K breaks exactly those rows. Guarding the quantizer is not enough: tensors already inside
// published models never pass through it.
//
// This exercises pre-existing PQ2_0 tensors directly (reference-encoded, as a model file holds
// them) for K in {128, 256, 384}, single-token and batched, and checks the numbers rather than
// just survival.
#include "ggml.h"
#include "ggml-cpu.h"
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>
static std::vector<float> make_weights(int64_t n, uint32_t seed) {
std::mt19937 rng(seed);
std::normal_distribution<float> dist(0.0f, 0.5f);
std::vector<float> v(n);
for (int64_t i = 0; i < n; i++) {
v[i] = dist(rng);
}
return v;
}
// rms(computed - reference) / rms(reference), with the reference taken from the dequantized
// weights and exact float activations. Activation quantization (Q8_K or Q8_0) is the only
// remaining error term, so this stays far below the threshold when the path is correct and
// explodes when it is not.
static double rel_rms(const std::vector<float> & got, const std::vector<float> & ref) {
double se = 0.0;
double sr = 0.0;
for (size_t i = 0; i < ref.size(); i++) {
const double d = (double) got[i] - (double) ref[i];
se += d * d;
sr += (double) ref[i] * (double) ref[i];
}
if (sr == 0.0) {
return se == 0.0 ? 0.0 : INFINITY;
}
return std::sqrt(se / sr);
}
static bool run_shape(int64_t K, int64_t M, int64_t N) {
const size_t mem_size = 64u * 1024 * 1024;
struct ggml_init_params ip = { mem_size, nullptr, false };
struct ggml_context * ctx = ggml_init(ip);
if (!ctx) {
printf(" K=%4lld M=%lld N=%lld : ggml_init failed\n", (long long) K, (long long) M, (long long) N);
return false;
}
// A pre-existing PQ2_0 weight tensor: encoded by the reference row quantizer, which is what
// a .gguf holds. It never sees llama-quant.cpp's type fallback.
struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_PQ2_0, K, M);
const std::vector<float> w = make_weights(K * M, 1234 + (uint32_t) K);
ggml_quantize_chunk(GGML_TYPE_PQ2_0, w.data(), a->data, 0, M, K, nullptr);
struct ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, N);
const std::vector<float> x = make_weights(K * N, 99);
memcpy(b->data, x.data(), x.size() * sizeof(float));
struct ggml_tensor * c = ggml_mul_mat(ctx, a, b);
struct ggml_cgraph * gf = ggml_new_graph(ctx);
ggml_build_forward_expand(gf, c);
ggml_graph_compute_with_ctx(ctx, gf, 4);
// reference: dequantized weights x exact float activations
std::vector<float> a_deq(K * M);
const ggml_type_traits * tt = ggml_get_type_traits(GGML_TYPE_PQ2_0);
tt->to_float(a->data, a_deq.data(), K * M);
std::vector<float> ref(M * N);
for (int64_t j = 0; j < N; j++) {
for (int64_t i = 0; i < M; i++) {
double s = 0.0;
for (int64_t k = 0; k < K; k++) {
s += (double) a_deq[i * K + k] * (double) x[j * K + k];
}
ref[j * M + i] = (float) s;
}
}
std::vector<float> got(M * N);
memcpy(got.data(), c->data, got.size() * sizeof(float));
const double err = rel_rms(got, ref);
const bool ok = err < 0.05;
printf(" K=%4lld M=%lld N=%2lld : %-4s rel_rms=%.4f\n",
(long long) K, (long long) M, (long long) N, ok ? "OK" : "FAIL", err);
ggml_free(ctx);
return ok;
}
int main() {
printf("test-pq2-row-shapes: PQ2_0 matmul over every legal row width\n");
bool all_ok = true;
// 128 and 384 are legal PQ2_0 row widths that are NOT multiples of the Q8_K block (256).
for (int64_t K : { (int64_t) 128, (int64_t) 256, (int64_t) 384, (int64_t) 512, (int64_t) 640 }) {
for (int64_t N : { (int64_t) 1, (int64_t) 4, (int64_t) 16 }) {
all_ok &= run_shape(K, 4, N);
}
}
printf("%s\n", all_ok ? "PASS" : "FAIL");
return all_ok ? 0 : 1;
}