[Cohere2Moe] Enable flashinfer_trtllm NVFP4 fused-MoE via SigmoidRenorm routing (#27401)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
zijiexia
2026-06-05 22:49:20 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent caeb449cd6
commit 9da88e32e0
2 changed files with 19 additions and 1 deletions
+7 -1
View File
@@ -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
+12
View File
@@ -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)