[GDN] [Backport] optimize: avoid fla l2-norm recompilation by token count (#31558)

python/sglang/srt/layers/attention/fla/chunk.py

@@ -14,13 +14,20 @@
 from sglang.srt.layers.attention.fla.index import (
     prepare_chunk_indices,
 )
-from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
+from sglang.srt.layers.attention.fla.l2norm import (
+    can_fuse_l2norm_qk,
+    fused_l2norm_qk,
+    l2norm_fwd,
+)
 from sglang.srt.layers.attention.fla.utils import (
     SUPPRESS_LEVEL,
     autocast_custom_fwd,
     input_guard,
     is_intel,
 )
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
 
 if is_intel:
     from sglang.srt.hardware_backend.xpu.kernels.fla.chunk_delta_h import (
@@ -102,12 +109,12 @@ def forward(
         cu_seqlens: Optional[torch.LongTensor] = None,
         use_qk_l2norm_in_kernel: bool = False,
     ):
-        q_orig = q
-        k_orig = k
-
         if use_qk_l2norm_in_kernel:
-            q = l2norm_fwd(q)
-            k = l2norm_fwd(k)
+            if _is_hip and can_fuse_l2norm_qk(q, k):
+                q, k = fused_l2norm_qk(q, k)
+            else:
+                q = l2norm_fwd(q)
+                k = l2norm_fwd(k)
 
         chunk_indices = (
             prepare_chunk_indices(cu_seqlens, CHUNK_SIZE)

python/sglang/srt/layers/attention/fla/l2norm.py

@@ -14,47 +14,89 @@
 BT_LIST = [8, 16, 32, 64, 128]
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["TQ", "TK"])
 def fused_l2norm_qk_kernel(
     q,
     k,
     q_out,
     k_out,
     eps,
-    T: tl.constexpr,
+    TQ,
+    TK,
     D: tl.constexpr,
     BT: tl.constexpr,
     BD: tl.constexpr,
 ):
     i_t = tl.program_id(0)
+
     rows = i_t * BT + tl.arange(0, BT)
     cols = tl.arange(0, BD)
-    mask = (rows[:, None] < T) & (cols[None, :] < D)
-    offsets = rows[:, None] * D + cols[None, :]
+    col_mask = cols < D
 
-    q_block = tl.load(q + offsets, mask=mask, other=0.0).to(tl.float32)
-    q_block = q_block / tl.sqrt(tl.sum(q_block * q_block, axis=1) + eps)[:, None]
-    tl.store(q_out + offsets, q_block.to(q_out.dtype.element_ty), mask=mask)
+    q_mask = (rows[:, None] < TQ) & col_mask[None, :]
+    q_offs = rows[:, None] * D + cols[None, :]
+    b_q = tl.load(q + q_offs, mask=q_mask, other=0.0).to(tl.float32)
+    q_var = tl.sum(b_q * b_q, axis=1)
+    b_q_out = b_q / tl.sqrt(q_var + eps)[:, None]
+    tl.store(q_out + q_offs, b_q_out.to(q_out.dtype.element_ty), mask=q_mask)
 
-    k_block = tl.load(k + offsets, mask=mask, other=0.0).to(tl.float32)
-    k_block = k_block / tl.sqrt(tl.sum(k_block * k_block, axis=1) + eps)[:, None]
-    tl.store(k_out + offsets, k_block.to(k_out.dtype.element_ty), mask=mask)
+    k_mask = (rows[:, None] < TK) & col_mask[None, :]
+    k_offs = rows[:, None] * D + cols[None, :]
+    b_k = tl.load(k + k_offs, mask=k_mask, other=0.0).to(tl.float32)
+    k_var = tl.sum(b_k * b_k, axis=1)
+    b_k_out = b_k / tl.sqrt(k_var + eps)[:, None]
+    tl.store(k_out + k_offs, b_k_out.to(k_out.dtype.element_ty), mask=k_mask)
 
 
-@triton.jit
-def fused_l2norm_qk_kernel1(q, k, q_out, k_out, D, BD: tl.constexpr, eps):
-    row = tl.program_id(0)
+@triton.jit(do_not_specialize=["TQ", "TK"])
+def fused_l2norm_qk_kernel1(
+    q,
+    k,
+    q_out,
+    k_out,
+    D,
+    BD: tl.constexpr,
+    eps,
+    TQ,
+    TK,
+):
+    i_t = tl.program_id(0)
     cols = tl.arange(0, BD)
     mask = cols < D
-    offsets = row * D + cols
-
-    q_block = tl.load(q + offsets, mask=mask, other=0.0).to(tl.float32)
-    q_block = q_block / tl.sqrt(tl.sum(q_block * q_block, axis=0) + eps)
-    tl.store(q_out + offsets, q_block.to(q_out.dtype.element_ty), mask=mask)
-
-    k_block = tl.load(k + offsets, mask=mask, other=0.0).to(tl.float32)
-    k_block = k_block / tl.sqrt(tl.sum(k_block * k_block, axis=0) + eps)
-    tl.store(k_out + offsets, k_block.to(k_out.dtype.element_ty), mask=mask)
+    offs = i_t * D + cols
+
+    if i_t < TQ:
+        b_q = tl.load(q + offs, mask=mask, other=0.0).to(tl.float32)
+        q_var = tl.sum(b_q * b_q, axis=0)
+        b_q_out = b_q / tl.sqrt(q_var + eps)
+        tl.store(q_out + offs, b_q_out.to(q_out.dtype.element_ty), mask=mask)
+
+    if i_t < TK:
+        b_k = tl.load(k + offs, mask=mask, other=0.0).to(tl.float32)
+        k_var = tl.sum(b_k * b_k, axis=0)
+        b_k_out = b_k / tl.sqrt(k_var + eps)
+        tl.store(k_out + offs, b_k_out.to(k_out.dtype.element_ty), mask=mask)
+
+
+def can_fuse_l2norm_qk(q: torch.Tensor, k: torch.Tensor) -> bool:
+    if q.device != k.device or q.dtype != k.dtype:
+        return False
+    if not q.is_cuda or q.dtype not in (torch.bfloat16, torch.float32):
+        return False
+    if q.shape != k.shape:
+        return False
+    q_rows = q.numel() // q.shape[-1]
+    # Wide-shape sweep is positive through D=512 except for the launch-bound
+    # D=512, rows<32 corner. Keep unbenchmarked larger head dims on the existing
+    # two-kernel path.
+    if q.shape[-1] > 512 or (q.shape[-1] == 512 and q_rows < 32):
+        return False
+    if q.stride(-1) != 1 or k.stride(-1) != 1:
+        return False
+    # `view(-1, D)` in the kernel wrapper must be valid without an implicit copy.
+    if not q.is_contiguous() or not k.is_contiguous():
+        return False
+    return True
 
 
 def fused_l2norm_qk(
@@ -63,55 +105,59 @@ def fused_l2norm_qk(
     eps: float = 1e-6,
     output_dtype: Optional[torch.dtype] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    q_shape, k_shape = q.shape, k.shape
+    if not can_fuse_l2norm_qk(q, k):
+        raise ValueError("Incompatible Q/K tensors for fused_l2norm_qk")
+
+    q_shape_og = q.shape
+    k_shape_og = k.shape
     q_flat = q.view(-1, q.shape[-1])
     k_flat = k.view(-1, k.shape[-1])
-    if q_flat.shape != k_flat.shape:
-        raise ValueError(
-            f"Fused Q/K L2Norm requires matching shapes, got {q_shape} and {k_shape}"
-        )
+    TQ, TK, D = q_flat.shape[0], k_flat.shape[0], q_flat.shape[-1]
 
     if output_dtype is None:
-        q_out, k_out = torch.empty_like(q_flat), torch.empty_like(k_flat)
+        q_out = torch.empty_like(q_flat)
+        k_out = torch.empty_like(k_flat)
     else:
         q_out = torch.empty_like(q_flat, dtype=output_dtype)
         k_out = torch.empty_like(k_flat, dtype=output_dtype)
 
-    tokens, dim = q_flat.shape
-    max_fused_size = 65536 // q.element_size()
-    block_dim = min(max_fused_size, triton.next_power_of_2(dim))
-    if dim > block_dim:
+    MAX_FUSED_SIZE = 65536 // q.element_size()
+    BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
+    if D > BD:
         raise RuntimeError("This layer doesn't support feature dim >= 64KB.")
 
-    if dim <= 512:
-        block_tokens = 16
-        fused_l2norm_qk_kernel[(triton.cdiv(tokens, block_tokens),)](
+    if D <= 512:
+        BT = 16
+        fused_l2norm_qk_kernel[(triton.cdiv(max(TQ, TK), BT),)](
             q_flat,
             k_flat,
             q_out,
             k_out,
             eps,
-            T=tokens,
-            D=dim,
-            BT=block_tokens,
-            BD=block_dim,
+            TQ=TQ,
+            TK=TK,
+            D=D,
+            BT=BT,
+            BD=BD,
             num_warps=8,
             num_stages=3,
         )
     else:
-        fused_l2norm_qk_kernel1[(tokens,)](
+        fused_l2norm_qk_kernel1[(max(TQ, TK),)](
             q_flat,
             k_flat,
             q_out,
             k_out,
-            D=dim,
-            BD=block_dim,
             eps=eps,
+            D=D,
+            BD=BD,
+            TQ=TQ,
+            TK=TK,
             num_warps=8,
             num_stages=3,
         )
 
-    return q_out.view(q_shape), k_out.view(k_shape)
+    return q_out.view(q_shape_og), k_out.view(k_shape_og)
 
 
 # @triton.autotune(
@@ -151,13 +197,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,
@@ -191,7 +236,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"]),)
@@ -200,7 +244,6 @@ def grid(meta):
             x,
             y,
             eps,
-            NB=NB,
             T=T,
             D=D,
             BD=BD,

test/registered/jit/benchmark/bench_gdn_l2norm.py

@@ -0,0 +1,56 @@
+import torch
+
+from sglang.jit_kernel.benchmark import marker
+from sglang.srt.layers.attention.fla.l2norm import fused_l2norm_qk, l2norm_fwd
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(
+    est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+register_amd_ci(est_time=8, stage="jit-kernel-benchmark", runner_config="amd")
+
+
+def _run_fused(q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return fused_l2norm_qk(q, k)
+
+
+def _run_separate(
+    q: torch.Tensor, k: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return l2norm_fwd(q), l2norm_fwd(k)
+
+
+FN_MAP = {
+    "fused": _run_fused,
+    "separate": _run_separate,
+}
+
+
+@marker.parametrize("tokens", [15, 16, 17, 257, 1024], [257])
+@marker.parametrize(
+    "local_heads,head_dim",
+    [
+        (2, 128),
+        (4, 128),
+        (8, 128),
+        (16, 128),
+        (8, 256),
+    ],
+)
+@marker.benchmark("impl", ["fused", "separate"])
+def benchmark(tokens: int, local_heads: int, head_dim: int, impl: str):
+    q = torch.randn(tokens, local_heads, head_dim, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(tokens, local_heads, head_dim, dtype=torch.bfloat16, device="cuda")
+    return marker.do_bench(
+        FN_MAP[impl],
+        input_args=(q, k),
+        warmup_iters=80,
+        replay_iters=1200,
+        graph_clone_args=(0, 1),
+        memory_args=(q, k),
+        memory_output=None,
+    )
+
+
+if __name__ == "__main__":
+    benchmark.run()