fix: inverse WHT in test-turbo-quant.c round-trip (#59)

Author: TheTom

ggml/src/ggml-turbo-quant.c

@@ -236,6 +236,37 @@ static void turbo_cpu_fwht(float * x, int group_size) {
     for (int i = 0; i < group_size; i++) x[i] *= inv_sqrt * s2[i];
 }
 
+/* ---------- CPU inverse WHT (in-place, group_size elements) ----------
+ *
+ * Forward is  y = D(s2) * N * H * D(s1) * x   (N = 1/sqrt(group_size))
+ * H is the unnormalized Hadamard butterfly with H*H = group_size * I, so
+ * (N*H) is self-inverse.  s1 and s2 are ±1 diagonals, also self-inverse.
+ * The inverse therefore has the same structure with s1 and s2 swapped:
+ *     x = D(s1) * N * H * D(s2) * y
+ */
+GGML_API void turbo_cpu_fwht_inverse(float * x, int group_size) {
+    const float * s1 = turbo_cpu_s1;
+    const float * s2 = turbo_cpu_s2;
+    const float inv_sqrt = (group_size == 128) ? 0.08838834764831845f : 0.125f;
+
+    // signs2 (undoes the s2 that was applied last in the forward pass)
+    for (int i = 0; i < group_size; i++) x[i] *= s2[i];
+
+    // butterfly stages (same as forward — self-inverse up to the inv_sqrt scaling below)
+    for (int h = 1; h < group_size; h *= 2) {
+        for (int i = 0; i < group_size; i += h * 2) {
+            for (int j = i; j < i + h; j++) {
+                float a = x[j], b = x[j + h];
+                x[j]     = a + b;
+                x[j + h] = a - b;
+            }
+        }
+    }
+
+    // normalize + signs1
+    for (int i = 0; i < group_size; i++) x[i] *= inv_sqrt * s1[i];
+}
+
 /* ---------- TURBO3_0: 3-bit PolarQuant with WHT rotation ---------- */
 
 void quantize_row_turbo3_0_ref(const float * GGML_RESTRICT x, block_turbo3_0 * GGML_RESTRICT y, int64_t k) {

tests/test-turbo-quant.c

@@ -6,6 +6,7 @@ extern void quantize_row_turbo3_0_ref(const float * x, void * y, long long k);
 extern void dequantize_row_turbo3_0(const void * x, float * y, long long k);
 extern void quantize_row_turbo4_0_ref(const float * x, void * y, long long k);
 extern void dequantize_row_turbo4_0(const void * x, float * y, long long k);
+extern void turbo_cpu_fwht_inverse(float * x, int group_size);
 
 int main(void) {
     const int d = 128;
@@ -15,11 +16,17 @@ int main(void) {
 
     printf("=== TurboQuant C Round-Trip Test ===\n\n");
 
-    /* Test 1: basis vector */
+    /* Test 1: basis vector
+     *
+     * dequantize_row_turbo3_0 leaves output in the WHT-rotated domain (Q is
+     * also rotated by the graph, so <Q_rot, K_rot> yields correct attention
+     * scores without an explicit inverse). To verify the round-trip, apply
+     * the inverse WHT before comparing against the original input. */
     memset(input, 0, sizeof(input));
     input[0] = 1.0f;
     quantize_row_turbo3_0_ref(input, buf, d);
     dequantize_row_turbo3_0(buf, output, d);
+    turbo_cpu_fwht_inverse(output, d);
     printf("Test 1 (turbo3): e0 = [1, 0, ...]\n");
     printf("  In:  [%.6f, %.6f, %.6f, %.6f]\n", input[0], input[1], input[2], input[3]);
     printf("  Out: [%.6f, %.6f, %.6f, %.6f]\n", output[0], output[1], output[2], output[3]);
@@ -31,17 +38,23 @@ int main(void) {
     for (int i = 0; i < d; i++) input[i] = sinf(i*0.1f+0.5f) * 10.0f;
     quantize_row_turbo3_0_ref(input, buf, d);
     dequantize_row_turbo3_0(buf, output, d);
+    turbo_cpu_fwht_inverse(output, d);
     printf("Test 2 (turbo3): sin*10\n");
     printf("  In:  [%.4f, %.4f, %.4f, %.4f]\n", input[0], input[1], input[2], input[3]);
     printf("  Out: [%.4f, %.4f, %.4f, %.4f]\n", output[0], output[1], output[2], output[3]);
     mse = cosv = ni = no = 0;
     for (int i = 0; i < d; i++) { mse += (input[i]-output[i])*(input[i]-output[i]); cosv += input[i]*output[i]; ni += input[i]*input[i]; no += output[i]*output[i]; }
     printf("  MSE=%.8f Cosine=%.6f InNorm=%.2f OutNorm=%.2f\n\n", mse/d, cosv/sqrtf(ni)/sqrtf(no), sqrtf(ni), sqrtf(no));
 
-    /* Test 3: turbo4 */
+    /* Test 3: turbo4
+     *
+     * Same convention as turbo3: dequant leaves output in the rotated domain
+     * (see comment in dequantize_row_turbo4_0 @ ggml-turbo-quant.c). Apply
+     * the inverse WHT before comparing. */
     for (int i = 0; i < d; i++) input[i] = cosf(i*0.2f) * 5.0f;
     quantize_row_turbo4_0_ref(input, buf, d);
     dequantize_row_turbo4_0(buf, output, d);
+    turbo_cpu_fwht_inverse(output, d);
     printf("Test 3 (turbo4): cos*5\n");
     printf("  In:  [%.4f, %.4f, %.4f, %.4f]\n", input[0], input[1], input[2], input[3]);
     printf("  Out: [%.4f, %.4f, %.4f, %.4f]\n", output[0], output[1], output[2], output[3]);