feat(v4): vendor mlx-lm PR #1189 mHC — HyperConnection + fused Sinkhorn

Two verbatim vendored modules from ml-explore/mlx-lm PR #1189 (MIT,
© 2025-2026 Apple Inc.):

  - _mlx_lm_sinkhorn_vendored.py — Sinkhorn projection onto the Birkhoff
    polytope with a fused register-resident Metal kernel. One thread per
    token, hc² floats in registers, fully unrolled row_softmax + 20
    iterations of (row_norm, col_norm). No threadgroup memory, no atomics,
    no global-memory traffic between iters. Pure-MLX fallback for hc>8 or
    when Metal isn't available.

  - _mlx_lm_hyper_connection_vendored.py — HyperConnection (per-block mHC)
    + HyperHead (final-layer head). hc_pre reduces [B,S,hc,D]→[B,S,D] via
    learned pre-gating + Sinkhorn; hc_post expands back via comb @ residual
    + post * f_out.

Complements our existing tilekernels_mhc.py (pure-MLX port of DeepSeek's
torch reference): the vendored version gives us the fused Sinkhorn kernel
that the upstream paper relies on for production speed.

Tests (6 new in tests/v4/test_mhc_metal_v4.py):
  - hc_split_sinkhorn shapes / doubly-stochastic property (rows, cols ≈ 1)
  - Metal kernel vs pure-MLX parity (atol 1e-3 on comb, 1e-6 on pre/post)
  - HyperConnection module hc_pre/hc_post shape contract
  - HyperHead head reduction shape
  - Sinkhorn kernel cache returns identity for repeated build

v4 suite: 184 passed / 2 skipped.

Author: apstenku123

cppmega_v4/nn/_external/_mlx_lm_hyper_connection_vendored.py

@@ -0,0 +1,125 @@
+# Vendored verbatim from ml-explore/mlx-lm PR #1189
+#   (mlx_lm/models/hyper_connection.py)
+# Upstream license: MIT, Copyright © 2026 Apple Inc.
+#
+# Public surface used by cppmega_v4:
+#   - HyperConnection  — per-block mHC (hc_pre + hc_post)
+#   - HyperHead        — final-layer head variant (sigmoid-weighted reduction)
+#
+# No edits — pure copy. cppmega_v4/nn/mhc_v4.py exposes the v4 wrapper
+# that integrates these into our UnifiedSuperblock.
+
+"""Manifold-constrained Hyper-Connections (mHC) — shared building blocks.
+
+DeepSeek-V4 introduced multi-HyperConnection (mHC) as a residual replacement:
+expand the hidden state into `hc_mult` parallel copies, mix them via a
+doubly-stochastic matrix on the Birkhoff polytope, apply the block, and
+recombine. This module hosts the `nn.Module` layers; the Sinkhorn projection
+that produces the doubly-stochastic mixing matrix lives in
+`_mlx_lm_sinkhorn_vendored`.
+
+Two layers:
+  - HyperConnection — per-block mHC: hc_pre reduces hc_mult -> 1; block F runs;
+    hc_post expands 1 -> hc_mult via (post * f_out + comb @ residual).
+  - HyperHead      — final-layer head variant: sigmoid-weighted reduction
+    hc_mult -> 1 with no Sinkhorn (simpler than HyperConnection).
+
+References:
+  - mHC: arXiv:2512.24880 (DeepSeek, Dec 2025)
+  - HC base: arXiv:2409.19606 (Sep 2024)
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ._mlx_lm_sinkhorn_vendored import hc_split_sinkhorn
+
+
+class HyperConnection(nn.Module):
+    """Per-block mHC parameters: projects x -> (pre, post, comb) used in hc_pre/hc_post.
+
+    Paper/ref stores the weights as:
+        hc_fn    : [(2+hc)*hc, hc*dim]
+        hc_scale : [3]
+        hc_base  : [(2+hc)*hc]
+
+    hc_pre reduces `hc_mult` parallel hidden states to 1 via `pre`.
+    Block F is applied to the reduced state. hc_post expands 1 -> hc via `post` (the new
+    contribution) added to `comb @ residual` (where `comb` is a doubly-stochastic mix
+    that recombines the input `hc_mult` copies to stay on the Birkhoff manifold).
+    """
+
+    def __init__(self, dim: int, hc_mult: int, norm_eps: float, sinkhorn_iters: int, hc_eps: float):
+        super().__init__()
+        self.dim = dim
+        self.hc_mult = hc_mult
+        self.norm_eps = norm_eps
+        self.sinkhorn_iters = sinkhorn_iters
+        self.hc_eps = hc_eps
+        mix_hc = (2 + hc_mult) * hc_mult
+        hc_dim = hc_mult * dim
+        # All mHC params are fp32 in the checkpoint.
+        self.fn = mx.zeros((mix_hc, hc_dim), dtype=mx.float32)
+        self.base = mx.zeros((mix_hc,), dtype=mx.float32)
+        self.scale = mx.zeros((3,), dtype=mx.float32)
+        self._fn_t = None  # lazy transpose cache (avoids 86 .T calls/token)
+
+    def hc_pre(self, x: mx.array):
+        B, S, hc, D = x.shape
+        dtype = x.dtype
+        xf = x.reshape(B, S, hc * D).astype(mx.float32)
+        xf_norm = mx.fast.rms_norm(xf, weight=None, eps=self.norm_eps)
+        if self._fn_t is None:
+            self._fn_t = self.fn.T
+        mixes = (xf_norm @ self._fn_t).reshape(B * S, -1)
+        pre, post, comb = hc_split_sinkhorn(
+            mixes, self.scale, self.base, hc, self.sinkhorn_iters, self.hc_eps
+        )
+        pre  = pre.reshape(B, S, hc)
+        post = post.reshape(B, S, hc)
+        comb = comb.reshape(B, S, hc, hc)
+        y = (pre[..., None] * x.astype(mx.float32)).sum(axis=2)
+        return y.astype(dtype), post, comb
+
+    def hc_post(self, f_out: mx.array, residual: mx.array, post: mx.array, comb: mx.array):
+        # f_out    [B,S,D] (block output, reduced state)
+        # residual [B,S,hc,D] (input to hc_pre)
+        # post     [B,S,hc]
+        # comb     [B,S,hc,hc]
+        # returns  [B,S,hc,D]
+        dtype = f_out.dtype
+        # post.unsqueeze(-1) * f_out.unsqueeze(-2)  -> [B,S,hc,D]
+        term_new = post[..., None] * f_out[:, :, None, :].astype(mx.float32)
+        # comb @ residual: [B,S,hc,hc] @ [B,S,hc,D] -> [B,S,hc,D]
+        term_res = comb.astype(mx.float32) @ residual.astype(mx.float32)
+        y = term_new + term_res
+        return y.astype(dtype)
+
+
+class HyperHead(nn.Module):
+    """Final (head) mHC projection: reduces [B,S,hc,D] -> [B,S,D] via sigmoid-weighted sum.
+    No Sinkhorn here — this is the simpler head variant from `ParallelHead.hc_head`.
+    """
+
+    def __init__(self, dim: int, hc_mult: int, norm_eps: float, hc_eps: float):
+        super().__init__()
+        self.dim = dim
+        self.hc_mult = hc_mult
+        self.norm_eps = norm_eps
+        self.hc_eps = hc_eps
+        self.fn = mx.zeros((hc_mult, hc_mult * dim), dtype=mx.float32)
+        self.base = mx.zeros((hc_mult,), dtype=mx.float32)
+        self.scale = mx.zeros((1,), dtype=mx.float32)
+        self._fn_t = None  # lazy transpose cache
+
+    def __call__(self, x: mx.array):
+        B, S, hc, D = x.shape
+        dtype = x.dtype
+        xf = x.reshape(B, S, hc * D).astype(mx.float32)
+        inv = mx.rsqrt((xf * xf).mean(axis=-1, keepdims=True) + self.norm_eps)
+        if self._fn_t is None:
+            self._fn_t = self.fn.T
+        mixes = (xf @ self._fn_t) * inv                     # [B,S,hc]
+        pre = mx.sigmoid(mixes * self.scale[0] + self.base) + self.hc_eps
+        y = (pre[..., None] * x.astype(mx.float32)).sum(axis=2)
+        return y.astype(dtype)

cppmega_v4/nn/_external/_mlx_lm_sinkhorn_vendored.py

@@ -0,0 +1,173 @@
+# Vendored verbatim from ml-explore/mlx-lm PR #1189 (mlx_lm/models/sinkhorn.py)
+# Upstream license: MIT, Copyright © 2025 Apple Inc.
+#
+# Public surface used by cppmega_v4:
+#   - hc_split_sinkhorn(mixes, hc_scale, hc_base, hc_mult, sinkhorn_iters, eps)
+#     → (pre, post, comb) where comb is doubly-stochastic on the Birkhoff
+#       polytope (Sinkhorn-normalized).
+#   - _make_sinkhorn_kernel(hc, iters, eps) — fused register-resident Metal
+#     kernel (one thread per token, hc^2 floats in registers, no atomics).
+#
+# No edits — pure copy. cppmega_v4/nn/_external/mlx_lm_hyper_connection.py
+# wraps our (pre, post, comb) API to this module's surface.
+
+"""Sinkhorn normalization for doubly-stochastic matrices (HyperConnection mHC).
+
+Shared module for models that use learned mixing weights on the Birkhoff
+polytope (e.g., DeepSeek-V4's multi-HyperConnection). Provides both a
+fused Metal kernel (register-resident, one thread per token) and a pure-MLX
+reference path.
+"""
+
+import mlx.core as mx
+
+_SINKHORN_KERNELS: dict = {}
+
+
+def _make_sinkhorn_kernel(hc: int, iters: int, eps: float):
+    """Build a fully-unrolled, register-resident Sinkhorn Metal kernel.
+
+    Each thread owns one token's [hc, hc] matrix in registers (hc^2 floats).
+    No threadgroup memory, no atomics, no global memory traffic between iters
+    — kernel reads input once and writes output once.
+    """
+    key = (hc, iters)
+    if key in _SINKHORN_KERNELS:
+        return _SINKHORN_KERNELS[key]
+
+    n_elem = hc * hc
+    eps_lit = f"{eps:.8e}f"
+
+    def row_softmax():
+        out = []
+        for r in range(hc):
+            base = r * hc
+            out.append(f"        {{ float mx = m[{base}];")
+            for c in range(1, hc):
+                out.append(f"          mx = metal::max(mx, m[{base + c}]);")
+            out.append(f"          float s = 0.0f;")
+            for c in range(hc):
+                out.append(f"          m[{base + c}] = metal::exp(m[{base + c}] - mx); s += m[{base + c}];")
+            out.append(f"          float inv = 1.0f / s;")
+            for c in range(hc):
+                out.append(f"          m[{base + c}] *= inv;")
+            out.append("        }")
+        return "\n".join(out)
+
+    def row_norm():
+        out = []
+        for r in range(hc):
+            base = r * hc
+            terms = " + ".join(f"m[{base + c}]" for c in range(hc))
+            out.append(f"        {{ float inv = 1.0f / ({terms} + {eps_lit});")
+            for c in range(hc):
+                out.append(f"          m[{base + c}] *= inv;")
+            out.append("        }")
+        return "\n".join(out)
+
+    def col_norm():
+        out = []
+        for c in range(hc):
+            terms = " + ".join(f"m[{r * hc + c}]" for r in range(hc))
+            out.append(f"        {{ float inv = 1.0f / ({terms} + {eps_lit});")
+            for r in range(hc):
+                out.append(f"          m[{r * hc + c}] *= inv;")
+            out.append("        }")
+        return "\n".join(out)
+
+    def add_eps():
+        return "\n".join(f"        m[{i}] += {eps_lit};" for i in range(n_elem))
+
+    iter_body = "\n".join([row_norm(), col_norm()])
+    inner_iters = "\n".join([iter_body] * (iters - 1))
+
+    source = f"""
+        uint n = thread_position_in_grid.x;
+        if (n >= n_tokens[0]) return;
+
+        uint base = n * {n_elem};
+        float m[{n_elem}];
+{chr(10).join(f"        m[{i}] = comb_log[base + {i}];" for i in range(n_elem))}
+
+        // Row softmax
+{row_softmax()}
+
+        // Add eps
+{add_eps()}
+
+        // Initial column normalization (matches reference: cols first after softmax)
+{col_norm()}
+
+        // Remaining (iters - 1) rounds of (row_norm, col_norm)
+{inner_iters}
+
+{chr(10).join(f"        comb[base + {i}] = m[{i}];" for i in range(n_elem))}
+    """
+
+    kernel = mx.fast.metal_kernel(
+        name=f"sinkhorn_hc{hc}_it{iters}",
+        input_names=["comb_log", "n_tokens"],
+        output_names=["comb"],
+        source=source,
+    )
+    _SINKHORN_KERNELS[key] = kernel
+    return kernel
+
+
+def hc_split_sinkhorn(
+    mixes: mx.array,        # [B*S, (2+hc)*hc] fp32
+    hc_scale: mx.array,     # [3] fp32
+    hc_base: mx.array,      # [(2+hc)*hc] fp32
+    hc_mult: int = 4,
+    sinkhorn_iters: int = 20,
+    eps: float = 1e-6,
+):
+    """Split `mixes` into (pre, post, comb_logits); Sinkhorn-normalize comb to doubly stochastic.
+
+    Returns:
+        pre  [N, hc]        — sigmoid(mixes[:,:hc] * s0 + base[:hc]) + eps
+        post [N, hc]        — 2*sigmoid(mixes[:,hc:2hc] * s1 + base[hc:2hc])
+        comb [N, hc, hc]    — Sinkhorn-normalized (rows & cols ~= 1) from the last hc*hc logits.
+    """
+    n = mixes.shape[0]
+    mix = mixes
+    s0, s1, s2 = hc_scale[0], hc_scale[1], hc_scale[2]
+
+    pre_log  = mix[:, :hc_mult]             * s0 + hc_base[:hc_mult]
+    post_log = mix[:, hc_mult:2 * hc_mult]  * s1 + hc_base[hc_mult:2 * hc_mult]
+    comb_log = (
+        mix[:, 2 * hc_mult:].reshape(n, hc_mult, hc_mult) * s2
+        + hc_base[2 * hc_mult:].reshape(hc_mult, hc_mult)
+    )
+
+    pre  = mx.sigmoid(pre_log) + eps
+    post = 2 * mx.sigmoid(post_log)
+
+    use_kernel = (
+        hc_mult <= 8
+        and mx.metal.is_available()
+        and comb_log.size > 0
+    )
+    if use_kernel:
+        kernel = _make_sinkhorn_kernel(hc_mult, sinkhorn_iters, eps)
+        flat = comb_log.reshape(n, hc_mult * hc_mult).astype(mx.float32)
+        n_tokens = mx.array([n], dtype=mx.uint32)
+        (comb_flat,) = kernel(
+            inputs=[flat, n_tokens],
+            output_shapes=[(n, hc_mult * hc_mult)],
+            output_dtypes=[mx.float32],
+            grid=(n, 1, 1),
+            threadgroup=(min(n, 256) or 1, 1, 1),
+        )
+        comb = comb_flat.reshape(n, hc_mult, hc_mult)
+    else:
+        comb = mx.softmax(comb_log, axis=-1, precise=True) + eps
+        col_sum = comb.sum(axis=1, keepdims=True) + eps
+        comb = comb / col_sum
+        for _ in range(sinkhorn_iters - 1):
+            row_sum = comb.sum(axis=2, keepdims=True) + eps
+            comb = comb / row_sum
+            col_sum = comb.sum(axis=1, keepdims=True) + eps
+            comb = comb / col_sum
+
+    return pre, post, comb