[AMD] Fuse sigmoid + mul into single Triton kernel for shared expert gating (#27636)

python/sglang/jit_kernel/triton/sigmoid_gate_mul.py

@@ -1,7 +1,22 @@
+"""Fused sigmoid-gate-multiply Triton kernels.
+
+Two variants:
+- ``sigmoid_gate_mul``: element-wise ``x * sigmoid(gate)`` when x and gate
+  have identical shapes.
+- ``sigmoid_gate_mul_broadcast``: broadcast ``x * sigmoid(gate)`` when gate
+  is ``(N, 1)`` and x is ``(N, D)``.
+"""
+
+from __future__ import annotations
+
 import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
 
 @triton.jit
 def _sigmoid_gate_mul_kernel(
@@ -21,9 +36,52 @@ def _sigmoid_gate_mul_kernel(
 
 
 def sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
-    """Compute x * sigmoid(gate) in a single fused kernel."""
+    """Compute ``x * sigmoid(gate)`` in a single fused kernel (same-shape)."""
     out = torch.empty_like(x)
     n = x.numel()
     grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
     _sigmoid_gate_mul_kernel[grid](x, gate, out, n, BLOCK_SIZE=1024)
     return out
+
+
+@triton.jit
+def _sigmoid_gate_mul_broadcast_kernel(
+    out_ptr,
+    gate_ptr,
+    x_ptr,
+    hidden_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    g = tl.load(gate_ptr + row).to(tl.float32)
+    g = tl.sigmoid(g)
+
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < hidden_dim
+    x = tl.load(x_ptr + row * hidden_dim + offs, mask=mask).to(tl.float32)
+    out = x * g
+    tl.store(
+        out_ptr + row * hidden_dim + offs,
+        out.to(x_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def sigmoid_gate_mul_broadcast(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    """Compute ``x * sigmoid(gate)`` where gate is (N, 1) and x is (N, D)."""
+    bs, hidden_dim = x.shape
+    out = torch.empty_like(x)
+    BLOCK_SIZE = triton.next_power_of_2(hidden_dim)
+    max_warps = 16 if _is_hip else 32
+    num_warps = max(
+        min(triton.next_power_of_2(triton.cdiv(hidden_dim, 8 * 32)), max_warps), 4
+    )
+    _sigmoid_gate_mul_broadcast_kernel[(bs,)](
+        out,
+        gate,
+        x,
+        hidden_dim=hidden_dim,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return out

python/sglang/srt/models/qwen2_moe.py

@@ -422,6 +422,13 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
                         True,
                         shared_output,
                     )
+                elif _is_hip:
+                    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+                        sigmoid_gate_mul_broadcast,
+                    )
+
+                    gate = self.shared_expert_gate(hidden_states)
+                    shared_output = sigmoid_gate_mul_broadcast(shared_output, gate)
                 else:
                     shared_output = (
                         F.sigmoid(self.shared_expert_gate(hidden_states))

test/registered/kernels/test_sigmoid_gate_mul.py

@@ -0,0 +1,123 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_amd_ci
+
+register_amd_ci(est_time=4, suite="jit-kernel-unit-test-amd")
+
+DEVICE = "cuda"
+
+
+def reference_sigmoid_gate_mul(x, gate):
+    return x * torch.sigmoid(gate)
+
+
+# ── element-wise variant ──
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (1, 4096),
+        (4, 4096),
+        (8, 4096),
+        (32, 8192),
+        (1, 128),
+        (16, 16384),
+    ],
+)
+def test_sigmoid_gate_mul_correctness(shape, dtype):
+    from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul
+
+    torch.manual_seed(42)
+    x = torch.randn(shape, dtype=dtype, device=DEVICE)
+    gate = torch.randn(shape, dtype=dtype, device=DEVICE)
+
+    ref = reference_sigmoid_gate_mul(x, gate)
+    out = sigmoid_gate_mul(x, gate)
+
+    rtol = 1e-2 if dtype == torch.bfloat16 else 1e-3
+    atol = 2e-2 if dtype == torch.bfloat16 else 1e-3
+    torch.testing.assert_close(out, ref, rtol=rtol, atol=atol)
+
+
+# ── broadcast variant ──
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (1, 4096),
+        (4, 4096),
+        (8, 4096),
+        (32, 8192),
+        (1, 128),
+        (16, 16384),
+    ],
+)
+def test_sigmoid_gate_mul_broadcast_correctness(shape, dtype):
+    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+        sigmoid_gate_mul_broadcast,
+    )
+
+    torch.manual_seed(42)
+    bs, hidden_dim = shape
+    x = torch.randn(shape, dtype=dtype, device=DEVICE)
+    gate = torch.randn(bs, 1, dtype=dtype, device=DEVICE)
+
+    ref = x * torch.sigmoid(gate.float()).to(dtype)
+    out = sigmoid_gate_mul_broadcast(x, gate)
+
+    rtol = 1e-2 if dtype == torch.bfloat16 else 1e-3
+    atol = 2e-2 if dtype == torch.bfloat16 else 1e-3
+    torch.testing.assert_close(out, ref, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("shape", [(4, 4096), (1, 128)])
+def test_sigmoid_gate_mul_broadcast_does_not_modify_inputs(shape):
+    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+        sigmoid_gate_mul_broadcast,
+    )
+
+    torch.manual_seed(42)
+    bs, hidden_dim = shape
+    x = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE)
+    gate = torch.randn(bs, 1, dtype=torch.bfloat16, device=DEVICE)
+    x_orig = x.clone()
+    gate_orig = gate.clone()
+
+    sigmoid_gate_mul_broadcast(x, gate)
+
+    torch.testing.assert_close(x, x_orig, rtol=0, atol=0)
+    torch.testing.assert_close(gate, gate_orig, rtol=0, atol=0)
+
+
+def test_sigmoid_gate_mul_broadcast_output_dtype():
+    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+        sigmoid_gate_mul_broadcast,
+    )
+
+    for dtype in [torch.bfloat16, torch.float16, torch.float32]:
+        x = torch.randn(4, 4096, dtype=dtype, device=DEVICE)
+        gate = torch.randn(4, 1, dtype=dtype, device=DEVICE)
+        out = sigmoid_gate_mul_broadcast(x, gate)
+        assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}"
+
+
+def test_sigmoid_gate_mul_broadcast_contiguous_output():
+    from sglang.jit_kernel.triton.sigmoid_gate_mul import (
+        sigmoid_gate_mul_broadcast,
+    )
+
+    x = torch.randn(4, 4096, dtype=torch.bfloat16, device=DEVICE)
+    gate = torch.randn(4, 1, dtype=torch.bfloat16, device=DEVICE)
+    out = sigmoid_gate_mul_broadcast(x, gate)
+    assert out.is_contiguous()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))