diff --git a/python/sglang/jit_kernel/triton/sigmoid_gate_mul.py b/python/sglang/jit_kernel/triton/sigmoid_gate_mul.py index e05cb801e..5696a54c7 100644 --- a/python/sglang/jit_kernel/triton/sigmoid_gate_mul.py +++ b/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 diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 40b2b108d..324e8678d 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/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)) diff --git a/test/registered/kernels/test_sigmoid_gate_mul.py b/test/registered/kernels/test_sigmoid_gate_mul.py new file mode 100644 index 000000000..2b74fd742 --- /dev/null +++ b/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"]))