fix(quantization): correct progressive GPTQ compensation

python/cactus/convert/quantization/cq.py

@@ -230,13 +230,15 @@ def _prepare_input_scale(input_scale: np.ndarray | None, k: int) -> np.ndarray:
     return np.clip(scale, 1e-8, 65504.0).astype(np.float32, copy=False)
 
 
-def _gptq_correct_group(work: np.ndarray, recon: np.ndarray, h_inv: np.ndarray | None, start: int, stop: int) -> None:
-    if h_inv is None or stop >= work.shape[1]:
+def _gptq_correct_group(work: np.ndarray, recon: np.ndarray, u_factor: np.ndarray | None, start: int, stop: int) -> None:
+    if u_factor is None or stop >= work.shape[1]:
         return
     try:
-        m_bb = h_inv[start:stop, start:stop].astype(np.float32, copy=False)
-        m_bs = h_inv[start:stop, stop:].astype(np.float32, copy=False)
-        update = np.linalg.solve(m_bb + np.eye(m_bb.shape[0], dtype=np.float32) * 1e-6, m_bs)
+        from scipy.linalg import solve_triangular
+
+        u_bb = u_factor[start:stop, start:stop].astype(np.float32, copy=False)
+        u_bs = u_factor[start:stop, stop:].astype(np.float32, copy=False)
+        update = solve_triangular(u_bb, u_bs, lower=False)
         work[:, stop:] -= (work[:, start:stop] - recon) @ update
     except Exception:
         return
@@ -266,7 +268,7 @@ def quantize_hadamard(
     groups = k // GROUP_SIZE
     indices = np.zeros((n, k), dtype=np.uint8)
     norms = np.zeros((n, groups), dtype=np.float16)
-    h_inv = None
+    u_factor = None
     if use_gptq and hessian is not None:
         try:
             h = np.asarray(hessian, dtype=np.float32)
@@ -277,9 +279,9 @@ def quantize_hadamard(
                 s = np.clip(input_scale.astype(np.float32), 1e-6, None)
                 h = h / (s[:, None] * s[None, :])
                 h = h + np.eye(h.shape[0], dtype=np.float32) * (0.01 * np.mean(np.diag(h)) + 1e-6)
-                h_inv = np.linalg.inv(h)
+                u_factor = np.linalg.cholesky(np.linalg.inv(h)).T
         except Exception:
-            h_inv = None
+            u_factor = None
     for g in range(groups):
         start, stop = g * GROUP_SIZE, (g + 1) * GROUP_SIZE
         group = work[:, start:stop]
@@ -289,8 +291,8 @@ def quantize_hadamard(
         recon = (codebook[idx] @ rot.T) * row_norms[:, None]
         indices[:, start:stop] = idx
         norms[:, g] = row_norms.astype(np.float16)
-        _gptq_correct_group(work, recon, h_inv, start, stop)
-    return CQTensor(indices=indices, norms=norms, input_scale=input_scale.astype(np.float16), bits=bits, gptq_used=bool(use_gptq and h_inv is not None))
+        _gptq_correct_group(work, recon, u_factor, start, stop)
+    return CQTensor(indices=indices, norms=norms, input_scale=input_scale.astype(np.float16), bits=bits, gptq_used=bool(use_gptq and u_factor is not None))
 
 
 def quantize_orthogonal(weight, bits: int = 4, seed: int = 1234, input_scale: np.ndarray | None = None) -> CQTensor:

python/cactus/convert/tests/test_cq.py

@@ -15,7 +15,15 @@
 )
 from cactus.convert.export.qdq import FLAG_INTERLEAVED_4ROW, FLAG_ORTHOGONAL_ROTATION, dequantize_cq_file, read_header
 from cactus.convert.interleave_orthogonal_cq4 import interleave_orthogonal_cq4_file
-from cactus.convert.quantization.cq import PRECISION_CQ, pack_indices_lsb, quantize_hadamard, quantize_orthogonal, write_cq_tensor
+from cactus.convert.quantization.cq import (
+    GROUP_SIZE as CQ_GROUP_SIZE,
+    PRECISION_CQ,
+    _gptq_correct_group,
+    pack_indices_lsb,
+    quantize_hadamard,
+    quantize_orthogonal,
+    write_cq_tensor,
+)
 
 
 def test_pack_indices_lsb_bits():
@@ -155,3 +163,76 @@ def test_pointwise_conv1d_int8_preserves_rank3_shape(tmp_path):
     assert scales_bytes == 3 * 2 * 2
     assert group_size == GROUP_SIZE
     assert num_groups == 2
+
+
+def test_gptq_multiblock_compensation_matches_progressive_obs():
+    rng = np.random.default_rng(12345)
+    k = 384
+    b = CQ_GROUP_SIZE
+    assert k == 3 * b
+
+    # Construct deterministic SPD correlated Hessian
+    x = rng.standard_normal((k * 2, k)).astype(np.float32)
+    h = (x.T @ x) / float(k * 2)
+    h = h + np.eye(k, dtype=np.float32) * (0.01 * np.mean(np.diag(h)) + 1e-6)
+
+    h_inv_global = np.linalg.inv(h)
+    u_factor = np.linalg.cholesky(h_inv_global).T
+
+    work_orig = rng.standard_normal((8, k)).astype(np.float32)
+    recon_orig = rng.standard_normal((8, k)).astype(np.float32)
+    tol = 1e-4
+
+    for g in range(k // b - 1):
+        start, stop = g * b, (g + 1) * b
+        work = work_orig.copy()
+        work_static = work_orig.copy()
+
+        # 1. Independent ground-truth OBS reference from remaining sub-Hessian
+        h_rem = h[start:, start:]
+        h_rem_inv = np.linalg.inv(h_rem)
+        ref_bb = h_rem_inv[0:b, 0:b]
+        ref_bs = h_rem_inv[0:b, b:]
+        ref_update = np.linalg.solve(ref_bb, ref_bs)
+
+        residual = work_orig[:, start:stop] - recon_orig[:, start:stop]
+        expected_work_rem = work_orig[:, stop:] - residual @ ref_update
+
+        # 2. Production Cholesky implementation passes OBS invariant
+        _gptq_correct_group(work, recon_orig[:, start:stop], u_factor, start, stop)
+        actual_work_rem = work[:, stop:]
+        actual_diff = float(np.max(np.abs(actual_work_rem - expected_work_rem)))
+        assert actual_diff <= tol, f"Block {g} failed progressive OBS invariant: {actual_diff} > {tol}"
+
+        # 3. Old static-inverse implementation violates the invariant on correlated blocks g >= 1
+        old_m_bb = h_inv_global[start:stop, start:stop]
+        old_m_bs = h_inv_global[start:stop, stop:]
+        old_update = np.linalg.solve(old_m_bb + np.eye(b, dtype=np.float32) * 1e-6, old_m_bs)
+        old_work_rem = work_static[:, stop:] - residual @ old_update
+        old_diff = float(np.max(np.abs(old_work_rem - expected_work_rem)))
+
+        if g == 0:
+            assert old_diff <= tol, f"Block 0 unexpectedly diverged on static inverse: {old_diff} > {tol}"
+        else:
+            assert old_diff > tol, f"Old static inverse unexpectedly passed tolerance on block {g}: {old_diff} <= {tol}"
+
+
+def test_quantize_hadamard_with_gptq_multiblock():
+    rng = np.random.default_rng(2026)
+    k = 384
+    n = 16
+    w = rng.standard_normal((n, k)).astype(np.float32)
+    x = rng.standard_normal((k * 2, k)).astype(np.float32)
+    h = (x.T @ x) / float(k * 2)
+
+    cq_gptq = quantize_hadamard(w, bits=4, hessian=h, use_gptq=True)
+    cq_rtn = quantize_hadamard(w, bits=4, hessian=h, use_gptq=False)
+
+    assert cq_gptq.gptq_used is True
+    assert cq_rtn.gptq_used is False
+    assert cq_gptq.indices.shape == (n, k)
+    assert cq_gptq.norms.shape == (n, k // CQ_GROUP_SIZE)
+    # Block 0 has no prior compensation, so index decisions match RTN
+    assert np.array_equal(cq_gptq.indices[:, :CQ_GROUP_SIZE], cq_rtn.indices[:, :CQ_GROUP_SIZE])
+    # Multi-block compensation actively modifies subsequent quantization decisions
+    assert not np.array_equal(cq_gptq.indices[:, CQ_GROUP_SIZE:], cq_rtn.indices[:, CQ_GROUP_SIZE:])