[Kernel] Fuse Qwen2/3-MoE shared-expert sigmoid gate into a Triton kernel

tests/kernels/moe/test_fused_shared_expert_gate.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness tests for `fused_shared_expert_gate`.
+
+Run `pytest tests/kernels/moe/test_fused_shared_expert_gate.py`.
+
+The Triton fusion replaces the three-kernel `F.sigmoid(linear(x)) * out`
+tail of `Qwen2MoeMLP.forward` / `Qwen3MoeMLP.forward`. This test
+parametrizes over real Qwen3-Next-style shapes (`K=2048`, hidden=2048)
+plus a smaller config and mask-boundary token counts (`N=1`, `N=7`) to
+exercise the within-block masking path.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm.model_executor.layers.fused_moe.shared_expert_gate import (
+    fused_shared_expert_gate,
+)
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="fused_shared_expert_gate requires a Triton-capable GPU (CUDA or ROCm).",
+)
+
+
+def _reference(
+    x: torch.Tensor, weight: torch.Tensor, out: torch.Tensor
+) -> torch.Tensor:
+    return F.sigmoid(F.linear(x, weight)) * out
+
+
+@pytest.mark.parametrize("num_tokens", [1, 7, 33, 1024, 7177, 8192])
+@pytest.mark.parametrize("hidden_size", [1024, 2048])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_fused_shared_expert_gate_matches_reference(
+    num_tokens: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+):
+    torch.manual_seed(0)
+    device = "cuda"
+    x = torch.randn((num_tokens, hidden_size), device=device, dtype=dtype)
+    weight = torch.randn((1, hidden_size), device=device, dtype=dtype)
+    out = torch.randn((num_tokens, hidden_size), device=device, dtype=dtype)
+
+    expected = _reference(x, weight, out)
+    actual = fused_shared_expert_gate(x, weight, out)
+
+    assert actual.dtype == expected.dtype
+    assert actual.shape == expected.shape
+    # Tolerance matches the existing pattern in tests/kernels/ for bf16
+    # row-fused ops; fp16 comfortably fits the same bound.
+    torch.testing.assert_close(actual, expected, atol=3.125e-2, rtol=2e-2)
+
+
+def test_fused_shared_expert_gate_fallback_on_unsupported_shape():
+    """A non-2D `x` must fall back to the PyTorch reference."""
+    torch.manual_seed(0)
+    device = "cuda"
+    dtype = torch.bfloat16
+    # 3D input -- the Triton kernel is 2D-only, so the wrapper must
+    # dispatch to the PyTorch reference path. `F.sigmoid(F.linear(...))`
+    # broadcasts the `[B, N, 1]` gate against the `[B, N, K]` output, so
+    # the reference expression is well-defined and we can compare equality.
+    x = torch.randn((2, 16, 2048), device=device, dtype=dtype)
+    out = torch.randn((2, 16, 2048), device=device, dtype=dtype)
+    weight = torch.randn((1, 2048), device=device, dtype=dtype)
+
+    expected = _reference(x, weight, out)
+    actual = fused_shared_expert_gate(x, weight, out)
+
+    assert actual.shape == expected.shape
+    torch.testing.assert_close(actual, expected, atol=3.125e-2, rtol=2e-2)

vllm/model_executor/layers/fused_moe/shared_expert_gate.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused Triton kernel for the Qwen2/3-MoE shared-expert sigmoid gate.
+
+Replaces the three-kernel `F.sigmoid(linear(x)) * out` tail of
+`Qwen2MoeMLP.forward` / `Qwen3MoeMLP.forward` with a single row-fused
+pass that removes the two HBM-resident intermediates.
+
+The wrapper is shape-guarded and silently falls back to the PyTorch
+reference (`F.sigmoid(F.linear(x, weight)) * out`) for any input shape
+this kernel does not handle, so it is safe to use behind the existing
+`expert_gate` call sites without further checks.
+"""
+
+import torch
+import torch.nn.functional as F
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _fused_shared_expert_gate_kernel(
+    x_ptr,
+    weight_ptr,
+    out_ptr,
+    y_ptr,
+    K: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_K)
+    mask = offsets < K
+
+    x = tl.load(x_ptr + row * K + offsets, mask=mask, other=0.0).to(tl.float32)
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+    gate = tl.sigmoid(tl.sum(x * weight, axis=0))
+
+    out = tl.load(out_ptr + row * K + offsets, mask=mask, other=0.0).to(tl.float32)
+    tl.store(y_ptr + row * K + offsets, out * gate, mask=mask)
+
+
+def fused_shared_expert_gate(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Compute ``F.sigmoid(F.linear(x, weight)) * out`` in a single pass.
+
+    Specialised for a one-row gate weight (``weight.shape == [1, K]``), as
+    produced by ``ReplicatedLinear(hidden_size, 1)`` in the Qwen2/3-MoE
+    shared-expert blocks. For any other shape, the function falls back to
+    the PyTorch reference so callers can use it unconditionally.
+
+    Args:
+        x: Shared-expert input, shape ``[N, K]``.
+        weight: Gate weight, shape ``[1, K]``.
+        out: Shared-expert MLP output, shape ``[N, K]``.
+
+    Returns:
+        ``[N, K]`` tensor equal to ``sigmoid(x @ weight.T) * out`` within
+        bf16/fp16 tolerance.
+    """
+    if (
+        x.ndim != 2
+        or out.ndim != 2
+        or weight.ndim != 2
+        or weight.shape[0] != 1
+        or x.shape != out.shape
+        or weight.shape[1] != x.shape[1]
+    ):
+        return F.sigmoid(F.linear(x, weight)) * out
+
+    y = torch.empty_like(out)
+    _fused_shared_expert_gate_kernel[(x.shape[0],)](
+        x,
+        weight,
+        out,
+        y,
+        K=x.shape[1],
+        BLOCK_K=triton.next_power_of_2(x.shape[1]),
+        num_warps=8,
+    )
+    return y

vllm/model_executor/models/qwen2_moe.py

@@ -30,7 +30,6 @@
 from typing import Any
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 from transformers import Qwen2MoeConfig
 
@@ -44,6 +43,9 @@
     FusedMoE,
     fused_moe_make_expert_params_mapping,
 )
+from vllm.model_executor.layers.fused_moe.shared_expert_gate import (
+    fused_shared_expert_gate,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -117,7 +119,7 @@ def forward(self, x):
         out, _ = self.down_proj(out)
 
         if self.expert_gate is not None:
-            out = F.sigmoid(self.expert_gate(x)[0]) * out
+            out = fused_shared_expert_gate(x, self.expert_gate.weight, out)
 
         return out
 

vllm/model_executor/models/qwen3_moe.py

@@ -29,7 +29,6 @@
 from typing import Any
 
 import torch
-import torch.nn.functional as F
 from torch import nn
 
 from vllm.compilation.decorators import support_torch_compile
@@ -47,6 +46,9 @@
     FusedMoE,
     fused_moe_make_expert_params_mapping,
 )
+from vllm.model_executor.layers.fused_moe.shared_expert_gate import (
+    fused_shared_expert_gate,
+)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
@@ -129,7 +131,7 @@ def forward(self, x):
         out, _ = self.down_proj(out)
 
         if self.expert_gate is not None:
-            out = F.sigmoid(self.expert_gate(x)[0]) * out
+            out = fused_shared_expert_gate(x, self.expert_gate.weight, out)
 
         return out