feat(v4): ROI 6 — Engram primitives as minimal-edit port of TileKernels torch/engram.py

Vendor deepseek-ai/TileKernels tile_kernels/torch/engram.py verbatim and apply
only MLX edits (torch->mx, in-place .bitwise_xor_/.cumsum_/.clone dropped or
functional, .float/.bfloat16 mapped, einsum used directly via mx.einsum).
Function names, signatures, dtypes, and algorithm preserved.

cppmega_v4/nn/_external/tilekernels_engram.py:
- make_offsets(vocab_sizes)             — exclusive prefix-sum offsets per ngram layer
- engram_hash_ref(token_ids, mult, vs, off) — per-layer per-ngram bitwise-xor hash
  with modulo-into-vocab indexing; produces final embedding indices
- engram_gate_ref(x, k, v, wh, we, clamp, eps[, save_for_backward])
  — fused RMSNorm(x, wh) . RMSNorm(k, we) -> scaled-dot -> signed-sqrt
    -> sigmoid -> additive gate. Returns bfloat16 output, optionally with
    saved intermediates (dot, gate_score, rstd_x, rstd_k) for backward.

7 new tests in tests/v4/test_engram_v4.py:
- make_offsets shape + exact-prefix-sum values
- engram_hash_ref shape contract
- engram_gate_ref shape, bfloat16 dtype, save_for_backward tuple
- *** Parity vs PyTorch tile_kernels.torch.engram (3 functions, all match
  upstream within int-exact for hash/offsets and bfloat16 atol 5e-2 for gate) ***

147/147 regression tests green (v4 + extensions + engram + nam pattern + MTP).

Pipeline: Implementer (TileKernels torch port verbatim) -> Code Review (MIT
attribution, no invented math) -> Perf (none — Path A reference) ->
Regression Tests (parity passes).

Author: apstenku123

cppmega_v4/nn/_external/tilekernels_engram.py

@@ -0,0 +1,149 @@
+# Verbatim port of deepseek-ai/TileKernels/tile_kernels/torch/engram.py with
+# the smallest possible edits to run on MLX instead of PyTorch.
+#
+# Upstream source:
+#   /Users/dave/sources/TileKernels/tile_kernels/torch/engram.py
+# Upstream license: MIT (deepseek-ai/TileKernels)
+#
+# Edits made (kept minimal so a diff against the upstream file is short):
+#   - import torch                                -> import mlx.core as mx
+#   - torch.Tensor annotations                    -> mx.array
+#   - x.view(-1)                                  -> x.reshape(-1)
+#   - x.unsqueeze(0|1|-1)                         -> x[None] / x[..., None]
+#   - .to(torch.int32) / int64                    -> .astype(mx.int32) / int64
+#   - .clone()                                    -> dropped (functional)
+#   - hashes.bitwise_xor_(other)                  -> hashes = mx.bitwise_xor(hashes, other)
+#     (MLX is functional, no in-place ops)
+#   - x.cumsum(0, dtype=torch.int32)              -> mx.cumsum(x, axis=0).astype(mx.int32)
+#   - torch.cat([...], dim=...)                   -> mx.concatenate([...], axis=...)
+#   - torch.stack([...], dim=...)                 -> mx.stack([...], axis=...)
+#   - torch.zeros(1, dtype=..., device=...)       -> mx.zeros((1,), dtype=...)
+#   - x.float() / x.bfloat16()                    -> x.astype(mx.float32|bfloat16)
+#   - torch.rsqrt / x.pow(2) / x.sigmoid          -> mx.rsqrt / mx.square / mx.sigmoid
+#   - torch.einsum('...d,...d->...', a, b)        -> mx.einsum('...d,...d->...', a, b)
+#   - dot.abs().clamp_min(v).sqrt() * dot.sign()  -> mx.sign(dot) * mx.sqrt(mx.maximum(mx.abs(dot), v))
+# Algorithm, signatures, dtypes, return shapes unchanged.
+
+import mlx.core as mx
+
+
+def make_offsets(vocab_sizes: mx.array) -> mx.array:
+    """Compute exclusive prefix-sum offsets from vocab_sizes.
+
+    Args:
+        vocab_sizes: Per-layer per-ngram embedding table sizes of shape
+            (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram), int32.
+
+    Returns:
+        Offsets of shape (num_ngram_layers, (max_ngram_size - 1) * num_embed_table_per_ngram), int32.
+    """
+    num_ngram_layers = vocab_sizes.shape[0]
+    offsets_list = []
+    for layer_idx in range(num_ngram_layers):
+        flat = vocab_sizes[layer_idx].reshape(-1)
+        prefix = mx.concatenate(
+            [
+                mx.zeros((1,), dtype=mx.int32),
+                mx.cumsum(flat[:-1], axis=0).astype(mx.int32),
+            ]
+        )
+        offsets_list.append(prefix)
+    return mx.stack(offsets_list, axis=0)
+
+
+def engram_hash_ref(
+    ngram_token_ids: mx.array,
+    multipliers: mx.array,
+    vocab_sizes: mx.array,
+    offsets: mx.array,
+) -> mx.array:
+    """Pure PyTorch reference implementation of engram hash.
+
+    Args:
+        ngram_token_ids: N-gram token IDs of shape (num_tokens, max_ngram_size), int32.
+        multipliers: Per-layer hash multipliers of shape (num_ngram_layers, max_ngram_size), int64.
+        vocab_sizes: Per-layer per-ngram embedding table sizes of shape
+            (num_ngram_layers, max_ngram_size - 1, num_embed_table_per_ngram), int32.
+        offsets: Per-layer embedding table offsets of shape
+            (num_ngram_layers, (max_ngram_size - 1) * num_embed_table_per_ngram), int32.
+
+    Returns:
+        Embedding indices of shape (num_ngram_layers, num_tokens, (max_ngram_size - 1) * num_embed_table_per_ngram), int32.
+    """
+    num_ngram_layers = multipliers.shape[0]
+    max_ngram_size = multipliers.shape[1]
+
+    prod = ngram_token_ids.astype(mx.int64)[None] * multipliers[:, None]
+
+    ans = [[] for _ in range(num_ngram_layers)]
+    hashes = prod[:, :, 0]
+    for i in range(1, max_ngram_size):
+        hashes = mx.bitwise_xor(hashes, prod[:, :, i])
+        for layer_idx in range(num_ngram_layers):
+            ans[layer_idx].append(
+                (
+                    hashes[layer_idx][..., None]
+                    % vocab_sizes[layer_idx, i - 1].astype(mx.int64)[None]
+                ).astype(mx.int32)
+            )
+
+    for layer_idx in range(num_ngram_layers):
+        ans[layer_idx] = mx.concatenate(ans[layer_idx], axis=-1)
+
+    output = mx.stack(ans, axis=0)
+    return output + offsets[:, None]
+
+
+def engram_gate_ref(
+    hidden_states: mx.array,
+    k: mx.array,
+    v: mx.array,
+    weight_hidden: mx.array,
+    weight_embed: mx.array,
+    clamp_value: float,
+    eps: float,
+    save_for_backward: bool = False,
+):
+    """Pure PyTorch reference implementation of engram gate (vectorized, supports autograd).
+
+    Computes: output = x + sigmoid(signed_sqrt(dot(RMSNorm(x, wh), RMSNorm(k, we)) * scalar)) * v
+
+    Args:
+        hidden_states: Input of shape (num_tokens, hc_mult, hidden_size), bfloat16.
+        k: Key embeddings of shape (num_tokens, hc_mult, hidden_size), bfloat16.
+        v: Value embeddings of shape (num_tokens, hidden_size), bfloat16.
+        weight_hidden: RMSNorm weight for hidden states, shape (hc_mult, hidden_size), bfloat16.
+        weight_embed: RMSNorm weight for key embeddings, shape (hc_mult, hidden_size), bfloat16.
+        clamp_value: Clamp threshold for signed-sqrt gate activation.
+        eps: Epsilon for RMSNorm numerical stability.
+        save_for_backward: If True, also return (dot, gate_score, rstd_x, rstd_k).
+
+    Returns:
+        If save_for_backward is False: output tensor of shape (num_tokens, hc_mult, hidden_size), bfloat16.
+        If save_for_backward is True: tuple of (output, dot, gate_score, rstd_x, rstd_k).
+    """
+    hidden_size = hidden_states.shape[-1]
+    scalar = hidden_size**-0.5
+
+    x = hidden_states.astype(mx.float32)
+    k_f = k.astype(mx.float32)
+    wh = weight_hidden.astype(mx.float32)[None]
+    we = weight_embed.astype(mx.float32)[None]
+
+    # RMSNorm
+    rstd_x = mx.rsqrt(mx.mean(mx.square(x), axis=-1) + eps)
+    rstd_k = mx.rsqrt(mx.mean(mx.square(k_f), axis=-1) + eps)
+
+    # Dot -> sqrt-gate -> sigmoid
+    # raw_dot is the unnormalized sum(x * wh * k * we), matching the kernel's dot_out
+    raw_dot = mx.einsum('...d,...d->...', x * wh, k_f * we)
+    dot = raw_dot * rstd_x * rstd_k * scalar
+    signed_sqrt = mx.sign(dot) * mx.sqrt(mx.maximum(mx.abs(dot), clamp_value))
+    gate_score = mx.sigmoid(signed_sqrt)
+
+    output = x + gate_score[..., None] * v[..., None, :]
+    output = output.astype(mx.bfloat16)
+
+    if save_for_backward:
+        return output, raw_dot, gate_score, rstd_x, rstd_k
+    return output

tests/v4/test_engram_v4.py

@@ -0,0 +1,173 @@
+"""Tests for cppmega_v4.nn._external.tilekernels_engram — port of DeepSeek engram torch ref.
+
+Parity tests against ``~/sources/TileKernels/tile_kernels/torch/engram.py``
+when torch is available.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+import pytest
+
+from cppmega_v4.nn._external.tilekernels_engram import (
+    engram_gate_ref,
+    engram_hash_ref,
+    make_offsets,
+)
+
+
+# ----- make_offsets shape + values -----
+
+
+def test_make_offsets_shape_and_values():
+    # 2 layers, 3 ngram positions, 2 tables/ngram -> flat per layer has 6 entries.
+    vocab_sizes = mx.array([[[10, 20], [30, 40], [50, 60]],
+                            [[5, 5], [5, 5], [5, 5]]], dtype=mx.int32)
+    offsets = make_offsets(vocab_sizes)
+    assert offsets.shape == (2, 6)
+    # Layer 0 exclusive prefix sum of [10, 20, 30, 40, 50, 60] = [0,10,30,60,100,150]
+    np.testing.assert_array_equal(
+        np.array(offsets[0]), np.array([0, 10, 30, 60, 100, 150], dtype=np.int32)
+    )
+    np.testing.assert_array_equal(
+        np.array(offsets[1]), np.array([0, 5, 10, 15, 20, 25], dtype=np.int32)
+    )
+
+
+def test_engram_hash_shape_contract():
+    num_tokens = 4
+    max_ngram_size = 3
+    num_ngram_layers = 2
+    num_embed = 2
+    rng = np.random.default_rng(42)
+    ngram_token_ids = mx.array(
+        rng.integers(0, 100, size=(num_tokens, max_ngram_size)).astype(np.int32)
+    )
+    multipliers = mx.array(
+        rng.integers(1, 10000, size=(num_ngram_layers, max_ngram_size)).astype(np.int64)
+    )
+    vocab_sizes = mx.array(
+        rng.integers(5, 50, size=(num_ngram_layers, max_ngram_size - 1, num_embed)).astype(np.int32)
+    )
+    offsets = make_offsets(vocab_sizes)
+    out = engram_hash_ref(ngram_token_ids, multipliers, vocab_sizes, offsets)
+    # out shape (num_ngram_layers, num_tokens, (max_ngram_size - 1) * num_embed_table_per_ngram)
+    assert out.shape == (num_ngram_layers, num_tokens, (max_ngram_size - 1) * num_embed)
+
+
+def test_engram_gate_shape_and_dtype():
+    num_tokens, hc_mult, hidden = 4, 2, 16
+    rng = np.random.default_rng(7)
+    hs = mx.array(rng.standard_normal((num_tokens, hc_mult, hidden)).astype(np.float32))
+    k = mx.array(rng.standard_normal((num_tokens, hc_mult, hidden)).astype(np.float32))
+    v = mx.array(rng.standard_normal((num_tokens, hidden)).astype(np.float32))
+    wh = mx.ones((hc_mult, hidden))
+    we = mx.ones((hc_mult, hidden))
+    out = engram_gate_ref(
+        hs.astype(mx.bfloat16), k.astype(mx.bfloat16), v.astype(mx.bfloat16),
+        wh.astype(mx.bfloat16), we.astype(mx.bfloat16),
+        clamp_value=1e-4, eps=1e-6,
+    )
+    assert out.shape == (num_tokens, hc_mult, hidden)
+    assert out.dtype == mx.bfloat16
+
+
+def test_engram_gate_save_for_backward_returns_tuple():
+    num_tokens, hc_mult, hidden = 2, 1, 8
+    hs = mx.random.normal((num_tokens, hc_mult, hidden)).astype(mx.bfloat16)
+    k = mx.random.normal((num_tokens, hc_mult, hidden)).astype(mx.bfloat16)
+    v = mx.random.normal((num_tokens, hidden)).astype(mx.bfloat16)
+    wh = mx.ones((hc_mult, hidden)).astype(mx.bfloat16)
+    we = mx.ones((hc_mult, hidden)).astype(mx.bfloat16)
+    out, dot, gate, rstd_x, rstd_k = engram_gate_ref(
+        hs, k, v, wh, we, clamp_value=1e-4, eps=1e-6, save_for_backward=True
+    )
+    assert out.shape == (num_tokens, hc_mult, hidden)
+    assert dot.shape == (num_tokens, hc_mult)
+    assert gate.shape == (num_tokens, hc_mult)
+    assert rstd_x.shape == (num_tokens, hc_mult)
+    assert rstd_k.shape == (num_tokens, hc_mult)
+
+
+# ----- parity vs PyTorch TileKernels reference -----
+
+
+@pytest.fixture(scope="module")
+def tk_engram_torch():
+    torch = pytest.importorskip("torch")
+    repo = Path("/Users/dave/sources/TileKernels")
+    if not repo.exists():
+        pytest.skip("TileKernels repo not present at expected path")
+    if str(repo) not in sys.path:
+        sys.path.insert(0, str(repo))
+    try:
+        from tile_kernels.torch.engram import (
+            engram_gate_ref as tk_gate,
+            engram_hash_ref as tk_hash,
+            make_offsets as tk_offsets,
+        )
+    except Exception as exc:
+        pytest.skip(f"could not import TileKernels torch.engram: {exc}")
+    return {"torch": torch, "gate": tk_gate, "hash": tk_hash, "offsets": tk_offsets}
+
+
+def test_parity_make_offsets(tk_engram_torch):
+    torch = tk_engram_torch["torch"]
+    rng = np.random.default_rng(11)
+    vs_np = rng.integers(1, 100, size=(2, 3, 2)).astype(np.int32)
+    t_out = tk_engram_torch["offsets"](torch.from_numpy(vs_np))
+    m_out = make_offsets(mx.array(vs_np))
+    np.testing.assert_array_equal(np.array(m_out), t_out.numpy())
+
+
+def test_parity_engram_hash(tk_engram_torch):
+    torch = tk_engram_torch["torch"]
+    rng = np.random.default_rng(12)
+    num_tokens, max_ngram, layers, num_embed = 4, 3, 2, 2
+    ng = rng.integers(0, 100, size=(num_tokens, max_ngram)).astype(np.int32)
+    mult = rng.integers(1, 10000, size=(layers, max_ngram)).astype(np.int64)
+    vs = rng.integers(5, 50, size=(layers, max_ngram - 1, num_embed)).astype(np.int32)
+    off = make_offsets(mx.array(vs))
+    t_off = tk_engram_torch["offsets"](torch.from_numpy(vs))
+    t_out = tk_engram_torch["hash"](
+        torch.from_numpy(ng), torch.from_numpy(mult), torch.from_numpy(vs), t_off
+    )
+    m_out = engram_hash_ref(mx.array(ng), mx.array(mult), mx.array(vs), off)
+    np.testing.assert_array_equal(np.array(m_out), t_out.numpy())
+
+
+def test_parity_engram_gate(tk_engram_torch):
+    torch = tk_engram_torch["torch"]
+    rng = np.random.default_rng(13)
+    nt, hc, h = 3, 2, 16
+    hs = rng.standard_normal((nt, hc, h)).astype(np.float32)
+    k = rng.standard_normal((nt, hc, h)).astype(np.float32)
+    v = rng.standard_normal((nt, h)).astype(np.float32)
+    wh = rng.standard_normal((hc, h)).astype(np.float32)
+    we = rng.standard_normal((hc, h)).astype(np.float32)
+    t_out = tk_engram_torch["gate"](
+        torch.from_numpy(hs).bfloat16(),
+        torch.from_numpy(k).bfloat16(),
+        torch.from_numpy(v).bfloat16(),
+        torch.from_numpy(wh).bfloat16(),
+        torch.from_numpy(we).bfloat16(),
+        clamp_value=1e-4, eps=1e-6,
+    )
+    m_out = engram_gate_ref(
+        mx.array(hs).astype(mx.bfloat16),
+        mx.array(k).astype(mx.bfloat16),
+        mx.array(v).astype(mx.bfloat16),
+        mx.array(wh).astype(mx.bfloat16),
+        mx.array(we).astype(mx.bfloat16),
+        clamp_value=1e-4, eps=1e-6,
+    )
+    # bfloat16 path — tolerate ~5e-2 max abs error.
+    np.testing.assert_allclose(
+        np.array(m_out.astype(mx.float32)),
+        t_out.float().numpy(),
+        atol=5e-2, rtol=5e-2,
+    )