Upload folder using huggingface_hub

README.md

@@ -0,0 +1,16 @@
+---
+license: apache-2.0
+---
+# Foundry-4
+
+Kernel-specialized hybrid LM: 5.02B total / ~0.58B active params.
+18 GDN (gated DeltaNet, fla) + 6 RoPE attention layers, 384-expert
+sigmoid-routed dropless MoE (top-8, packed Triton segmm), KernelGram
+n-gram tables, 4-iteration shared-weight latent ponder loop, JEV-style
+judge head.
+
+Trained 2000 steps on a curated kernel-optimization corpus
+(CUDA/Triton episodes with measured RESULT/BEST/HISTORY feedback) plus
+GPU-verified synthetic optimization ladders (synth4.py).
+
+Code: model4.py / gdn.py / segmm.py in this repo.

config.json

@@ -0,0 +1,38 @@
+{
+  "architecture": "Foundry4",
+  "config": {
+    "vocab_size": 65536,
+    "n_specials": 16,
+    "d_model": 1280,
+    "n_layers": 24,
+    "n_heads": 10,
+    "head_dim": 128,
+    "seq_len": 4096,
+    "dropout": 0.0,
+    "rope_theta": 10000.0,
+    "qk_norm": true,
+    "attn_every": 4,
+    "gdn_heads": 10,
+    "gdn_head_k": 128,
+    "gdn_head_v": 128,
+    "gdn_chunk": 32,
+    "n_experts": 384,
+    "expert_hidden": 128,
+    "moe_topk": 8,
+    "shared_hidden": 640,
+    "lb_weight": 0.01,
+    "kg_buckets": 262144,
+    "kg_dim": 80,
+    "kg_max_n": 4,
+    "kg_probes": 2,
+    "ponder_iters": 4,
+    "ponder_hidden": 2560,
+    "ponder_iter_weight_decay": 1.0,
+    "judge_hidden": 512,
+    "use_mtp": true,
+    "mtp_weight": 0.05,
+    "moe_ckpt": false
+  },
+  "trained_steps": 2000,
+  "dtype": "bfloat16"
+}

gdn.py

@@ -0,0 +1,273 @@
+"""Gated DeltaNet mixer for Foundry-4.
+
+Recurrence (per head), GDN convention:
+
+    S_t = exp(g_t) * (I - beta_t k_t k_t^T) S_{t-1} + beta_t k_t v_t^T
+    o_t = q_t^T S_t
+
+with S in R^{dk x dv} — keys address, values retrieved (o ~ v when
+q ~ k).  g <= 0 is log-decay per head per step, beta in (0,1) the write
+strength.
+
+Implementations:
+
+- ``gdn_sequential`` — per-token recurrence.  Obviously correct oracle
+  for the unit test; too slow for training.
+- ``gdn_chunk`` — chunkwise WY form (Yang et al. 2024, "Gated Delta
+  Networks").  All within-chunk work is batched bmms + one triangular
+  solve; the only sequential loop is the inter-chunk state scan over
+  T/C chunks of cheap [dk x dv] ops.
+
+Derivation (verified against the N=2 expansion by hand, and against
+gdn_sequential in test_gdn.py):
+
+    within chunk, G_i = sum_{j<=i} g_j,  Ghat_i = G_i - G_N  (<= 0)
+    W = beta * K,  U = beta * V
+    M[i,j]   = exp(G_i - G_j)          for i >= j, else 0
+    A_k[i,l] = beta_l <k_i, k_l>       for i > l   (strict lower)
+    A_v[i,j] = exp(G_i - G_j) beta_i <k_i, k_j>   for i > j (strict lower)
+    Kt  = (I + A_k)^{-1} K             transformed keys (forward solve —
+                                     prefix property: leading submatrix
+                                     is the restricted solve)
+    Ut  = (I + A_v)^{-1} U             transformed values (decayed)
+    Kd  = K * exp(G_N - G_j)           decay-safe anti-decayed keys
+
+    removal product:  prod_{j<=i}(I - w_j k_j^T) = I - W^T Kt  (prefix
+    property lets the same Kt serve every intra-chunk position)
+
+    o_i = e^{G_i} q_i S_in
+          - e^{G_i} sum_{l<=i} <q_i, W_l> (Kt_l S_in)
+          + sum_{j<=i} M[i,j] <q_i, k_j> Ut_j
+
+    S_out = e^{G_N} ( S_in - W^T (Kt S_in) )  +  Kd^T Ut
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# sequential oracle
+# ---------------------------------------------------------------------------
+
+def gdn_sequential(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                   g: torch.Tensor, beta: torch.Tensor,
+                   S0: torch.Tensor | None = None):
+    """q,k [B,T,H,dk]; v [B,T,H,dv]; g,beta [B,T,H] -> (o [B,T,H,dv], S).
+
+    fp32 oracle — the unit test compares gdn_chunk against this.
+    """
+    B, T, H, dk = q.shape
+    dv = v.shape[-1]
+    S = (S0 if S0 is not None else
+         q.new_zeros(B, H, dk, dv)).float()
+    o = q.new_zeros(B, T, H, dv, dtype=torch.float32)
+    qf, kf, vf = q.float(), k.float(), v.float()
+    gf, bf = g.float(), beta.float()
+    for t in range(T):
+        gt = gf[:, t].view(B, H, 1, 1).exp()
+        bt = bf[:, t].view(B, H, 1, 1)
+        kt, vt, qt = kf[:, t], vf[:, t], qf[:, t]
+        kTS = torch.einsum("bhd,bhdv->bhv", kt, S)          # k^T S
+        S = gt * (S - bt * torch.einsum("bhd,bhv->bhdv", kt, kTS)) \
+            + bt * torch.einsum("bhd,bhv->bhdv", kt, vt)
+        o[:, t] = torch.einsum("bhd,bhdv->bhv", qt, S)
+    return o, S
+
+
+# ---------------------------------------------------------------------------
+# chunked WY implementation
+# ---------------------------------------------------------------------------
+
+def gdn_chunk(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+              g: torch.Tensor, beta: torch.Tensor,
+              S0: torch.Tensor | None = None, chunk_size: int = 32,
+              matmul_dtype: torch.dtype | None = None):
+    """Chunkwise gated delta rule, vectorized across chunks.
+
+    Only the [dk x dv] state scan is sequential (T/C cheap iterations).
+    Gram/output matmuls run in ``matmul_dtype`` (bf16 on GPU = tensor
+    cores); the triangular solves and the state recurrence stay fp32.
+    """
+    B, T, H, dk = q.shape
+    dv = v.shape[-1]
+    C = chunk_size
+    NC = (T + C - 1) // C
+    pad = NC * C - T
+    if pad:                                        # pad tail chunk
+        zp = lambda t: F.pad(t, (0, 0, 0, 0, 0, pad))      # [B,T,H,*]
+        q, k, v = zp(q), zp(k), zp(v)
+        g = F.pad(g, (0, 0, 0, pad))                        # g=0 no decay
+        beta = F.pad(beta, (0, 0, 0, pad))
+
+    mdt = matmul_dtype or torch.float32            # big matmuls
+    dt = torch.float32                             # solves + state
+    dev = q.device
+    # -> [B,H,NC,C,*]
+    qc = q.float().view(B, NC, C, H, dk).permute(0, 3, 1, 2, 4)
+    kc = k.float().view(B, NC, C, H, dk).permute(0, 3, 1, 2, 4)
+    vc = v.float().view(B, NC, C, H, dv).permute(0, 3, 1, 2, 4)
+    gc = g.float().view(B, NC, C, H).permute(0, 3, 1, 2)
+    bc = beta.float().view(B, NC, C, H).permute(0, 3, 1, 2)
+
+    G = gc.cumsum(-1)                                     # [B,H,NC,C]
+    Gh = G - G[..., -1:]                                  # <=0
+    EG = Gh.exp()                                         # e^{Ghat_i}
+    EGN = G[..., -1].exp()                                # e^{G_N}  [B,H,NC]
+
+    tril = torch.tril(torch.ones(C, C, device=dev, dtype=torch.bool))
+    tril_m1 = torch.tril(torch.ones(C, C, device=dev, dtype=torch.bool),
+                         diagonal=-1)
+
+    # M[i,j] = e^{G_i - G_j} on i >= j  (<=1, stable)
+    M = torch.where(tril, (G.unsqueeze(-1) - G.unsqueeze(-2)).exp(),
+                    G.new_zeros(()))
+    # ---- fp32 triangular solves (bounded by chunk size C) --------------
+    kbc32 = kc * bc.unsqueeze(-1)                         # W = beta k
+    vbc32 = vc * bc.unsqueeze(-1)                         # U = beta v
+    A_k = torch.where(tril_m1, kc @ kbc32.transpose(-1, -2),
+                      kc.new_zeros(()))                   # beta_l<k_i,k_l>
+    A_v = torch.where(tril_m1, (kbc32 @ kc.transpose(-1, -2)) * M,
+                      kc.new_zeros(()))                   # e^{Gi-Gj} b_i<k_i,k_j>
+    eye = torch.eye(C, device=dev, dtype=dt).view(
+        1, 1, 1, C, C).expand(B, H, NC, C, C)
+    Kt = torch.linalg.solve_triangular(eye + A_k, kc, upper=False)
+    Ut = torch.linalg.solve_triangular(eye + A_v, vbc32, upper=False)
+
+    # ---- batch the matmul dtype: Gram rows, outputs, Kd----------------
+    kbc = kbc32.to(mdt)
+    Ut_m = Ut.to(mdt)
+    qc_m, kc_m = qc.to(mdt), kc.to(mdt)
+    M_m = M.to(mdt)
+    Kd = kc * (G[..., -1:] - G).exp().unsqueeze(-1)       # e^{G_N-G_j} k
+    Kd_m = Kd.to(mdt)
+    EG_m, EGN_m = EG.to(mdt), EGN.to(mdt)
+    QW = torch.where(tril, qc_m @ kbc.transpose(-1, -2),
+                     qc_m.new_zeros(()))                  # <q_i,W_l> l<=i
+    QWd = QW * EG_m.unsqueeze(-1)                         # e^{Ghat_i} rows
+    QK_m = (qc_m @ kc_m.transpose(-1, -2)) * M_m          # tril by M
+    o_loc_all = QK_m @ Ut_m                               # batchable term
+
+    # scan states over chunks:  o_inter needs S_in, so sequential here.
+    S = (S0.float() if S0 is not None else
+         q.new_zeros(B, H, dk, dv, dtype=dt))
+    o = q.new_zeros(B, H, NC, C, dv, dtype=dt)
+    for c in range(NC):
+        KS_c = Kt[:, :, c].to(mdt) @ S.to(mdt)            # Kt_l S_in [C,dv]
+        o_in = (qc_m[:, :, c] * EG_m[:, :, c].unsqueeze(-1)) @ S.to(mdt) \
+            - QWd[:, :, c] @ KS_c
+        o_in = o_in * EGN_m[:, :, c].view(B, H, 1, 1)
+        o[:, :, c] = (o_in + o_loc_all[:, :, c]).float()
+        S = EGN[:, :, c].view(B, H, 1, 1) * (
+            S - kbc[:, :, c].transpose(-1, -2).float() @ KS_c.float()) \
+            + (Kd_m[:, :, c].transpose(-1, -2) @ Ut_m[:, :, c]).float()
+
+    o = o.permute(0, 2, 3, 1, 4).reshape(B, NC * C, H, dv)[:, :T]
+    return o.to(q.dtype), S
+
+
+# ---------------------------------------------------------------------------
+# nn.Module: causal conv1d + projections + gated delta + output gate
+# (the FLA/Qwen3-Next layer shape: short conv on qkv, per-head decay/beta,
+# normalized output gate)
+# ---------------------------------------------------------------------------
+
+# fla provides the production Triton kernel (verified vs gdn_sequential
+# in test_gdn: rel ~1e-6 on ROCm).  Fall back to gdn_chunk elsewhere.
+try:
+    from fla.ops.gated_delta_rule import chunk_gated_delta_rule as _fla_gdr
+except Exception:
+    _fla_gdr = None
+
+
+class GatedDeltaNet(nn.Module):
+    """Linear-attention mixer.  h: [B,T,d_model] -> [B,T,d_model]."""
+
+    def __init__(self, d_model: int, n_heads: int, head_k: int = 128,
+                 head_v: int = 128, conv_k: int = 4,
+                 chunk_size: int = 32):
+        super().__init__()
+        self.h = n_heads
+        self.dk, self.dv = head_k, head_v
+        self.chunk_size = chunk_size
+        dk, dv = n_heads * head_k, n_heads * head_v
+        self.conv_k = conv_k
+        # q,k share dk side; v on dv side; conv covers all three
+        self.qkv = nn.Linear(d_model, 2 * dk + dv, bias=False)
+        self.conv = nn.Conv1d(2 * dk + dv, 2 * dk + dv, conv_k,
+                              groups=2 * dk + dv, padding=conv_k - 1,
+                              bias=True)
+        self.a_proj = nn.Linear(d_model, n_heads, bias=True)   # decay pre-activ
+        self.b_proj = nn.Linear(d_model, n_heads, bias=True)   # beta
+        self.g_out = nn.Linear(d_model, dv, bias=False)        # output gate
+        self.out_norm = nn.GroupNorm(n_heads, dv)              # per-head rms-ish
+        self.o_proj = nn.Linear(dv, d_model, bias=False)
+        nn.init.normal_(self.o_proj.weight, std=0.02)
+
+    def forward(self, x: torch.Tensor, S=None, use_cache=False,
+                conv_prev: torch.Tensor | None = None):
+        B, T, D = x.shape
+        qkv = self.qkv(x).transpose(1, 2)                     # [B,C,T]
+        if use_cache:
+            # decode path: the causal conv needs the previous k-1
+            # columns of qkv-projected input as its left context.
+            # conv_tail must be taken from the CONCATENATED context —
+            # a short chunk can carry fewer than k-1 new columns.
+            k1 = self.conv_k - 1
+            p_len = 0 if conv_prev is None else conv_prev.size(-1)
+            if conv_prev is not None:
+                qkv = torch.cat([conv_prev, qkv], -1)
+            conv_tail = qkv[..., -k1:]
+            # out[i] covers cat_in[i-k1 .. i]; new token t sits at cat
+            # index p_len+t -> want out[p_len : p_len+T]
+            qkv = self.conv(qkv)[..., p_len:p_len + T]
+            qkv = F.silu(qkv).transpose(1, 2)                 # [B,T,C]
+        else:
+            conv_tail = None
+            qkv = self.conv(qkv)[..., :T]                     # causal conv
+            qkv = F.silu(qkv).transpose(1, 2)                 # [B,T,C]
+        dk_t = self.h * self.dk
+        q = qkv[..., :dk_t].view(B, T, self.h, self.dk)
+        kk = qkv[..., dk_t:2 * dk_t].view(B, T, self.h, self.dk)
+        vv = qkv[..., 2 * dk_t:].view(B, T, self.h, self.dv)
+        # F.normalize upcasts to fp32 under autocast — fla's triton
+        # kernels need all main operands in ONE dtype, so normalize in
+        # fp32 for accuracy then cast back to the input dtype.
+        kk = F.normalize(kk.float(), p=2, dim=-1).to(vv.dtype)
+        a = self.a_proj(x)                                    # [B,T,H]
+        b = self.b_proj(x)
+        g_log = -F.softplus(a).clamp(1e-4, 8.0)               # log decay <=0
+        beta = torch.sigmoid(b).clamp(1e-4, 1 - 1e-4)
+        if _fla_gdr is not None and x.is_cuda:
+            # fla accepts bf16 operands + fp32 gates on ROCm (verified —
+            # the earlier fp32-only workaround was stale).  GDN_FP32=1
+            # forces the whole call fp32 — fallback if bf16 fla
+            # instability is ever suspected in a training divergence.
+            if os.environ.get("GDN_FP32") == "1":
+                q, kk, vv, beta = (q.float(), kk.float(), vv.float(),
+                                   beta.float())
+            res = _fla_gdr(q, kk, vv,
+                           g=g_log.float(), beta=beta, scale=1.0,
+                           initial_state=(S.to(q.dtype) if S is not None
+                                          else None),
+                           output_final_state=use_cache)
+            o = res[0].to(x.dtype)
+            S = res[1] if (use_cache and len(res) > 1) else None
+        else:
+            mdt = x.dtype if x.dtype in (torch.bfloat16,
+                                         torch.float16) else None
+            o, S = gdn_chunk(q, kk, vv, g_log, beta, S0=S,
+                             chunk_size=self.chunk_size,
+                             matmul_dtype=mdt)
+        # per-head GroupNorm over dv: reshape so channels = H*dv
+        o = self.out_norm(o.reshape(B * T, self.h * self.dv, 1))
+        o = o.reshape(B, T, -1)
+        o = o * torch.sigmoid(self.g_out(x))
+        out = self.o_proj(o)
+        if use_cache:
+            return out, S, conv_tail
+        return out

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e9798dedae530810fc606a77ad8cbb31763700e3927d20b660d69398af55049
+size 10202500036

model4.py

@@ -0,0 +1,597 @@
+"""Foundry-4: sparse hybrid kernel-optimization model.
+
+Architecture (see DESIGN notes at bottom):
+
+- Embedding + KernelGram n-gram memory (hashed bigram..5-gram tables,
+  zero-init gates — proven on Foundry-2/3 to carry ~half the val loss).
+- Hybrid trunk: GatedDeltaNet linear-attention blocks everywhere except
+  every ``attn_every``-th layer which is full causal attention (RoPE +
+  QK-norm + SDPA).  Linear layers handle the stream; attention layers
+  give exact long-range recall for pointer names / block sizes.
+- Sparse MoE FFN on every block: 1 shared expert + N fine-grained
+  routed experts, top-k sigmoid router (DeepSeekMoE-style), Switch-style
+  load-balance aux loss.
+- PonderLoop: a recurrent block (shared weights, iteration embedding)
+  applied K times to the trunk output.  Per-iteration deep supervision
+  (weighted CE on each iteration's logits) so every loop step is a
+  usable state — the substrate for latent "try many paths" reasoning
+  and later RL-driven non-lazy search.
+- JudgmentHead (Jev-like): small MLP over the pooled latent state that
+  emits typed verdicts — p(compiles), p(correct), expected speedup,
+  halt.  In pretrain it is self-supervised on a signal we already have:
+  per-iteration LM improvement.  In RL it is re-pointed at real harness
+  outcomes and becomes the best-of-loop selector / idea filter.
+- Tied 65k-token embedding + chunked CE scoring (never materialise
+  [B,T,V] logits at once), optional MTP head.
+
+Invariants (same discipline as Foundry-1/2 — do not break):
+  1. Trunk is causal.  Attention uses the explicit -inf mask; GDN is
+     causal by construction; the ponder block uses the same mask.
+  2. Rewards/judgments are measured or self-supervised, never faked.
+  3. save_checkpoint stage->verify->copy->verify on every write.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+from gdn import GatedDeltaNet
+from segmm import seg_mm, wscatter
+
+
+# ---------------------------------------------------------------------------
+# config
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Foundry4Config:
+    vocab_size: int = 65536          # bpe2_65k_v2 + specials 0..15
+    n_specials: int = 16
+    d_model: int = 1280
+    n_layers: int = 24
+    n_heads: int = 10                # attn heads, head_dim 128
+    head_dim: int = 128
+    seq_len: int = 4096
+    dropout: float = 0.0
+    rope_theta: float = 10000.0
+    qk_norm: bool = True
+    attn_every: int = 4              # layer i%4==3 -> full attention
+    # gated deltanet
+    gdn_heads: int = 10
+    gdn_head_k: int = 128
+    gdn_head_v: int = 128
+    gdn_chunk: int = 32          # C=64 can overflow the WY solve (2^C growth)
+    # sparse MoE
+    n_experts: int = 384
+    expert_hidden: int = 128
+    moe_topk: int = 8
+    shared_hidden: int = 640         # always-on shared expert
+    lb_weight: float = 0.01
+    # kernelgram
+    kg_buckets: int = 262144
+    kg_dim: int = 80
+    kg_max_n: int = 4                # n-gram orders 2..4
+    kg_probes: int = 2               # salted probes per order
+    # ponder loop (latent reasoning)
+    ponder_iters: int = 4
+    ponder_hidden: int = 2560        # dense SwiGLU inside the loop block
+    ponder_iter_weight_decay: float = 1.0   # loss weight ramp exponent
+    # judgment head
+    judge_hidden: int = 512
+    # mtp
+    use_mtp: bool = True
+    mtp_weight: float = 0.05
+    # memory
+    moe_ckpt: bool = False           # recompute MoE dispatch in backward
+
+
+def beast_config() -> Foundry4Config:
+    """~5.1B total / ~570M active.  The target."""
+    return Foundry4Config()
+
+
+def dev_config() -> Foundry4Config:
+    """~700M total / ~120M active — smoke + kernel-tuning config."""
+    return Foundry4Config(
+        d_model=768, n_layers=12, n_heads=6, head_dim=128,
+        gdn_heads=6, gdn_head_k=128, gdn_head_v=128,
+        n_experts=32, expert_hidden=192, moe_topk=4,
+        shared_hidden=512, seq_len=2048,
+        kg_buckets=65536, kg_dim=64, kg_max_n=4, kg_probes=2,
+        ponder_iters=3, ponder_hidden=1536,
+        judge_hidden=256)
+
+
+# ---------------------------------------------------------------------------
+# KernelGram — identical semantics to model2 (proven); kept self-contained
+# ---------------------------------------------------------------------------
+
+class KernelGram(nn.Module):
+    _C1 = 0x9E3779B1
+    _C2 = 0x85EBCA77
+    _SEEDS = (0x0, 0x9E3779B1, 0x85EBCA77, 0xC2B2AE3D)
+
+    def __init__(self, cfg: Foundry4Config):
+        super().__init__()
+        self.cfg = cfg
+        self.n_grams = cfg.kg_max_n - 1
+        self.n_probes = cfg.kg_probes
+        self.n_slots = self.n_grams * self.n_probes
+        self.tables = nn.ParameterList([
+            nn.Parameter(torch.randn(cfg.kg_buckets, cfg.kg_dim)
+                         / math.sqrt(cfg.kg_dim))
+            for _ in range(self.n_slots)])
+        self.proj = nn.ParameterList([
+            nn.Parameter(torch.empty(cfg.kg_dim, cfg.d_model))
+            for _ in range(self.n_slots)])
+        for p in self.proj:
+            nn.init.normal_(p, std=0.02)
+        self.gates = nn.Parameter(torch.zeros(self.n_slots))
+
+    def _hashes(self, ids: torch.Tensor) -> list[torch.Tensor]:
+        B, T = ids.shape
+        out = []
+        for k_ in range(self.n_grams):
+            n = k_ + 2
+            for p in range(self.n_probes):
+                h = torch.zeros(B, T, dtype=torch.long, device=ids.device)
+                if p:
+                    h[:, n - 1:] = self._SEEDS[p % len(self._SEEDS)]
+                for j in range(n):
+                    tok = ids[:, j:T - n + 1 + j]
+                    h[:, n - 1:] = (h[:, n - 1:] * self._C1
+                                    + tok * self._C2)
+                out.append(h % self.cfg.kg_buckets)
+        return out
+
+    def forward(self, ids, x_emb):
+        for s, (h, tab, pr) in enumerate(
+                zip(self._hashes(ids), self.tables, self.proj)):
+            g = torch.tanh(self.gates[s])
+            x_emb = x_emb + g * (F.embedding(h, tab) @ pr)
+        return x_emb
+
+
+# ---------------------------------------------------------------------------
+# Sparse MoE — shared expert + fine-grained routed experts (sigmoid router,
+# DeepSeek-style top-k with weight renormalisation; bucketed bmm dispatch)
+# ---------------------------------------------------------------------------
+
+class SparseMoE(nn.Module):
+    def __init__(self, cfg: Foundry4Config):
+        super().__init__()
+        d, h = cfg.d_model, cfg.expert_hidden
+        E = cfg.n_experts
+        self.cfg = cfg
+        self.E, self.k = E, cfg.moe_topk
+        self.router = nn.Linear(d, E, bias=False)
+        self.w_g = nn.Parameter(torch.empty(E, d, h))
+        self.w_w = nn.Parameter(torch.empty(E, d, h))
+        self.w_o = nn.Parameter(torch.empty(E, h, d))
+        for w in (self.w_g, self.w_w, self.w_o):
+            nn.init.normal_(w, std=0.02)
+        # shared expert (always on)
+        self.sh_g = nn.Linear(d, cfg.shared_hidden, bias=False)
+        self.sh_w = nn.Linear(d, cfg.shared_hidden, bias=False)
+        self.sh_o = nn.Linear(cfg.shared_hidden, d, bias=False)
+
+    def _swiglu(self, x, wg, ww, wo):
+        return (F.silu(x @ wg) * (x @ ww)) @ wo
+
+    def _dispatch(self, xf, top_i, w):
+        """Packed-segment expert compute (segmm.py Triton grouped GEMM).
+        DROPLESS (causality invariant): every routed token keeps all k
+        experts — a capacity mask would couple position i's output to
+        future tokens' routing and break KV-cache semantics.  Packed
+        layout: memory stays at Nk rows regardless of routing skew, and
+        there is no host<->device sync at all."""
+        N, d = xf.shape
+        counts = torch.bincount(
+            top_i.reshape(-1), minlength=self.E)
+        order = torch.argsort(top_i.reshape(-1))          # stable sort
+        flat_t = (order // self.k)                        # token ids
+        flat_w = w.reshape(-1)[order]
+        offs = torch.zeros(self.E + 1, dtype=torch.int64,
+                           device=xf.device)
+        offs[1:] = torch.cumsum(counts, 0)
+        xs = xf[flat_t]                                   # [Nk,d] packed
+        dt = xs.dtype
+        wg = self.w_g.to(dt)                              # cast: same-dtype
+        ww = self.w_w.to(dt)                              # tl.dot operands
+        wo = self.w_o.to(dt)
+        h1 = seg_mm(xs, wg, offs)
+        h2 = seg_mm(xs, ww, offs)
+        yp = seg_mm(F.silu(h1) * h2, wo, offs)            # [Nk,d]
+        out = xf.new_zeros(N, d)
+        out.index_add_(0, flat_t,
+                       (yp * flat_w.unsqueeze(-1)).to(out.dtype))
+        return out
+
+    def forward(self, x: torch.Tensor):
+        """x [B,T,d] -> out [B,T,d], lb_loss scalar."""
+        B, T, d = x.shape
+        xf = x.reshape(-1, d)
+        N = xf.size(0)
+
+        logits = self.router(xf)                       # [N,E]
+        scores = torch.sigmoid(logits)
+        top_s, top_i = scores.topk(self.k, dim=-1)     # [N,k]
+        w = top_s / top_s.sum(-1, keepdim=True).clamp_min(1e-6)
+
+        if self.cfg.moe_ckpt and self.training and xf.requires_grad:
+            out = torch.utils.checkpoint.checkpoint(
+                self._dispatch, xf, top_i, w,
+                use_reentrant=False)
+        else:
+            out = self._dispatch(xf, top_i, w)
+
+        out = out + self.sh_o(F.silu(self.sh_g(xf)) * self.sh_w(xf))
+        out = out.view(B, T, d)
+
+        # Switch-style load balance on the full softmax
+        p = F.softmax(logits, -1)
+        f = torch.zeros_like(p).scatter_(
+            1, top_i, 1.0 / self.k).mean(0)
+        lb = self.E * (f * p.mean(0)).sum()
+        return out, lb
+
+
+# ---------------------------------------------------------------------------
+# attention mixer + block
+# ---------------------------------------------------------------------------
+
+class RMSNorm(nn.Module):
+    def __init__(self, d: int, eps: float = 1e-6):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True)
+                               + self.eps) * self.w
+
+
+def _rope_cos_sin(T, hd, theta, device):
+    half = hd // 2
+    fr = 1.0 / (theta ** (torch.arange(0, half, device=device,
+                                       dtype=torch.float32) / half))
+    t = torch.arange(T, device=device, dtype=torch.float32)
+    ang = torch.outer(t, fr)
+    return ang.cos(), ang.sin()
+
+
+def _rope(x, cos, sin):
+    half = x.size(-1) // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    c = cos[None, None]
+    s = sin[None, None]
+    return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], -1)
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg: Foundry4Config):
+        super().__init__()
+        d = cfg.d_model
+        self.h, self.hd = cfg.n_heads, cfg.head_dim
+        self.wq = nn.Linear(d, self.h * self.hd, bias=False)
+        self.wk = nn.Linear(d, self.h * self.hd, bias=False)
+        self.wv = nn.Linear(d, self.h * self.hd, bias=False)
+        self.wo = nn.Linear(self.h * self.hd, d, bias=False)
+        self.qn = RMSNorm(self.hd) if cfg.qk_norm else None
+        self.kn = RMSNorm(self.hd) if cfg.qk_norm else None
+
+    def forward(self, x, cos, sin, causal):
+        B, T, _ = x.shape
+        q = self.wq(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        k = self.wk(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        v = self.wv(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        if self.qn is not None:
+            q, k = self.qn(q), self.kn(k)
+        q, k = _rope(q, cos, sin), _rope(k, cos, sin)
+        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.wo(a.transpose(1, 2).reshape(B, T, -1))
+
+    def forward_cached(self, x, cos, sin, past=None):
+        """Decode-time path.  past = (k, v) post-rope from earlier calls.
+        Returns (out, (k_all, v_all)) — the returned tensors become the
+        next call's `past`.  cos/sin must cover only the NEW positions."""
+        B, T, _ = x.shape
+        q = self.wq(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        k = self.wk(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        v = self.wv(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        if self.qn is not None:
+            q, k = self.qn(q), self.kn(k)
+        q, k = _rope(q, cos, sin), _rope(k, cos, sin)
+        if past is not None:
+            k = torch.cat([past[0], k], 2)
+            v = torch.cat([past[1], v], 2)
+        if past is None or T == k.size(2):
+            a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            P = k.size(2) - T
+            m = torch.ones(T, k.size(2), dtype=torch.bool,
+                           device=x.device)
+            m[:, P:] = torch.tril(torch.ones(T, T, dtype=torch.bool,
+                                             device=x.device))
+            a = F.scaled_dot_product_attention(q, k, v, attn_mask=m)
+        return self.wo(a.transpose(1, 2).reshape(B, T, -1)), (k, v)
+
+
+class Block(nn.Module):
+    """mixer (GDN or attention) + sparse MoE, pre-norm."""
+
+    def __init__(self, cfg: Foundry4Config, use_attn: bool):
+        super().__init__()
+        self.use_attn = use_attn
+        d = cfg.d_model
+        self.norm1 = RMSNorm(d)
+        self.norm2 = RMSNorm(d)
+        if use_attn:
+            self.mixer = Attention(cfg)
+        else:
+            self.mixer = GatedDeltaNet(
+                d, cfg.gdn_heads, cfg.gdn_head_k, cfg.gdn_head_v,
+                chunk_size=cfg.gdn_chunk)
+        self.ffn = SparseMoE(cfg)
+
+    def forward(self, h, cos, sin, causal):
+        if self.use_attn:
+            h = h + self.mixer(self.norm1(h), cos, sin, causal)
+        else:
+            h = h + self.mixer(self.norm1(h))
+        ff, lb = self.ffn(self.norm2(h))
+        h = h + ff
+        return h, lb
+
+    def forward_cached(self, h, cos, sin, cache):
+        """cache: dict — {'kv': (k,v)} for attn, {'S': state} for GDN."""
+        if self.use_attn:
+            o, kv = self.mixer.forward_cached(self.norm1(h), cos, sin,
+                                              cache.get("kv"))
+            cache["kv"] = kv
+            h = h + o
+        else:
+            o, S, ct = self.mixer(self.norm1(h), S=cache.get("S"),
+                                  use_cache=True,
+                                  conv_prev=cache.get("conv"))
+            cache["S"] = S
+            cache["conv"] = ct
+            h = h + o
+        ff, _ = self.ffn(self.norm2(h))
+        return h + ff
+
+
+# ---------------------------------------------------------------------------
+# PonderLoop — the latent reasoning core
+# ---------------------------------------------------------------------------
+
+class PonderBlock(nn.Module):
+    """One shared-weight recurrent step: causal attn + dense SwiGLU."""
+
+    def __init__(self, cfg: Foundry4Config):
+        super().__init__()
+        d = cfg.d_model
+        self.norm1 = RMSNorm(d)
+        self.norm2 = RMSNorm(d)
+        self.attn = Attention(cfg)
+        self.ff_g = nn.Linear(d, cfg.ponder_hidden, bias=False)
+        self.ff_w = nn.Linear(d, cfg.ponder_hidden, bias=False)
+        self.ff_o = nn.Linear(cfg.ponder_hidden, d, bias=False)
+        self.iter_emb = nn.Parameter(torch.zeros(cfg.ponder_iters, d))
+
+    def forward(self, h, it, cos, sin, causal):
+        h = h + self.iter_emb[it]
+        h = h + self.attn(self.norm1(h), cos, sin, causal)
+        x = self.norm2(h)
+        h = h + self.ff_o(F.silu(self.ff_g(x)) * self.ff_w(x))
+        return h
+
+    def forward_cached(self, h, it, cos, sin, past):
+        o, kv = self.attn.forward_cached(
+            self.norm1(h + self.iter_emb[it]), cos, sin, past)
+        h = h + self.iter_emb[it] + o
+        x = self.norm2(h)
+        return h + self.ff_o(F.silu(self.ff_g(x)) * self.ff_w(x)), kv
+
+
+class JudgmentHead(nn.Module):
+    """Typed verdicts over the pooled latent state.
+
+    Outputs: p_compile, p_correct logits, log-speedup regression,
+    halt logit.  Trained self-supervised in pretrain (per-iter LM
+    improvement); re-pointed at harness rewards in SFT/RL.
+    """
+
+    def __init__(self, cfg: Foundry4Config):
+        super().__init__()
+        d = cfg.d_model
+        self.net = nn.Sequential(
+            nn.Linear(d, cfg.judge_hidden), nn.SiLU(),
+            nn.Linear(cfg.judge_hidden, 4))
+
+    def forward(self, h):                      # h [B,T,d] -> [B,4]
+        pooled = h.mean(1)
+        return self.net(pooled)
+
+
+# ---------------------------------------------------------------------------
+# model
+# ---------------------------------------------------------------------------
+
+class Foundry4Model(nn.Module):
+    def __init__(self, cfg: Foundry4Config | None = None):
+        super().__init__()
+        self.cfg = cfg = cfg or Foundry4Config()
+        d = cfg.d_model
+        self.tok_emb = nn.Embedding(cfg.vocab_size, d)
+        self.kg = KernelGram(cfg)
+        self.blocks = nn.ModuleList([
+            Block(cfg, use_attn=((i + 1) % cfg.attn_every == 0))
+            for i in range(cfg.n_layers)])
+        self.norm = RMSNorm(d)
+        self.ponder = PonderBlock(cfg)
+        self.judge = JudgmentHead(cfg)
+        self.head = nn.Linear(d, cfg.vocab_size, bias=False)
+        self.head.weight = self.tok_emb.weight          # tied
+        self.mtp = (nn.Linear(d, d, bias=False)
+                    if cfg.use_mtp else None)
+        self.apply(self._init)
+
+    @staticmethod
+    def _init(m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.02)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, std=0.02)
+
+    # -- trunk -------------------------------------------------------------
+
+    def trunk(self, ids):
+        x = self.kg(ids, self.tok_emb(ids))
+        dev = ids.device
+        T = ids.size(1)
+        cos, sin = _rope_cos_sin(T, self.cfg.head_dim,
+                                 self.cfg.rope_theta, dev)
+        h, lb = x, x.new_zeros(())
+        for blk in self.blocks:
+            h, l = blk(h, cos, sin, None)
+            lb = lb + l
+        return self.norm(h), lb
+
+    # -- forward: trunk -> ponder loop -> per-iter hidden states -----------
+
+    def forward(self, ids, return_aux=False):
+        cfg = self.cfg
+        h, lb = self.trunk(ids)
+        T = ids.size(1)
+        cos, sin = _rope_cos_sin(T, cfg.head_dim, cfg.rope_theta,
+                                 ids.device)
+        h_iters, judge_per_iter = [h], []
+        for it in range(cfg.ponder_iters):
+            h = self.ponder(h, it, cos, sin, None)
+            h_iters.append(h)
+            if return_aux:
+                judge_per_iter.append(self.judge(h))
+        if return_aux:
+            # training path: h_iters drive chunked/fused CE — never
+            # materialise the [B,T,V] logits here (saves ~4GB + backward)
+            aux = {"lb_loss": lb * cfg.lb_weight,
+                   "judge": torch.stack(judge_per_iter, 1),
+                   "h_iters": h_iters}
+            if self.mtp is not None:
+                aux["mtp_h"] = self.mtp(h)
+            return None, aux
+        return self.head(h)                              # final iter
+
+    # -- cached decode ------------------------------------------------------
+
+    def new_cache(self):
+        return {"blocks": [{} for _ in self.blocks],
+                "ponder": [None] * self.cfg.ponder_iters,
+                "tail": None}                     # last kg_max_n-1 ids
+
+    def forward_cached(self, ids, pos0, cache):
+        """Decode/prefill for generation.  ids [B,Tn] at absolute
+        positions pos0..pos0+Tn-1.  Returns logits [B,Tn,V]; cache is
+        updated in place.  Matches forward() token-for-token (causality
+        invariant), just incremental."""
+        cfg = self.cfg
+        B, Tn = ids.shape
+        dev = ids.device
+        # KernelGram needs the trailing n-1 ids for the n-gram window
+        tail = cache["tail"]
+        kg_in = ids if tail is None else torch.cat([tail, ids], 1)
+        x = self.kg(kg_in, self.tok_emb(kg_in))[:, -Tn:]
+        cache["tail"] = kg_in[:, -(cfg.kg_max_n - 1):]
+        cos, sin = _rope_cos_sin(pos0 + Tn, cfg.head_dim,
+                                 cfg.rope_theta, dev)
+        cos, sin = cos[pos0:], sin[pos0:]
+        h = x
+        for blk, bc in zip(self.blocks, cache["blocks"]):
+            h = blk.forward_cached(h, cos, sin, bc)
+        h = self.norm(h)
+        for it in range(cfg.ponder_iters):
+            h, cache["ponder"][it] = self.ponder.forward_cached(
+                h, it, cos, sin, cache["ponder"][it])
+        return self.head(h)
+
+
+# ---------------------------------------------------------------------------
+# parameter accounting + checkpointing (same discipline as model2)
+# ---------------------------------------------------------------------------
+
+def param_report(model: Foundry4Model) -> dict:
+    cfg = model.cfg
+    total = sum(p.numel() for p in model.parameters())
+
+    emb = cfg.vocab_size * cfg.d_model                       # tied
+    kg = sum(p.numel() for p in model.kg.parameters())
+    experts = cfg.n_layers * cfg.n_experts * 3 * cfg.d_model \
+        * cfg.expert_hidden
+
+    # active = everything except the cold (non-selected) routed-expert
+    # capacity.  Embedding lookup counts once (tied head reuses it for
+    # logits, so it's active anyway).
+    routed = cfg.n_layers * cfg.moe_topk * 3 * cfg.d_model \
+        * cfg.expert_hidden
+    active = total - experts + routed
+    n_attn = cfg.n_layers // cfg.attn_every
+    n_gdn = cfg.n_layers - n_attn
+    return {"total": total, "active": active,
+            "experts": experts, "kg": kg, "emb": emb,
+            "n_attn": n_attn, "n_gdn": n_gdn}
+
+
+def save_checkpoint(obj: dict, out_path) -> None:
+    """Stage-write-verify-copy-verify; survives flaky filesystems."""
+    import io
+    import shutil
+    import tempfile
+    from pathlib import Path
+
+    out_path = Path(out_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    buf = io.BytesIO()
+    torch.save(obj, buf)
+    data = buf.getvalue()
+    with tempfile.NamedTemporaryFile(
+            dir=out_path.parent, delete=False) as f:
+        f.write(data)
+        tmp = Path(f.name)
+    # verify the staged file loads before it replaces anything
+    torch.load(tmp, map_location="cpu", weights_only=False)
+    shutil.copy2(tmp, out_path)
+    torch.load(out_path, map_location="cpu", weights_only=False)
+    tmp.unlink()
+
+
+# ---------------------------------------------------------------------------
+# FLOP model for MFU accounting
+# ---------------------------------------------------------------------------
+
+def flops_per_token(cfg: Foundry4Config) -> float:
+    """Forward+backward FLOPs per token (~6N_active + attention terms)."""
+    d = cfg.d_model
+    n_attn = cfg.n_layers // cfg.attn_every
+    n_gdn = cfg.n_layers - n_attn
+    # dense per-token params: mixers + shared expert + topk experts +
+    # router + ponder(iters) + heads — times 6
+    mixer = n_attn * 4 * d * cfg.n_heads * cfg.head_dim \
+        + n_gdn * d * (4 * cfg.gdn_heads * cfg.gdn_head_k)
+    ffn = cfg.n_layers * 3 * d * (cfg.shared_hidden
+                                + cfg.moe_topk * cfg.expert_hidden)
+    ponder = cfg.ponder_iters * (4 * d * cfg.n_heads * cfg.head_dim
+                                 + 3 * d * cfg.ponder_hidden)
+    heads = cfg.vocab_size * d                       # tied head proj
+    n_active = mixer + ffn + ponder + heads
+    # attention QK/AV term: ~4 * T * d * n_attn_heads-ish; and GDN chunk
+    # work ~= O(C * dk * dv) per token — lump as seq term:
+    seq_term = n_attn * 12 * cfg.seq_len * d \
+        + n_gdn * 2 * cfg.gdn_chunk * cfg.gdn_heads \
+        * cfg.gdn_head_v
+    return 6.0 * n_active + seq_term

segmm.py

@@ -0,0 +1,393 @@
+"""segmm.py — segmented/grouped GEMM for dropless MoE dispatch.
+
+Tokens are gathered into expert-sorted order once: xs [M, K] where the
+rows for expert e occupy the contiguous slice [offs[e], offs[e+1]).
+Each expert multiplies its slice by its own weight w[e] [K, N].
+
+Why not torch.bmm + padding?  Dropless padding materialises
+[E, max_count, *] — on skewed routing max_count spikes (we measured
+cap ~7600 on real data -> multi-GB buffers per layer, OOM).  Packed
+segments keep memory at exactly M rows regardless of skew, and the
+same slices power a dropless (causality-safe) dispatch.
+
+Why not torch._grouped_mm?  On this ROCm build it measured ~10 TFLOP/s
+(a scalar fallback path) vs ~300 TFLOP/s for bmm.  This kernel is a
+straight tile loop and hits bmm-class throughput.
+
+Forward:  c[rows_e] = a[rows_e] @ w[e]
+Backward: da[rows_e] = dc[rows_e] @ w[e]^T          (same kernel, w
+          loaded transposed via strides)
+          dw[e]      = a[rows_e]^T @ dc[rows_e]     (reduction kernel)
+"""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        # BM must stay 64 — _tile_meta partitions tiles with this size
+        triton.Config({"BM": 64, "BK": 64, "BN": 64}, num_warps=4,
+                      num_stages=3),
+        triton.Config({"BM": 64, "BK": 128, "BN": 64}, num_warps=4,
+                      num_stages=4),
+        triton.Config({"BM": 64, "BK": 64, "BN": 128}, num_warps=8,
+                      num_stages=3),
+        triton.Config({"BM": 64, "BK": 128, "BN": 128}, num_warps=8,
+                      num_stages=4),
+        triton.Config({"BM": 64, "BK": 256, "BN": 64}, num_warps=8,
+                      num_stages=3),
+    ],
+    key=["K", "N", "E"],
+)
+@triton.jit
+def _seg_mm_fwd(a_ptr, w_ptr, c_ptr, offs_ptr, mt_starts_ptr, n_tiles_ptr,
+                gidx_ptr,
+                K, N,
+                sa_m, sa_k, sw_e, sw_k, sw_n, sc_m, sc_n,
+                E: tl.constexpr, BE: tl.constexpr,
+                HAS_IDX: tl.constexpr,
+                BM: tl.constexpr, BK: tl.constexpr, BN: tl.constexpr):
+    pid_m = tl.program_id(0)          # global m-tile across all experts
+    pid_n = tl.program_id(1)
+    total = tl.load(n_tiles_ptr)      # scalar: total m-tiles
+    if pid_m >= total:
+        return
+    # expert = last e with mt_starts[e] <= pid_m  (linear scan, E small)
+    e_off = tl.arange(0, BE)
+    starts = tl.load(mt_starts_ptr + e_off, mask=e_off < E, other=1 << 30)
+    e = tl.sum((starts <= pid_m).to(tl.int32)) - 1
+    mt_local = pid_m - tl.load(mt_starts_ptr + e)
+    row0 = tl.load(offs_ptr + e)
+    row1 = tl.load(offs_ptr + e + 1)
+    rows = row0 + mt_local * BM + tl.arange(0, BM)
+    rmask = rows < row1
+    if HAS_IDX:
+        a_rows = tl.load(gidx_ptr + rows, mask=rmask, other=0)
+    else:
+        a_rows = rows
+    cols = pid_n * BN + tl.arange(0, BN)
+    cmask = cols < N
+    acc = tl.zeros((BM, BN), dtype=tl.float32)
+    for k0 in range(0, K, BK):
+        kk = k0 + tl.arange(0, BK)
+        a = tl.load(a_ptr + a_rows[:, None] * sa_m + kk[None, :] * sa_k,
+                    mask=rmask[:, None] & (kk < K)[None, :], other=0.0)
+        w = tl.load(w_ptr + e * sw_e + kk[:, None] * sw_k
+                    + cols[None, :] * sw_n,
+                    mask=(kk < K)[:, None] & cmask[None, :], other=0.0)
+        acc += tl.dot(a, w)
+    tl.store(c_ptr + rows[:, None] * sc_m + cols[None, :] * sc_n,
+             acc.to(c_ptr.dtype.element_ty),
+             mask=rmask[:, None] & cmask[None, :])
+
+
+@triton.jit
+def _seg_mm_dw(a_ptr, dy_ptr, dw_ptr, offs_ptr, gidx_ptr,
+               K, N,
+               sa_m, sa_k, sd_m, sd_n, sw_e, sw_k, sw_n,
+               HAS_IDX: tl.constexpr,
+               BR: tl.constexpr, BK: tl.constexpr, BN: tl.constexpr):
+    # grid: (E, ceil(K/BK), ceil(N/BN)) flattened into pid via program_id
+    e = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    pid_n = tl.program_id(2)
+    row0 = tl.load(offs_ptr + e)
+    row1 = tl.load(offs_ptr + e + 1)
+    kk = pid_k * BK + tl.arange(0, BK)
+    nn = pid_n * BN + tl.arange(0, BN)
+    kmask = kk < K
+    nmask = nn < N
+    acc = tl.zeros((BK, BN), dtype=tl.float32)
+    for r0 in range(row0, row1, BR):
+        rows = r0 + tl.arange(0, BR)
+        rmask = rows < row1
+        if HAS_IDX:
+            a_rows = tl.load(gidx_ptr + rows, mask=rmask, other=0)
+        else:
+            a_rows = rows
+        a = tl.load(a_ptr + a_rows[:, None] * sa_m + kk[None, :] * sa_k,
+                    mask=rmask[:, None] & kmask[None, :], other=0.0)
+        dy = tl.load(dy_ptr + rows[:, None] * sd_m + nn[None, :] * sd_n,
+                     mask=rmask[:, None] & nmask[None, :], other=0.0)
+        acc += tl.dot(tl.trans(a), dy)
+    tl.store(dw_ptr + e * sw_e + kk[:, None] * sw_k + nn[None, :] * sw_n,
+             acc.to(dw_ptr.dtype.element_ty),
+             mask=kmask[:, None] & nmask[None, :])
+
+
+# ---------------------------------------------------------------------------
+# fused weighted scatter: out[gidx[r]] += s[r] * yp[r]   (atomics, fp32 out)
+# backward:   dyp[r] = s[r] * dout[gidx[r]]
+#             ds[r]  = <dout[gidx[r]], yp[r]>
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def _wscatter_fwd(yp_ptr, gidx_ptr, s_ptr, out_ptr,
+                  D, sm, sd_m,
+                  BR: tl.constexpr, BD: tl.constexpr):
+    r = tl.program_id(0)
+    dn = tl.program_id(1)
+    t = tl.load(gidx_ptr + r)
+    s = tl.load(s_ptr + r).to(tl.float32)
+    cols = dn * BD + tl.arange(0, BD)
+    m = cols < D
+    v = tl.load(yp_ptr + r * sm + cols, mask=m, other=0.).to(tl.float32)
+    tl.atomic_add(out_ptr + t * sd_m + cols, s * v, mask=m)
+
+
+@triton.jit
+def _wscatter_bwd(dout_ptr, gidx_ptr, s_ptr, yp_ptr, dyp_ptr, ds_ptr,
+                  D, sd_m, sm,
+                  BR: tl.constexpr, BD: tl.constexpr):
+    r = tl.program_id(0)
+    dn = tl.program_id(1)
+    t = tl.load(gidx_ptr + r)
+    s = tl.load(s_ptr + r).to(tl.float32)
+    cols = dn * BD + tl.arange(0, BD)
+    m = cols < D
+    do = tl.load(dout_ptr + t * sd_m + cols, mask=m, other=0.).to(tl.float32)
+    yp = tl.load(yp_ptr + r * sm + cols, mask=m, other=0.).to(tl.float32)
+    tl.store(dyp_ptr + r * sm + cols, (s * do).to(dyp_ptr.dtype.element_ty),
+             mask=m)
+    tl.atomic_add(ds_ptr + r, tl.sum(do * yp, 0))
+
+
+def _next_pow2(x):
+    return 1 << (x - 1).bit_length()
+
+
+class SegMM(torch.autograd.Function):
+    """c[offs[e]:offs[e+1]] = a[gidx[offs[e]:offs[e+1]]] @ w[e]  for each e.
+
+    a: [M,K] rows sorted by segment (or an arbitrarily-ordered dense
+    matrix when gidx is given — the kernel gathers rows itself);
+    w: [E,K,N];  offs: [E+1] int32;  gidx (optional): [M] int32 row
+    indices into a (gather fused into the GEMM read).
+    Grid is a static upper bound — no host<->device sync anywhere.
+    """
+
+    BM, BK, BN, BR = 64, 64, 64, 32
+
+    @staticmethod
+    def _tile_meta(offs, E, M, dev):
+        counts = (offs[1:] - offs[:-1])
+        tiles = (counts + SegMM.BM - 1) // SegMM.BM
+        mt_starts = torch.cumsum(tiles, 0) - tiles
+        total = tiles.sum()
+        return mt_starts.to(torch.int32), total[None].to(torch.int32)
+
+    @staticmethod
+    def forward(ctx, a, w, offs, gidx=None):
+        if gidx is not None:
+            M = gidx.shape[0]
+        else:
+            M = a.shape[0]
+        _, K = a.shape
+        E, _, N = w.shape
+        dev = a.device
+        offs = offs.to(torch.int32)
+        has_idx = gidx is not None
+        if has_idx:
+            gidx = gidx.to(torch.int32).contiguous()
+        else:
+            gidx = offs  # dummy ptr — kernel ignores when HAS_IDX=False
+        mt_starts, total = SegMM._tile_meta(offs, E, M, dev)
+        c = torch.empty(M, N, device=dev, dtype=a.dtype)
+        max_tiles = (M + SegMM.BM - 1) // SegMM.BM + E
+        grid_n = (N + SegMM.BN - 1) // SegMM.BN
+        _seg_mm_fwd[(max_tiles, grid_n)](
+            a, w, c, offs, mt_starts, total, gidx, K, N,
+            a.stride(0), a.stride(1), w.stride(0), w.stride(1),
+            w.stride(2), c.stride(0), c.stride(1),
+            E=E, BE=_next_pow2(E + 1), HAS_IDX=has_idx)
+        ctx.save_for_backward(a, w, offs, gidx if has_idx else None)
+        ctx.has_idx = has_idx
+        return c
+
+    @staticmethod
+    def backward(ctx, dc):
+        a, w, offs, gidx = ctx.saved_tensors
+        has_idx = ctx.has_idx
+        M = dc.shape[0]
+        _, K = a.shape
+        E, _, N = w.shape
+        dev = a.device
+        dc = dc.contiguous()
+        # da = seg_mm(dc, w^T) — packed-space grad (M rows either way)
+        mt_starts, total = SegMM._tile_meta(offs, E, M, dev)
+        da_packed = torch.empty(M, K, device=dev, dtype=a.dtype)
+        max_tiles = (M + SegMM.BM - 1) // SegMM.BM + E
+        grid_k = (K + SegMM.BN - 1) // SegMM.BN
+        dummy = gidx if has_idx else offs
+        _seg_mm_fwd[(max_tiles, grid_k)](
+            dc, w, da_packed, offs, mt_starts, total, dummy, N, K,
+            dc.stride(0), dc.stride(1), w.stride(0), w.stride(2),
+            w.stride(1), da_packed.stride(0), da_packed.stride(1),
+            E=E, BE=_next_pow2(E + 1), HAS_IDX=False)
+        if has_idx:
+            # scatter packed grads back to the dense a-rows
+            da = torch.zeros_like(a)
+            da.index_add_(0, gidx.long(), da_packed)
+        else:
+            da = da_packed
+        # dw[e] = a_e^T @ dc_e — reduction kernel (gathers via gidx)
+        dw = torch.empty_like(w)
+        _seg_mm_dw[(E, (K + SegMM.BK - 1) // SegMM.BK,
+                    (N + SegMM.BN - 1) // SegMM.BN)](
+            a, dc, dw, offs, dummy, K, N,
+            a.stride(0), a.stride(1), dc.stride(0), dc.stride(1),
+            dw.stride(0), dw.stride(1), dw.stride(2),
+            HAS_IDX=has_idx,
+            BR=SegMM.BR, BK=SegMM.BK, BN=SegMM.BN)
+        return da, dw, None, None
+
+
+def seg_mm(a, w, offs, gidx=None):
+    """Grouped matmul: a [M,K] sorted rows, w [E,K,N], offs [E+1].
+    With gidx [M] int32: a stays dense, packed row r reads a[gidx[r]]."""
+    return SegMM.apply(a, w, offs, gidx)
+
+
+class WScatter(torch.autograd.Function):
+    """out[t[r]] += s[r] * yp[r]  — fused weighted scatter-add.
+
+    out: [N,D] fp32 (must be pre-zeroed);  yp: [M,D];  s: [M];
+    t: [M] int32 destination rows (duplicates allowed — atomics).
+    Backward: dyp = s * dout[t];  ds = <dout[t], yp>.
+    """
+
+    BD = 256
+
+    @staticmethod
+    def forward(ctx, out, yp, s, t):
+        M, D = yp.shape
+        t32 = t.to(torch.int32).contiguous()
+        grid_d = (D + WScatter.BD - 1) // WScatter.BD
+        _wscatter_fwd[(M, grid_d)](
+            yp, t32, s, out, D, yp.stride(0), out.stride(0),
+            BR=1, BD=WScatter.BD)
+        ctx.mark_dirty(out)
+        ctx.save_for_backward(yp, s, t32)
+        ctx.D, ctx.sm = D, yp.stride(0)
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        yp, s, t32 = ctx.saved_tensors
+        M, D = yp.shape
+        dout = dout.contiguous()
+        dyp = torch.empty_like(yp)
+        ds = torch.zeros(M, device=yp.device, dtype=torch.float32)
+        grid_d = (D + WScatter.BD - 1) // WScatter.BD
+        _wscatter_bwd[(M, grid_d)](
+            dout, t32, s, yp, dyp, ds, D, dout.stride(0),
+            yp.stride(0), BR=1, BD=WScatter.BD)
+        return dout, dyp, ds.to(s.dtype), None
+
+
+def wscatter(out, yp, s, t):
+    """out [N,D] fp32 pre-zeroed; yp [M,D]; s [M]; t [M] int rows."""
+    return WScatter.apply(out, yp, s, t)
+
+
+# ---------------------------------------------------------------------------
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    dev = "cuda"
+    E, M, K, N = 32, 9000, 256, 192
+    # ragged segment sizes
+    counts = torch.multinomial(torch.rand(E), E, replacement=True) * 7
+    counts = (counts.float() / counts.sum() * M).long()
+    counts[0] += M - counts.sum()
+    offs = torch.zeros(E + 1, dtype=torch.int64)
+    offs[1:] = torch.cumsum(counts, 0)
+    offs = offs.to(dev)
+    a = torch.randn(M, K, device=dev, dtype=torch.bfloat16,
+                    requires_grad=True)
+    w = (torch.randn(E, K, N, device=dev, dtype=torch.bfloat16)
+         / (K ** 0.5)).requires_grad_(True)
+
+    a_ref = a.detach().clone().requires_grad_(True)
+    w_ref = w.detach().clone().requires_grad_(True)
+    c2 = torch.zeros(M, N, device=dev, dtype=torch.bfloat16)
+    for e in range(E):
+        s, t = int(offs[e]), int(offs[e + 1])
+        if t > s:
+            c2[s:t] = a_ref[s:t] @ w_ref[e]
+
+    c = seg_mm(a, w, offs)
+    err = (c.float() - c2.float()).abs().max().item()
+    rel = err / (c2.float().abs().max().item() + 1e-9)
+    print(f"fwd rel err {rel:.2e}")
+
+    g = torch.randn_like(c)
+    c.backward(g); c2.backward(g)
+    da_err = ((a.grad - a_ref.grad).abs().max()
+              / a_ref.grad.abs().max()).item()
+    dw_err = ((w.grad - w_ref.grad).abs().max()
+              / w_ref.grad.abs().max()).item()
+    print(f"bwd da rel err {da_err:.2e}  dw rel err {dw_err:.2e}")
+
+    # --- gidx gather variant: dense a_src[N,K], gidx[M] picks rows ---
+    Nsrc = 4000
+    gidx = torch.randint(0, Nsrc, (M,), device=dev)
+    a_src = torch.randn(Nsrc, K, device=dev, dtype=torch.bfloat16,
+                        requires_grad=True)
+    w3 = (torch.randn(E, K, N, device=dev, dtype=torch.bfloat16)
+          / (K ** 0.5)).requires_grad_(True)
+    a_src_r = a_src.detach().clone().requires_grad_(True)
+    w3_r = w3.detach().clone().requires_grad_(True)
+    c3 = torch.zeros(M, N, device=dev, dtype=torch.bfloat16)
+    for e in range(E):
+        s, t = int(offs[e]), int(offs[e + 1])
+        if t > s:
+            c3[s:t] = a_src_r[gidx[s:t]] @ w3_r[e]
+    c4 = seg_mm(a_src, w3, offs, gidx=gidx)
+    rel = ((c4.float() - c3.float()).abs().max()
+           / (c3.float().abs().max() + 1e-9)).item()
+    print(f"gidx fwd rel err {rel:.2e}")
+    g4 = torch.randn_like(c4)
+    c4.backward(g4); c3.backward(g4)
+    da_err = ((a_src.grad - a_src_r.grad).abs().max()
+              / a_src_r.grad.abs().max()).item()
+    dw_err = ((w3.grad - w3_r.grad).abs().max()
+              / w3_r.grad.abs().max()).item()
+    print(f"gidx bwd da rel {da_err:.2e}  dw rel {dw_err:.2e}")
+
+    # --- wscatter: out[t[r]] += s[r]*yp[r] vs index_add ---
+    Nout = 3000
+    yp = torch.randn(M, N, device=dev, dtype=torch.bfloat16,
+                     requires_grad=True)
+    tt = torch.randint(0, Nout, (M,), device=dev)
+    sw = torch.rand(M, device=dev, dtype=torch.float32,
+                     requires_grad=True)
+    yp_r = yp.detach().clone().requires_grad_(True)
+    sw_r = sw.detach().clone().requires_grad_(True)
+    ref = torch.zeros(Nout, N, device=dev, dtype=torch.float32)
+    ref.index_add_(0, tt, (yp_r * sw_r[:, None]).float())
+    out5 = wscatter(torch.zeros(Nout, N, device=dev, dtype=torch.float32),
+                    yp, sw, tt)
+    rel = ((out5 - ref).abs().max()
+           / (ref.abs().max() + 1e-9)).item()
+    print(f"wscatter fwd rel err {rel:.2e}")
+    g5 = torch.randn(Nout, N, device=dev, dtype=torch.float32)
+    out5.backward(g5); ref.backward(g5)
+    dyp_err = ((yp.grad.float() - yp_r.grad.float()).abs().max()
+               / yp_r.grad.float().abs().max()).item()
+    ds_err = ((sw.grad - sw_r.grad).abs().max()
+              / (sw_r.grad.abs().max() + 1e-9)).item()
+    print(f"wscatter bwd dyp rel {dyp_err:.2e}  ds rel {ds_err:.2e}")
+
+    # speed vs padded bmm
+    import time
+    def bench(fn, n=50):
+        for _ in range(5): fn()
+        torch.cuda.synchronize(); t0 = time.time()
+        for _ in range(n): fn()
+        torch.cuda.synchronize(); return (time.time() - t0) / n * 1e3
+    print(f"seg_mm fwd {bench(lambda: seg_mm(a, w, offs)):.3f}ms")
+    cap = int(counts.max())
+    xp = torch.zeros(E, cap, K, device=dev, dtype=torch.bfloat16)
+    print(f"padded bmm {bench(lambda: torch.bmm(xp, w)):.3f}ms")
+    print("segmm self-test done")