Exponential speed-up

hexstate_quantize.c

@@ -5466,7 +5466,7 @@ static double hex_quad(const float *e, const float *H, int64_t n)
 
 /* Optimise one output row of length n_in (a whole number of Q2_K blocks). */
 static void hex_analog_row(const float *w, int64_t n_in, const float *H,
-                           BlockQ2K *out, int n_sweeps,
+                           int h_diag, BlockQ2K *out, int n_sweeps,
                            float *E, float *Gd, uint8_t *q,
                            uint8_t *Ls, uint8_t *Lm, float *dv, float *mv)
 {
@@ -5519,11 +5519,15 @@ static void hex_analog_row(const float *w, int64_t n_in, const float *H,
         float ds = dv[b] * (float)Ls[j], ms = mv[b] * (float)Lm[j];
         E[i] = w[i] - (ds * (float)q[i] - ms);
     }
-    for (int64_t i = 0; i < n_in; i++) {
-        const float *Hi = H + i * n_in;
-        double s = 0.0;
-        for (int64_t k = 0; k < n_in; k++) s += (double)Hi[k] * (double)E[k];
-        Gd[i] = (float)s;
+    if (h_diag) {
+        for (int64_t i = 0; i < n_in; i++) Gd[i] = H[i * n_in + i] * E[i];
+    } else {
+        for (int64_t i = 0; i < n_in; i++) {
+            const float *Hi = H + i * n_in;
+            double s = 0.0;
+            for (int64_t k = 0; k < n_in; k++) s += (double)Hi[k] * (double)E[k];
+            Gd[i] = (float)s;
+        }
     }
 
     for (int sweep = 0; sweep < n_sweeps; sweep++) {
@@ -5550,71 +5554,83 @@ static void hex_analog_row(const float *w, int64_t n_in, const float *H,
                     q[j] = (uint8_t)bestq;
                     E[j] -= bestd;
                     changed++;
-                    if (HEX_HESS_FLUSH <= 1) {
+                    if (h_diag) {
+                        /* diagonal H: only Gd[j] depends on E[j]. O(1). */
+                        Gd[j] = Hjj * E[j];
+                    } else {
                         const float *Hj = H + j * n_in;
                         for (int64_t k = 0; k < n_in; k++) Gd[k] -= bestd * Hj[k];
-                    } else {
-                        pend++;
-                    }
-                }
-                if (HEX_HESS_FLUSH > 1 && pend && ((j + 1) % HEX_HESS_FLUSH == 0)) {
-                    for (int64_t i = 0; i < n_in; i++) {
-                        const float *Hi = H + i * n_in;
-                        double s = 0.0;
-                        for (int64_t k = 0; k < n_in; k++) s += (double)Hi[k] * (double)E[k];
-                        Gd[i] = (float)s;
                     }
-                    pend = 0;
-                }
-            }
-            if (HEX_HESS_FLUSH > 1) {
-                for (int64_t i = 0; i < n_in; i++) {
-                    const float *Hi = H + i * n_in;
-                    double s = 0.0;
-                    for (int64_t k = 0; k < n_in; k++) s += (double)Hi[k] * (double)E[k];
-                    Gd[i] = (float)s;
                 }
             }
+            (void)pend;
         }
 
         /* ── move B: joint (ls, lm) per sub-block, full 16x16 ── */
         for (int64_t sj = 0; sj < n_in / 16; sj++) {
             int64_t b = (sj * 16) / QK_K, c0 = sj * 16;
             float dd = dv[b], mm = mv[b];
-            float old[16], del[16], bestdel[16];
-            for (int k = 0; k < 16; k++)
-                old[k] = dd * (float)Ls[sj] * (float)q[c0 + k] - mm * (float)Lm[sj];
+            /* del(tl,tm) is AFFINE in the candidate:
+             *     del = a*u - b*e,   a = tl-Ls,  b = tm-Lm
+             *     u[k] = d*q[k],     e[k] = dmin  (constant vector)
+             * so the gain expands with five precomputed scalars
+             *     gain = 2(a*Lu - b*Le) - (a^2*Quu - 2ab*Que + b^2*Qee)
+             * and each of the 256 candidates costs O(1) instead of the
+             * O(256) 16x16 quadratic form. This inner loop dominated the
+             * whole routine: ~65k ops per sub-block became ~3k. */
+            float u[16], del[16], bestdel[16];
+            for (int k = 0; k < 16; k++) u[k] = dd * (float)q[c0 + k];
+            const float ev = mm;
+
+            double Quu = 0, Que = 0, Qee = 0, Lu = 0, Le = 0;
+            if (h_diag) {
+                for (int k = 0; k < 16; k++) {
+                    double h = H[(c0 + k) * n_in + c0 + k];
+                    Quu += h * u[k] * u[k];
+                    Que += h * u[k] * ev;
+                    Qee += h * ev * ev;
+                }
+            } else {
+                for (int a = 0; a < 16; a++) {
+                    const float *Ha = H + (c0 + a) * n_in + c0;
+                    double su = 0, se = 0;
+                    for (int c = 0; c < 16; c++) { su += Ha[c] * u[c]; se += Ha[c] * ev; }
+                    Quu += u[a] * su;  Que += u[a] * se;  Qee += ev * se;
+                }
+            }
+            for (int k = 0; k < 16; k++) { Lu += u[k] * Gd[c0 + k]; Le += ev * Gd[c0 + k]; }
+
             float bestg = 0.0f; int bl = Ls[sj], bm = Lm[sj]; int found = 0;
+            double ba = 0.0, bb2 = 0.0;
             for (int tl = 0; tl <= 15; tl++) {
+                double a = (double)tl - (double)Ls[sj];
                 for (int tm = 0; tm <= 15; tm++) {
-                    float lin = 0.0f, quad = 0.0f;
-                    for (int k = 0; k < 16; k++) {
-                        float nw = dd * (float)tl * (float)q[c0 + k] - mm * (float)tm;
-                        del[k] = nw - old[k];
-                        lin += del[k] * Gd[c0 + k];
-                    }
-                    for (int a = 0; a < 16; a++) {
-                        if (del[a] == 0.0f) continue;
-                        const float *Ha = H + (c0 + a) * n_in + c0;
-                        float s = 0.0f;
-                        for (int c = 0; c < 16; c++) s += Ha[c] * del[c];
-                        quad += del[a] * s;
-                    }
-                    float gain = 2.0f * lin - quad;
-                    if (gain > bestg + 1e-20f) {
-                        bestg = gain; bl = tl; bm = tm; found = 1;
-                        memcpy(bestdel, del, sizeof(del));
+                    double b = (double)tm - (double)Lm[sj];
+                    double gain = 2.0 * (a * Lu - b * Le)
+                                - (a*a*Quu - 2.0*a*b*Que + b*b*Qee);
+                    if (gain > bestg + 1e-20) {
+                        bestg = (float)gain; bl = tl; bm = tm; found = 1;
+                        ba = a; bb2 = b;
                     }
                 }
             }
+            if (found)
+                for (int k = 0; k < 16; k++)
+                    bestdel[k] = (float)(ba * u[k] - bb2 * ev);
+            (void)del;
             if (found) {
                 Ls[sj] = (uint8_t)bl; Lm[sj] = (uint8_t)bm;
                 for (int k = 0; k < 16; k++) E[c0 + k] -= bestdel[k];
-                for (int64_t i = 0; i < n_in; i++) {
-                    const float *Hi = H + i * n_in + c0;
-                    float s = 0.0f;
-                    for (int k = 0; k < 16; k++) s += Hi[k] * bestdel[k];
-                    Gd[i] -= s;
+                if (h_diag) {
+                    for (int k = 0; k < 16; k++)
+                        Gd[c0 + k] = H[(c0 + k) * n_in + c0 + k] * E[c0 + k];
+                } else {
+                    for (int64_t i = 0; i < n_in; i++) {
+                        const float *Hi = H + i * n_in + c0;
+                        float s = 0.0f;
+                        for (int k = 0; k < 16; k++) s += Hi[k] * bestdel[k];
+                        Gd[i] -= s;
+                    }
                 }
                 changed++;
             }
@@ -5634,16 +5650,25 @@ static void hex_analog_row(const float *w, int64_t n_in, const float *H,
                     bvv[i] = (float)Lm[b * NSl + j];
                     tv[i]  = w[c0 + i];
                 }
-            for (int i = 0; i < QK_K; i++) {
-                const float *Hi = H + (c0 + i) * n_in + c0;
-                double sa=0, sb=0, st=0;
-                for (int k = 0; k < QK_K; k++) {
-                    sa += (double)Hi[k] * av[k];
-                    sb += (double)Hi[k] * bvv[k];
-                    st += (double)Hi[k] * tv[k];
+            if (h_diag) {
+                for (int i = 0; i < QK_K; i++) {
+                    double h = H[(c0 + i) * n_in + c0 + i];
+                    Saa += h*av[i]*av[i];  Sab += h*av[i]*bvv[i];
+                    Sbb += h*bvv[i]*bvv[i];
+                    Sat += h*av[i]*tv[i];  Sbt += h*bvv[i]*tv[i];
+                }
+            } else {
+                for (int i = 0; i < QK_K; i++) {
+                    const float *Hi = H + (c0 + i) * n_in + c0;
+                    double sa=0, sb=0, st=0;
+                    for (int k = 0; k < QK_K; k++) {
+                        sa += (double)Hi[k] * av[k];
+                        sb += (double)Hi[k] * bvv[k];
+                        st += (double)Hi[k] * tv[k];
+                    }
+                    Saa += av[i]*sa; Sab += av[i]*sb; Sbb += bvv[i]*sb;
+                    Sat += av[i]*st; Sbt += bvv[i]*st;
                 }
-                Saa += av[i]*sa; Sab += av[i]*sb; Sbb += bvv[i]*sb;
-                Sat += av[i]*st; Sbt += bvv[i]*st;
             }
             double det = Saa * Sbb - Sab * Sab;
             if (fabs(det) < 1e-30) continue;
@@ -5653,31 +5678,49 @@ static void hex_analog_row(const float *w, int64_t n_in, const float *H,
             if (mn < 0.0) mn = 0.0;
             float dtry = gguf_fp16_to_fp32(gguf_fp32_to_fp16((float)dn));
             float mtry = gguf_fp16_to_fp32(gguf_fp32_to_fp16((float)mn));
-            /* exact accept test on the block's contribution */
-            float eo[QK_K], en[QK_K];
+            /* Accept test. E changes only inside this block, by dl = en - eo,
+             * so the exact objective change is
+             *     dL = 2*dl^T Gd + dl^T H_blk dl
+             * -- the SAME gain form used by moves A and B. The previous
+             * version compared eo^T H_blk eo against en^T H_blk en, which
+             * silently dropped the cross-block coupling carried in Gd and so
+             * could accept a move that raised the true objective. */
+            float eo[QK_K], en[QK_K], dl[QK_K];
             for (int i = 0; i < QK_K; i++) {
                 eo[i] = E[c0 + i];
                 en[i] = w[c0 + i] - (dtry * av[i] - mtry * bvv[i]);
+                dl[i] = en[i] - eo[i];
             }
-            double Lo=0, Ln=0;
-            for (int i = 0; i < QK_K; i++) {
-                const float *Hi = H + (c0 + i) * n_in + c0;
-                double so=0, sn=0;
-                for (int k = 0; k < QK_K; k++) { so += Hi[k]*eo[k]; sn += Hi[k]*en[k]; }
-                Lo += eo[i]*so; Ln += en[i]*sn;
-            }
-            if (Ln < Lo) {
-                dv[b] = dtry; mv[b] = mtry;
+            double lin = 0.0, quad = 0.0;
+            for (int i = 0; i < QK_K; i++) lin += (double)dl[i] * Gd[c0 + i];
+            if (h_diag) {
+                for (int i = 0; i < QK_K; i++)
+                    quad += (double)dl[i]*dl[i]*H[(c0 + i) * n_in + c0 + i];
+            } else {
                 for (int i = 0; i < QK_K; i++) {
-                    float d0 = en[i] - E[c0 + i];
-                    E[c0 + i] = en[i];
-                    (void)d0;
+                    if (dl[i] == 0.0f) continue;
+                    const float *Hi = H + (c0 + i) * n_in + c0;
+                    double sn = 0.0;
+                    for (int k = 0; k < QK_K; k++) sn += (double)Hi[k] * dl[k];
+                    quad += (double)dl[i] * sn;
                 }
-                for (int64_t i = 0; i < n_in; i++) {
-                    const float *Hi = H + i * n_in;
-                    double s = 0.0;
-                    for (int64_t k = 0; k < n_in; k++) s += (double)Hi[k] * (double)E[k];
-                    Gd[i] = (float)s;
+            }
+            if (2.0 * lin + quad < 0.0) {          /* strict decrease */
+                dv[b] = dtry; mv[b] = mtry;
+                for (int i = 0; i < QK_K; i++) E[c0 + i] = en[i];
+                /* Banded gradient update: Gd += H[:, blk] @ dl.
+                 * Was a full O(n_in^2) recompute -- 10x waste at n_in=2560,
+                 * 24x at 6144, and it dominated the whole runtime. */
+                if (h_diag) {
+                    for (int i = 0; i < QK_K; i++)
+                        Gd[c0 + i] = H[(c0 + i) * n_in + c0 + i] * E[c0 + i];
+                } else {
+                    for (int64_t i = 0; i < n_in; i++) {
+                        const float *Hi = H + i * n_in + c0;
+                        double sm2 = 0.0;
+                        for (int k = 0; k < QK_K; k++) sm2 += (double)Hi[k] * dl[k];
+                        Gd[i] += (float)sm2;
+                    }
                 }
                 changed++;
             }
@@ -5713,6 +5756,20 @@ static void quantize_tensor_q2k_analog(const float *weights, int64_t n_elements,
     int64_t nb_row = n_in / QK_K;
     double  total_err = 0.0;
 
+    /* Detect a diagonal Hessian once. An imatrix-derived H is diagonal, and
+     * then Gd = diag(H) .* E has NO coupling: every gradient update collapses
+     * from O(n_in) or O(n_in^2) to O(1). Same output, orders of magnitude
+     * less work -- worth the one-time O(n_in^2) scan. */
+    int h_diag = 1;
+    for (int64_t i = 0; i < n_in && h_diag; i++) {
+        const float *Hi = H + i * n_in;
+        for (int64_t k = 0; k < n_in; k++)
+            if (k != i && Hi[k] != 0.0f) { h_diag = 0; break; }
+    }
+    if (verbose)
+        printf("  [analog] Hessian is %s -> %s path\n",
+               h_diag ? "DIAGONAL" : "dense", h_diag ? "fast" : "full");
+
     #pragma omp parallel for schedule(dynamic, 1) reduction(+:total_err)
     for (int64_t r = 0; r < n_rows; r++) {
         float *E  = (float *)malloc((size_t)n_in * sizeof(float));
@@ -5725,7 +5782,7 @@ static void quantize_tensor_q2k_analog(const float *weights, int64_t n_elements,
 
         const float *w = weights + r * n_in;
         BlockQ2K   *o = output + r * nb_row;
-        hex_analog_row(w, n_in, H, o, n_sweeps, E, Gd, q, Ls, Lm, dv, mv);
+        hex_analog_row(w, n_in, H, h_diag, o, n_sweeps, E, Gd, q, Ls, Lm, dv, mv);
 
         for (int64_t b = 0; b < nb_row; b++)
             total_err += (double)gguf_q2_k_block_error(w + b * QK_K, &o[b]);