refactor(layers): share one GatedDeltaNet op in layers/linear_attn

python/freetoken/attention/linear.py

@@ -38,7 +38,7 @@ class FLAMetadata:
     # --- hybrid-radix track-checkpoint (extra_buffer) fields; all None when not caching ---
     # For each request crossing a chunk-aligned (×CHUNK) boundary this forward, snapshot its
     # recurrent + conv state into a donatable pool slot, written on the forward stream by the
-    # GDN op (see Qwen3_5GatedDeltaNet._write_track_snapshot). Built by the scheduler in P2;
+    # GDN op (see GatedDeltaNet._write_track_snapshot). Built by the scheduler in P2;
     # left None by build_fla_metadata so the existing path is unchanged.
     track_dst: torch.Tensor | None = None        # [nt] int64 dst pool slot per tracked req
     track_h_row: torch.Tensor | None = None      # [nt] int64 row into h (boh_i + aligned//CHUNK)

python/freetoken/layers/linear_attn/__init__.py

@@ -0,0 +1,3 @@
+from .gdn import GatedDeltaNet
+
+__all__ = ["GatedDeltaNet"]

python/freetoken/models/qwen4_exp/gdn.py → python/freetoken/layers/linear_attn/gdn.py

@@ -6,7 +6,6 @@
 from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen
 from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearReplicated
 from freetoken.layers.quantization import QuantConfig
-from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla
 
 
 class _DepthwiseConv1d(BaseOP):
@@ -16,7 +15,7 @@ def __init__(self, conv_dim: int, kernel: int):
         self.weight = torch.empty(conv_dim, 1, kernel)
 
 
-class Qwen4ExpGatedDeltaNet(BaseOP):
+class GatedDeltaNet(BaseOP):
     """GatedDeltaNet op using the vendored flash-linear-attention triton kernels
     (``freetoken.kernel.fla``) for the recurrence and a per-request
     recurrent + conv state held in ``ctx.linear_state_pool`` (keyed by ``Req.table_idx``).
@@ -31,7 +30,7 @@ class Qwen4ExpGatedDeltaNet(BaseOP):
 
     def __init__(
         self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim,
-        conv_kernel_size, rms_norm_eps, layer_id, output_gate: str = "sigmoid",
+        conv_kernel_size, rms_norm_eps, layer_id, output_gate: str = "silu",
         *, quant_config: QuantConfig | None = None, prefix: str = "", out_heads_tiled: bool = False,
     ):
         self.layer_id = layer_id
@@ -75,7 +74,7 @@ def __init__(
         self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size)
         # Recurrence-gating params kept in fp32 (exp/softplus is precision-sensitive,
         # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a
-        # per-call .float() upcast in the decode wrapper. The weight loader exempts
+        # per-call .float() upcast before the decode kernel. The weight loader exempts
         # *.A_log / *.dt_bias from the model-dtype downcast.
         self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32)
         self.A_log = torch.empty(num_v_heads, dtype=torch.float32)
@@ -156,18 +155,25 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             # Fused fla decode kernel: gating + in-kernel l2norm + recurrent update +
             # per-request state read/write-by-index, all in one kernel (no gather/scatter,
             # no clone, no external l2norm). q/k stay at num_k_heads (kernel handles GQA).
+            from freetoken.kernel.fla import fused_sigmoid_gating_delta_rule_update
+
             mixed = self._conv_decode(conv_in, fla.cache_indices, pool)  # [B, conv_dim]
             B = mixed.shape[0]
             qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
             q = qf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype)
             k = kf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype)
             v = vf.reshape(1, B, self.num_v_heads, self.head_v_dim).to(dtype)
-            core_out = gdn_decode_fla(
-                q, k, v, a, b, A_log=self.A_log, dt_bias=self.dt_bias,
-                state_source=pool.recurrent_states[li], indices=fla.cache_indices,
-                cu_seqlens=fla.cu_seqlens, scale=self.head_k_dim ** -0.5,
-            )
+            core_out = fused_sigmoid_gating_delta_rule_update(
+                A_log=self.A_log, a=a, dt_bias=self.dt_bias,  # already fp32 (stored fp32)
+                softplus_beta=1.0, softplus_threshold=20.0,
+                q=q, k=k, v=v, b=b,
+                initial_state_source=pool.recurrent_states[li],
+                initial_state_indices=fla.cache_indices,  # already int32 (built int32 by build_fla_metadata)
+                scale=self.head_k_dim ** -0.5, use_qk_l2norm_in_kernel=True, cu_seqlens=fla.cu_seqlens,
+            )[0]  # [B, num_v, V]
         else:
+            from freetoken.kernel.fla import chunk_gated_delta_rule
+
             mixed = self._conv_prefill(
                 conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state)
             # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads.
@@ -182,18 +188,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             # fresh sequences (cached_len==0) must start from a zeroed slot.
             if fla.fresh_state_indices is not None:
                 pool.recurrent_states[li].index_fill_(0, fla.fresh_state_indices, 0.0)
-            track = fla.track_dst is not None
-            result = gdn_prefill_chunk_fla(
-                q, k, v, g, beta,
-                state_source=pool.recurrent_states[li], indices=fla.cache_indices,
-                cu_seqlens=fla.cu_seqlens, scale=self.head_k_dim ** -0.5,
-                return_h=track,
+            o, _, h = chunk_gated_delta_rule(
+                q=q, k=k, v=v, g=g, beta=beta, scale=self.head_k_dim ** -0.5,
+                initial_state=pool.recurrent_states[li], initial_state_indices=fla.cache_indices.to(torch.int32),
+                cu_seqlens=fla.cu_seqlens.to(torch.int64), head_first=False,
+                use_qk_l2norm_in_kernel=True,
             )
-            if track:
-                core_out, h = result
+            core_out = o[0]  # [total, num_v_heads, head_v_dim]
+            if fla.track_dst is not None:
                 self._write_track_snapshot(pool, li, conv_in, h, fla)
-            else:
-                core_out = result
 
         core_out = core_out.reshape(-1, self.head_v_dim)
         z = z.reshape(-1, self.head_v_dim)
@@ -205,4 +208,4 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return self.out_proj.forward(out)
 
 
-__all__ = ["Qwen4ExpGatedDeltaNet"]
+__all__ = ["GatedDeltaNet"]

python/freetoken/models/qwen3_5_moe/gdn_reference.py → python/freetoken/layers/linear_attn/gdn_reference.py

@@ -1,10 +1,11 @@
 """Pure-torch Gated DeltaNet reference (text-only, no cache).
 
-Correctness oracle for the kernel-backed GDN op (``gdn.Qwen3_5GatedDeltaNet``).
-The recurrence and forward are transcribed from
-``transformers.models.qwen3_5_moe.modeling_qwen3_5_moe`` (``torch_recurrent_gated_delta_rule``
-and ``Qwen3_5MoeGatedDeltaNet.forward`` no-cache path) and were validated against the
-HF module bit-for-bit when ported.
+Correctness oracle for the kernel-backed GDN op (``freetoken.layers.linear_attn.GatedDeltaNet``).
+The delta rule and the forward are transcribed from
+``transformers.models.qwen4_exp.modeling_qwen4_exp`` (``torch_recurrent_gated_delta_rule``
+and ``Qwen4ExpTextGatedDeltaNet.forward`` no-cache path).
+Qwen3.8-Flash-Next gates the output norm with ``config.output_gate_type`` (sigmoid) where
+Qwen3.5 hardcodes silu; the conv keeps ``config.hidden_act`` (silu).
 """
 
 from __future__ import annotations
@@ -13,6 +14,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+_GATE_ACTS = {"silu": F.silu, "swish": F.silu, "sigmoid": torch.sigmoid}
+
 
 def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
     return x * torch.rsqrt(x.pow(2).sum(dim=-1, keepdim=True) + eps)
@@ -65,26 +68,27 @@ def recurrent_gated_delta_rule(
 
 
 class _RMSNormGated(nn.Module):
-    """RMSNorm of x followed by a silu(z) gate (norm_before_gate=True).
+    """RMSNorm of x followed by an ``activation(z)`` gate (norm_before_gate=True).
 
-    Mirrors ``Qwen3_5MoeRMSNormGated`` over head_v_dim groups.
+    Mirrors ``Qwen4ExpTextRMSNormGated`` over head_v_dim groups.
     """
 
-    def __init__(self, dim: int, eps: float):
+    def __init__(self, dim: int, eps: float, activation: str = "silu"):
         super().__init__()
         self.weight = nn.Parameter(torch.ones(dim))
         self.eps = eps
+        self.act = _GATE_ACTS[activation]
 
     def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
         in_dtype = x.dtype
         x = x.float()
         x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
         x = x * self.weight.float()
-        x = x * F.silu(z.float())
+        x = x * self.act(z.float())
         return x.to(in_dtype)
 
 
-class Qwen3_5GatedDeltaNetReference(nn.Module):
+class GatedDeltaNetReference(nn.Module):
     """Pure-torch Gated DeltaNet (text-only, no cache)."""
 
     def __init__(
@@ -97,10 +101,13 @@ def __init__(
         conv_kernel_size: int,
         rms_norm_eps: float,
         hidden_act: str = "silu",
+        output_gate: str = "silu",
     ):
         super().__init__()
         if hidden_act != "silu":
-            raise ValueError(f"GDN reference only supports silu activation, got {hidden_act!r}")
+            raise ValueError(f"GDN reference only supports silu conv activation, got {hidden_act!r}")
+        if output_gate not in _GATE_ACTS:
+            raise ValueError(f"unsupported GDN output gate {output_gate!r}")
         self.num_k_heads = num_k_heads
         self.num_v_heads = num_v_heads
         self.head_k_dim = head_k_dim
@@ -120,12 +127,12 @@ def __init__(
         )
         self.dt_bias = nn.Parameter(torch.zeros(num_v_heads))
         self.A_log = nn.Parameter(torch.zeros(num_v_heads))
-        self.norm = _RMSNormGated(head_v_dim, eps=rms_norm_eps)
+        self.norm = _RMSNormGated(head_v_dim, eps=rms_norm_eps, activation=output_gate)
         self.out_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
 
     @torch.no_grad()
     def load_from_hf(self, hf_gdn) -> None:
-        """Copy weights from a transformers ``Qwen3_5MoeGatedDeltaNet``."""
+        """Copy weights from a transformers ``Qwen4ExpTextGatedDeltaNet``."""
         self.in_proj_qkv.weight.copy_(hf_gdn.in_proj_qkv.weight)
         self.in_proj_z.weight.copy_(hf_gdn.in_proj_z.weight)
         self.in_proj_b.weight.copy_(hf_gdn.in_proj_b.weight)
@@ -173,4 +180,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return self.out_proj(core)
 
 
-__all__ = ["Qwen3_5GatedDeltaNetReference", "recurrent_gated_delta_rule"]
+__all__ = [
+    "GatedDeltaNetReference",
+    "recurrent_gated_delta_rule",
+]

python/freetoken/models/glm5_next/kda.py

@@ -11,7 +11,7 @@
 merged q|k|v stream (width 3*H*D == the pool's ``2*K*dk + V*dv``) and the
 recurrent state is one [D, D] matrix per head, stored in the KERNEL's [V, K]
 layout (coincides with the pool's [K, V] declaration because D_k == D_v; same
-convention as GDN, see qwen3_5_moe/gdn.py).
+convention as GDN, see layers/linear_attn/gdn.py).
 
 Kernels: ``fused_recurrent_kda`` decodes with in-kernel gate/beta/l2norm and
 per-slot state read/write (slot 0 is its NULL sentinel == the pool's padding
@@ -84,7 +84,7 @@ def _conv_weight(self) -> torch.Tensor:
     def _write_track_snapshot(self, pool, li, conv_in, h, fla) -> None:
         """Hybrid-radix: snapshot recurrent + conv state at the chunk-aligned track
         boundary into a donatable pool slot (same contract as GDN, see
-        qwen3_5_moe/gdn.py). h rows are the kernel's per-chunk [V, K] states --
+        layers/linear_attn/gdn.py). h rows are the kernel's per-chunk [V, K] states --
         a direct copy into the pool's [K, V] slots (D_k == D_v)."""
         rec = pool.recurrent_states[li]
         rec.index_copy_(0, fla.track_dst, h[0, fla.track_h_row].to(rec.dtype))