Repository: huggingface/transformers
Pull request: Align torch implementation of Gated DeltaNet in Qwen3-Next with fla library.

This pull request aims to align the PyTorch implementation of Gated DeltaNet in Qwen3-Next with the corresponding implementation in the fla library, ensuring consistency between the two implementations.

Changed files:
filename: src/transformers/models/qwen3_next/modeling_qwen3_next.py
patch: @@ -446,8 +446,10 @@ def torch_chunk_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        query = F.normalize(query, p=2, dim=-1)
-        key = F.normalize(key, p=2, dim=-1)
+        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, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]
@@ -518,8 +520,10 @@ def torch_recurrent_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        query = F.normalize(query, p=2, dim=-1)
-        key = F.normalize(key, p=2, dim=-1)
+        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, 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: @@ -282,8 +282,10 @@ def torch_chunk_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        query = F.normalize(query, p=2, dim=-1)
-        key = F.normalize(key, p=2, dim=-1)
+        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, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]
@@ -354,8 +356,10 @@ def torch_recurrent_gated_delta_rule(
 ):
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        query = F.normalize(query, p=2, dim=-1)
-        key = F.normalize(key, p=2, dim=-1)
+        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, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]

filename: src/transformers/utils/import_utils.py
patch: @@ -592,12 +592,14 @@ def is_flash_linear_attention_available():
 
         if not torch.cuda.is_available():
             return False
-        else:
-            if _is_package_available("fla"):
-                import fla
 
-                if version.parse(fla.__version__) >= version.parse("0.2.2"):
-                    return True
+        try:
+            import fla
+
+            if version.parse(fla.__version__) >= version.parse("0.2.2"):
+                return True
+        except Exception:
+            pass
     return False