From 9da88e32e09619d6fb5b8bd97438270382082aae Mon Sep 17 00:00:00 2001
From: zijiexia <37504505+zijiexia@users.noreply.github.com>
Date: Fri, 5 Jun 2026 22:49:20 -0700
Subject: [PATCH] [Cohere2Moe] Enable flashinfer_trtllm NVFP4 fused-MoE via
 SigmoidRenorm routing (#27401)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 python/sglang/srt/layers/moe/utils.py   |  8 +++++++-
 python/sglang/srt/models/cohere2_moe.py | 12 ++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py
index 37f1f0f63..e29f7efec 100644
--- a/python/sglang/srt/layers/moe/utils.py
+++ b/python/sglang/srt/layers/moe/utils.py
@@ -515,8 +515,14 @@ class RoutingMethodType(IntEnum):
     RenormalizeNaive = (4,)
     # TopK only (no softmax)
     TopK = (5,)
+    # SigmoidRenorm: Sigmoid -> TopK -> Renormalize
+    SigmoidRenorm = (6,)
+    # MiniMax2
+    MiniMax2 = (7,)
+    # Sigmoid: Sigmoid -> TopK (no renormalize)
+    Sigmoid = (8,)
     # Unspecified
-    Unspecified = 6
+    Unspecified = 9
 
 
 AITER_PADDING_SIZE = 128
diff --git a/python/sglang/srt/models/cohere2_moe.py b/python/sglang/srt/models/cohere2_moe.py
index a9924a338..8326c35ef 100644
--- a/python/sglang/srt/models/cohere2_moe.py
+++ b/python/sglang/srt/models/cohere2_moe.py
@@ -24,6 +24,7 @@ from sglang.srt.layers.linear import (
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
 from sglang.srt.layers.moe.topk import TopK
+from sglang.srt.layers.moe.utils import RoutingMethodType
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.rotary_embedding import get_rope
@@ -249,9 +250,19 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
         if self.expert_selection_fn == "sigmoid":
             custom_routing_function = cohere2_sigmoid_topk
             scoring_func = "sigmoid"
+            routing_method_type = (
+                RoutingMethodType.SigmoidRenorm
+                if self.norm_topk_prob
+                else RoutingMethodType.Sigmoid
+            )
         else:
             custom_routing_function = None
             scoring_func = "softmax"
+            routing_method_type = (
+                RoutingMethodType.RenormalizeNaive
+                if self.norm_topk_prob
+                else RoutingMethodType.Default
+            )
 
         self.gate = ReplicatedLinear(
             config.hidden_size,
@@ -278,6 +289,7 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
             quant_config=quant_config,
             layer_id=layer_id,
             prefix=add_prefix("experts", prefix),
+            routing_method_type=routing_method_type,
         )
 
         num_shared_experts = getattr(config, "num_shared_experts", 0)