From 484b4e7c65f7776b9afb87a081a9f1ecd5bd064e Mon Sep 17 00:00:00 2001
From: Kamleecoder <501362869@qq.com>
Date: Wed, 26 Aug 2026 22:14:27 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20L2Norm=20Triton=20kernel?=
 =?UTF-8?q?=20=E5=9C=A8=E5=8A=A8=E6=80=81=20token=20=E6=95=B0=E5=9C=BA?=
 =?UTF-8?q?=E6=99=AF=E4=B8=8B=E5=9B=A0=E8=BE=93=E5=85=A5=20T/NB=20?=
 =?UTF-8?q?=E5=8F=98=E5=8C=96=E5=8F=8D=E5=A4=8D=E8=A7=A6=E5=8F=91=20autotu?=
 =?UTF-8?q?ne=EF=BC=8C=E5=AF=BC=E8=87=B4=E8=AE=AD=E7=BB=83=20iteration=20?=
 =?UTF-8?q?=E8=80=97=E6=97=B6=E5=BC=82=E5=B8=B8=E5=A2=9E=E9=95=BF=E7=9A=84?=
 =?UTF-8?q?=E9=97=AE=E9=A2=98=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 mindspeed/ops/triton/l2norm.py | 20 ++++++--------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/mindspeed/ops/triton/l2norm.py b/mindspeed/ops/triton/l2norm.py
index 0050bce3..280aee6a 100644
--- a/mindspeed/ops/triton/l2norm.py
+++ b/mindspeed/ops/triton/l2norm.py
@@ -82,18 +82,17 @@ def l2norm_bwd_kernel1(
         for num_warps in [1, 2, 4, 8, 16]
         for BT in BT_LIST
     ],
-    key=['D', 'NB']
+    key=['D']
 )
-@triton.jit
+@triton.jit(do_not_specialize=['T'])
 def l2norm_fwd_kernel(
     x,
     y,
     rstd,
     eps,
-    T: tl.constexpr,
+    T,
     D: tl.constexpr,
     BD: tl.constexpr,
-    NB: tl.constexpr,
     BT: tl.constexpr,
     bt_size,
 ):
@@ -119,19 +118,18 @@ def l2norm_fwd_kernel(
         for num_warps in [1, 2, 4, 8, 16]
         for BT in BT_LIST
     ],
-    key=['D', 'NB']
+    key=['D']
 )
-@triton.jit
+@triton.jit(do_not_specialize=['T'])
 def l2norm_bwd_kernel(
     y,
     rstd,
     dy,
     dx,
     eps,
-    T: tl.constexpr,
+    T,
     D: tl.constexpr,
     BD: tl.constexpr,
-    NB: tl.constexpr,
     BT: tl.constexpr,
     bt_size,
 ):
@@ -187,7 +185,6 @@ def l2norm_fwd(
 
     rstd = torch.empty((T,), dtype=torch.float32, device=x.device)
     if D <= 512:
-        NB = triton.cdiv(T, 2048)
         bt_size = 32
 
         def grid(meta):
@@ -202,7 +199,6 @@ def grid(meta):
             T=T,
             D=D,
             BD=BD,
-            NB=NB,
             bt_size=bt_size,
         )
     else:
@@ -235,9 +231,7 @@ def l2norm_bwd(
     BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
     if D > BD:
         raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
-
     if D <= 512:
-        NB = triton.cdiv(T, 2048)
         bt_size = 40
         l2norm_bwd_kernel[(bt_size,)](
             y=y,
@@ -248,7 +242,6 @@ def l2norm_bwd(
             T=T,
             D=D,
             BD=BD,
-            NB=NB,
             bt_size=bt_size,
         )
     else:
@@ -261,7 +254,6 @@ def l2norm_bwd(
             D=D,
             BD=BD,
         )
-
     return dx.view(y_shape_og)