minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-07-18 07:58:36 +08:00
committed by GitHub
parent 01a96720c6
commit 42a058c760
1 changed files with 2 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fla/l2norm.py
+2 -5
View File
@@ -51,13 +51,12 @@ def l2norm_fwd_kernel1(
# ], # ],
# key=["D", "NB"], # key=["D", "NB"],
# ) # )
@triton.jit @triton.jit(do_not_specialize=["T"])
def l2norm_fwd_kernel( def l2norm_fwd_kernel(
x, x,
y, y,
eps, eps,
NB: tl.constexpr, T,
T: tl.constexpr,
D: tl.constexpr, D: tl.constexpr,
BT: tl.constexpr, BT: tl.constexpr,
BD: tl.constexpr, BD: tl.constexpr,
@@ -91,7 +90,6 @@ def l2norm_fwd(
raise RuntimeError("This layer doesn't support feature dim >= 64KB.") raise RuntimeError("This layer doesn't support feature dim >= 64KB.")
if D <= 512: if D <= 512:
NB = triton.cdiv(T, 2048)
def grid(meta): def grid(meta):
return (triton.cdiv(T, meta["BT"]),) return (triton.cdiv(T, meta["BT"]),)
@@ -100,7 +98,6 @@ def l2norm_fwd(
x, x,
y, y,
eps, eps,
NB=NB,
T=T, T=T,
D=D, D=D,
BD=BD, BD=BD,