Repository: huggingface/transformers
Pull request: Fix the misalignment between the l2norm in GDN of Qwen3-Next and the implementation in the FLA library.

Fix the misalignment between the l2norm in GDN of Qwen3-Next and the implementation in the FLA library.

Changed files:
filename: src/transformers/models/qwen3_next/modeling_qwen3_next.py
patch: @@ -433,6 +433,12 @@ def torch_causal_conv1d_update(
     return out
 
 
+def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
+    """This function is intended to align with the l2norm implementation in the FLA library."""
+    inv_norm = 1 / torch.sqrt((x * x).sum(dim=dim, keepdim=True) + eps)
+    return x * inv_norm
+
+
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -446,10 +452,8 @@ def torch_chunk_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        head_dim = query.size(-1)
-        inv_scale = head_dim**-0.5
-        query = F.rms_norm(query, (head_dim,), eps=1e-6) * inv_scale
-        key = F.rms_norm(key, (head_dim,), eps=1e-6) * inv_scale
+        query = l2norm(query, dim=-1, eps=1e-6)
+        key = l2norm(key, dim=-1, eps=1e-6)
     query, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]
@@ -520,10 +524,8 @@ def torch_recurrent_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        head_dim = query.size(-1)
-        inv_scale = head_dim**-0.5
-        query = F.rms_norm(query, (head_dim,), eps=1e-6) * inv_scale
-        key = F.rms_norm(key, (head_dim,), eps=1e-6) * inv_scale
+        query = l2norm(query, dim=-1, eps=1e-6)
+        key = l2norm(key, dim=-1, eps=1e-6)
     query, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]

filename: src/transformers/models/qwen3_next/modular_qwen3_next.py
patch: @@ -269,6 +269,12 @@ def torch_causal_conv1d_update(
     return out
 
 
+def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
+    """This function is intended to align with the l2norm implementation in the FLA library."""
+    inv_norm = 1 / torch.sqrt((x * x).sum(dim=dim, keepdim=True) + eps)
+    return x * inv_norm
+
+
 def torch_chunk_gated_delta_rule(
     query,
     key,
@@ -282,10 +288,8 @@ def torch_chunk_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        head_dim = query.size(-1)
-        inv_scale = head_dim**-0.5
-        query = F.rms_norm(query, (head_dim,), eps=1e-6) * inv_scale
-        key = F.rms_norm(key, (head_dim,), eps=1e-6) * inv_scale
+        query = l2norm(query, dim=-1, eps=1e-6)
+        key = l2norm(key, dim=-1, eps=1e-6)
     query, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]
@@ -356,10 +360,8 @@ def torch_recurrent_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        head_dim = query.size(-1)
-        inv_scale = head_dim**-0.5
-        query = F.rms_norm(query, (head_dim,), eps=1e-6) * inv_scale
-        key = F.rms_norm(key, (head_dim,), eps=1e-6) * inv_scale
+        query = l2norm(query, dim=-1, eps=1e-6)
+        key = l2norm(key, dim=-1, eps=1e-6)
     query, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]