Create pytorch_compiled_kernel.py

pytorch_compiled_kernel.py

@@ -0,0 +1,1049 @@
+"""
+CompiledEigh: torch.compile(fullgraph=True) drop-in for torch.linalg.eigh.
+
+Eliminates all graph breaks, device-host syncs, and dynamic allocation.
+Output contract matches torch.linalg.eigh:
+  eigenvalues:  [*, n] real, ascending
+  eigenvectors: [*, n, n] orthonormal columns
+
+Author: AbstractPhil / GeoLIP project
+"""
+
+import math
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Tuple, Optional
+
+
+# =============================================================================
+# Constants
+# =============================================================================
+
+DEFAULT_MAX_NEWTON: int = 8
+DEFAULT_MAX_JACOBI_SWEEPS: int = 10  # 10 sweeps saturates for n <= 16
+JACOBI_THRESHOLD: int = 16
+
+
+# =============================================================================
+# Atom: 2x2 Symmetric Eigenproblem
+# =============================================================================
+
+def eigh_2x2(a: Tensor, b: Tensor, c: Tensor, eps: float = 1e-30
+             ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """
+    Closed-form eigendecomposition of batched 2x2 symmetric matrices.
+    Returns: (lambda1, lambda2, cos_theta, sin_theta), lambda1 <= lambda2.
+    """
+    trace = a + c
+    diff = a - c
+    two_b = 2.0 * b
+    hyp = torch.sqrt(diff * diff + two_b * two_b + eps)
+
+    lambda1 = 0.5 * (trace - hyp)
+    lambda2 = 0.5 * (trace + hyp)
+
+    vx = two_b
+    vy = lambda2 - a
+    norm_v = torch.sqrt(vx * vx + vy * vy + eps)
+    cos_theta = vy / norm_v
+    sin_theta = vx / norm_v
+
+    return lambda1, lambda2, cos_theta, sin_theta
+
+
+# =============================================================================
+# Utility: Newton-Schulz Orthogonalization (all-bmm, GPU-native)
+# =============================================================================
+
+def orthogonalize_ns(V: Tensor, n_iter: int = 2) -> Tensor:
+    """
+    Re-orthogonalize columns of V via Newton-Schulz iteration.
+
+    Computes V @ (V^T V)^{-1/2} using the coupled iteration:
+        X_0 = I,  Y_0 = V^T V
+        X_{k+1} = 0.5 * X_k @ (3I - Y_k)
+        Y_{k+1} = 0.5 * (3I - Y_k) @ Y_k
+
+    X converges to (V^T V)^{-1/2}, Y converges to I.
+    Cubically convergent when V^T V ≈ I.
+
+    Convergence from ||V^T V - I|| = ε:
+        1 iteration: error → O(ε²)  ≈ 1e-6 from 1e-3
+        2 iterations: error → O(ε⁴) ≈ 1e-12 from 1e-3
+
+    All ops are bmm — fully compiled, no sequential column processing.
+
+    V: [B, n, n] (square, columns are approximate eigenvectors)
+    Returns: [B, n, n] with orthonormal columns
+    """
+    B, n, m = V.shape
+    I_n = torch.eye(m, device=V.device, dtype=V.dtype).unsqueeze(0).expand(B, -1, -1)
+
+    # Z = V^T V ≈ I
+    Y = torch.bmm(V.transpose(-2, -1), V)
+    X = I_n.clone()
+
+    for _ in range(n_iter):
+        T = 3.0 * I_n - Y           # (3I - Y_k)
+        X = 0.5 * torch.bmm(X, T)   # X_{k+1}
+        Y = 0.5 * torch.bmm(T, Y)   # Y_{k+1}
+
+    return torch.bmm(V, X)
+
+
+# =============================================================================
+# Phase 1: Householder Tridiagonalization
+# =============================================================================
+
+class HouseholderTridiagonalizer(nn.Module):
+    """
+    Reduces batched symmetric A to tridiagonal T = Q^T A Q via
+    Householder reflections. Fixed loop bounds, compilable.
+    """
+
+    def __init__(self, max_n: int, eps: float = 1e-30):
+        super().__init__()
+        self.max_n = max_n
+        self.eps = eps
+
+    def forward(self, A: Tensor, d: Tensor, e: Tensor,
+                reflectors: Tensor) -> None:
+        B, n, _ = A.shape
+        eps = self.eps
+
+        for k in range(n - 2):
+            tail_len = n - k - 1
+            x = A[:, k + 1:, k].clone()
+
+            sigma = torch.sqrt((x * x).sum(dim=-1, keepdim=True) + eps)
+            sign_x0 = torch.where(x[:, 0:1] >= 0,
+                                  torch.ones_like(sigma),
+                                  -torch.ones_like(sigma))
+            alpha = -sign_x0 * sigma
+
+            v = x.clone()
+            v[:, 0:1] = v[:, 0:1] - alpha
+            v_norm = torch.sqrt((v * v).sum(dim=-1, keepdim=True) + eps)
+            v = v / v_norm
+
+            reflectors[k, :, :tail_len] = v
+            if tail_len < n:
+                reflectors[k, :, tail_len:] = 0.0
+
+            sub_A = A[:, k + 1:, k + 1:]
+            v_col = v.unsqueeze(-1)
+
+            p = torch.bmm(sub_A, v_col).squeeze(-1)
+            vtp = (v * p).sum(dim=-1, keepdim=True)
+            q = p - vtp * v
+
+            q_col = q.unsqueeze(-1)
+            q_row = q.unsqueeze(-2)
+            v_row = v.unsqueeze(-2)
+
+            A[:, k + 1:, k + 1:] -= 2.0 * (v_col @ q_row + q_col @ v_row)
+            A[:, k, k + 1] = alpha.squeeze(-1)
+            A[:, k + 1, k] = alpha.squeeze(-1)
+
+        for i in range(n):
+            d[:, i] = A[:, i, i]
+        for i in range(n - 1):
+            e[:, i] = A[:, i, i + 1]
+
+
+# =============================================================================
+# Phase 2a: Secular Equation Newton Solver (Fixed Budget)
+# =============================================================================
+
+class SecularNewtonSolver(nn.Module):
+
+    def __init__(self, max_newton: int = DEFAULT_MAX_NEWTON,
+                 eps: float = 1e-30, tol: float = 1e-7):
+        super().__init__()
+        self.max_newton = max_newton
+        self.eps = eps
+        self.tol = tol
+
+    def forward(self, delta: Tensor, z_sq: Tensor,
+                rho: Tensor, mask: Tensor) -> Tensor:
+        B, m = delta.shape
+        eps = self.eps
+        tol = self.tol
+
+        z_sq_sum = (z_sq * mask).sum(dim=-1, keepdim=True)
+        rho_abs = rho.abs().unsqueeze(-1)
+        upper_bound = delta[:, -1:] + z_sq_sum * rho_abs + 1.0
+
+        lo = delta + eps
+        hi = torch.cat([delta[:, 1:], upper_bound], dim=-1) - eps
+        lam = 0.5 * (lo + hi)
+        rho_exp = rho.unsqueeze(-1)
+
+        for _step in range(self.max_newton):
+            delta_exp = delta.unsqueeze(-1)
+            lam_exp = lam.unsqueeze(-2)
+            denom = delta_exp - lam_exp
+
+            denom_safe = torch.where(
+                denom.abs() < eps,
+                torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),
+                denom
+            )
+
+            z_sq_exp = z_sq.unsqueeze(-1)
+            mask_exp = mask.unsqueeze(-1)
+            masked_z = z_sq_exp * mask_exp
+
+            terms = masked_z / denom_safe
+            f = 1.0 + rho_exp * terms.sum(dim=-2)
+            f_prime = rho_exp * (masked_z / (denom_safe * denom_safe)).sum(dim=-2)
+
+            f_prime_safe = torch.where(
+                f_prime.abs() < eps,
+                torch.full_like(f_prime, eps),
+                f_prime
+            )
+            delta_lam = -f / f_prime_safe
+            lam_new = torch.clamp(lam + delta_lam, lo, hi)
+
+            f_pos = f > 0
+            lo = torch.where(f_pos & mask.bool(), lam, lo)
+            hi = torch.where(~f_pos & mask.bool(), lam, hi)
+
+            converged = (f.abs() < tol) | ~mask.bool()
+            lam = torch.where(converged, lam, lam_new)
+
+        return lam
+
+
+# =============================================================================
+# Phase 2b: Eigenvectors from Secular Equation
+# =============================================================================
+
+def secular_eigenvectors(delta: Tensor, lam: Tensor, z: Tensor,
+                         mask: Tensor, eps: float = 1e-30) -> Tensor:
+    delta_exp = delta.unsqueeze(-1)
+    lam_exp = lam.unsqueeze(-2)
+    denom = delta_exp - lam_exp
+
+    denom_safe = torch.where(
+        denom.abs() < eps,
+        torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),
+        denom
+    )
+
+    z_exp = z.unsqueeze(-1)
+    mask_exp = mask.unsqueeze(-1)
+    V = (z_exp * mask_exp) / denom_safe
+    col_norms = torch.sqrt((V * V).sum(dim=-2, keepdim=True) + eps)
+    V = V / col_norms
+    return V
+
+
+# =============================================================================
+# Phase 2c: Fixed-Depth Tensor Tree D&C
+# =============================================================================
+
+class TensorTreeDC(nn.Module):
+
+    def __init__(self, max_n: int,
+                 max_newton: int = DEFAULT_MAX_NEWTON,
+                 eps: float = 1e-30, tol: float = 1e-7):
+        super().__init__()
+        self.padded_n = 1 << math.ceil(math.log2(max(max_n, 2)))
+        self.depth = int(math.log2(self.padded_n))
+        self.max_n = max_n
+        self.eps = eps
+        self.secular_solver = SecularNewtonSolver(
+            max_newton=max_newton, eps=eps, tol=tol
+        )
+
+    def forward(self, d: Tensor, e: Tensor) -> Tuple[Tensor, Tensor]:
+        B, n = d.shape
+        pn = self.padded_n
+        eps = self.eps
+        device = d.device
+        dtype = d.dtype
+
+        if n < pn:
+            d_max = d.abs().max(dim=-1, keepdim=True).values + 1.0
+            pad_diag = d_max + torch.arange(1, pn - n + 1, device=device, dtype=dtype).unsqueeze(0)
+            d_padded = torch.cat([d, pad_diag], dim=-1)
+            e_padded = torch.zeros(B, pn - 1, device=device, dtype=dtype)
+            e_padded[:, :n - 1] = e
+        else:
+            d_padded = d.clone()
+            e_padded = e.clone()
+
+        # DOWNWARD PASS
+        coupling_rho = []
+        current_d = d_padded.clone()
+        current_e = e_padded.clone()
+
+        for level in range(self.depth):
+            num_sub = 2 ** level
+            sub_size = pn // num_sub
+            half = sub_size // 2
+
+            cd = current_d.reshape(B, num_sub, sub_size)
+            ce = current_e.reshape(B, num_sub, sub_size - 1)
+
+            rho = ce[:, :, half - 1].clone()
+            coupling_rho.append(rho)
+
+            cd[:, :, half - 1] = cd[:, :, half - 1] - rho.abs()
+            cd[:, :, half] = cd[:, :, half] - rho.abs()
+            ce[:, :, half - 1] = 0.0
+
+            left_d = cd[:, :, :half].reshape(B, num_sub * half)
+            right_d = cd[:, :, half:].reshape(B, num_sub * half)
+            current_d = torch.stack([left_d.reshape(B, num_sub, half),
+                                     right_d.reshape(B, num_sub, half)],
+                                    dim=2).reshape(B, pn)
+
+            left_e = ce[:, :, :half - 1].reshape(B, num_sub, half - 1)
+            right_e = ce[:, :, half:].reshape(B, num_sub, half - 1)
+            current_e = torch.stack([left_e, right_e], dim=2).reshape(
+                B, num_sub * 2 * (half - 1))
+            expected_e_len = pn - 1
+            if current_e.shape[-1] < expected_e_len:
+                current_e = torch.nn.functional.pad(
+                    current_e, (0, expected_e_len - current_e.shape[-1]))
+
+        # BASE
+        base_evals = current_d
+        V_current = torch.ones(B, pn, 1, 1, device=device, dtype=dtype)
+
+        # UPWARD PASS
+        current_evals = base_evals
+
+        for level in range(self.depth - 1, -1, -1):
+            num_sub = 2 ** level
+            sub_size = pn // num_sub
+            half = sub_size // 2
+            child_size = half
+
+            evals_grouped = current_evals.reshape(B, num_sub, 2, child_size)
+            left_evals = evals_grouped[:, :, 0, :]
+            right_evals = evals_grouped[:, :, 1, :]
+            delta = torch.cat([left_evals, right_evals], dim=-1)
+
+            V_grouped = V_current.reshape(B, num_sub, 2, child_size, child_size)
+            V_left = V_grouped[:, :, 0, :, :]
+            V_right = V_grouped[:, :, 1, :, :]
+
+            z_left = V_left[:, :, -1, :]
+            z_right = V_right[:, :, 0, :]
+            z_cat = torch.cat([z_left, z_right], dim=-1)
+
+            delta_sorted, sort_idx = delta.sort(dim=-1)
+            z_sorted = z_cat.gather(-1, sort_idx)
+
+            rho = coupling_rho[level]
+            mask = torch.ones(B, num_sub, sub_size, device=device, dtype=dtype)
+
+            gaps = (delta_sorted[:, :, 1:] - delta_sorted[:, :, :-1]).abs()
+            degenerate = gaps < (eps * 100)
+            avg = 0.5 * (delta_sorted[:, :, :-1] + delta_sorted[:, :, 1:])
+
+            delta_defl = delta_sorted.clone()
+            delta_defl[:, :, :-1] = torch.where(degenerate, avg, delta_sorted[:, :, :-1])
+            delta_defl[:, :, 1:] = torch.where(degenerate, avg, delta_sorted[:, :, 1:])
+
+            z_defl = z_sorted.clone()
+            defl_kill = torch.ones_like(z_sorted)
+            defl_kill[:, :, 1:] = torch.where(
+                degenerate, torch.zeros_like(gaps), torch.ones_like(gaps))
+            z_defl = z_defl * defl_kill
+
+            z_sq = z_defl * z_defl
+            Bns = B * num_sub
+            new_evals_flat = self.secular_solver(
+                delta_defl.reshape(Bns, sub_size),
+                z_sq.reshape(Bns, sub_size),
+                rho.reshape(Bns),
+                mask.reshape(Bns, sub_size),
+            )
+            new_evals = new_evals_flat.reshape(B, num_sub, sub_size)
+
+            V_secular_flat = secular_eigenvectors(
+                delta_defl.reshape(Bns, sub_size),
+                new_evals_flat,
+                z_defl.reshape(Bns, sub_size),
+                mask.reshape(Bns, sub_size),
+                eps=eps
+            )
+            V_secular = V_secular_flat.reshape(B, num_sub, sub_size, sub_size)
+
+            inv_sort = sort_idx.argsort(dim=-1)
+            inv_exp = inv_sort.unsqueeze(-1).expand_as(V_secular)
+            V_unsorted = V_secular.gather(-2, inv_exp)
+
+            V_block = torch.zeros(B, num_sub, sub_size, sub_size,
+                                  device=device, dtype=dtype)
+            V_block[:, :, :half, :half] = V_left
+            V_block[:, :, half:, half:] = V_right
+
+            V_merged = torch.bmm(
+                V_block.reshape(Bns, sub_size, sub_size),
+                V_unsorted.reshape(Bns, sub_size, sub_size)
+            ).reshape(B, num_sub, sub_size, sub_size)
+
+            current_evals = new_evals.reshape(B, pn)
+            V_current = V_merged
+
+        eigenvalues = current_evals
+        eigenvectors = V_current.squeeze(1)
+
+        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)
+        sort_exp = sort_perm.unsqueeze(-2).expand_as(eigenvectors)
+        sorted_evecs = eigenvectors.gather(-1, sort_exp)
+
+        if n < pn:
+            sorted_evals = sorted_evals[:, :n]
+            sorted_evecs = sorted_evecs[:, :n, :n]
+
+        return sorted_evals, sorted_evecs
+
+
+# =============================================================================
+# Phase 2 (alternate): Jacobi for small n
+# =============================================================================
+
+class JacobiEigh(nn.Module):
+    """
+    Jacobi eigenvalue algorithm for small symmetric matrices.
+    Fixed sweep count, fully vectorized, zero branches.
+
+    COMPILE FIX: Pair indices stored as plain Python lists (not tensors).
+    Dynamo sees these as constants — no SymInt issues.
+    """
+
+    def __init__(self, max_n: int,
+                 max_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,
+                 eps: float = 1e-30):
+        super().__init__()
+        self.max_n = max_n
+        self.max_sweeps = max_sweeps
+        self.eps = eps
+
+        # CRITICAL: plain Python lists, NOT registered buffers.
+        # Dynamo traces these as compile-time constants.
+        pairs = []
+        for p in range(max_n):
+            for q in range(p + 1, max_n):
+                pairs.append((p, q))
+        self._pairs_p: list[int] = [p for p, q in pairs]
+        self._pairs_q: list[int] = [q for p, q in pairs]
+        self._n_pairs: int = len(pairs)
+
+    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:
+        """
+        A: [B, n, n] symmetric
+        Returns: (eigenvalues [B, n] ascending, eigenvectors [B, n, n])
+        """
+        B, n, _ = A.shape
+        eps = self.eps
+
+        W = A.clone()
+        V = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0).expand(B, -1, -1).clone()
+
+        for _sweep in range(self.max_sweeps):
+            for idx in range(self._n_pairs):
+                # Plain Python ints — Dynamo sees these as constants
+                p: int = self._pairs_p[idx]
+                q: int = self._pairs_q[idx]
+
+                app = W[:, p, p]
+                aqq = W[:, q, q]
+                apq = W[:, p, q]
+
+                # Givens rotation angle
+                two_apq = 2.0 * apq
+                diff = aqq - app
+
+                # Safe division: sign-preserving eps guard
+                abs_two_apq = two_apq.abs().clamp(min=eps)
+                sign_two_apq = torch.where(two_apq >= 0,
+                                           torch.ones_like(two_apq),
+                                           -torch.ones_like(two_apq))
+                tau = diff / (abs_two_apq * sign_two_apq)
+
+                tau_sign = torch.where(tau >= 0,
+                                       torch.ones_like(tau),
+                                       -torch.ones_like(tau))
+                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))
+
+                # Zero rotation when off-diagonal is already negligible
+                skip = (apq.abs() < eps).float()
+                t = t * (1.0 - skip)
+
+                c = 1.0 / torch.sqrt(1.0 + t * t)
+                s = t * c
+
+                # ── Rotate W columns p, q ──
+                Wp = W[:, :, p].clone()
+                Wq = W[:, :, q].clone()
+                c_col = c.unsqueeze(-1)
+                s_col = s.unsqueeze(-1)
+                W[:, :, p] = c_col * Wp - s_col * Wq
+                W[:, :, q] = s_col * Wp + c_col * Wq
+
+                # ── Rotate W rows p, q ──
+                Wp = W[:, p, :].clone()
+                Wq = W[:, q, :].clone()
+                W[:, p, :] = c_col * Wp - s_col * Wq
+                W[:, q, :] = s_col * Wp + c_col * Wq
+
+                # ── Exact diagonal repair (prevents accumulation drift) ──
+                W[:, p, q] = 0.0
+                W[:, q, p] = 0.0
+                W[:, p, p] = app - t * apq
+                W[:, q, q] = aqq + t * apq
+
+                # ── Accumulate eigenvectors ──
+                Vp = V[:, :, p].clone()
+                Vq = V[:, :, q].clone()
+                V[:, :, p] = c_col * Vp - s_col * Vq
+                V[:, :, q] = s_col * Vp + c_col * Vq
+
+        # ── Newton-Schulz re-orthogonalization ──
+        # 2 iterations: orth error 1e-3 → ~1e-12 via bmm (GPU-native)
+        V = orthogonalize_ns(V, n_iter=2)
+
+        # ── Extract and sort ──
+        eigenvalues = torch.diagonal(W, dim1=-2, dim2=-1)
+        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)
+        sort_exp = sort_perm.unsqueeze(-2).expand_as(V)
+        sorted_evecs = V.gather(-1, sort_exp)
+
+        return sorted_evals, sorted_evecs
+
+
+# =============================================================================
+# Phase 3: Householder Back-Accumulation
+# =============================================================================
+
+class HouseholderBackAccumulate(nn.Module):
+
+    def __init__(self, max_n: int, eps: float = 1e-30):
+        super().__init__()
+        self.max_n = max_n
+        self.eps = eps
+
+    def forward(self, reflectors: Tensor, Z: Tensor, n: int) -> Tensor:
+        V = Z.clone()
+        for k in range(n - 3, -1, -1):
+            tail_len = n - k - 1
+            v = reflectors[k, :, :tail_len]
+            v_col = v.unsqueeze(-1)
+            V_sub = V[:, k + 1:, :]
+            vtV = torch.bmm(v_col.transpose(-2, -1), V_sub)
+            V[:, k + 1:, :] = V_sub - 2.0 * v_col @ vtV
+        return V
+
+
+# =============================================================================
+# Validation
+# =============================================================================
+
+class EighValidator(nn.Module):
+
+    def forward(self, A: Tensor, eigenvalues: Tensor,
+                eigenvectors: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        B, n, _ = A.shape
+        AV = torch.bmm(A, eigenvectors)
+        VL = eigenvectors * eigenvalues.unsqueeze(-2)
+        residual = AV - VL
+        A_norm = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)
+        residual_norm = torch.linalg.norm(residual.reshape(B, -1), dim=-1) / A_norm
+
+        VtV = torch.bmm(eigenvectors.transpose(-2, -1), eigenvectors)
+        I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)
+        orth_err = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)
+
+        return residual_norm, orth_err, residual_norm.max()
+
+
+# =============================================================================
+# Top-Level: CompiledEigh
+# =============================================================================
+
+class CompiledEigh(nn.Module):
+    """
+    Drop-in replacement for torch.linalg.eigh.
+
+    Usage:
+        solver = CompiledEigh(max_n=6)
+        solver = torch.compile(solver, fullgraph=True)
+        eigenvalues, eigenvectors = solver(A)
+    """
+
+    def __init__(self, max_n: int,
+                 use_jacobi: Optional[bool] = None,
+                 max_newton: int = DEFAULT_MAX_NEWTON,
+                 max_jacobi_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,
+                 eps: float = 1e-30, tol: float = 1e-7):
+        super().__init__()
+        self.max_n = max_n
+        self.eps = eps
+
+        if use_jacobi is None:
+            use_jacobi = (max_n <= JACOBI_THRESHOLD)
+        self.use_jacobi = use_jacobi
+
+        if use_jacobi:
+            self.jacobi = JacobiEigh(
+                max_n=max_n, max_sweeps=max_jacobi_sweeps, eps=eps)
+        else:
+            self.tridiag = HouseholderTridiagonalizer(max_n=max_n, eps=eps)
+            self.dc = TensorTreeDC(
+                max_n=max_n, max_newton=max_newton, eps=eps, tol=tol)
+            self.back_accum = HouseholderBackAccumulate(max_n=max_n, eps=eps)
+
+        self.validator = EighValidator()
+
+    def forward(self, A: Tensor, validate: bool = False
+                ) -> Tuple[Tensor, Tensor]:
+        B, n, _ = A.shape
+
+        if self.use_jacobi:
+            eigenvalues, eigenvectors = self.jacobi(A)
+        else:
+            A_work = A.clone()
+            d = torch.empty(B, n, device=A.device, dtype=A.dtype)
+            e = torch.empty(B, n - 1, device=A.device, dtype=A.dtype)
+            reflectors = torch.zeros(max(n - 2, 1), B, n,
+                                     device=A.device, dtype=A.dtype)
+            self.tridiag(A_work, d, e, reflectors)
+            eigenvalues, Z = self.dc(d, e)
+            eigenvectors = self.back_accum(reflectors, Z, n)
+            # Newton-Schulz re-orthogonalization for D&C path
+            eigenvectors = orthogonalize_ns(eigenvectors, n_iter=2)
+
+        if validate:
+            res_norm, orth_err, max_err = self.validator(A, eigenvalues, eigenvectors)
+            print(f"[CompiledEigh] max residual: {max_err.item():.2e}, "
+                  f"mean orth err: {orth_err.mean().item():.2e}")
+
+        return eigenvalues, eigenvectors
+
+
+# =============================================================================
+# Functional API
+# =============================================================================
+
+_cached_solvers = {}
+
+def compiled_eigh(A: Tensor, validate: bool = False) -> Tuple[Tensor, Tensor]:
+    B, n, _ = A.shape
+    key = (n, A.device, A.dtype)
+    if key not in _cached_solvers:
+        _cached_solvers[key] = CompiledEigh(max_n=n).to(A.device)
+    return _cached_solvers[key](A, validate=validate)
+
+
+"""
+CompiledEigh — Colab GPU Benchmark v3
+Fixes:
+  v2: Jacobi pairs as plain Python lists (Dynamo compile fix), sweeps 6→10
+  v3: Replaced Gram-Schmidt with Newton-Schulz orthogonalization (all-bmm),
+      disabled TF32 to ensure fp32 precision on Blackwell
+"""
+
+import torch
+import time
+import gc
+import sys
+
+# ── Ensure full fp32 precision on Ampere/Hopper/Blackwell ──
+# TF32 uses 10-bit mantissa for matmul which can degrade orthogonality
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+torch.set_float32_matmul_precision('highest')
+
+
+def sync():
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def gpu_timer(fn, warmup=10, repeats=200):
+    for _ in range(warmup):
+        fn()
+    sync()
+    start = time.perf_counter()
+    for _ in range(repeats):
+        fn()
+    sync()
+    return (time.perf_counter() - start) / repeats
+
+
+def make_symmetric_batch(B, n, device, dtype=torch.float32):
+    R = torch.randn(B, n, n, device=device, dtype=dtype)
+    return (R + R.transpose(-2, -1)) / 2.0
+
+
+def make_cm_like_batch(B, n, device, dtype=torch.float32):
+    points = torch.randn(B, n, n, device=device, dtype=dtype)
+    points = points / (points.norm(dim=-1, keepdim=True) + 1e-8)
+    return torch.bmm(points, points.transpose(-2, -1)) * 0.3
+
+
+def fmt_time(seconds):
+    if seconds < 1e-3:
+        return f"{seconds*1e6:.1f} us"
+    elif seconds < 1.0:
+        return f"{seconds*1e3:.2f} ms"
+    return f"{seconds:.3f} s"
+
+
+# ─── Test 0: Newton-Schulz Diagnostic ───
+
+def test_ns_diagnostic(device):
+    """Verify Newton-Schulz orthogonalization works on GPU independently."""
+    print("\n" + "=" * 70)
+    print("  TEST 0: NEWTON-SCHULZ DIAGNOSTIC")
+    print("=" * 70)
+
+    for n in [5, 6, 8]:
+        B = 1024
+        # Create nearly-orthogonal matrix (simulating Jacobi output)
+        Q, _ = torch.linalg.qr(torch.randn(B, n, n, device=device))
+        # Perturb to ~1e-3 orthogonality error
+        noise = torch.randn(B, n, n, device=device) * 1e-3
+        V_dirty = Q + noise
+
+        I_n = torch.eye(n, device=device).unsqueeze(0)
+
+        # Before NS
+        VtV_before = torch.bmm(V_dirty.transpose(-2, -1), V_dirty)
+        orth_before = torch.linalg.norm((VtV_before - I_n).reshape(B, -1), dim=-1).max().item()
+
+        # After NS (2 iterations)
+        V_clean = orthogonalize_ns(V_dirty, n_iter=2)
+        VtV_after = torch.bmm(V_clean.transpose(-2, -1), V_clean)
+        orth_after = torch.linalg.norm((VtV_after - I_n).reshape(B, -1), dim=-1).max().item()
+
+        # After NS (3 iterations for comparison)
+        V_clean3 = orthogonalize_ns(V_dirty, n_iter=3)
+        VtV_after3 = torch.bmm(V_clean3.transpose(-2, -1), V_clean3)
+        orth_after3 = torch.linalg.norm((VtV_after3 - I_n).reshape(B, -1), dim=-1).max().item()
+
+        print(f"  n={n}: before={orth_before:.2e}  "
+              f"after(2iter)={orth_after:.2e}  "
+              f"after(3iter)={orth_after3:.2e}")
+
+    # Also test with actual Jacobi output
+    print(f"\n  --- With actual Jacobi output ---")
+    for n in [5, 6]:
+        B = 2048
+        A = make_symmetric_batch(B, n, device)
+        solver = JacobiEigh(max_n=n, max_sweeps=10).to(device)
+
+        # Run Jacobi WITHOUT the NS cleanup
+        W = A.clone()
+        V = torch.eye(n, device=device).unsqueeze(0).expand(B, -1, -1).clone()
+        for _sweep in range(solver.max_sweeps):
+            for idx in range(solver._n_pairs):
+                p, q = solver._pairs_p[idx], solver._pairs_q[idx]
+                app, aqq, apq = W[:, p, p], W[:, q, q], W[:, p, q]
+                two_apq = 2.0 * apq
+                diff = aqq - app
+                abs_2apq = two_apq.abs().clamp(min=1e-30)
+                sign_2apq = torch.where(two_apq >= 0,
+                    torch.ones_like(two_apq), -torch.ones_like(two_apq))
+                tau = diff / (abs_2apq * sign_2apq)
+                tau_sign = torch.where(tau >= 0,
+                    torch.ones_like(tau), -torch.ones_like(tau))
+                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))
+                skip = (apq.abs() < 1e-30).float()
+                t = t * (1.0 - skip)
+                c = 1.0 / torch.sqrt(1.0 + t * t)
+                s = t * c
+                c_col, s_col = c.unsqueeze(-1), s.unsqueeze(-1)
+                Wp = W[:, :, p].clone(); Wq = W[:, :, q].clone()
+                W[:, :, p] = c_col * Wp - s_col * Wq
+                W[:, :, q] = s_col * Wp + c_col * Wq
+                Wp = W[:, p, :].clone(); Wq = W[:, q, :].clone()
+                W[:, p, :] = c_col * Wp - s_col * Wq
+                W[:, q, :] = s_col * Wp + c_col * Wq
+                W[:, p, q] = 0.0; W[:, q, p] = 0.0
+                W[:, p, p] = app - t * apq
+                W[:, q, q] = aqq + t * apq
+                Vp = V[:, :, p].clone(); Vq = V[:, :, q].clone()
+                V[:, :, p] = c_col * Vp - s_col * Vq
+                V[:, :, q] = s_col * Vp + c_col * Vq
+
+        I_n = torch.eye(n, device=device).unsqueeze(0)
+        VtV = torch.bmm(V.transpose(-2, -1), V)
+        orth_raw = torch.linalg.norm((VtV - I_n).reshape(B, -1), dim=-1).max().item()
+
+        V_ns = orthogonalize_ns(V, n_iter=2)
+        VtV_ns = torch.bmm(V_ns.transpose(-2, -1), V_ns)
+        orth_ns = torch.linalg.norm((VtV_ns - I_n).reshape(B, -1), dim=-1).max().item()
+
+        print(f"  Jacobi raw n={n}: orth={orth_raw:.2e}  after NS(2)={orth_ns:.2e}")
+
+
+# ─── Test 1: Accuracy ───
+
+def test_accuracy(device):
+    print("\n" + "=" * 70)
+    print("  TEST 1: ACCURACY vs torch.linalg.eigh")
+    print("=" * 70)
+
+    validator = EighValidator()
+    configs = [
+        (3,   4096, "3x3   small"),
+        (5,   4096, "5x5   CM matrix size"),
+        (6,   4096, "6x6   pentachoron bordered"),
+        (8,   2048, "8x8   padded CM"),
+        (12,  1024, "12x12 medium"),
+        (16,  512,  "16x16 Jacobi boundary"),
+    ]
+
+    all_pass = True
+    for n, B, label in configs:
+        A = make_symmetric_batch(B, n, device)
+        ref_vals, ref_vecs = torch.linalg.eigh(A)
+
+        solver = CompiledEigh(max_n=n).to(device)
+        our_vals, our_vecs = solver(A)
+
+        val_err = (our_vals - ref_vals).abs().max().item()
+        val_mean = (our_vals - ref_vals).abs().mean().item()
+
+        dots = torch.bmm(ref_vecs.transpose(-2, -1), our_vecs)
+        alignment = dots.abs().max(dim=-1).values.min().item()
+
+        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)
+        max_orth = orth_err.max().item()
+
+        # Thresholds: eigenval 1e-3, alignment 0.999, orth 1e-4
+        ok = val_err < 1e-3 and alignment > 0.999 and max_orth < 1e-4
+        if not ok:
+            all_pass = False
+
+        print(f"\n  [{'PASS' if ok else 'FAIL'}] {label}  (n={n}, B={B})")
+        print(f"    eigenvalue err   max={val_err:.2e}  mean={val_mean:.2e}")
+        print(f"    eigvec alignment min={alignment:.8f}")
+        print(f"    residual norm    max={max_res.item():.2e}")
+        print(f"    orthogonality    max={max_orth:.2e}")
+
+    print(f"\n  --- CM-like spectral distribution ---")
+    for n in [5, 6]:
+        A = make_cm_like_batch(2048, n, device)
+        ref_vals, _ = torch.linalg.eigh(A)
+        solver = CompiledEigh(max_n=n).to(device)
+        our_vals, our_vecs = solver(A)
+        val_err = (our_vals - ref_vals).abs().max().item()
+        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)
+        print(f"    CM-like n={n}: val_err={val_err:.2e}  "
+              f"res={max_res.item():.2e}  orth={orth_err.max().item():.2e}")
+
+    return all_pass
+
+
+# ─── Test 2: torch.compile fullgraph ───
+
+def test_compile(device):
+    print("\n" + "=" * 70)
+    print("  TEST 2: torch.compile(fullgraph=True)")
+    print("=" * 70)
+
+    results = {}
+    for n, B, label in [(5, 1024, "5x5"), (6, 1024, "6x6"), (8, 512, "8x8")]:
+        A = make_symmetric_batch(B, n, device)
+        solver = CompiledEigh(max_n=n).to(device)
+
+        try:
+            compiled_solver = torch.compile(solver, fullgraph=True)
+            vals, vecs = compiled_solver(A)
+            sync()
+            ref_vals, _ = torch.linalg.eigh(A)
+            err = (vals - ref_vals).abs().max().item()
+            results[label] = ("PASS", err)
+            print(f"  [{label}] fullgraph=True SUCCESS  (val_err={err:.2e})")
+        except Exception as e:
+            results[label] = ("FAIL", str(e)[:200])
+            print(f"  [{label}] COMPILE FAILED: {str(e)[:200]}")
+
+    return all(v[0] == "PASS" for v in results.values())
+
+
+# ─── Test 3: Throughput ───
+
+def test_benchmark(device):
+    print("\n" + "=" * 70)
+    print("  TEST 3: GPU THROUGHPUT BENCHMARK")
+    print("=" * 70)
+    print(f"  Device: {torch.cuda.get_device_name(0)}")
+    print(f"  Timing: 10 warmup + 200 repeats\n")
+
+    configs = [
+        (5,  1024,  "CM 5x5  B=1024"),
+        (5,  4096,  "CM 5x5  B=4096"),
+        (5,  8192,  "CM 5x5  B=8192"),
+        (6,  1024,  "CM 6x6  B=1024"),
+        (6,  4096,  "CM 6x6  B=4096"),
+        (6,  8192,  "CM 6x6  B=8192"),
+        (8,  2048,  "8x8     B=2048"),
+        (16, 1024,  "16x16   B=1024"),
+    ]
+
+    print(f"  {'Config':<22} {'eigh ref':>10} {'ours eager':>12} "
+          f"{'ours compiled':>14} {'vs ref':>8}")
+    print(f"  {'-'*22} {'-'*10} {'-'*12} {'-'*14} {'-'*8}")
+
+    for n, B, label in configs:
+        A = make_symmetric_batch(B, n, device)
+
+        ref_time = gpu_timer(lambda: torch.linalg.eigh(A))
+
+        solver = CompiledEigh(max_n=n).to(device)
+        eager_time = gpu_timer(lambda: solver(A))
+
+        try:
+            compiled_solver = torch.compile(solver, fullgraph=True)
+            for _ in range(5):
+                compiled_solver(A)
+            sync()
+            compiled_time = gpu_timer(lambda: compiled_solver(A))
+            compiled_str = fmt_time(compiled_time)
+            speedup = ref_time / compiled_time
+            speedup_str = f"{speedup:.2f}x"
+        except Exception:
+            compiled_str = "FAIL"
+            speedup_str = "N/A"
+
+        print(f"  {label:<22} {fmt_time(ref_time):>10} "
+              f"{fmt_time(eager_time):>12} {compiled_str:>14} {speedup_str:>8}")
+
+    print(f"\n  --- High batch stress test ---")
+    for n in [5, 6]:
+        for B in [16384, 32768]:
+            try:
+                A = make_symmetric_batch(B, n, device)
+                solver = CompiledEigh(max_n=n).to(device)
+                compiled_solver = torch.compile(solver, fullgraph=True)
+                for _ in range(3):
+                    compiled_solver(A)
+                sync()
+                t = gpu_timer(lambda: compiled_solver(A), warmup=5, repeats=100)
+                ref_t = gpu_timer(lambda: torch.linalg.eigh(A), warmup=5, repeats=100)
+                print(f"    n={n} B={B}: compiled={fmt_time(t)}  ref={fmt_time(ref_t)}  "
+                      f"ratio={ref_t/t:.2f}x  throughput={B/t:.0f}/sec")
+            except RuntimeError as e:
+                if "out of memory" in str(e).lower():
+                    print(f"    n={n} B={B}: OOM")
+                    torch.cuda.empty_cache()
+                else:
+                    raise
+
+
+# ─── Test 4: Autograd ───
+
+def test_autograd(device):
+    print("\n" + "=" * 70)
+    print("  TEST 4: AUTOGRAD BACKWARD")
+    print("=" * 70)
+
+    for n, B in [(5, 512), (6, 512)]:
+        A_ref = make_symmetric_batch(B, n, device).requires_grad_(True)
+        vals_ref, vecs_ref = torch.linalg.eigh(A_ref)
+        (vals_ref.sum() + (vecs_ref ** 2).sum()).backward()
+        grad_ref = A_ref.grad.clone()
+
+        # Eager backward
+        A_e = A_ref.detach().clone().requires_grad_(True)
+        solver = CompiledEigh(max_n=n).to(device)
+        try:
+            vals_e, vecs_e = solver(A_e)
+            (vals_e.sum() + (vecs_e ** 2).sum()).backward()
+            err_e = (A_e.grad - grad_ref).abs().max().item()
+            rel_e = err_e / (grad_ref.abs().max().item() + 1e-30)
+            print(f"  [{'PASS' if rel_e < 0.1 else 'WARN'}] n={n} eager backward: "
+                  f"grad_err={err_e:.2e}  rel={rel_e:.2e}")
+        except Exception as e:
+            print(f"  [FAIL] n={n} eager backward: {e}")
+
+        # Compiled backward (may break — forward fullgraph is the key win)
+        A_c = A_ref.detach().clone().requires_grad_(True)
+        try:
+            compiled_solver = torch.compile(solver)
+            vals_c, vecs_c = compiled_solver(A_c)
+            (vals_c.sum() + (vecs_c ** 2).sum()).backward()
+            err_c = (A_c.grad - grad_ref).abs().max().item()
+            rel_c = err_c / (grad_ref.abs().max().item() + 1e-30)
+            print(f"  [{'PASS' if rel_c < 0.1 else 'WARN'}] n={n} compiled backward: "
+                  f"grad_err={err_c:.2e}  rel={rel_c:.2e}")
+        except Exception as e:
+            print(f"  [INFO] n={n} compiled backward: {str(e)[:150]}")
+            print(f"         (forward fullgraph is the main win)")
+
+
+# ─── Test 5: VRAM ───
+
+def test_vram(device):
+    print("\n" + "=" * 70)
+    print("  TEST 5: VRAM USAGE")
+    print("=" * 70)
+
+    for n, B in [(5, 4096), (6, 4096), (6, 8192), (5, 8192)]:
+        torch.cuda.empty_cache()
+        gc.collect()
+        torch.cuda.reset_peak_memory_stats()
+        base_mem = torch.cuda.memory_allocated()
+
+        A = make_symmetric_batch(B, n, device)
+        solver = CompiledEigh(max_n=n).to(device)
+        vals, vecs = solver(A)
+
+        peak_mem = torch.cuda.max_memory_allocated()
+        delta_mb = (peak_mem - base_mem) / (1024 ** 2)
+        print(f"  n={n} B={B}: peak delta = {delta_mb:.1f} MB")
+
+        del A, solver, vals, vecs
+        torch.cuda.empty_cache()
+        gc.collect()
+
+
+# ─── Main ───
+
+def main():
+    print("=" * 70)
+    print("  CompiledEigh v3 — GPU Benchmark Suite")
+    print("=" * 70)
+
+    if not torch.cuda.is_available():
+        print("\n  No CUDA. Run on Colab with A100/H100.")
+        sys.exit(1)
+
+    device = torch.device('cuda')
+    print(f"\n  GPU:     {torch.cuda.get_device_name(0)}")
+    print(f"  CUDA:    {torch.version.cuda}")
+    print(f"  PyTorch: {torch.__version__}")
+    mem_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+    print(f"  VRAM:    {mem_gb:.1f} GB")
+    print(f"  TF32 matmul: {torch.backends.cuda.matmul.allow_tf32}")
+    print(f"  float32 precision: {torch.get_float32_matmul_precision()}")
+
+    test_ns_diagnostic(device)
+    acc_ok = test_accuracy(device)
+    compile_ok = test_compile(device)
+    test_benchmark(device)
+    test_autograd(device)
+    test_vram(device)
+
+    print("\n" + "=" * 70)
+    print("  SUMMARY")
+    print("=" * 70)
+    print(f"  Accuracy:  {'PASS' if acc_ok else 'FAIL'}")
+    print(f"  Compile:   {'PASS' if compile_ok else 'FAIL'}")
+    print("=" * 70)
+
+
+if __name__ == '__main__':
+    main()