"""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 for i > l (strict lower) A_v[i,j] = exp(G_i - G_j) beta_i 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} (Kt_l S_in) + sum_{j<=i} M[i,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 A_v = torch.where(tril_m1, (kbc32 @ kc.transpose(-1, -2)) * M, kc.new_zeros(())) # e^{Gi-Gj} b_i 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(())) # 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