From 531d85597e8b5eceec339e901999fe15a533a41d Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 28 Oct 2026 22:08:16 +0000 Subject: [PATCH 1/2] refactor(layers): share one GatedDeltaNet op in layers/linear_attn --- python/freetoken/attention/linear.py | 2 +- .../freetoken/layers/linear_attn/__init__.py | 3 + .../qwen4_exp => layers/linear_attn}/gdn.py | 41 +-- .../linear_attn}/gdn_reference.py | 38 ++- python/freetoken/models/glm5_next/kda.py | 4 +- python/freetoken/models/qwen3_5_moe/gdn.py | 206 ------------- .../models/qwen3_5_moe/gdn_kernels.py | 76 ----- python/freetoken/models/qwen3_5_moe/model.py | 5 +- .../models/qwen4_exp/gdn_reference.py | 279 ------------------ python/freetoken/models/qwen4_exp/model.py | 4 +- tests/README.md | 2 +- .../{models/qwen4_exp => layers}/test_gdn.py | 65 +--- 12 files changed, 73 insertions(+), 652 deletions(-) create mode 100644 python/freetoken/layers/linear_attn/__init__.py rename python/freetoken/{models/qwen4_exp => layers/linear_attn}/gdn.py (89%) rename python/freetoken/{models/qwen3_5_moe => layers/linear_attn}/gdn_reference.py (82%) delete mode 100644 python/freetoken/models/qwen3_5_moe/gdn.py delete mode 100644 python/freetoken/models/qwen3_5_moe/gdn_kernels.py delete mode 100644 python/freetoken/models/qwen4_exp/gdn_reference.py rename tests/{models/qwen4_exp => layers}/test_gdn.py (65%) diff --git a/python/freetoken/attention/linear.py b/python/freetoken/attention/linear.py index 1c4c34a6af..3c63f2b3d7 100644 --- a/python/freetoken/attention/linear.py +++ b/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) diff --git a/python/freetoken/layers/linear_attn/__init__.py b/python/freetoken/layers/linear_attn/__init__.py new file mode 100644 index 0000000000..de8f28e700 --- /dev/null +++ b/python/freetoken/layers/linear_attn/__init__.py @@ -0,0 +1,3 @@ +from .gdn import GatedDeltaNet + +__all__ = ["GatedDeltaNet"] diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/layers/linear_attn/gdn.py similarity index 89% rename from python/freetoken/models/qwen4_exp/gdn.py rename to python/freetoken/layers/linear_attn/gdn.py index 48c92a8e8a..631a9a83c1 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/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 @@ -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 in the scheduler) + 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"] diff --git a/python/freetoken/models/qwen3_5_moe/gdn_reference.py b/python/freetoken/layers/linear_attn/gdn_reference.py similarity index 82% rename from python/freetoken/models/qwen3_5_moe/gdn_reference.py rename to python/freetoken/layers/linear_attn/gdn_reference.py index ba57a91d60..466be6f51a 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn_reference.py +++ b/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", +] diff --git a/python/freetoken/models/glm5_next/kda.py b/python/freetoken/models/glm5_next/kda.py index ba722ae287..7b5251c718 100644 --- a/python/freetoken/models/glm5_next/kda.py +++ b/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)) diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py deleted file mode 100644 index 0785ad4968..0000000000 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ /dev/null @@ -1,206 +0,0 @@ -from __future__ import annotations - -import torch -import torch.nn.functional as F -from freetoken.core import get_global_ctx -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 .gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla - - -class _DepthwiseConv1d(BaseOP): - """Holds the depthwise conv weight ``[conv_dim, 1, K]`` (key ``conv1d.weight``).""" - - def __init__(self, conv_dim: int, kernel: int): - self.weight = torch.empty(conv_dim, 1, kernel) - - -class Qwen3_5GatedDeltaNet(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``). - - Parameter names match HF (``in_proj_qkv``/``in_proj_z``/``in_proj_b``/``in_proj_a``/ - ``conv1d``/``A_log``/``dt_bias``/``norm``/``out_proj``). Handles prefill (incl. chunked - continuation) and single-token decode; state is fresh when ``req.cached_len == 0``. - """ - - 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, *, quant_config: QuantConfig | None = None, - prefix: str = "", out_heads_tiled: bool = False, - ): - self.layer_id = layer_id - self._out_heads_tiled = out_heads_tiled - # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as - # [V, K] while the LinearStatePool declares it [K, V]; these coincide (and the - # hybrid-radix snapshot scatter h[h_row]->slot is a plain copy) only when the two head - # dims are equal. Qwen3.5/3.6 satisfy this (128/128); guard any future config. - assert head_k_dim == head_v_dim, ( - f"GatedDeltaNet requires head_k_dim == head_v_dim, got {head_k_dim} != {head_v_dim}" - ) - self.num_k_heads = num_k_heads - self.num_v_heads = num_v_heads - self.head_k_dim = head_k_dim - self.head_v_dim = head_v_dim - self.key_dim = num_k_heads * head_k_dim - self.value_dim = num_v_heads * head_v_dim - self.conv_dim = 2 * self.key_dim + self.value_dim - self.conv_kernel_size = conv_kernel_size - # quantized checkpoints quantize qkv|z but not b|a, so the fusion splits into a qkvz GEMM and a ba GEMM with their own schemes (matches sglang / vLLM) - self._split_in_proj = ( - quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None - ) - - self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._split_in_proj: - self.in_proj_qkvz = LinearColParallelMerged( - hidden_size, [self.conv_dim, self.value_dim], has_bias=False, - quant_config=quant_config, prefix=f"{prefix}.in_proj_qkvz", - ) - self.in_proj_ba = LinearColParallelMerged( - hidden_size, [num_v_heads, num_v_heads], has_bias=False, - quant_config=quant_config, prefix=f"{prefix}.in_proj_ba", - ) - else: - # Fused input projection (one GEMM instead of four): qkv | z | b | a. - self.in_proj = LinearColParallelMerged( - hidden_size, self._in_proj_split, has_bias=False, - quant_config=quant_config, prefix=f"{prefix}.in_proj", - ) - 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 - # *.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) - self.norm = GatedRMSNorm(head_v_dim, eps=rms_norm_eps) - self.out_proj = LinearReplicated( - self.value_dim, hidden_size, has_bias=False, - quant_config=quant_config, prefix=f"{prefix}.out_proj", - ) - - def _gate_params(self, a: torch.Tensor, b: torch.Tensor): - beta = b.sigmoid() - g = -self.A_log.exp() * F.softplus(a.float() + self.dt_bias) - return g, beta - - def _conv_weight(self) -> torch.Tensor: - return self.conv1d.weight.squeeze(1) # [conv_dim, kernel] for the fused kernel - - def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state) -> torch.Tensor: - """Varlen causal conv (fused sgl_kernel) with silu; reads/updates each request's - conv state in place by ``cache_indices`` slot. ``conv_in`` [total, conv_dim]. - ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from FLAMetadata.""" - li = pool.local_index(self.layer_id) - x = conv_in.transpose(0, 1).contiguous() # [conv_dim, total] - out = causal_conv1d_varlen(x, self._conv_weight(), pool.conv_states[li], - cu_seqlens, cache_indices, has_initial_state) - return out.transpose(0, 1) # [total, conv_dim] - - def _conv_decode(self, conv_in: torch.Tensor, table_idx: torch.Tensor, pool) -> torch.Tensor: - """Single-token causal conv update (fused sgl_kernel) by ``table_idx`` slot; - updates conv state in place, no host loop -> CUDA-graph capturable. - ``conv_in`` [B, conv_dim] -> silu(conv) [B, conv_dim].""" - li = pool.local_index(self.layer_id) - return causal_conv1d_decode(conv_in, pool.conv_states[li], self._conv_weight(), table_idx) - - def _write_track_snapshot(self, pool, li: int, conv_in: torch.Tensor, - h: torch.Tensor, fla) -> None: - """Snapshot this layer's recurrent + conv state at the chunk-aligned track boundary - into a donatable pool slot, on the forward stream (hybrid-radix extra_buffer path). - SSM: ``recurrent_states[li, dst] = h[0, h_row]`` -- a DIRECT copy (h is [V,K], the - state pool is [K,V]; they coincide because GDN requires head_k_dim == head_v_dim). - Conv: the last (kernel-1) raw conv-input timesteps ending at the boundary.""" - rec = pool.recurrent_states[li] - rec.index_copy_(0, fla.track_dst, h[0, fla.track_h_row].to(rec.dtype)) - cv = pool.conv_states[li] - # conv_in [total, conv_dim]; gather the (kernel-1) window per tracked req. - conv_win = conv_in[fla.track_conv_src].transpose(-1, -2).contiguous() # [nt, conv_dim, K-1] - cv.index_copy_(0, fla.track_dst, conv_win.to(cv.dtype)) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - ctx = get_global_ctx() - batch = ctx.batch - pool = ctx.linear_state_pool - total = hidden_states.shape[0] - dtype = hidden_states.dtype - - # Per-forward GDN metadata (cu_seqlens / cache_indices / continuation flags), - # built once and shared by all GDN layers. The scheduler/graph set it; build it - # lazily here (cached on the batch) for direct-op callers (tests). - fla = batch.fla_metadata - if fla is None: - from freetoken.attention.linear import build_fla_metadata - - fla = build_fla_metadata(batch, hidden_states.device) - batch.fla_metadata = fla - - if self._split_in_proj: - qkvz = self.in_proj_qkvz.forward(hidden_states) - conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) - ba = self.in_proj_ba.forward(hidden_states) - b, a = torch.split(ba, [self.num_v_heads, self.num_v_heads], dim=-1) - else: - proj = self.in_proj.forward(hidden_states) - conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1) - z = z.reshape(total, self.num_v_heads, self.head_v_dim) - li = pool.local_index(self.layer_id) - - if batch.is_decode: - # 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). - 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, - ) - else: - 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. - qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) - q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - k = kf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - v = vf.reshape(1, total, self.num_v_heads, self.head_v_dim).to(dtype) - g, beta = self._gate_params(a, b) - g = g.reshape(1, total, self.num_v_heads) - beta = beta.float().reshape(1, total, self.num_v_heads) - # The chunk kernel reads + writes back initial_state[cache_indices] in place; - # 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, - ) - if track: - core_out, h = result - 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) - out = self.norm.forward(core_out, z).reshape(total, -1) - if self._out_heads_tiled: - # grouped value heads (k-head major) -> tiled (k-head minor), the column order of a GGUF ssm_out - r = self.num_v_heads // self.num_k_heads - out = out.view(total, self.num_k_heads, r, self.head_v_dim).transpose(1, 2).reshape(total, -1) - return self.out_proj.forward(out) - - -__all__ = ["Qwen3_5GatedDeltaNet"] diff --git a/python/freetoken/models/qwen3_5_moe/gdn_kernels.py b/python/freetoken/models/qwen3_5_moe/gdn_kernels.py deleted file mode 100644 index 061a6d5499..0000000000 --- a/python/freetoken/models/qwen3_5_moe/gdn_kernels.py +++ /dev/null @@ -1,76 +0,0 @@ -from __future__ import annotations - -import torch - - -def gdn_prefill_chunk_fla( - q: torch.Tensor, # [1, total, num_k_heads, head_k_dim] bf16 (NOT GQA-expanded) - k: torch.Tensor, # [1, total, num_k_heads, head_k_dim] bf16 - v: torch.Tensor, # [1, total, num_v_heads, head_v_dim] bf16 - g: torch.Tensor, # [1, total, num_v_heads] log-decay (<=0), fp32 - beta: torch.Tensor, # [1, total, num_v_heads] fp32 - *, - state_source: torch.Tensor, # [num_slots, num_v_heads, head_k_dim, head_v_dim] fp32 (in place) - indices: torch.Tensor, # [num_seqs] slot id per sequence - cu_seqlens: torch.Tensor, # [num_seqs+1] int64 - scale: float, - return_h: bool = False, -) -> torch.Tensor: - """Chunked gated-delta-rule prefill via the vendored fla kernel. GQA is handled - in-kernel (q/k at num_k_heads), q/k l2norm is done in-kernel, and the per-sequence - recurrent state is read from and written back to ``state_source[indices]`` IN PLACE - (no external l2norm, no Python stack of initial states, no copy_ writeback loop). - Fresh sequences must have their ``state_source`` slot pre-zeroed by the caller. - Returns ``o`` of shape ``[total, num_v_heads, head_v_dim]`` (bf16). - - When ``return_h=True`` also returns the per-chunk hidden-state buffer ``h`` of shape - ``[1, NT_total, num_v_heads, head_v_dim, head_k_dim]`` (bf16). ``h[0, boh_i + c]`` is the - recurrent state after ``c*64`` tokens of packed sequence ``i`` (chunk granularity 64), where - ``boh_i = prepare_chunk_offsets(cu_seqlens, 64)[i]``. Note the last two dims are ``[V, K]`` -- - transposed vs ``state_source``'s ``[K, V]``. Used by the hybrid-radix track-checkpoint path.""" - from freetoken.kernel.fla import chunk_gated_delta_rule - - o, _, h = chunk_gated_delta_rule( - q=q, k=k, v=v, g=g, beta=beta, scale=scale, - initial_state=state_source, initial_state_indices=indices.to(torch.int32), - cu_seqlens=cu_seqlens.to(torch.int64), head_first=False, - use_qk_l2norm_in_kernel=True, - ) - if return_h: - return o[0], h # h: [1, NT_total, num_v_heads, head_v_dim, head_k_dim] - return o[0] # [total, num_v_heads, head_v_dim] - - -def gdn_decode_fla( - q: torch.Tensor, # [1, B, num_k_heads, head_k_dim] bf16 (NOT GQA-expanded) - k: torch.Tensor, # [1, B, num_k_heads, head_k_dim] bf16 - v: torch.Tensor, # [1, B, num_v_heads, head_v_dim] bf16 - a: torch.Tensor, # [B, num_v_heads] raw - b: torch.Tensor, # [B, num_v_heads] raw - *, - A_log: torch.Tensor, # [num_v_heads] - dt_bias: torch.Tensor, # [num_v_heads] - state_source: torch.Tensor, # [num_slots, num_v_heads, head_k_dim, head_v_dim] fp32 (in place) - indices: torch.Tensor, # [B] int32 slot id per request - cu_seqlens: torch.Tensor, # [B+1] query indptr (arange) from FLAMetadata - scale: float, -) -> torch.Tensor: - """Fused sigmoid-gating gated-delta-rule decode (vendored fla triton kernel): gating + - in-kernel l2norm + recurrent update + state read/write-by-index in one kernel, with no - external gating or gather/scatter/clone glue. Returns [B, num_v, V].""" - from freetoken.kernel.fla import fused_sigmoid_gating_delta_rule_update - - o = fused_sigmoid_gating_delta_rule_update( - A_log=A_log, a=a, dt_bias=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=state_source, - initial_state_indices=indices, # already int32 (built int32 in the scheduler) - scale=scale, use_qk_l2norm_in_kernel=True, cu_seqlens=cu_seqlens, - ) - # kernel returns o = [NK, *v.shape] then squeeze(NK) -> [1, B, num_v, V]. - # o[0] -> [B, num_v, V] (all B decode tokens; o[0,0] would drop B>1). - return o[0] - - -__all__ = ["gdn_prefill_chunk_fla", "gdn_decode_fla"] diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index 88c6ae1345..fc6d066364 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -11,13 +11,13 @@ ParallelLMHead, VocabParallelEmbedding, ) +from freetoken.layers.linear_attn import GatedDeltaNet from freetoken.models.blocks import BaseLLMModel from freetoken.models.blocks import embed_input_ids from freetoken.models.qwen3_vl.vision import Qwen3VLVisionModel, QwenVLVisionMixin from freetoken.utils import nvtx_annotate from .attention import Qwen3_5Attention -from .gdn import Qwen3_5GatedDeltaNet from .moe import Qwen3_5DenseMLP, Qwen3_5MoE if TYPE_CHECKING: @@ -35,7 +35,7 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): if self._is_linear: g = config.linear_attention_group() assert g is not None - self.linear_attn = Qwen3_5GatedDeltaNet( + self.linear_attn = GatedDeltaNet( hidden_size=config.hidden_size, num_k_heads=g.num_key_heads, num_v_heads=g.num_value_heads, @@ -44,6 +44,7 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): conv_kernel_size=g.conv_kernel_dim, rms_norm_eps=config.rms_norm_eps, layer_id=layer_id, + output_gate=g.output_gate, quant_config=config.quant, prefix=f"{prefix}.linear_attn", # llama.cpp stores every qwen35 ssm_out with tiled value-head columns diff --git a/python/freetoken/models/qwen4_exp/gdn_reference.py b/python/freetoken/models/qwen4_exp/gdn_reference.py deleted file mode 100644 index fbd9647bcb..0000000000 --- a/python/freetoken/models/qwen4_exp/gdn_reference.py +++ /dev/null @@ -1,279 +0,0 @@ -"""Pure-torch Gated DeltaNet reference (text-only, no cache). - -Correctness oracle for the kernel-backed GDN op (``gdn.Qwen4ExpGatedDeltaNet``). -The two delta rules and the forward are transcribed from -``transformers.models.qwen4_exp.modeling_qwen4_exp`` (``torch_chunk_gated_delta_rule``, -``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 - -import torch -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) - - -def recurrent_gated_delta_rule( - query: torch.Tensor, # [B, T, Hv, Dk] - key: torch.Tensor, # [B, T, Hv, Dk] - value: torch.Tensor, # [B, T, Hv, Dv] - g: torch.Tensor, # [B, T, Hv] (log-decay; per-step decay = exp(g)) - beta: torch.Tensor, # [B, T, Hv] - *, - initial_state: torch.Tensor | None = None, - use_qk_l2norm: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - """Verbatim port of HF ``torch_recurrent_gated_delta_rule`` (output_final_state=True).""" - initial_dtype = query.dtype - if use_qk_l2norm: - query = _l2norm(query, eps=1e-6) - key = _l2norm(key, eps=1e-6) - query, key, value, beta, g = [ - t.transpose(1, 2).contiguous().to(torch.float32) - for t in (query, key, value, beta, g) - ] - b, h, t_len, dk = key.shape - dv = value.shape[-1] - scale = 1.0 / (dk ** 0.5) - query = query * scale - - out = torch.zeros(b, h, t_len, dv, dtype=value.dtype, device=value.device) - state = ( - torch.zeros(b, h, dk, dv, dtype=value.dtype, device=value.device) - if initial_state is None - else initial_state.to(value) - ) - for i in range(t_len): - q_t = query[:, :, i] - k_t = key[:, :, i] - v_t = value[:, :, i] - g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1) - beta_t = beta[:, :, i].unsqueeze(-1) - state = state * g_t - kv_mem = (state * k_t.unsqueeze(-1)).sum(dim=-2) - delta = (v_t - kv_mem) * beta_t - state = state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) - out[:, :, i] = (state * q_t.unsqueeze(-1)).sum(dim=-2) - - out = out.transpose(1, 2).contiguous().to(initial_dtype) # [B, T, Hv, Dv] - return out, state - - -def chunk_gated_delta_rule( - query: torch.Tensor, # [B, T, Hv, Dk] - key: torch.Tensor, # [B, T, Hv, Dk] - value: torch.Tensor, # [B, T, Hv, Dv] - g: torch.Tensor, # [B, T, Hv] - beta: torch.Tensor, # [B, T, Hv] - *, - chunk_size: int = 64, - initial_state: torch.Tensor | None = None, - use_qk_l2norm: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - """Verbatim port of HF ``torch_chunk_gated_delta_rule`` (output_final_state=True). - - Same recurrence as ``recurrent_gated_delta_rule`` in exact arithmetic; it is the form the - fla chunk kernel implements, so it is the closer oracle for the prefill path.""" - initial_dtype = query.dtype - if use_qk_l2norm: - query = _l2norm(query, eps=1e-6) - key = _l2norm(key, eps=1e-6) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) - for x in (query, key, value, beta, g) - ] - - batch_size, num_heads, sequence_length, k_head_dim = key.shape - v_head_dim = value.shape[-1] - pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size - query = F.pad(query, (0, 0, 0, pad_size)) - key = F.pad(key, (0, 0, 0, pad_size)) - value = F.pad(value, (0, 0, 0, pad_size)) - beta = F.pad(beta, (0, pad_size)) - g = F.pad(g, (0, pad_size)) - total_sequence_length = sequence_length + pad_size - scale = 1 / (query.shape[-1] ** 0.5) - query = query * scale - - v_beta = value * beta.unsqueeze(-1) - k_beta = key * beta.unsqueeze(-1) - # reshape to chunks - query, key, value, k_beta, v_beta = [ - x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) - for x in (query, key, value, k_beta, v_beta) - ] - g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) - mask = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0 - ) - - # chunk decay - g = g.cumsum(dim=-1) - decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) - for i in range(1, chunk_size): - row = attn[..., i, :i].clone() - sub = attn[..., :i, :i].clone() - attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) - attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - value = attn @ v_beta - k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) - last_recurrent_state = ( - torch.zeros( - batch_size, num_heads, k_head_dim, v_head_dim, - dtype=value.dtype, device=value.device, - ) - if initial_state is None - else initial_state.to(value) - ) - core_attn_out = torch.zeros_like(value) - - # for each chunk; decay_mask is already lower-triangular so no extra causal mask is needed - for i in range(0, total_sequence_length // chunk_size): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] - attn = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i] - v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state - v_new = v_i - v_prime - attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - core_attn_out[:, :, i] = attn_inter + attn @ v_new - last_recurrent_state = ( - last_recurrent_state * g[:, :, i, -1, None, None].exp() - + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) - @ v_new - ) - - core_attn_out = core_attn_out.reshape( - core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1] - ) - core_attn_out = core_attn_out[:, :, :sequence_length] - core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) - return core_attn_out, last_recurrent_state - - -class _RMSNormGated(nn.Module): - """RMSNorm of x followed by an ``activation(z)`` gate (norm_before_gate=True). - - Mirrors ``Qwen4ExpTextRMSNormGated`` over head_v_dim groups. - """ - - def __init__(self, dim: int, eps: float, activation: str = "sigmoid"): - 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 * self.act(z.float()) - return x.to(in_dtype) - - -class Qwen4ExpGatedDeltaNetReference(nn.Module): - """Pure-torch Gated DeltaNet (text-only, no cache).""" - - def __init__( - self, - hidden_size: int, - num_k_heads: int, - num_v_heads: int, - head_k_dim: int, - head_v_dim: int, - conv_kernel_size: int, - rms_norm_eps: float, - hidden_act: str = "silu", - output_gate: str = "sigmoid", - ): - super().__init__() - if hidden_act != "silu": - 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 - self.head_v_dim = head_v_dim - self.key_dim = num_k_heads * head_k_dim - self.value_dim = num_v_heads * head_v_dim - self.conv_dim = self.key_dim * 2 + self.value_dim - self.conv_kernel_size = conv_kernel_size - - self.in_proj_qkv = nn.Linear(hidden_size, self.conv_dim, bias=False) - self.in_proj_z = nn.Linear(hidden_size, self.value_dim, bias=False) - self.in_proj_b = nn.Linear(hidden_size, num_v_heads, bias=False) - self.in_proj_a = nn.Linear(hidden_size, num_v_heads, bias=False) - self.conv1d = nn.Conv1d( - self.conv_dim, self.conv_dim, kernel_size=conv_kernel_size, - groups=self.conv_dim, padding=conv_kernel_size - 1, bias=False, - ) - 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, 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 ``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) - self.in_proj_a.weight.copy_(hf_gdn.in_proj_a.weight) - # HF stores conv1d weight as [conv_dim, 1, K]; our depthwise Conv1d matches. - self.conv1d.weight.copy_(hf_gdn.conv1d.weight.view_as(self.conv1d.weight)) - if hf_gdn.conv1d.bias is not None and self.conv1d.bias is not None: - self.conv1d.bias.copy_(hf_gdn.conv1d.bias) - self.dt_bias.copy_(hf_gdn.dt_bias) - self.A_log.copy_(hf_gdn.A_log) - self.norm.weight.copy_(hf_gdn.norm.weight) - self.out_proj.weight.copy_(hf_gdn.out_proj.weight) - - def forward(self, hidden_states: torch.Tensor, *, use_chunk_rule: bool = False) -> torch.Tensor: - b, t_len, _ = hidden_states.shape - - mixed_qkv = self.in_proj_qkv(hidden_states).transpose(1, 2) # [B, conv_dim, T] - z = self.in_proj_z(hidden_states).reshape(b, t_len, -1, self.head_v_dim) - a = self.in_proj_a(hidden_states) - bb = self.in_proj_b(hidden_states) - - # causal depthwise conv + silu (drop the right padding back to T) - mixed_qkv = F.silu(self.conv1d(mixed_qkv)[..., :t_len]).transpose(1, 2) # [B, T, conv_dim] - query, key, value = torch.split( - mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1 - ) - query = query.reshape(b, t_len, -1, self.head_k_dim) - key = key.reshape(b, t_len, -1, self.head_k_dim) - value = value.reshape(b, t_len, -1, self.head_v_dim) - - beta = bb.sigmoid() - g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) - - # GQA expand: replicate q/k heads up to num_v_heads - rep = self.num_v_heads // self.num_k_heads - if rep > 1: - query = query.repeat_interleave(rep, dim=2) - key = key.repeat_interleave(rep, dim=2) - - rule = chunk_gated_delta_rule if use_chunk_rule else recurrent_gated_delta_rule - core, _ = rule(query, key, value, g, beta, use_qk_l2norm=True) - - core = core.reshape(-1, self.head_v_dim) - z = z.reshape(-1, self.head_v_dim) - core = self.norm(core, z).reshape(b, t_len, -1) - return self.out_proj(core) - - -__all__ = [ - "Qwen4ExpGatedDeltaNetReference", - "chunk_gated_delta_rule", - "recurrent_gated_delta_rule", -] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 4eb69165d5..cca354d369 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -37,10 +37,10 @@ def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseOP: """GDN mixer of a linear_attention layer (Qwen3.5's GDN with a configurable output gate).""" - from .gdn import Qwen4ExpGatedDeltaNet + from freetoken.layers.linear_attn import GatedDeltaNet g = config.linear_attention_group() - return Qwen4ExpGatedDeltaNet( + return GatedDeltaNet( hidden_size=config.hidden_size, num_k_heads=g.num_key_heads, num_v_heads=g.num_value_heads, diff --git a/tests/README.md b/tests/README.md index 3cd4cc9045..4f202047b6 100644 --- a/tests/README.md +++ b/tests/README.md @@ -17,7 +17,7 @@ directories (`kvcache/test_dsv4_pool.py`) or in `dsv4/`, and will move there. | `engine/` | `freetoken.engine` — cache budget planning, config resolution, backend gating | | `kernels/` | `freetoken.kernel` / `freetoken.layers` — attention, rope, fused ops, JIT cache, pinned memory | | `attention/` | `freetoken.attention` — attention backends against eager references | -| `layers/` | `freetoken.layers` — whole layer modules (mHC); a layer's kernels and fused ops go in `kernels/` | +| `layers/` | `freetoken.layers` — whole layer modules (mHC, GatedDeltaNet); a layer's kernels and fused ops go in `kernels/` | | `moe/` | `freetoken.moe` — offload cache, CPU/GPU expert kernels, quantized backends (fp8/nvfp4/mxfp4/gguf) | | `models/` | `freetoken.models` — the registry, and the loading machinery every model shares (sharding, qkv/expert merge, streaming layers into banks); one directory per model for tests specific to it | | `mm/` | `freetoken.mm` — multimodal processor, placeholder expansion, encoder cache | diff --git a/tests/models/qwen4_exp/test_gdn.py b/tests/layers/test_gdn.py similarity index 65% rename from tests/models/qwen4_exp/test_gdn.py rename to tests/layers/test_gdn.py index a854626a80..ea8c13052e 100644 --- a/tests/models/qwen4_exp/test_gdn.py +++ b/tests/layers/test_gdn.py @@ -1,10 +1,4 @@ -"""qwen4_exp GatedDeltaNet op vs the pure-torch HF reference math. - -The oracle is ``models/qwen4_exp/gdn_reference.py``, whose two delta rules and forward are -transcribed from the ``modeling_qwen4_exp.py`` snapshot, so no transformers build carrying -qwen4_exp is needed here. Covered: prefill at 128 and 1000 tokens, a decode step continuing -from the prefill state, ragged bs=3, both GQA head ratios, and the sigmoid output gate. -""" +"""The GatedDeltaNet layer against the pure-torch HF reference math (``layers/linear_attn/gdn_reference.py``): a ragged bs=3 prefill (up to 1000 tokens) and one decode step off the carried state, for both GQA head ratios and both output gates.""" from __future__ import annotations @@ -12,9 +6,9 @@ import torch from freetoken.core import Batch, Context, Req, SamplingParams +from freetoken.layers.linear_attn import GatedDeltaNet +from freetoken.layers.linear_attn.gdn_reference import GatedDeltaNetReference from freetoken.models.config import LinearGatedDeltaGroupConfig -from freetoken.models.qwen4_exp.gdn import Qwen4ExpGatedDeltaNet -from freetoken.models.qwen4_exp.gdn_reference import Qwen4ExpGatedDeltaNetReference from freetoken.utils import torch_dtype pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -47,9 +41,13 @@ def _state_dict(ref) -> dict[str, torch.Tensor]: def _make_layer(ratio: int, output_gate: str = "sigmoid", seed: int = 0): """fp32 reference + bf16 kernel op over one set of weights. The op is built on meta under the serving dtype, the way the engine builds a model, so load_state_dict's dtype check bites.""" + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) num_k, num_v = HEADS[ratio] torch.manual_seed(seed) - ref = Qwen4ExpGatedDeltaNetReference( + ref = GatedDeltaNetReference( hidden_size=HIDDEN, num_k_heads=num_k, num_v_heads=num_v, head_k_dim=HEAD_DIM, head_v_dim=HEAD_DIM, conv_kernel_size=CONV_K, rms_norm_eps=EPS, output_gate=output_gate, ).to(DEV).float().eval() @@ -59,7 +57,7 @@ def _make_layer(ratio: int, output_gate: str = "sigmoid", seed: int = 0): ref.dt_bias.uniform_(-1.0, 1.0) ref.norm.weight.normal_(1.0, 0.1) with torch.device("meta"), torch_dtype(torch.bfloat16): - op = Qwen4ExpGatedDeltaNet( + op = GatedDeltaNet( hidden_size=HIDDEN, num_k_heads=num_k, num_v_heads=num_v, head_k_dim=HEAD_DIM, head_v_dim=HEAD_DIM, conv_kernel_size=CONV_K, rms_norm_eps=EPS, layer_id=0, output_gate=output_gate, @@ -110,24 +108,17 @@ def _decode(op, ctx: Context, reqs, hidden: torch.Tensor) -> torch.Tensor: @torch.no_grad() -def _ref_out(ref, hidden: torch.Tensor, use_chunk_rule: bool = False) -> torch.Tensor: - return ref(hidden.float().unsqueeze(0), use_chunk_rule=use_chunk_rule)[0] - +def _ref_out(ref, hidden: torch.Tensor) -> torch.Tensor: + return ref(hidden.float().unsqueeze(0))[0] -@pytest.mark.parametrize("length", (1000,)) -@pytest.mark.parametrize("ratio", (2, 3)) -def test_prefill_matches_reference(ratio, length): - op, ref = _make_layer(ratio, seed=ratio) - hidden, _, out = _prefill(op, _ctx(ratio), [length], seed=11) - torch.testing.assert_close(out.float(), _ref_out(ref, hidden[0]), rtol=RTOL, atol=ATOL) - -@pytest.mark.parametrize("ratio", (2, 3)) -def test_ragged_prefill_then_decode(ratio): +# 3:1 with sigmoid is Qwen3.8-Flash-Next; the silu case catches a gate that ignores the config +@pytest.mark.parametrize("ratio,output_gate", [(2, "silu"), (3, "sigmoid")]) +def test_ragged_prefill_then_decode(ratio, output_gate): """bs=3 ragged prefill, then one decode step per request off the carried conv + recurrent state. The decode oracle is the whole (prefill + 1) sequence in one reference pass, so a state that did not survive the prefill shows up immediately.""" - op, ref = _make_layer(ratio, seed=ratio) + op, ref = _make_layer(ratio, output_gate, seed=ratio) ctx = _ctx(ratio) lengths = [128, 1000, 37] hidden, reqs, out = _prefill(op, ctx, lengths, seed=13) @@ -144,29 +135,3 @@ def test_ragged_prefill_then_decode(ratio): for i, h in enumerate(hidden): full = _ref_out(ref, torch.cat([h, nxt[i:i + 1]], dim=0)) torch.testing.assert_close(dec[i].float(), full[-1], rtol=RTOL, atol=ATOL) - - -def test_chunk_and_recurrent_rules_agree(): - """The chunked form (what the fla prefill kernel implements) against the sequential - definition, both fp32: the chunk oracle is only worth anything if it reproduces the - recurrence to fp32 precision.""" - _, ref = _make_layer(3, seed=1) - torch.manual_seed(17) - hidden = torch.randn(1000, HIDDEN, device=DEV, dtype=torch.bfloat16) - torch.testing.assert_close( - _ref_out(ref, hidden, use_chunk_rule=True), _ref_out(ref, hidden), rtol=1e-4, atol=1e-4 - ) - - -def test_output_gate_comes_from_the_config(): - """The gate activation is the group config's string, not a hardcoded silu. Both gates track - their own reference, and the two are far apart -- so a stuck activation cannot pass.""" - op_silu, ref_silu = _make_layer(3, output_gate="silu", seed=2) - hidden, _, out_silu = _prefill(op_silu, _ctx(3), [128], seed=19) - torch.testing.assert_close(out_silu.float(), _ref_out(ref_silu, hidden[0]), rtol=RTOL, atol=ATOL) - - op_sig, ref_sig = _make_layer(3, output_gate="sigmoid", seed=2) - _, _, out_sig = _prefill(op_sig, _ctx(3), [128], seed=19) - torch.testing.assert_close(out_sig.float(), _ref_out(ref_sig, hidden[0]), rtol=RTOL, atol=ATOL) - - assert (out_sig.float() - out_silu.float()).abs().max().item() > 10 * ATOL From 84cbd873ce4e8929bce864d580d87a947d935266 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Thu, 29 Oct 2026 00:11:39 +0000 Subject: [PATCH 2/2] test(layers): restore the global context in the GDN test and fix stale comments --- python/freetoken/layers/linear_attn/gdn.py | 4 ++-- python/freetoken/models/qwen3_5_moe/weight.py | 2 +- python/freetoken/models/qwen4_exp/weight.py | 2 +- tests/layers/test_gdn.py | 13 ++++++------- tests/models/qwen4_exp/test_config.py | 2 +- 5 files changed, 11 insertions(+), 12 deletions(-) diff --git a/python/freetoken/layers/linear_attn/gdn.py b/python/freetoken/layers/linear_attn/gdn.py index 631a9a83c1..78a0cddfe6 100644 --- a/python/freetoken/layers/linear_attn/gdn.py +++ b/python/freetoken/layers/linear_attn/gdn.py @@ -74,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) @@ -168,7 +168,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 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 in the scheduler) + 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: diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index afd32eb2c0..947e1b4833 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -125,7 +125,7 @@ def target(self, module: str) -> tuple[str, int, int]: if not candidates: return module, 0, 1 if len(candidates) > 1: - # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as gdn.py + # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as layers/linear_attn/gdn.py split = self.scheme(f"{parent}.in_proj_qkvz") is not None keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"} candidates = [c for c in candidates if c[0] in keep] diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 6c62d141cb..734eb80a83 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -137,7 +137,7 @@ def target(self, module: str) -> tuple[str, int, int]: if not candidates: return module, 0, 1 if len(candidates) > 1: - # GDN: quantized checkpoints split qkv|z from b|a; same test as gdn.py + # GDN: quantized checkpoints split qkv|z from b|a; same test as layers/linear_attn/gdn.py split = self.scheme(f"{parent}.in_proj_qkvz") is not None keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"} candidates = [c for c in candidates if c[0] in keep] diff --git a/tests/layers/test_gdn.py b/tests/layers/test_gdn.py index ea8c13052e..b1aa2b58fd 100644 --- a/tests/layers/test_gdn.py +++ b/tests/layers/test_gdn.py @@ -38,7 +38,7 @@ def _state_dict(ref) -> dict[str, torch.Tensor]: } -def _make_layer(ratio: int, output_gate: str = "sigmoid", seed: int = 0): +def _make_layer(ratio: int, output_gate: str, seed: int = 0): """fp32 reference + bf16 kernel op over one set of weights. The op is built on meta under the serving dtype, the way the engine builds a model, so load_state_dict's dtype check bites.""" from freetoken.distributed import set_tp_info, try_get_tp_info @@ -66,7 +66,7 @@ def _make_layer(ratio: int, output_gate: str = "sigmoid", seed: int = 0): return op, ref -def _ctx(ratio: int, num_slots: int = 8) -> Context: +def _ctx(monkeypatch, ratio: int, num_slots: int = 8) -> Context: import freetoken.core as core from freetoken.kvcache.linear_state_pool import LinearStatePool @@ -76,10 +76,9 @@ def _ctx(ratio: int, num_slots: int = 8) -> Context: key_head_dim=HEAD_DIM, value_head_dim=HEAD_DIM, conv_kernel_dim=CONV_K, output_gate="sigmoid", ) - core._GLOBAL_CTX = None ctx = Context(page_size=64) ctx.linear_state_pool = LinearStatePool(group, num_slots, torch.bfloat16, DEV, tp_size=1) - core.set_global_ctx(ctx) + monkeypatch.setattr(core, "_GLOBAL_CTX", ctx) return ctx @@ -112,14 +111,14 @@ def _ref_out(ref, hidden: torch.Tensor) -> torch.Tensor: return ref(hidden.float().unsqueeze(0))[0] -# 3:1 with sigmoid is Qwen3.8-Flash-Next; the silu case catches a gate that ignores the config +# 3:1 with sigmoid is Qwen3.8-Flash-Next; the sigmoid case also catches a layer that ignores output_gate (the default is silu) @pytest.mark.parametrize("ratio,output_gate", [(2, "silu"), (3, "sigmoid")]) -def test_ragged_prefill_then_decode(ratio, output_gate): +def test_ragged_prefill_then_decode(monkeypatch, ratio, output_gate): """bs=3 ragged prefill, then one decode step per request off the carried conv + recurrent state. The decode oracle is the whole (prefill + 1) sequence in one reference pass, so a state that did not survive the prefill shows up immediately.""" op, ref = _make_layer(ratio, output_gate, seed=ratio) - ctx = _ctx(ratio) + ctx = _ctx(monkeypatch, ratio) lengths = [128, 1000, 37] hidden, reqs, out = _prefill(op, ctx, lengths, seed=13) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index f850927b99..d11ce243fa 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -182,7 +182,7 @@ def test_vision_turns_on_mrope_and_the_tower(): def test_text_only_keeps_the_1d_rope(): config = parse_config(_hf_config()) assert not config.is_multimodal and not config.model_is_mrope and config.rotary_config.mrope_section is None -# the merged-projection prefixes the model asks the QuantConfig about (attention.py / gdn.py) +# the merged-projection prefixes the model asks the QuantConfig about (attention.py / layers/linear_attn/gdn.py) DENSE_PREFIXES = ( "model.layers.3.self_attn.qkv_proj", "model.layers.3.self_attn.o_proj", "model.layers.0.linear_attn.in_proj_qkvz", "model.layers.0.linear_attn.out_proj",