optimize: avoid fla l2-norm recompilation by token count (#31558)

This commit is contained in:
Mick
2026-07-18 07:58:36 +08:00
committed by GitHub
parent 01a96720c6
commit 42a058c760
@@ -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,