[4/N] Qwen3.5Opt: Overlap mamba verify update with draft extend (#26924)

benchmark/kernels/bench_fused_gate_sigmoid_mul_add.py

@@ -0,0 +1,60 @@
+"""Benchmark fused_gate_sigmoid_mul_add: Triton kernel vs PyTorch eager.
+
+Compares the fused Triton kernel against a plain PyTorch implementation
+over the Qwen3.5 MoE target hidden size.
+"""
+
+import torch
+import triton
+
+from sglang.srt.layers.elementwise import fused_gate_sigmoid_mul_add
+
+HIDDEN_DIMS = [4096]
+
+
+def _pytorch_reference(hidden_states, gate_weight, shared_output, final_hidden_states):
+    gate = hidden_states @ gate_weight
+    final_hidden_states += torch.sigmoid(gate).unsqueeze(1) * shared_output
+
+
+def make_bench(hidden_dim):
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["num_tokens"],
+            x_vals=[1, 2, 4, 8, 16, 1024, 2048, 4096, 8192],
+            line_arg="impl",
+            line_vals=["triton", "pytorch"],
+            line_names=["Triton fused", "PyTorch eager"],
+            styles=[("blue", "-"), ("orange", "--")],
+            ylabel="us",
+            plot_name=f"fused_gate_sigmoid_mul_add-hidden{hidden_dim}",
+            args={"hidden_dim": hidden_dim},
+        )
+    )
+    def bench(num_tokens, impl, hidden_dim, dtype=torch.bfloat16):
+        hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
+        gate_weight = torch.randn(hidden_dim, dtype=dtype, device="cuda")
+        shared_output = torch.randn(num_tokens, hidden_dim, dtype=dtype, device="cuda")
+        final_hidden_states = torch.randn(
+            num_tokens, hidden_dim, dtype=dtype, device="cuda"
+        )
+
+        if impl == "triton":
+            fn = lambda: fused_gate_sigmoid_mul_add(
+                hidden_states, gate_weight, shared_output, final_hidden_states
+            )
+        else:
+            fn = lambda: _pytorch_reference(
+                hidden_states, gate_weight, shared_output, final_hidden_states
+            )
+
+        ms = triton.testing.do_bench(fn, warmup=100, rep=200)
+        return ms * 1000  # convert to us
+
+    return bench
+
+
+if __name__ == "__main__":
+    for d in HIDDEN_DIMS:
+        print(f"\n===== hidden_dim={d} =====")
+        make_bench(d).run(print_data=True)

benchmark/kernels/bench_fused_sigmoid_mul.py

@@ -0,0 +1,59 @@
+"""Benchmark fused_sigmoid_mul: auto-dispatch vs PyTorch eager.
+
+The auto-dispatch path uses a strided Triton kernel for Qwen3.5 MoE attention
+output gates.
+
+Both paths start from a strided 3D gate (from torch.chunk) to ensure
+a fair comparison — the reshape/contiguous cost is included.
+"""
+
+import torch
+import triton
+
+from sglang.srt.layers.elementwise import fused_sigmoid_mul
+
+NUM_HEADS = 32
+HEAD_DIM = 256
+HIDDEN_DIM = NUM_HEADS * HEAD_DIM  # 8192
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["num_tokens"],
+        x_vals=[1, 2, 4, 8, 16, 1024, 2048, 4096, 8192],
+        line_arg="impl",
+        line_vals=["auto", "auto_inplace", "pytorch_from_strided"],
+        line_names=[
+            "fused_sigmoid_mul (auto)",
+            "fused_sigmoid_mul (auto, inplace)",
+            "PyTorch eager (incl. reshape)",
+        ],
+        styles=[("blue", "-"), ("green", "-"), ("orange", "--")],
+        ylabel="us",
+        plot_name="fused_sigmoid_mul_qwen3_5_moe_target",
+        args={},
+    )
+)
+def bench(num_tokens, impl, dtype=torch.bfloat16):
+    q_gate = torch.randn(
+        num_tokens, NUM_HEADS, 2 * HEAD_DIM, dtype=dtype, device="cuda"
+    )
+    _, gate_strided = torch.chunk(q_gate, 2, dim=-1)
+    attn_output = torch.randn(num_tokens, HIDDEN_DIM, dtype=dtype, device="cuda")
+
+    if impl == "auto":
+        fn = lambda: fused_sigmoid_mul(attn_output, gate_strided, inplace=False)
+    elif impl == "auto_inplace":
+        fn = lambda: fused_sigmoid_mul(attn_output, gate_strided, inplace=True)
+    else:
+        # Fair comparison: include reshape cost in every iteration
+        def fn():
+            g = gate_strided.contiguous().view(num_tokens, HIDDEN_DIM)
+            return attn_output * torch.sigmoid(g)
+
+    ms = triton.testing.do_bench(fn, warmup=100, rep=200)
+    return ms * 1000
+
+
+if __name__ == "__main__":
+    bench.run(print_data=True)