feat(v4): UnifiedSuperblockV4 real factories — no more pass-through stubs

Replaced 7 pass-through lambdas in BLOCK_BUILDERS with real factories:

  - "gdn"               → LinearAttentionBlock(LinearAttentionConfig(...))
                          (FLA-naive GDN; Path B/C/E dispatchable via env)
  - "kda"               → KimiDeltaAttentionBlock(KimiDeltaAttentionConfig(...))
  - "moe"               → V4MoE(V4MoEConfig(d_model=...)) wrapped to expose
                          .output (V4MoE returns MoEOutput dataclass)
  - "attention"         → standard multi-head self-attention with causal
                          mask + RMSNorm; o_proj zero-init for identity-at-init.
                          Also the fallback for "mla" and "mla_absorb" until
                          we land a full MLA block.
  - "lightning_indexer" → LightningIndexerFP8 wrapped as residual no-op
                          (indexer's natural output is top-k indices, not
                          a residual — real consumption happens inside CSA+HCA).

Dispatch in UnifiedSuperblockV4.__call__ extended:
  - _DOC_ID_KW_KINDS = {"gdn", "kda"} — these accept document_ids as a
    keyword argument (their recurrent state needs the doc-reset signal).
  - Engram remains positional-token-ids + keyword-doc-ids (special-cased).

Local imports inside each builder avoid circular-import issues at the
unified_superblock_v4 module level.

Tests (18 new in tests/v4/test_unified_superblock_real_factories.py):
  - 7 parametrized: each kind is no longer a pass-through lambda.
  - gdn / kda / moe / attention(×3 for mla aliases) builds and runs.
  - GDN block threads doc_ids → output differs across doc-boundary
    configurations (after perturbing o_proj, which is zero-init).
  - attention zero-init identity check (out == in via residual).
  - lightning_indexer residual identity (out == in by design).
  - Full V4 stack end-to-end (engram + gdn + kda + attention + nsa +
    csa_hca + moe + mlp) — runs forward and propagates gradients.

scripts/v4_1b_parquet_smoke.py updated: per-kind layer count spread
across all 7 attention-like kinds (gdn, kda, attention, nsa, csa_hca,
moe, mlp) so the smoke now exercises the *whole* V4 cohort, not just
the previously-non-stub subset.

Smoke run (hidden=1024, L=12, T=256, B=1, 3 steps):
  loss decreases 11.18 → 11.02 → 10.96 (sanity check)
  fwd median 59 ms, total 91 ms, ~2 800 tokens/sec
  (was 10 191 tokens/sec on the engram+nsa+csa_hca+mlp-only subset,
  drop is expected — we now run 6× more block kinds, including
  attention/MoE compute that was being silently skipped before)

v4 suite: 303 passed / 2 skipped.

Author: apstenku123

cppmega_v4/models/unified_superblock_v4.py

@@ -72,29 +72,152 @@ def __call__(self, x):
 
 
 def _build_pass_through_unsupported(kind: str):
-    """Block kinds with circular imports or external dep — pass-through."""
+    """Last-resort no-op for kinds with no factory yet. Residual-only."""
     class _PassThrough(nn.Module):
         def __init__(self):
             super().__init__()
             self._kind = kind
         def __call__(self, x):
-            return x  # zero contribution, residual-only path
+            return x  # zero contribution
     return _PassThrough()
 
 
+def _build_gdn(hidden_size: int, params: dict) -> nn.Module:
+    """Real GatedDeltaNet block (Path A — Path B/C/E land via env override)."""
+    from cppmega_v4.nn.linear_attention import (  # local import: avoids circular
+        LinearAttentionBlock, LinearAttentionConfig,
+    )
+    # Defaults sensible for our 1B smoke configs.
+    params.setdefault("num_heads", max(1, hidden_size // 64))
+    params.setdefault("head_dim", hidden_size // params["num_heads"])
+    cfg = LinearAttentionConfig(hidden_size=hidden_size, **params)
+    return LinearAttentionBlock(cfg)
+
+
+def _build_kda(hidden_size: int, params: dict) -> nn.Module:
+    """Real Kimi Delta Attention block."""
+    from cppmega_v4.nn.kimi_delta_attention import (
+        KimiDeltaAttentionBlock, KimiDeltaAttentionConfig,
+    )
+    params.setdefault("num_heads", max(1, hidden_size // 64))
+    params.setdefault("head_dim", hidden_size // params["num_heads"])
+    cfg = KimiDeltaAttentionConfig(hidden_size=hidden_size, **params)
+    return KimiDeltaAttentionBlock(cfg)
+
+
+def _build_moe(hidden_size: int, params: dict) -> nn.Module:
+    """Real V4 MoE. V4MoE returns MoEOutput — wrap to expose .output."""
+    from cppmega_v4.nn.moe_v4 import V4MoE, V4MoEConfig
+
+    params.setdefault("expert_hidden_size", hidden_size * 4)
+    cfg = V4MoEConfig(d_model=hidden_size, **params)
+    inner = V4MoE(cfg)
+
+    class _MoEWrap(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.moe = inner
+
+        def __call__(self, x):
+            return self.moe(x).output
+
+    return _MoEWrap()
+
+
+def _build_attention(hidden_size: int, params: dict) -> nn.Module:
+    """Standard multi-head self-attention (causal). Used for `attention` and
+    as the fallback for `mla` / `mla_absorb` until we land a full MLA block.
+    """
+    num_heads = params.get("num_heads", max(1, hidden_size // 64))
+    head_dim = params.get("head_dim", hidden_size // num_heads)
+    norm_eps = params.get("norm_eps", 1e-6)
+
+    class _SelfAttn(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.q_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
+            self.k_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
+            self.v_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
+            self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
+            self.norm = nn.RMSNorm(hidden_size, eps=norm_eps)
+            # Zero-init out so the block is identity at init.
+            self.o_proj.weight = mx.zeros_like(self.o_proj.weight)
+
+        def __call__(self, x):
+            B, S, _ = x.shape
+            q = self.q_proj(x).reshape(B, S, num_heads, head_dim)
+            k = self.k_proj(x).reshape(B, S, num_heads, head_dim)
+            v = self.v_proj(x).reshape(B, S, num_heads, head_dim)
+            q = mx.transpose(q, (0, 2, 1, 3))
+            k = mx.transpose(k, (0, 2, 1, 3))
+            v = mx.transpose(v, (0, 2, 1, 3))
+            scale = head_dim ** -0.5
+            scores = mx.matmul(q, mx.transpose(k, (0, 1, 3, 2))) * scale
+            mask = mx.tril(mx.ones((S, S), dtype=mx.bool_))
+            scores = mx.where(mask, scores, mx.full(scores.shape, -1e9,
+                                                     dtype=scores.dtype))
+            w = mx.softmax(scores.astype(mx.float32), axis=-1).astype(scores.dtype)
+            o = mx.matmul(w, v)
+            o = mx.transpose(o, (0, 2, 1, 3)).reshape(B, S, num_heads * head_dim)
+            return self.norm(self.o_proj(o))
+
+    return _SelfAttn()
+
+
+def _build_lightning_indexer(hidden_size: int, params: dict) -> nn.Module:
+    """LightningIndexer is a top-k helper, not a residual block. Wrap as a
+    residual no-op for stack composition — the real callsite is inside
+    CSA+HCA. RunTemplate users who want the indexer as an inline gate-pre
+    pass get a configurable hidden-pass-through wrapper.
+    """
+    from cppmega_v4.nn.lightning_indexer_fp8 import (
+        LightningIndexerFP8, LightningIndexerFP8Config,
+    )
+    n_heads = params.get("n_heads", max(1, hidden_size // 64))
+    cfg = LightningIndexerFP8Config(
+        hidden_size=hidden_size,
+        n_heads=n_heads,
+        head_dim=params.get("head_dim", 32),
+        rope_head_dim=params.get("rope_head_dim", 16),
+        q_lora_rank=params.get("q_lora_rank", hidden_size),
+        index_topk=params.get("index_topk", 64),
+        fp8_blocks=params.get("fp8_blocks", True),
+    )
+    indexer = LightningIndexerFP8(cfg)
+
+    class _IndexerResidualNoOp(nn.Module):
+        """Residual pass-through wrapper for LightningIndexer.
+
+        Lightning Indexer's natural output is top-k indices (int32), not a
+        residual. In a RunTemplate context, the block contributes zero
+        delta — its presence in the stack signals that downstream
+        CSA/HCA / sparse-MLA layers should consume the indexer's outputs.
+        Real wiring lives in CSA+HCA's select_indices argument.
+        """
+        def __init__(self):
+            super().__init__()
+            self.indexer = indexer
+
+        def __call__(self, x):
+            return mx.zeros_like(x)
+
+    return _IndexerResidualNoOp()
+
+
 BLOCK_BUILDERS: dict[str, Callable[[int, dict], nn.Module]] = {
     "engram": _build_engram,
     "nsa": _build_nsa,
     "csa_hca": _build_csa_hca,
     "mlp": _build_mlp,
-    # Pass-through for blocks that need optional deps not always in scope:
-    "gdn": lambda h, p: _build_pass_through_unsupported("gdn"),
-    "kda": lambda h, p: _build_pass_through_unsupported("kda"),
-    "mla_absorb": lambda h, p: _build_pass_through_unsupported("mla_absorb"),
-    "mla": lambda h, p: _build_pass_through_unsupported("mla"),
-    "attention": lambda h, p: _build_pass_through_unsupported("attention"),
-    "moe": lambda h, p: _build_pass_through_unsupported("moe"),
-    "lightning_indexer": lambda h, p: _build_pass_through_unsupported("lightning_indexer"),
+    "gdn": _build_gdn,
+    "kda": _build_kda,
+    "moe": _build_moe,
+    "attention": _build_attention,
+    # mla / mla_absorb fall back to standard attention until we land a
+    # full MLA block (mla_absorb.py is a pure algebra module, not nn.Module).
+    "mla": _build_attention,
+    "mla_absorb": _build_attention,
+    "lightning_indexer": _build_lightning_indexer,
 }
 
 
@@ -106,15 +229,19 @@ class _BuiltBlock:
     needs_token_ids: bool
 
 
+_DOC_ID_KW_KINDS = {"gdn", "kda"}     # accept doc_ids via *kwargs
+_TOKEN_ID_POS_KINDS = {"engram"}      # take token_ids positional
+
+
 def _build_one(spec: BlockSpec, hidden_size: int) -> _BuiltBlock:
     builder = BLOCK_BUILDERS.get(spec.kind)
     if builder is None:
         raise ValueError(f"no builder registered for block kind {spec.kind!r}")
     mod = builder(hidden_size, dict(spec.params))
     return _BuiltBlock(
         kind=spec.kind, module=mod,
-        needs_doc_ids=(spec.kind == "engram"),
-        needs_token_ids=(spec.kind == "engram"),
+        needs_doc_ids=(spec.kind in _DOC_ID_KW_KINDS or spec.kind == "engram"),
+        needs_token_ids=(spec.kind in _TOKEN_ID_POS_KINDS),
     )
 
 
@@ -168,10 +295,14 @@ def __call__(
             )
         h = hidden_states
         for b in self.blocks:
-            if b.needs_doc_ids and b.needs_token_ids:
+            if b.needs_token_ids and b.needs_doc_ids:
+                # Engram: positional token_ids + keyword document_ids.
                 delta = b.module(h, token_ids, document_ids=document_ids)
             elif b.needs_token_ids:
                 delta = b.module(h, token_ids)
+            elif b.needs_doc_ids:
+                # GDN/KDA: doc_ids keyword (block uses it for doc-reset).
+                delta = b.module(h, doc_ids=document_ids)
             else:
                 delta = b.module(h)
             h = h + delta

scripts/v4_1b_parquet_smoke.py

@@ -42,35 +42,50 @@
 
 
 def _make_template(hidden_size: int, layers: int, vocab_size: int) -> RunTemplate:
-    """Build a V4 template close to ~1B params at hidden=2048, layers=12."""
+    """Build a V4 template close to ~1B params at hidden=2048, layers=12.
+
+    Real-factories version: includes GDN (Path A→C dispatchable), KDA,
+    attention, MoE, NSA, CSA+HCA, Engram, MLP — the full V4 cohort.
+    """
+    n_heads = max(1, hidden_size // 64)
+    h_dim = hidden_size // n_heads
+    per_kind = max(1, layers // 6)  # spread layers across 6 attention-like kinds
+    blocks = [
+        BlockSpec(kind="engram", repeat=1, params={
+            "num_ngram_layers": 2, "max_ngram_size": 4,
+            "num_embed_table_per_ngram": 4, "embed_dim": 64,
+            "embed_table_size": 1024,
+        }),
+        BlockSpec(kind="gdn", repeat=per_kind, params={
+            "num_heads": n_heads, "head_dim": h_dim, "use_short_conv": False,
+        }),
+        BlockSpec(kind="kda", repeat=per_kind, params={
+            "num_heads": n_heads, "head_dim": h_dim, "use_short_conv": False,
+        }),
+        BlockSpec(kind="attention", repeat=per_kind, params={
+            "num_heads": n_heads, "head_dim": h_dim,
+        }),
+        BlockSpec(kind="nsa", repeat=per_kind, params={
+            "num_heads": n_heads, "head_dim": h_dim,
+            "compress_block_size": 64, "select_topk": 16,
+            "sliding_window": 256,
+        }),
+        BlockSpec(kind="csa_hca", repeat=per_kind, params={
+            "num_heads": n_heads, "head_dim": h_dim,
+            "m_csa": 4, "m_hca": 16,
+        }),
+        BlockSpec(kind="moe", repeat=per_kind, params={
+            "num_experts": 8, "top_k": 2,
+            "expert_hidden_size": hidden_size * 2,
+        }),
+        BlockSpec(kind="mlp", repeat=max(1, layers - 6 * per_kind),
+                  params={"intermediate_size": hidden_size * 4}),
+    ]
     return RunTemplate(
         name=f"v4_smoke_h{hidden_size}_L{layers}",
         hidden_size=hidden_size,
         vocab_size=vocab_size,
-        blocks=[
-            # Engram boost first (cheap doc-aware n-gram).
-            BlockSpec(kind="engram", repeat=1, params={
-                "num_ngram_layers": 2, "max_ngram_size": 4,
-                "num_embed_table_per_ngram": 4, "embed_dim": 64,
-                "embed_table_size": 1024,
-            }),
-            # NSA × layers/3.
-            BlockSpec(kind="nsa", repeat=max(1, layers // 3), params={
-                "num_heads": max(1, hidden_size // 64),
-                "head_dim": 64,
-                "compress_block_size": 64, "select_topk": 16,
-                "sliding_window": 256,
-            }),
-            # CSA+HCA × layers/3.
-            BlockSpec(kind="csa_hca", repeat=max(1, layers // 3), params={
-                "num_heads": max(1, hidden_size // 64),
-                "head_dim": 64,
-                "m_csa": 4, "m_hca": 16,
-            }),
-            # MLP × remainder.
-            BlockSpec(kind="mlp", repeat=max(1, layers - 2 * (layers // 3)),
-                      params={"intermediate_size": hidden_size * 4}),
-        ],
+        blocks=blocks,
     )