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

python/sglang/srt/layers/elementwise.py

@@ -4,6 +4,7 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.srt.layers.triton_ops.softcap import softcap_out as fused_softcap
 from sglang.srt.utils import is_hip
 from sglang.srt.utils.custom_op import register_custom_op
@@ -528,3 +529,174 @@ def silu_and_mul_triton(
         return out_hidden_states, out_scales
     else:
         return out_hidden_states, None
+
+
+@triton.jit
+def _fused_sigmoid_mul_kernel(
+    output_ptr,
+    attn_output_ptr,
+    gate_ptr,
+    gate_stride_row,
+    gate_stride_head,
+    hidden_dim: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    """Fuse sigmoid(gate) * attn_output into a single kernel."""
+    pid_row = tl.program_id(0).to(tl.int64)
+    pid_block = tl.program_id(1)
+
+    offsets = pid_block * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = offsets < hidden_dim
+    head = offsets // HEAD_DIM
+    d = offsets - head * HEAD_DIM
+
+    attn_off = pid_row * hidden_dim + offsets
+    attn = tl.load(attn_output_ptr + attn_off, mask=mask, other=0.0).to(tl.float32)
+
+    gate_off = pid_row * gate_stride_row + head * gate_stride_head + d
+    g = tl.load(gate_ptr + gate_off, mask=mask, other=0.0).to(tl.float32)
+
+    result = attn * tl.sigmoid(g)
+    tl.store(output_ptr + attn_off, result, mask=mask)
+
+
+def fused_sigmoid_mul(
+    attn_output: torch.Tensor,
+    gate: torch.Tensor,
+    inplace: bool = False,
+) -> torch.Tensor:
+    """
+    Fused sigmoid-mul for attention output gating.
+
+    Equivalent to: attn_output * sigmoid(gate)
+
+    The production Qwen3.5 path passes a 3D strided gate. A single hidden-block
+    Triton kernel handles both that path and flat contiguous inputs.
+
+    When inplace=True, writes result back to attn_output and returns it.
+
+    Supports strided gate: if gate is 3D (num_tokens, num_heads, head_dim)
+    and attn_output is 2D (num_tokens, hidden_dim), the kernel reads gate
+    via explicit strides without requiring a contiguous copy.
+    """
+    if gate.ndim == 3 and attn_output.ndim == 2:
+        # Strided gate path: gate is 3D (num_tokens, num_heads, head_dim)
+        num_tokens, num_heads, head_dim = gate.shape
+        hidden_dim = num_heads * head_dim
+        assert attn_output.shape == (num_tokens, hidden_dim)
+        gate_stride_row = gate.stride(0)
+        gate_stride_head = gate.stride(1)
+    else:
+        # Flat path: both tensors have the same shape
+        assert (
+            attn_output.shape == gate.shape
+        ), "attn_output and gate must have the same shape"
+        hidden_dim = attn_output.shape[-1]
+        num_tokens = attn_output.numel() // hidden_dim
+        head_dim = hidden_dim
+        gate_stride_row = hidden_dim
+        gate_stride_head = hidden_dim
+
+    out = attn_output if inplace else torch.empty_like(attn_output)
+    block_h = 1024 if num_tokens < 1024 else 2048
+    grid = (num_tokens, triton.cdiv(hidden_dim, block_h))
+    _fused_sigmoid_mul_kernel[grid](
+        out,
+        attn_output,
+        gate,
+        gate_stride_row,
+        gate_stride_head,
+        hidden_dim,
+        HEAD_DIM=head_dim,
+        BLOCK_H=block_h,
+        num_warps=4,
+    )
+    return out
+
+
+@triton.jit
+def _fused_gate_sigmoid_mul_add_kernel(
+    hidden_states_ptr,  # [num_tokens, hidden_dim]
+    gate_weight_ptr,  # [hidden_dim]
+    shared_output_ptr,  # [num_tokens, hidden_dim]
+    final_hidden_states_ptr,  # [num_tokens, hidden_dim]
+    hidden_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    USE_PDL: tl.constexpr = False,
+):
+    pid = tl.program_id(axis=0).to(tl.int64)
+    row_offset = pid * hidden_dim
+
+    offsets = tl.arange(0, BLOCK_SIZE)
+    mask = offsets < hidden_dim
+
+    w = tl.load(gate_weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    h = tl.load(hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    s = tl.load(shared_output_ptr + row_offset + offsets, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    f = tl.load(
+        final_hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0
+    ).to(tl.float32)
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    gate_val = tl.sigmoid(tl.sum(h * w, axis=0))
+    result = f + gate_val * s
+
+    tl.store(final_hidden_states_ptr + row_offset + offsets, result, mask=mask)
+
+
+def fused_gate_sigmoid_mul_add(
+    hidden_states: torch.Tensor,
+    gate_weight: torch.Tensor,
+    shared_output: torch.Tensor,
+    final_hidden_states: torch.Tensor,
+) -> None:
+    """
+    Fused gate-sigmoid-mul-add for MoE shared expert gating.
+
+    Equivalent to:
+        gate = hidden_states @ gate_weight
+        final_hidden_states += sigmoid(gate).unsqueeze(1) * shared_output
+    """
+    assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
+    assert gate_weight.is_contiguous(), "gate_weight must be contiguous"
+    assert shared_output.is_contiguous(), "shared_output must be contiguous"
+    assert final_hidden_states.is_contiguous(), "final_hidden_states must be contiguous"
+
+    num_tokens, hidden_dim = hidden_states.shape
+    assert gate_weight.shape == (hidden_dim,)
+    assert shared_output.shape == (num_tokens, hidden_dim)
+    assert final_hidden_states.shape == (num_tokens, hidden_dim)
+
+    max_warps = 16 if _is_hip else 32
+    config = {
+        "BLOCK_SIZE": triton.next_power_of_2(hidden_dim),
+        "num_warps": max(
+            min(triton.next_power_of_2(triton.cdiv(hidden_dim, 256)), max_warps), 4
+        ),
+    }
+
+    if num_tokens >= 1024:
+        config["num_warps"] = min(config["num_warps"], 8)
+
+    pdl_kwargs = {"USE_PDL": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+
+    _fused_gate_sigmoid_mul_add_kernel[(num_tokens,)](
+        hidden_states,
+        gate_weight,
+        shared_output,
+        final_hidden_states,
+        hidden_dim=hidden_dim,
+        **config,
+        **pdl_kwargs,
+    )

python/sglang/srt/models/qwen2_moe.py

@@ -53,6 +53,7 @@ from sglang.srt.layers.dp_attention import (
     get_attention_tp_size,
     is_dp_attention_enabled,
 )
+from sglang.srt.layers.elementwise import fused_gate_sigmoid_mul_add
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     MergedColumnParallelLinear,
@@ -265,6 +266,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             layer_id=layer_id,
         )
 
+        # Disable inplace MoE when fused gate will need hidden_states after experts
+        _needs_hidden_after_experts = (
+            config.shared_expert_intermediate_size > 0
+            and not self.enable_shared_expert_fusion
+        )
         self.experts = get_moe_impl_class(quant_config)(
             layer_id=self.layer_id,
             top_k=(
@@ -285,6 +291,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             prefix=add_prefix("experts", prefix),
             routing_method_type=RoutingMethodType.RenormalizeNaive,
             num_fused_shared_experts=self.num_fused_shared_experts,
+            inplace=not _needs_hidden_after_experts,
         )
 
         self.gate = ReplicatedLinear(
@@ -400,11 +407,13 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             router_logits=topk_output.router_logits,
         )
 
-    def _forward_shared_experts(self, hidden_states: torch.Tensor):
+    def _forward_shared_experts(
+        self, hidden_states: torch.Tensor, apply_gate: bool = True
+    ):
         shared_output = None
         if self.shared_expert is not None:
             shared_output = self.shared_expert(hidden_states)
-            if self.shared_expert_gate is not None:
+            if self.shared_expert_gate is not None and apply_gate:
                 if use_intel_amx_backend(self.shared_expert_gate):
                     shared_output = torch.ops.sgl_kernel.fused_linear_sigmoid_mul(
                         hidden_states,
@@ -476,11 +485,14 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
     def forward_normal_dual_stream(
         self,
         hidden_states: torch.Tensor,
+        use_fused_gate: bool = False,
     ) -> torch.Tensor:
         current_stream = torch.cuda.current_stream()
         self.alt_stream.wait_stream(current_stream)
         shared_output = (
-            self._forward_shared_experts(hidden_states.clone())
+            self._forward_shared_experts(
+                hidden_states.clone(), apply_gate=not use_fused_gate
+            )
             if self.shared_expert is not None
             else None
         )
@@ -525,6 +537,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
         if get_moe_a2a_backend().is_deepep():
             return self._forward_deepep(hidden_states, forward_batch)
 
+        use_fused_gate = (
+            self.shared_expert_gate is not None
+            and not use_intel_amx_backend(self.shared_expert_gate)
+        )
+
         if hidden_states.shape[0] == 0:
             # M=0 guard for idle DP ranks: skip shared_experts and gate
             # (which crash on empty tensors in FP4 GEMM), but still call
@@ -534,14 +551,24 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             final_hidden_states = self.experts(hidden_states, topk_output)
         elif self.alt_stream is not None and get_is_capture_mode():
             final_hidden_states, shared_output = self.forward_normal_dual_stream(
-                hidden_states
+                hidden_states, use_fused_gate=use_fused_gate
             )
         else:
-            shared_output = self._forward_shared_experts(hidden_states)
+            shared_output = self._forward_shared_experts(
+                hidden_states, apply_gate=not use_fused_gate
+            )
             final_hidden_states = self._forward_router_experts(hidden_states)
 
         if shared_output is not None:
-            final_hidden_states += shared_output
+            if use_fused_gate:
+                fused_gate_sigmoid_mul_add(
+                    hidden_states,
+                    self.shared_expert_gate.weight.squeeze(),
+                    shared_output,
+                    final_hidden_states,
+                )
+            else:
+                final_hidden_states += shared_output
         if (
             self.tp_size > 1
             and not should_skip_post_experts_all_reduce(

python/sglang/srt/models/qwen3_5.py

@@ -47,6 +47,7 @@ from sglang.srt.layers.dp_attention import (
     get_attention_tp_size,
     is_dp_attention_enabled,
 )
+from sglang.srt.layers.elementwise import fused_sigmoid_mul
 
 # Layers - Others
 from sglang.srt.layers.layernorm import GemmaRMSNorm
@@ -887,7 +888,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
             q_gate = q_gate.view(*orig_shape, self.num_heads, -1)
             q, gate = torch.chunk(q_gate, 2, dim=-1)
             q = q.reshape(*orig_shape, -1)
-            gate = gate.reshape(*orig_shape, -1)
+            # gate stays as 3D strided view; fused_sigmoid_mul handles it directly
         else:
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
             gate = None
@@ -974,15 +975,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
         attn_output = self.attn(q, k, v, forward_batch)
 
         if self.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
+            attn_output = fused_sigmoid_mul(attn_output, gate, inplace=True)
 
         output, _ = self.o_proj(attn_output)
         return output