feat(v4): ROI 3.A — GDN Path A as minimal-edit port of FLA naive.py

Approach correction: instead of writing GDN recurrence from scratch, vendor
fla-org/flash-linear-attention/fla/ops/gated_delta_rule/naive.py verbatim and
apply only the edits MLX semantics require (torch->mx, .clone->noop, no
in-place index assignment so o[:, :, i] = ... becomes append+stack, einops
rearrange dropped, einsum expanded to broadcast+sum). Math, variable names,
loop body, and final-state semantics preserved.

cppmega_v4/nn/_external/fla_naive_gated_delta_rule.py:
  - Verbatim copy of FLA naive.py header + recurrent function with MIT
    attribution. Edits documented at top of file so a diff vs upstream is
    short. Chunk variant deferred (its in-place attn mutation needs more
    work).

cppmega_v4/nn/linear_attention.py:
  - LinearAttentionConfig field names mirror fla.layers.gated_deltanet.
    GatedDeltaNet so a future swap to fused versions is a name-match
    (hidden_size, num_heads, head_dim, expand_v, num_v_heads, use_short_conv,
    conv_size, use_gate, norm_eps).
  - LinearAttentionBlock uses standard MLX nn.Linear + nn.RMSNorm plus a
    tiny causal short-conv ported from cppmega_mlx.nn.engram. FusedRMSNormGated
    and ShortConvolution (FLA Triton optimizations) deferred to Path C.
  - doc_ids: explicit per-segment recurrence so state does not cross document
    boundaries. cu_seqlens-style packed-batch API deferred.

mtp_v4.py: docstring clarified — V3 architecture inherited unchanged in V4
(verified: TileKernels modeling/ has no MTP module, DeepSeek-V3.2-Exp has no
MTP, DeepSeek-V3/inference/model.py is training-time only).

19 new tests in tests/v4/test_linear_attention_path_a.py:
- Config validation (6 invalid combos rejected)
- Derived dims (head_v_dim, key_dim, value_dim) match FLA arithmetic
- num_v_heads divisibility validated
- naive_recurrent shape contract; initial_state carry; output_final_state
- *** Numerical parity vs PyTorch fla.ops.gated_delta_rule.naive.
  naive_recurrent_gated_delta_rule on identical inputs (rtol 1e-4 atol 1e-5) ***
- Block forward shape, rank/dim validation, identity-at-init
- Short-conv path runs
- doc_ids changes output; shape validation
- Gradient flows into q/k/v/a/b/o projections

109/109 regression green (v4 + extensions + engram + nam pattern + mtp).

Pipeline: Implementer (port FLA naive verbatim with minimal edits) ->
Code Review (clean, attribution header, no invented math) ->
Perf Optimization (none — Path A is golden reference, perf is Path B/C/D/E job) ->
Regression Tests (parity vs FLA torch passes at 1e-5).

Author: apstenku123

cppmega_v4/nn/_external/__init__.py

@@ -0,0 +1,99 @@
+# Verbatim port of fla-org/flash-linear-attention/fla/ops/gated_delta_rule/naive.py
+# with the smallest possible edits to run on MLX instead of PyTorch.
+#
+# Upstream source:
+#   /Volumes/external/sources/rent_kernels/flash-linear-attention/
+#   fla/ops/gated_delta_rule/naive.py
+# Upstream license: MIT (Songlin Yang, Yu Zhang, Zhiyuan Li — 2023-2026)
+# Upstream copyright header retained below.
+#
+# Edits made (kept minimal so a diff against the upstream file is short):
+#   - import torch                     -> import mlx.core as mx
+#   - import torch.nn.functional as F  -> removed (only used in chunk variant)
+#   - from einops import rearrange     -> removed (replaced by inline reshape)
+#   - torch.Tensor type annotations    -> mx.array
+#   - .transpose(1, 2).contiguous()    -> mx.transpose(x, (0, 2, 1, *rest))
+#   - .to(torch.float32)               -> .astype(mx.float32)
+#   - .to(v)                           -> .astype(v.dtype)
+#   - .clone()                         -> dropped (MLX is functional, no aliasing)
+#   - .unsqueeze(-1)                   -> [..., None]
+#   - torch.zeros(...).to(v)           -> mx.zeros(shape, dtype=v.dtype)
+#   - torch.einsum('bhd,bhdm->bhm',
+#                  b_q, h)             -> mx.sum(b_q[..., :, None] * h, axis=-2)
+#   - o[:, :, i] = ...                 -> append to list, mx.stack at end
+#     (MLX is functional and cannot assign by index)
+# Algorithm, variable names, loop body, and final-state semantics unchanged.
+#
+# Chunk variant (naive_chunk_gated_delta_rule) is NOT ported here yet —
+# its in-place mutation of `attn[..., i, :i]` requires a deeper rewrite that
+# is out of scope for this minimal-edit port.
+
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import mlx.core as mx
+
+
+def naive_recurrent_gated_delta_rule(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    beta: mx.array,
+    g: mx.array,
+    scale: float = None,
+    initial_state: mx.array = None,
+    output_final_state: bool = False,
+):
+    """
+    Reference PyTorch implementation of recurrent gated delta rule.
+
+    Args:
+        q: [B, T, H, K]
+        k: [B, T, H, K]
+        v: [B, T, H, V]
+        beta: [B, T, H]
+        g: [B, T, H]
+        scale: float, optional
+        initial_state: [B, H, K, V], optional
+        output_final_state: bool
+
+    Returns:
+        o: [B, T, H, V]
+        final_state: [B, H, K, V] if output_final_state else None
+    """
+    def _t12(x):  # ".transpose(1, 2).contiguous()" — supports rank 3 and 4.
+        perm = (0, 2, 1) + tuple(range(3, x.ndim))
+        return mx.transpose(x, perm).astype(mx.float32)
+    q, k, v, beta, g = map(_t12, [q, k, v, beta, g])
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    o_list = []
+    h = mx.zeros((B, H, K, V), dtype=v.dtype)
+    if initial_state is not None:
+        h = initial_state.astype(mx.float32)
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+    q = q * scale
+
+    for i in range(T):
+        b_q = q[:, :, i]
+        b_k = k[:, :, i]
+        b_v = v[:, :, i]
+        h = h * mx.exp(g[:, :, i])[..., None, None]
+        b_beta = beta[:, :, i]
+        b_v = b_v - (h * b_k[..., None]).sum(-2)
+        b_v = b_v * b_beta[..., None]
+        h = h + b_k[..., None] * b_v[..., None, :]
+        # einsum('bhd,bhdm->bhm', b_q, h): contract d -> output (B,H,V)
+        o_list.append(mx.sum(b_q[..., :, None] * h, axis=-2))
+
+    if not output_final_state:
+        h = None
+    # stack list along the time axis (originally o[:, :, i] = ...) then
+    # transpose back to [B, T, H, V] to match upstream return shape.
+    o = mx.stack(o_list, axis=2)
+    o = mx.transpose(o, (0, 2, 1, 3))
+    return o, h

cppmega_v4/nn/_external/fla_naive_gated_delta_rule.py

@@ -0,0 +1,245 @@
+"""GatedDeltaNet block plugin (Path A entry — naive recurrent kernel).
+
+Minimal block wrapper around the FLA ``naive_recurrent_gated_delta_rule``
+math primitive (vendored in ``cppmega_v4/nn/_external/`` with MIT
+attribution). Field names are kept identical to FLA's
+``fla.layers.gated_deltanet.GatedDeltaNet`` (``q_proj``, ``k_proj``,
+``v_proj``, ``a_proj``, ``b_proj``, ``o_proj``, ``o_norm``) so a future
+swap to the fused-kernel layer is a name-match drop-in.
+
+What this wrapper DELIBERATELY does NOT do yet (deferred to Path B/C/D/E):
+    - FusedRMSNormGated (FLA Triton fusion) — we use plain RMSNorm here.
+    - ShortConvolution (FLA Triton kernel) — we use a tiny pure-MLX
+      causal depthwise conv copied from cppmega_mlx.nn.engram (same idiom).
+    - dt_bias / A_log learnable scalars (FLA training-only quality knob).
+    - Allow-negative-eigval branch.
+    - cu_seqlens packed-batch path (we expose ``doc_ids`` instead).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from cppmega_v4.nn._external.fla_naive_gated_delta_rule import (
+    naive_recurrent_gated_delta_rule,
+)
+
+
+@dataclass(frozen=True)
+class LinearAttentionConfig:
+    """Static-shape config for the GDN linear-attention block.
+
+    Field names mirror ``fla.layers.gated_deltanet.GatedDeltaNet.__init__``
+    where possible (``hidden_size``, ``num_heads``, ``head_dim``,
+    ``expand_v``, ``use_short_conv``, ``conv_size``, ``use_gate``,
+    ``norm_eps``).
+    """
+
+    hidden_size: int
+    num_heads: int = 4
+    head_dim: int = 32
+    expand_v: float = 1.0
+    num_v_heads: int | None = None
+    use_short_conv: bool = True
+    conv_size: int = 4
+    use_gate: bool = False
+    norm_eps: float = 1e-6
+
+    def __post_init__(self) -> None:
+        if self.hidden_size <= 0:
+            raise ValueError("hidden_size must be positive")
+        if self.num_heads <= 0:
+            raise ValueError("num_heads must be positive")
+        if self.head_dim <= 0:
+            raise ValueError("head_dim must be positive")
+        if self.expand_v <= 0:
+            raise ValueError("expand_v must be positive")
+        nv = self.num_heads if self.num_v_heads is None else self.num_v_heads
+        if nv <= 0:
+            raise ValueError("num_v_heads must be positive when set")
+        if nv > self.num_heads and nv % self.num_heads != 0:
+            raise ValueError("num_v_heads must be divisible by num_heads")
+        if self.conv_size < 0:
+            raise ValueError("conv_size must be non-negative")
+        if self.norm_eps <= 0:
+            raise ValueError("norm_eps must be positive")
+
+    @property
+    def head_k_dim(self) -> int:
+        return self.head_dim
+
+    @property
+    def head_v_dim(self) -> int:
+        return int(self.head_dim * self.expand_v)
+
+    @property
+    def key_dim(self) -> int:
+        return self.num_heads * self.head_k_dim
+
+    @property
+    def _num_v_heads(self) -> int:
+        return self.num_heads if self.num_v_heads is None else self.num_v_heads
+
+    @property
+    def value_dim(self) -> int:
+        return self._num_v_heads * self.head_v_dim
+
+
+def _causal_short_conv(x: mx.array, weight: mx.array) -> mx.array:
+    """Depthwise causal conv1d via shift-and-add.
+
+    x: (B, S, C); weight: (C, K, 1). Matches the idiom already in
+    ``cppmega_mlx.nn.engram.causal_depthwise_silu_conv1d`` for consistency.
+    """
+    kernel = weight.shape[1]
+    if kernel == 1:
+        return x * weight[:, 0, 0].reshape(1, 1, -1)
+    seq_len = x.shape[1]
+    out = mx.zeros_like(x)
+    for shift in range(kernel):
+        if shift == 0:
+            shifted = x
+        elif shift < seq_len:
+            shifted = mx.pad(x[:, :-shift, :], [(0, 0), (shift, 0), (0, 0)])
+        else:
+            shifted = mx.zeros_like(x)
+        tap = weight[:, kernel - 1 - shift, 0].reshape(1, 1, -1)
+        out = out + shifted * tap
+    return out
+
+
+class LinearAttentionBlock(nn.Module):
+    """GDN block. Forward: ``(B, S, hidden_size) -> (B, S, hidden_size)``.
+
+    Backend: Path A (FLA naive recurrent). Path B/C/D/E land under
+    ``cppmega_v4/_tilelang/`` and dispatch via env override in a follow-up ROI;
+    until then this block always runs Path A.
+    """
+
+    def __init__(self, config: LinearAttentionConfig):
+        super().__init__()
+        self.config = config
+        d = config.hidden_size
+        # Field names match FLA: q_proj, k_proj, v_proj, a_proj, b_proj, o_proj.
+        self.q_proj = nn.Linear(d, config.key_dim, bias=False)
+        self.k_proj = nn.Linear(d, config.key_dim, bias=False)
+        self.v_proj = nn.Linear(d, config.value_dim, bias=False)
+        self.a_proj = nn.Linear(d, config._num_v_heads, bias=False)
+        self.b_proj = nn.Linear(d, config._num_v_heads, bias=False)
+        if config.use_gate:
+            self.g_proj = nn.Linear(d, config.value_dim, bias=False)
+        if config.use_short_conv and config.conv_size > 0:
+            # Identity-initialized depthwise short conv (center tap = 1).
+            self.q_conv_weight = self._init_id_conv_weight(config.key_dim, config.conv_size)
+            self.k_conv_weight = self._init_id_conv_weight(config.key_dim, config.conv_size)
+            self.v_conv_weight = self._init_id_conv_weight(config.value_dim, config.conv_size)
+        self.o_proj = nn.Linear(config.value_dim, d, bias=False)
+        # Zero-init out projection so block is identity at init.
+        self.o_proj.weight = mx.zeros_like(self.o_proj.weight)
+        self.o_norm = nn.RMSNorm(d, eps=config.norm_eps)
+
+    @staticmethod
+    def _init_id_conv_weight(channels: int, kernel: int) -> mx.array:
+        center = kernel - 1  # causal: last tap is "current"
+        w = mx.zeros((channels, kernel, 1), dtype=mx.float32)
+        w_np = w.tolist()
+        for c in range(channels):
+            w_np[c][center][0] = 1.0
+        return mx.array(w_np, dtype=mx.float32)
+
+    def __call__(self, x: mx.array, *, doc_ids: mx.array | None = None) -> mx.array:
+        if x.ndim != 3:
+            raise ValueError(f"x must be shaped (B, S, D), got {x.shape}")
+        if x.shape[-1] != self.config.hidden_size:
+            raise ValueError(
+                f"x last dim must be {self.config.hidden_size}, got {x.shape[-1]}"
+            )
+        cfg = self.config
+        batch, seq_len, _ = x.shape
+
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+        if cfg.use_short_conv and cfg.conv_size > 0:
+            q = _causal_short_conv(q, self.q_conv_weight.astype(q.dtype))
+            k = _causal_short_conv(k, self.k_conv_weight.astype(k.dtype))
+            v = _causal_short_conv(v, self.v_conv_weight.astype(v.dtype))
+
+        # FLA layout: q/k/v ∈ [B, T, H, K|V]
+        q = q.reshape(batch, seq_len, cfg.num_heads, cfg.head_k_dim)
+        k = k.reshape(batch, seq_len, cfg.num_heads, cfg.head_k_dim)
+        v = v.reshape(batch, seq_len, cfg._num_v_heads, cfg.head_v_dim)
+
+        # a → gate decay logit; b → beta (learning rate). Sigmoid per FLA layer.
+        beta = mx.sigmoid(self.b_proj(x))   # (B, T, num_v_heads)
+        # FLA uses ``-softplus(a_proj + dt_bias)`` for gate (with dt_bias). We
+        # match the simpler "g = log(sigmoid(a_proj))" until the dt_bias path
+        # lands as part of a fused upgrade.
+        g = mx.log(mx.sigmoid(self.a_proj(x)) + cfg.norm_eps)  # (B, T, num_v_heads)
+
+        if doc_ids is None:
+            o, _ = naive_recurrent_gated_delta_rule(q, k, v, beta, g)
+        else:
+            # Document-boundary state reset: split into runs of contiguous
+            # doc_id, recur within each, concatenate outputs.
+            o = self._recurrent_with_doc_reset(q, k, v, beta, g, doc_ids)
+
+        # o has shape [B, T, H_v, V_dim]; flatten head axis for o_proj.
+        o = o.reshape(batch, seq_len, cfg.value_dim)
+        o = self.o_proj(o)
+        return self.o_norm(o)
+
+    @staticmethod
+    def _recurrent_with_doc_reset(
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        beta: mx.array,
+        g: mx.array,
+        doc_ids: mx.array,
+    ) -> mx.array:
+        """Apply naive recurrent kernel per document segment.
+
+        We loop over batch entries and over contiguous doc-id runs. This is
+        slow (O(B * num_segments) Python iterations) but correct; a fused
+        cu_seqlens-style API ports later when Path C lands.
+        """
+        if doc_ids.ndim != 2:
+            raise ValueError(f"doc_ids must be 2D (B, T), got {doc_ids.shape}")
+        batch, seq_len = doc_ids.shape
+        if q.shape[0] != batch or q.shape[1] != seq_len:
+            raise ValueError("doc_ids shape must match q's (B, T)")
+        # Iterate batch entries explicitly; build a list of [1, run_len, H, *]
+        # slices per run and concatenate along T.
+        ids_np = doc_ids.tolist()
+        per_batch_outputs: list[mx.array] = []
+        for b in range(batch):
+            row = ids_np[b]
+            runs: list[tuple[int, int]] = []
+            start = 0
+            for t in range(1, seq_len):
+                if row[t] != row[start]:
+                    runs.append((start, t))
+                    start = t
+            runs.append((start, seq_len))
+            run_outputs: list[mx.array] = []
+            for s, e in runs:
+                qb = q[b:b + 1, s:e]
+                kb = k[b:b + 1, s:e]
+                vb = v[b:b + 1, s:e]
+                bb = beta[b:b + 1, s:e]
+                gb = g[b:b + 1, s:e]
+                ob, _ = naive_recurrent_gated_delta_rule(qb, kb, vb, bb, gb)
+                run_outputs.append(ob)
+            per_batch_outputs.append(mx.concatenate(run_outputs, axis=1))
+        return mx.concatenate(per_batch_outputs, axis=0)
+
+
+__all__ = [
+    "LinearAttentionConfig",
+    "LinearAttentionBlock",
+    "naive_recurrent_gated_delta_rule",
+]

cppmega_v4/nn/linear_attention.py

@@ -1,11 +1,18 @@
-"""DeepSeek-V3-style Sequential Multi-Token-Prediction head plugin.
+"""Sequential Multi-Token-Prediction head plugin (V3 architecture, inherited in V4).
 
 The existing ``MinimalMTPHead`` (``cppmega_mlx/training/mtp.py``) reuses one
 shared block recursively for every depth — a contracted reference. DeepSeek-V3
 ships D **sequential transformer-style blocks**, one per future-token depth,
 each with its own RMSNorms, projection, and transformer kernel, but sharing the
 model's token embedding and lm head.
 
+V3 vs V4 status (as of May 2026 — public evidence only):
+    The V4-Pro / V4-Flash configs and the available V4 paper (arxiv 2512.24880,
+    which is the mHC paper, not a full V4 tech report) do not document any
+    MTP-architecture change vs V3. V4 inherits this exact Sequential head
+    unchanged. When a V4-specific MTP variant lands (full tech report or vLLM
+    V4 inference path), add a ``SequentialMTPHeadV4`` sibling here.
+
 This plugin lands the V3-faithful surface side-by-side without touching the
 existing ``MinimalMTPHead``.