[AMD] Fuse shared_expert_gate GEMV into the MoE append kernel (HIP/aiter) (#28666)

Co-authored-by: sogalin_codegen <39478626+sogalin@users.noreply.github.com>

python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py

@@ -1342,6 +1342,8 @@ def _fused_append_shared_experts_with_weights_kernel(
     shared_weights_ptr,
     out_ids_ptr,
     out_weights_ptr,
+    hidden_ptr,
+    wgate_ptr,
     N_BASE,
     scale,
     K: tl.constexpr,
@@ -1349,6 +1351,9 @@ def _fused_append_shared_experts_with_weights_kernel(
     BLOCK_K: tl.constexpr,
     BLOCK_S: tl.constexpr,
     APPLY_SIGMOID: tl.constexpr,
+    FUSE_GATE: tl.constexpr,
+    HIDDEN: tl.constexpr,
+    BLOCK_H: tl.constexpr,
 ):
     pid = tl.program_id(0)
 
@@ -1366,11 +1371,19 @@ def _fused_append_shared_experts_with_weights_kernel(
     offs_s = tl.arange(0, BLOCK_S)
     mask_s = offs_s < S
     shared_ids = tl.cast(N_BASE + offs_s, ids.dtype)
+    if FUSE_GATE:
+        offs_h = tl.arange(0, BLOCK_H)
+        mask_h = offs_h < HIDDEN
+        h = tl.load(hidden_ptr + pid * HIDDEN + offs_h, mask=mask_h, other=0.0).to(
+            tl.float32
+        )
+        w = tl.load(wgate_ptr + offs_h, mask=mask_h, other=0.0).to(tl.float32)
+        logit = tl.sum(h * w)
+        shared_val = tl.sigmoid(logit) * scale
+        shared_ws = tl.zeros((BLOCK_S,), dtype=tl.float32) + shared_val
+    else:
         shared_ws = tl.load(shared_weights_ptr + pid * S + offs_s, mask=mask_s)
         if APPLY_SIGMOID:
-        # Fuse sigmoid(shared_gate) + dtype upcast (+ optional 1/ep_size scale)
-        # in-register so the raw bf16 logits stream straight into the fp32
-        # output, eliminating the standalone sigmoid and bf16->fp32 copy kernels.
             shared_ws = tl.sigmoid(shared_ws.to(tl.float32)) * scale
 
     tl.store(out_ids_ptr + out_row_ptr + K + offs_s, shared_ids, mask=mask_s)
@@ -1384,23 +1397,51 @@ def fused_append_shared_experts_with_weights(
     num_fused_shared_experts,
     N=None,
     apply_sigmoid=False,
+    fuse_gate=False,
+    hidden_states=None,
+    gate_weight=None,
     scale=1.0,
 ):
     """Like fused_append_shared_experts but accepts per-token shared weights tensor.
 
-    When ``apply_sigmoid`` is True, ``shared_weights`` are treated as raw gate
+    Two optional in-kernel fusions are supported (both default off → legacy
-    logits: the kernel applies ``sigmoid`` (in fp32) and the optional ``scale``
+    behavior is preserved byte-for-byte):
-    in-register, so the caller can skip the separate ``sigmoid`` activation and
+
-    the bf16->fp32 cast. When False the legacy behavior is preserved exactly.
+    - ``apply_sigmoid=True``: ``shared_weights`` are treated as raw gate logits;
+      the kernel applies ``sigmoid`` (in fp32) and the optional ``scale``
+      in-register, so the caller can skip the separate ``sigmoid`` activation
+      and the bf16->fp32 cast.
+    - ``fuse_gate=True``: the shared_expert_gate GEMV
+      (``hidden_states @ gate_weight.T``) + sigmoid + ``scale`` are computed
+      *inside* the kernel, eliminating the standalone gate GEMM launch.
+      ``shared_weights`` is ignored; ``hidden_states`` ([M, HIDDEN]) and
+      ``gate_weight`` ([1, HIDDEN] or [HIDDEN]) must be provided. This subsumes
+      ``apply_sigmoid`` (the sigmoid is intrinsic), so the two are mutually
+      exclusive.
     """
+    assert not (
+        fuse_gate and apply_sigmoid
+    ), "fuse_gate already applies sigmoid in-kernel; do not also set apply_sigmoid"
     assert N is not None, "N (shared expert base id) must be provided"
     m, k = topk_ids.shape
     s = int(num_fused_shared_experts)
     if s <= 0:
         return topk_ids, topk_weights
 
-    # When fusing sigmoid in-kernel, keep the raw logits dtype (the kernel emits
+    if fuse_gate:
-    # fp32 directly); otherwise match the output weight dtype as before.
+        assert (
+            hidden_states is not None and gate_weight is not None
+        ), "fuse_gate=True requires hidden_states and gate_weight"
+        hidden_arg = hidden_states.contiguous()
+        wgate_arg = gate_weight.reshape(-1).contiguous()
+        hidden_dim = hidden_arg.shape[1]
+        block_h = triton.next_power_of_2(hidden_dim)
+        shared_arg = topk_weights
+        num_warps = 8
+    else:
+        # When fusing sigmoid in-kernel (apply_sigmoid), keep the raw logits
+        # dtype (the kernel emits fp32 directly); otherwise match the output
+        # weight dtype as before.
         shared_weights_2d = (
             shared_weights if apply_sigmoid else shared_weights.to(topk_weights.dtype)
         )
@@ -1408,7 +1449,13 @@ def fused_append_shared_experts_with_weights(
             shared_weights_2d = shared_weights_2d.unsqueeze(-1)
         if shared_weights_2d.shape[1] < s:
             shared_weights_2d = shared_weights_2d.expand(m, s)
-    shared_weights_2d = shared_weights_2d.contiguous()
+        shared_arg = shared_weights_2d.contiguous()
+        # hidden_ptr / wgate_ptr are unused; pass placeholders.
+        hidden_arg = topk_weights
+        wgate_arg = topk_weights
+        hidden_dim = 1
+        block_h = 1
+        num_warps = 1
 
     out_ids = torch.empty((m, k + s), dtype=topk_ids.dtype, device=topk_ids.device)
     out_weights = torch.empty(
@@ -1421,9 +1468,11 @@ def fused_append_shared_experts_with_weights(
     _fused_append_shared_experts_with_weights_kernel[(m,)](
         topk_ids,
         topk_weights,
-        shared_weights_2d,
+        shared_arg,
         out_ids,
         out_weights,
+        hidden_arg,
+        wgate_arg,
         N_BASE=N,
         scale=scale,
         K=k,
@@ -1431,6 +1480,9 @@ def fused_append_shared_experts_with_weights(
         BLOCK_K=block_k,
         BLOCK_S=block_s,
         APPLY_SIGMOID=apply_sigmoid,
-        num_warps=1,
+        FUSE_GATE=fuse_gate,
+        HIDDEN=hidden_dim,
+        BLOCK_H=block_h,
+        num_warps=num_warps,
     )
     return out_ids, out_weights

python/sglang/srt/models/qwen2_moe.py

@@ -394,33 +394,63 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             return F.sigmoid(shared_logits) * scale, 1.0
         return shared_logits, scale
 
+    def _shared_expert_scale(self) -> float:
+        """1/ep_size pre-scale for the fused shared-expert routing weight.
+
+        Mirrors the scaling applied in _get_shared_expert_weights; see that
+        method for the allreduce-EP rationale.
+        """
+        moe_ep_size = get_parallel().moe_ep_size
+        if moe_ep_size > 1 and not is_deepep_class_backend():
+            return 1.0 / float(moe_ep_size)
+        return 1.0
+
     def _append_shared_to_topk_output(
         self,
         topk_output: StandardTopKOutput,
         hidden_states: torch.Tensor,
     ) -> StandardTopKOutput:
         """Append shared expert ids and weights to topk output before fused MoE."""
-        if not self.enable_shared_expert_fusion:
+        if not self.enable_shared_expert_fusion or self.shared_expert_gate is None:
             return topk_output
-        shared = self._get_shared_expert_weights(hidden_states)
-        if shared is None:
-            return topk_output
-        shared_weights, shared_scale = shared
 
         from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
             fused_append_shared_experts_with_weights,
         )
 
-        # AITER returns raw logits + scale for in-kernel sigmoid fusion; CUDA
+        if _use_aiter:
-        # returns pre-activated weights (scale already folded in) → no fusion.
+            # HIP/aiter: fuse the shared_expert_gate GEMV + sigmoid + scale into
-        fused_topk_ids, fused_topk_weights = fused_append_shared_experts_with_weights(
+            # the append kernel, eliminating the standalone gate GEMM launch.
+            # This subsumes the sigmoid-only fusion: there is no separate gate
+            # GEMM and no _get_shared_expert_weights call on this path.
+            fused_topk_ids, fused_topk_weights = (
+                fused_append_shared_experts_with_weights(
+                    topk_output.topk_ids,
+                    topk_output.topk_weights,
+                    None,
+                    self.num_fused_shared_experts,
+                    N=self.num_experts,
+                    fuse_gate=True,
+                    hidden_states=hidden_states,
+                    gate_weight=self.shared_expert_gate.weight,
+                    scale=self._shared_expert_scale(),
+                )
+            )
+        else:
+            # CUDA: _get_shared_expert_weights returns pre-activated weights
+            # (sigmoid + scale already folded in) → legacy append, no fusion.
+            shared = self._get_shared_expert_weights(hidden_states)
+            if shared is None:
+                return topk_output
+            shared_weights, _ = shared
+            fused_topk_ids, fused_topk_weights = (
+                fused_append_shared_experts_with_weights(
                     topk_output.topk_ids,
                     topk_output.topk_weights,
                     shared_weights,
                     self.num_fused_shared_experts,
                     N=self.num_experts,
-            apply_sigmoid=_use_aiter,
+                )
-            scale=shared_scale,
             )
         return StandardTopKOutput(
             topk_weights=fused_topk_weights,

test/registered/amd/accuracy/mi35x/test_qwen35_mxfp4_eval_mi35x.py

@@ -0,0 +1,249 @@
+"""MI35x Qwen3.5-397B-A17B MXFP4 GSM8K accuracy + serving-perf test (nightly).
+
+Covers the two AMD MXFP4 checkpoints, which differ ONLY in whether the shared
+expert is quantized -- and therefore in whether the shared-expert fusion
+``fuse_gate`` serving path (qwen2_moe.py ``_use_aiter`` branch + the fuse_gate
+append kernel) actually runs:
+
+  * ``Qwen3.5-397B-A17B-MXFP4``     -- shared expert EXCLUDED from quant (BF16),
+    so ``can_fuse_shared_expert`` returns False -> fusion OFF -> eager path.
+    This is the reference-accuracy checkpoint (gate: gsm8k > 0.91).
+  * ``Qwen3.5-397B-A17B-MoE-MXFP4`` -- shared expert IS quantized, so fusion is
+    ON and the aiter ``fuse_gate`` GEMV-in-kernel path runs end-to-end. No
+    per-PR CI exercises this path (the kernel unit test only checks the kernel
+    in isolation); here we require it to *run through* cleanly e2e.
+
+Launch config mirrors ~/run_qwen3.5_mxfp4_perf.sh (tp=2, aiter + unified attn +
+flydsl, allreduce fusion). MXFP4 is MI35x-only, so this is not registered on
+mi30x.
+
+Registry: nightly-amd-accuracy-8-gpu-mi35x-qwen35-mxfp4 suite
+"""
+
+import os
+import unittest
+from types import SimpleNamespace
+from typing import List
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.nightly_bench_utils import BenchmarkResult
+from sglang.test.nightly_utils import NightlyBenchmarkRunner
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    _parse_int_list_env,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+register_amd_ci(
+    est_time=14400, suite="nightly-amd-accuracy-8-gpu-mi35x-qwen35-mxfp4", nightly=True
+)
+
+# Local checkpoints on the MI35x runner; overridable for other hosts.
+MXFP4_MODEL_PATH = os.environ.get(
+    "QWEN35_MXFP4_MODEL_PATH", "/data/amd/Qwen3.5-397B-A17B-MXFP4"
+)
+MOE_MXFP4_MODEL_PATH = os.environ.get(
+    "QWEN35_MOE_MXFP4_MODEL_PATH", "/data/amd/Qwen3.5-397B-A17B-MoE-MXFP4"
+)
+
+SERVER_LAUNCH_TIMEOUT = 3600
+BENCH_TIMEOUT = 5400
+
+# Accuracy gate for the reference (non-fused, BF16-shared-expert) checkpoint.
+MXFP4_ACC_THRESHOLD = 0.91
+
+# Server flags from ~/run_qwen3.5_mxfp4_perf.sh (the reference dev config).
+COMMON_ARGS = [
+    "--tp",
+    "2",
+    "--attention-backend",
+    "aiter",
+    "--trust-remote-code",
+    "--chunked-prefill-size",
+    "32768",
+    "--model-loader-extra-config",
+    '{"enable_multithread_load": true}',
+    "--watchdog-timeout",
+    "1200",
+    "--mem-fraction-static",
+    "0.9",
+    "--disable-radix-cache",
+    "--enable-aiter-allreduce-fusion",
+    "--max-running-requests",
+    "512",
+    "--page-size",
+    "16",
+]
+
+COMMON_ENV = {
+    "SGLANG_USE_AITER": "1",
+    "SGLANG_USE_AITER_UNIFIED_ATTN": "1",
+    "AITER_FLYDSL_FORCE": "1",
+}
+
+
+def _run_gsm8k(
+    base_url: str, model: str, num_examples: int, max_tokens: int = 2048
+) -> dict:
+    """Few-shot GSM8K against a running server; returns run_eval metrics.
+
+    ``metrics["score"]`` is the accuracy in [0, 1].
+    """
+    requests.get(base_url + "/flush_cache")
+    args = SimpleNamespace(
+        base_url=base_url,
+        model=model,
+        eval_name="gsm8k",
+        api="completion",
+        max_tokens=max_tokens,
+        num_examples=num_examples,
+        num_threads=256,
+    )
+    return run_eval(args)
+
+
+def _generate_perf_report(results: List[BenchmarkResult]) -> str:
+    """Compact markdown perf table (skips a leading warmup duplicate)."""
+    header = results[0].model_path
+    if results[0].run_name and results[0].run_name != "default":
+        header += f" ({results[0].run_name})"
+    header += f" [{os.getenv('GPU_CONFIG', 'MI35x')}]"
+
+    summary = f"### {header}\n"
+    summary += "| batch size | input len | latency (s) | input tput (tok/s) | output tput (tok/s) | ITL (ms) |\n"
+    summary += "| ---------- | --------- | ----------- | ------------------ | ------------------- | -------- |\n"
+
+    report = (
+        results[1:]
+        if len(results) > 1 and results[0].batch_size == results[1].batch_size
+        else results
+    )
+    for r in report:
+        itl = 1 / (r.output_throughput / r.batch_size) * 1000
+        summary += (
+            f"| {r.batch_size} | {r.input_len} | {r.latency:.2f} | "
+            f"{r.input_throughput:.2f} | {r.output_throughput:.2f} | {itl:.2f} |\n"
+        )
+    return summary
+
+
+def _run_perf(model_path: str, variant: str, base_url: str) -> None:
+    """Run the nightly serving benchmark for ``model_path`` with COMMON_ENV set."""
+    batch_sizes = _parse_int_list_env("NIGHTLY_BATCH_SIZES", "1,8,16,64")
+    input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
+    output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
+
+    runner = NightlyBenchmarkRunner(
+        f"performance_profiles_{variant}", variant, base_url
+    )
+    runner.setup_profile_directory()
+    runner.full_report = f"## {variant}\n"
+
+    old_env = {}
+    for key, value in COMMON_ENV.items():
+        old_env[key] = os.environ.get(key)
+        os.environ[key] = value
+    try:
+        results, success = runner.run_benchmark_for_model(
+            model_path=model_path,
+            batch_sizes=batch_sizes,
+            input_lens=input_lens,
+            output_lens=output_lens,
+            other_args=COMMON_ARGS,
+            variant=variant,
+            extra_bench_args=["--trust-remote-code"],
+            enable_profile=False,
+            timeout=BENCH_TIMEOUT,
+        )[:2]
+        if results:
+            runner.full_report += _generate_perf_report(results) + "\n"
+        assert success, f"Perf benchmark failed for {model_path} on MI35x"
+    finally:
+        for key, value in old_env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+        runner.write_final_report()
+
+
+class TestQwen35Mxfp4MI35x(CustomTestCase):
+    """Non-fused reference checkpoint (shared expert BF16): accuracy-gated."""
+
+    base_url = DEFAULT_URL_FOR_TEST
+
+    def test_a_gsm8k(self):
+        """GSM8K accuracy must clear the reference gate (> 0.91)."""
+        env = os.environ.copy()
+        env.update(COMMON_ENV)
+        process = popen_launch_server(
+            MXFP4_MODEL_PATH,
+            self.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=COMMON_ARGS,
+            env=env,
+        )
+        try:
+            metrics = _run_gsm8k(self.base_url, MXFP4_MODEL_PATH, num_examples=1319)
+            print(f"[{MXFP4_MODEL_PATH}] {metrics=}")
+            if is_in_ci():
+                write_github_step_summary(
+                    f"### gsm8k accuracy ({MXFP4_MODEL_PATH})\n"
+                    f'score={metrics["score"]:.3f} '
+                    f"(threshold {MXFP4_ACC_THRESHOLD})\n"
+                )
+            self.assertGreater(metrics["score"], MXFP4_ACC_THRESHOLD)
+        finally:
+            kill_process_tree(process.pid)
+
+    def test_b_perf(self):
+        """Serving performance benchmark."""
+        _run_perf(MXFP4_MODEL_PATH, "qwen35-mxfp4-mi35x", self.base_url)
+
+
+class TestQwen35MoeMxfp4MI35x(CustomTestCase):
+    """Fused checkpoint (shared expert MXFP4 -> fuse_gate path): run-through smoke."""
+
+    base_url = DEFAULT_URL_FOR_TEST
+
+    def test_a_gsm8k_runthrough(self):
+        """Exercise the aiter fuse_gate serving path e2e; require it to run and
+        produce valid (non-degenerate) output. No strict accuracy gate -- the
+        point is that the fuse_gate path serves cleanly end-to-end (the coverage
+        gap no per-PR CI test fills)."""
+        env = os.environ.copy()
+        env.update(COMMON_ENV)
+        process = popen_launch_server(
+            MOE_MXFP4_MODEL_PATH,
+            self.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=COMMON_ARGS,
+            env=env,
+        )
+        try:
+            metrics = _run_gsm8k(self.base_url, MOE_MXFP4_MODEL_PATH, num_examples=200)
+            print(f"[{MOE_MXFP4_MODEL_PATH}] {metrics=}")
+            if is_in_ci():
+                write_github_step_summary(
+                    f"### gsm8k run-through ({MOE_MXFP4_MODEL_PATH}, fuse_gate)\n"
+                    f'score={metrics["score"]:.3f} (run-through, no gate)\n'
+                )
+            # Ran e2e and returned parseable answers -> fuse_gate path is healthy.
+            self.assertGreater(metrics["score"], 0.0)
+        finally:
+            kill_process_tree(process.pid)
+
+    def test_b_perf(self):
+        """Serving performance benchmark for the fused checkpoint."""
+        _run_perf(MOE_MXFP4_MODEL_PATH, "qwen35-moe-mxfp4-mi35x", self.base_url)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/moe/test_fused_append_shared_experts.py

@@ -0,0 +1,301 @@
+"""Unit tests for ``fused_append_shared_experts_with_weights``.
+
+Covers PR #28658, which folds the shared-expert ``sigmoid`` activation and the
+bf16->fp32 cast into the append kernel's prologue (``apply_sigmoid=True``), and
+PR #28666, which additionally folds the ``shared_expert_gate`` GEMV into the
+kernel (``fuse_gate=True``). The old (``apply_sigmoid=False, fuse_gate=False``)
+path must stay byte-for-byte identical; the ``apply_sigmoid`` path must equal
+the eager ``sigmoid(logits.float()) * scale`` it replaces; and the ``fuse_gate``
+path must equal the eager ``sigmoid((hidden @ W_gate).float()) * scale`` GEMV it
+replaces. These paths only run on the AITER shared-expert-fusion route at
+serving time, so they are otherwise uncovered by CI.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
+    fused_append_shared_experts_with_weights,
+)
+from sglang.srt.utils import get_device
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase, empty_gpu_cache
+
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-large")
+register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-small-amd")
+
+
+def _eager_append(
+    topk_ids, topk_weights, shared_weights, s, n_base, apply_sigmoid, scale
+):
+    """Reference: eager equivalent of the kernel (matches the pre-PR two-step path).
+
+    The last ``s`` columns hold the shared expert(s); ids are ``n_base + i`` and
+    weights are either the raw shared value cast to the output dtype (legacy) or
+    ``sigmoid(value.float()) * scale`` cast to the output dtype (fused).
+    """
+    m, k = topk_ids.shape
+    sw = shared_weights
+    if sw.ndim == 1:
+        sw = sw.unsqueeze(-1)
+    if sw.shape[1] < s:
+        sw = sw.expand(m, s)
+    if apply_sigmoid:
+        col = (torch.sigmoid(sw.float()) * scale).to(topk_weights.dtype)
+    else:
+        col = sw.to(topk_weights.dtype)
+
+    shared_ids = (
+        (n_base + torch.arange(s, device=topk_ids.device))
+        .view(1, s)
+        .expand(m, s)
+        .to(topk_ids.dtype)
+    )
+
+    out_ids = torch.cat([topk_ids, shared_ids], dim=1)
+    out_weights = torch.cat([topk_weights, col.contiguous()], dim=1)
+    return out_ids, out_weights
+
+
+def _eager_gate(topk_ids, topk_weights, hidden, gate_weight, s, n_base, scale):
+    """Reference for fuse_gate=True: in-kernel GEMV + sigmoid + scale.
+
+    ``logit[m] = hidden[m, :] . W_gate[:]`` (fp32), broadcast
+    ``sigmoid(logit) * scale`` to all ``s`` shared slots.
+    """
+    m, k = topk_ids.shape
+    logit = (hidden.float() @ gate_weight.float().reshape(-1)).view(m, 1)
+    col = (torch.sigmoid(logit) * scale).expand(m, s).to(topk_weights.dtype)
+
+    shared_ids = (
+        (n_base + torch.arange(s, device=topk_ids.device))
+        .view(1, s)
+        .expand(m, s)
+        .to(topk_ids.dtype)
+    )
+
+    out_ids = torch.cat([topk_ids, shared_ids], dim=1)
+    out_weights = torch.cat([topk_weights, col.contiguous()], dim=1)
+    return out_ids, out_weights
+
+
+class TestFusedAppendSharedExperts(CustomTestCase):
+    MS = [1, 3, 128]
+    KS = [2, 6, 8]
+    SS = [1, 2]
+    SCALES = [1.0, 0.5]
+
+    @staticmethod
+    def _rand_inputs(m, k, s, dtype, n_base):
+        device = get_device()
+        topk_ids = torch.randint(0, n_base, (m, k), dtype=torch.int32, device=device)
+        topk_weights = torch.rand(m, k, dtype=torch.float32, device=device).to(dtype)
+        # shared gate logits: wide range so sigmoid is meaningfully exercised
+        shared_logits = (
+            torch.randn(m, s, dtype=torch.float32, device=device) * 4.0
+        ).to(dtype)
+        return topk_ids, topk_weights, shared_logits
+
+    def _check(self, out_ids, out_w, ref_ids, ref_w, dtype, msg):
+        # ids are integers -> must match exactly
+        self.assertTrue(torch.equal(out_ids, ref_ids), f"ids mismatch: {msg}")
+        # triton tl.sigmoid vs torch.sigmoid differ by a few ULPs, so the fused
+        # activation path is only bitexact-free: tolerance, tight in fp32.
+        if dtype == torch.float32:
+            torch.testing.assert_close(
+                out_w, ref_w, rtol=1e-5, atol=1e-6, msg=f"weights mismatch: {msg}"
+            )
+        else:
+            torch.testing.assert_close(
+                out_w, ref_w, rtol=2e-3, atol=2e-3, msg=f"weights mismatch: {msg}"
+            )
+
+    def test_fused_sigmoid_matches_eager(self):
+        """apply_sigmoid=True must equal eager sigmoid(logits.float())*scale."""
+        n_base = 64
+        for dtype in [torch.float32, torch.bfloat16]:
+            for m in self.MS:
+                for k in self.KS:
+                    for s in self.SS:
+                        for scale in self.SCALES:
+                            topk_ids, topk_w, logits = self._rand_inputs(
+                                m, k, s, dtype, n_base
+                            )
+                            out_ids, out_w = fused_append_shared_experts_with_weights(
+                                topk_ids,
+                                topk_w,
+                                logits,
+                                s,
+                                N=n_base,
+                                apply_sigmoid=True,
+                                scale=scale,
+                            )
+                            ref_ids, ref_w = _eager_append(
+                                topk_ids, topk_w, logits, s, n_base, True, scale
+                            )
+                            self.assertEqual(out_ids.shape, (m, k + s))
+                            self._check(
+                                out_ids,
+                                out_w,
+                                ref_ids,
+                                ref_w,
+                                dtype,
+                                f"m={m} k={k} s={s} scale={scale} dtype={dtype}",
+                            )
+                            empty_gpu_cache()
+
+    def test_legacy_path_unchanged(self):
+        """apply_sigmoid=False must reproduce the legacy cast-and-append exactly."""
+        n_base = 64
+        for dtype in [torch.float32, torch.bfloat16]:
+            for m in self.MS:
+                for k in self.KS:
+                    for s in self.SS:
+                        topk_ids, topk_w, shared = self._rand_inputs(
+                            m, k, s, dtype, n_base
+                        )
+                        out_ids, out_w = fused_append_shared_experts_with_weights(
+                            topk_ids, topk_w, shared, s, N=n_base, apply_sigmoid=False
+                        )
+                        ref_ids, ref_w = _eager_append(
+                            topk_ids, topk_w, shared, s, n_base, False, 1.0
+                        )
+                        # legacy path is a pure copy/cast -> exact for all dtypes
+                        self.assertTrue(torch.equal(out_ids, ref_ids))
+                        self.assertTrue(torch.equal(out_w, ref_w))
+                        empty_gpu_cache()
+
+    def test_routed_columns_preserved(self):
+        """First k columns must be an untouched copy of the routed topk output."""
+        n_base = 32
+        m, k, s = 5, 6, 1
+        topk_ids, topk_w, logits = self._rand_inputs(m, k, s, torch.bfloat16, n_base)
+        out_ids, out_w = fused_append_shared_experts_with_weights(
+            topk_ids, topk_w, logits, s, N=n_base, apply_sigmoid=True, scale=1.0
+        )
+        self.assertTrue(torch.equal(out_ids[:, :k], topk_ids))
+        self.assertTrue(torch.equal(out_w[:, :k], topk_w))
+        # shared ids are exactly n_base .. n_base+s-1
+        self.assertTrue(
+            torch.equal(
+                out_ids[:, k],
+                torch.full((m,), n_base, dtype=topk_ids.dtype, device=topk_ids.device),
+            )
+        )
+
+    def test_fuse_gate_matches_eager(self):
+        """fuse_gate=True must equal eager sigmoid((hidden @ W_gate).float())*scale."""
+        n_base = 64
+        device = get_device()
+        for dtype in [torch.float32, torch.bfloat16]:
+            for hidden_dim in [512, 2048]:
+                for m in self.MS:
+                    for k in self.KS:
+                        for s in self.SS:
+                            for scale in self.SCALES:
+                                topk_ids = torch.randint(
+                                    0, n_base, (m, k), dtype=torch.int32, device=device
+                                )
+                                topk_w = torch.rand(
+                                    m, k, dtype=torch.float32, device=device
+                                ).to(dtype)
+                                hidden = (
+                                    torch.randn(
+                                        m,
+                                        hidden_dim,
+                                        dtype=torch.float32,
+                                        device=device,
+                                    )
+                                    * 0.1
+                                ).to(dtype)
+                                # shared_expert_gate weight: Linear(hidden, 1)
+                                gate_w = (
+                                    torch.randn(
+                                        1,
+                                        hidden_dim,
+                                        dtype=torch.float32,
+                                        device=device,
+                                    )
+                                    * 0.1
+                                ).to(dtype)
+
+                                out_ids, out_w = (
+                                    fused_append_shared_experts_with_weights(
+                                        topk_ids,
+                                        topk_w,
+                                        None,
+                                        s,
+                                        N=n_base,
+                                        fuse_gate=True,
+                                        hidden_states=hidden,
+                                        gate_weight=gate_w,
+                                        scale=scale,
+                                    )
+                                )
+                                ref_ids, ref_w = _eager_gate(
+                                    topk_ids, topk_w, hidden, gate_w, s, n_base, scale
+                                )
+                                self.assertEqual(out_ids.shape, (m, k + s))
+                                msg = (
+                                    f"m={m} k={k} s={s} h={hidden_dim} "
+                                    f"scale={scale} dtype={dtype}"
+                                )
+                                self.assertTrue(
+                                    torch.equal(out_ids, ref_ids), f"ids: {msg}"
+                                )
+                                # fp32 GEMV reduction order differs from torch
+                                # matmul -> tolerance (tight in fp32).
+                                if dtype == torch.float32:
+                                    torch.testing.assert_close(
+                                        out_w, ref_w, rtol=1e-4, atol=1e-5, msg=msg
+                                    )
+                                else:
+                                    torch.testing.assert_close(
+                                        out_w, ref_w, rtol=2e-3, atol=2e-3, msg=msg
+                                    )
+                                empty_gpu_cache()
+
+    def test_fuse_gate_and_apply_sigmoid_mutually_exclusive(self):
+        """fuse_gate already applies sigmoid; combining with apply_sigmoid asserts."""
+        device = get_device()
+        topk_ids = torch.randint(0, 8, (2, 2), dtype=torch.int32, device=device)
+        topk_w = torch.rand(2, 2, dtype=torch.float32, device=device)
+        hidden = torch.randn(2, 16, dtype=torch.float32, device=device)
+        gate_w = torch.randn(1, 16, dtype=torch.float32, device=device)
+        with self.assertRaises(AssertionError):
+            fused_append_shared_experts_with_weights(
+                topk_ids,
+                topk_w,
+                None,
+                1,
+                N=8,
+                apply_sigmoid=True,
+                fuse_gate=True,
+                hidden_states=hidden,
+                gate_weight=gate_w,
+            )
+
+    def test_fuse_gate_requires_hidden_and_gate(self):
+        """fuse_gate=True without hidden_states/gate_weight asserts."""
+        device = get_device()
+        topk_ids = torch.randint(0, 8, (2, 2), dtype=torch.int32, device=device)
+        topk_w = torch.rand(2, 2, dtype=torch.float32, device=device)
+        with self.assertRaises(AssertionError):
+            fused_append_shared_experts_with_weights(
+                topk_ids, topk_w, None, 1, N=8, fuse_gate=True
+            )
+
+    def test_zero_shared_is_noop(self):
+        """num_fused_shared_experts <= 0 returns the inputs unchanged."""
+        n_base = 8
+        topk_ids, topk_w, logits = self._rand_inputs(4, 2, 1, torch.float32, n_base)
+        out_ids, out_w = fused_append_shared_experts_with_weights(
+            topk_ids, topk_w, logits, 0, N=n_base, apply_sigmoid=True
+        )
+        self.assertIs(out_ids, topk_ids)
+        self.assertIs(out_w, topk_w)
+
+
+if __name__ == "__main__":
+    unittest.main()