Upload folder using huggingface_hub

fe7c66d verified 19 days ago

12.7 kB

"""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

	"""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