[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:
co-authored by
Claude Opus 4.8
parent
caeb449cd6
commit
9da88e32e0
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user