From 22c7285a2680946afae6e867bf623b68bbba5886 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Thu, 11 Jun 2026 17:05:05 +0800 Subject: [PATCH] [AMD] Fuse sigmoid + mul attention output gate into single Triton kernel (#27630) --- .../jit_kernel/tests/test_sigmoid_gate_mul.py | 81 +++++++++++++++++++ .../jit_kernel/triton/sigmoid_gate_mul.py | 29 +++++++ python/sglang/srt/models/qwen3_5.py | 11 ++- python/sglang/srt/models/qwen3_next.py | 13 ++- 4 files changed, 130 insertions(+), 4 deletions(-) create mode 100644 python/sglang/jit_kernel/tests/test_sigmoid_gate_mul.py create mode 100644 python/sglang/jit_kernel/triton/sigmoid_gate_mul.py diff --git a/python/sglang/jit_kernel/tests/test_sigmoid_gate_mul.py b/python/sglang/jit_kernel/tests/test_sigmoid_gate_mul.py new file mode 100644 index 000000000..83ca811b8 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_sigmoid_gate_mul.py @@ -0,0 +1,81 @@ +import sys + +import pytest +import torch + +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=4, suite="base-b-kernel-unit-1-gpu-large") +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) + + +@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) + + +@pytest.mark.parametrize("shape", [(4, 4096), (1, 128)]) +def test_sigmoid_gate_mul_does_not_modify_inputs(shape): + from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul + + torch.manual_seed(42) + x = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE) + gate = torch.randn(shape, dtype=torch.bfloat16, device=DEVICE) + x_orig = x.clone() + gate_orig = gate.clone() + + sigmoid_gate_mul(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_output_dtype(): + from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul + + for dtype in [torch.bfloat16, torch.float16, torch.float32]: + x = torch.randn(4, 4096, dtype=dtype, device=DEVICE) + gate = torch.randn(4, 4096, dtype=dtype, device=DEVICE) + out = sigmoid_gate_mul(x, gate) + assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}" + + +def test_sigmoid_gate_mul_contiguous_output(): + from sglang.jit_kernel.triton.sigmoid_gate_mul import sigmoid_gate_mul + + x = torch.randn(4, 4096, dtype=torch.bfloat16, device=DEVICE) + gate = torch.randn(4, 4096, dtype=torch.bfloat16, device=DEVICE) + out = sigmoid_gate_mul(x, gate) + assert out.is_contiguous() + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/python/sglang/jit_kernel/triton/sigmoid_gate_mul.py b/python/sglang/jit_kernel/triton/sigmoid_gate_mul.py new file mode 100644 index 000000000..e05cb801e --- /dev/null +++ b/python/sglang/jit_kernel/triton/sigmoid_gate_mul.py @@ -0,0 +1,29 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _sigmoid_gate_mul_kernel( + x_ptr, + gate_ptr, + out_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32) + g = tl.load(gate_ptr + offsets, mask=mask).to(tl.float32) + out = x * tl.sigmoid(g) + tl.store(out_ptr + offsets, out.to(x_ptr.dtype.element_ty), mask=mask) + + +def sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Compute x * sigmoid(gate) in a single fused kernel.""" + 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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 00801ec1b..2c059217c 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -970,8 +970,15 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): attn_output = self.attn(q, k, v, forward_batch) if self.attn_output_gate: - gate = torch.sigmoid(gate) - attn_output = attn_output * gate + if _is_hip: + from sglang.jit_kernel.triton.sigmoid_gate_mul import ( + sigmoid_gate_mul, + ) + + attn_output = sigmoid_gate_mul(attn_output, gate) + else: + gate = torch.sigmoid(gate) + attn_output = attn_output * gate output, _ = self.o_proj(attn_output) return output diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index f8fc8be00..3d1231e4a 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -49,6 +49,7 @@ from sglang.srt.utils import ( cpu_has_amx_support, is_cpu, is_cuda, + is_hip, is_npu, make_layers, set_weight_attrs, @@ -65,6 +66,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( ) _is_cuda = is_cuda() +_is_hip = is_hip() _is_npu = is_npu() _is_cpu = is_cpu() _is_amx_available = cpu_has_amx_support() @@ -824,8 +826,15 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): attn_output = self.attn(q, k, v, forward_batch) if self.attn_output_gate: - gate = torch.sigmoid(gate) - attn_output = attn_output * gate + if _is_hip: + from sglang.jit_kernel.triton.sigmoid_gate_mul import ( + sigmoid_gate_mul, + ) + + attn_output = sigmoid_gate_mul(attn_output, gate) + else: + gate = torch.sigmoid(gate) + attn_output = attn_output * gate output, _ = self.o_proj(attn_output) return output