From cf9bf7964151702cbfdef3b13745ff5db5974c90 Mon Sep 17 00:00:00 2001 From: bozheng-hit Date: Thu, 11 Sep 2025 13:12:26 +0800 Subject: [PATCH 1/6] align torch implementation of gdn with fla. --- .../models/qwen3_next/modeling_qwen3_next.py | 12 ++++++++---- .../models/qwen3_next/modular_qwen3_next.py | 12 ++++++++---- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 55b23945488b..08798fc0aa8d 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -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 = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + 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 = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + 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) ] diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py index b4c179c50d72..a0e6add1735f 100644 --- a/src/transformers/models/qwen3_next/modular_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py @@ -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 = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + 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 = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + 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) ] From 1af3e4043bdb6cc319c961c17124ff2ca5de823e Mon Sep 17 00:00:00 2001 From: bozheng-hit Date: Thu, 11 Sep 2025 18:09:43 +0800 Subject: [PATCH 2/6] fix fla import. --- src/transformers/utils/import_utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index 8155a5f1188e..71a99a46c92a 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -86,6 +86,12 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[ package_version = importlib.metadata.version("pytorch-triton") # pytorch-triton except Exception: package_exists = False + elif pkg_name == "fla": + try: + package = importlib.import_module(pkg_name) + package_version = getattr(package, "__version__", "N/A") + except Exception: + package_exists = False else: # For packages other than "torch", don't attempt the fallback and set as not available package_exists = False From e0c4db9e8105a78725bcf5928ab6f68ebfc561c4 Mon Sep 17 00:00:00 2001 From: bozheng-hit Date: Thu, 11 Sep 2025 18:18:01 +0800 Subject: [PATCH 3/6] fix --- src/transformers/models/qwen3_next/modeling_qwen3_next.py | 4 ++-- src/transformers/models/qwen3_next/modular_qwen3_next.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 08798fc0aa8d..5560276a20cc 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -447,7 +447,7 @@ def torch_chunk_gated_delta_rule( initial_dtype = query.dtype if use_qk_l2norm_in_kernel: head_dim = query.size(-1) - inv_scale = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + inv_scale = query.size(-1) ** -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 = [ @@ -521,7 +521,7 @@ def torch_recurrent_gated_delta_rule( initial_dtype = query.dtype if use_qk_l2norm_in_kernel: head_dim = query.size(-1) - inv_scale = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + inv_scale = query.size(-1) ** -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 = [ diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py index a0e6add1735f..e31ad9d81699 100644 --- a/src/transformers/models/qwen3_next/modular_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py @@ -283,7 +283,7 @@ def torch_chunk_gated_delta_rule( initial_dtype = query.dtype if use_qk_l2norm_in_kernel: head_dim = query.size(-1) - inv_scale = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + inv_scale = query.size(-1)**-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 = [ @@ -357,7 +357,7 @@ def torch_recurrent_gated_delta_rule( initial_dtype = query.dtype if use_qk_l2norm_in_kernel: head_dim = query.size(-1) - inv_scale = torch.rsqrt(torch.tensor(head_dim, device=query.device, dtype=query.dtype)) + inv_scale = query.size(-1)**-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 = [ From 7cac277ed13aa1fa67cf18f0a81fc10a74e46470 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 11 Sep 2025 12:34:11 +0200 Subject: [PATCH 4/6] remove unused attr --- src/transformers/models/qwen3_next/modular_qwen3_next.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py index e31ad9d81699..f42966982bb0 100644 --- a/src/transformers/models/qwen3_next/modular_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py @@ -282,7 +282,6 @@ def torch_chunk_gated_delta_rule( ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: - head_dim = query.size(-1) inv_scale = query.size(-1)**-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 @@ -356,7 +355,6 @@ def torch_recurrent_gated_delta_rule( ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: - head_dim = query.size(-1) inv_scale = query.size(-1)**-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 From ab9f5380e02f26e9a62f14f0038bd3c1794abdcf Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Thu, 11 Sep 2025 12:37:04 +0200 Subject: [PATCH 5/6] fixes --- .../models/qwen3_next/modeling_qwen3_next.py | 4 ++-- .../models/qwen3_next/modular_qwen3_next.py | 6 ++++-- src/transformers/utils/import_utils.py | 18 +++++++----------- 3 files changed, 13 insertions(+), 15 deletions(-) diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 5560276a20cc..990b776f0d57 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -447,7 +447,7 @@ def torch_chunk_gated_delta_rule( initial_dtype = query.dtype if use_qk_l2norm_in_kernel: head_dim = query.size(-1) - inv_scale = query.size(-1) ** -0.5 + 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 = [ @@ -521,7 +521,7 @@ def torch_recurrent_gated_delta_rule( initial_dtype = query.dtype if use_qk_l2norm_in_kernel: head_dim = query.size(-1) - inv_scale = query.size(-1) ** -0.5 + 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 = [ diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py index f42966982bb0..7512f4a3060a 100644 --- a/src/transformers/models/qwen3_next/modular_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py @@ -282,7 +282,8 @@ def torch_chunk_gated_delta_rule( ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: - inv_scale = query.size(-1)**-0.5 + 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 = [ @@ -355,7 +356,8 @@ def torch_recurrent_gated_delta_rule( ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: - inv_scale = query.size(-1)**-0.5 + 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 = [ diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index 71a99a46c92a..d3c08eea5971 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -86,12 +86,6 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[ package_version = importlib.metadata.version("pytorch-triton") # pytorch-triton except Exception: package_exists = False - elif pkg_name == "fla": - try: - package = importlib.import_module(pkg_name) - package_version = getattr(package, "__version__", "N/A") - except Exception: - package_exists = False else: # For packages other than "torch", don't attempt the fallback and set as not available package_exists = False @@ -598,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 From a0bf6f2fbf355cf953406ae705498e173af6ba55 Mon Sep 17 00:00:00 2001 From: bozheng-hit Date: Fri, 12 Sep 2025 17:06:27 +0800 Subject: [PATCH 6/6] strictly align l2norm in Qwen3-Next with FLA implementation. --- .../models/qwen3_next/modeling_qwen3_next.py | 18 ++++++++++-------- .../models/qwen3_next/modular_qwen3_next.py | 18 ++++++++++-------- 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 990b776f0d57..ae2d3664e10e 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -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) ] diff --git a/src/transformers/models/qwen3_next/modular_qwen3_next.py b/src/transformers/models/qwen3_next/modular_qwen3_next.py index 7512f4a3060a..1854a391cfb6 100644 --- a/src/transformers/models/qwen3_next/modular_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modular_qwen3_next.py @@ -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) ]