diff --git a/python/sglang/kernels/ops/attention/fla/l2norm.py b/python/sglang/kernels/ops/attention/fla/l2norm.py index dd1605cd3..0527324a8 100644 --- a/python/sglang/kernels/ops/attention/fla/l2norm.py +++ b/python/sglang/kernels/ops/attention/fla/l2norm.py @@ -51,13 +51,12 @@ def l2norm_fwd_kernel1( # ], # key=["D", "NB"], # ) -@triton.jit +@triton.jit(do_not_specialize=["T"]) def l2norm_fwd_kernel( x, y, eps, - NB: tl.constexpr, - T: tl.constexpr, + T, D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, @@ -91,7 +90,6 @@ def l2norm_fwd( raise RuntimeError("This layer doesn't support feature dim >= 64KB.") if D <= 512: - NB = triton.cdiv(T, 2048) def grid(meta): return (triton.cdiv(T, meta["BT"]),) @@ -100,7 +98,6 @@ def l2norm_fwd( x, y, eps, - NB=NB, T=T, D=D, BD=BD,