[AMD] Fuse sigmoid + mul attention output gate into single Triton kernel (#27630)

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"]))

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

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

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