[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,)
|
RenormalizeNaive = (4,)
|
||||||
# TopK only (no softmax)
|
# TopK only (no softmax)
|
||||||
TopK = (5,)
|
TopK = (5,)
|
||||||
|
# SigmoidRenorm: Sigmoid -> TopK -> Renormalize
|
||||||
|
SigmoidRenorm = (6,)
|
||||||
|
# MiniMax2
|
||||||
|
MiniMax2 = (7,)
|
||||||
|
# Sigmoid: Sigmoid -> TopK (no renormalize)
|
||||||
|
Sigmoid = (8,)
|
||||||
# Unspecified
|
# Unspecified
|
||||||
Unspecified = 6
|
Unspecified = 9
|
||||||
|
|
||||||
|
|
||||||
AITER_PADDING_SIZE = 128
|
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.logits_processor import LogitsProcessor
|
||||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||||
from sglang.srt.layers.moe.topk import TopK
|
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.quantization.base_config import QuantizationConfig
|
||||||
from sglang.srt.layers.radix_attention import RadixAttention
|
from sglang.srt.layers.radix_attention import RadixAttention
|
||||||
from sglang.srt.layers.rotary_embedding import get_rope
|
from sglang.srt.layers.rotary_embedding import get_rope
|
||||||
@@ -249,9 +250,19 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
|
|||||||
if self.expert_selection_fn == "sigmoid":
|
if self.expert_selection_fn == "sigmoid":
|
||||||
custom_routing_function = cohere2_sigmoid_topk
|
custom_routing_function = cohere2_sigmoid_topk
|
||||||
scoring_func = "sigmoid"
|
scoring_func = "sigmoid"
|
||||||
|
routing_method_type = (
|
||||||
|
RoutingMethodType.SigmoidRenorm
|
||||||
|
if self.norm_topk_prob
|
||||||
|
else RoutingMethodType.Sigmoid
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
custom_routing_function = None
|
custom_routing_function = None
|
||||||
scoring_func = "softmax"
|
scoring_func = "softmax"
|
||||||
|
routing_method_type = (
|
||||||
|
RoutingMethodType.RenormalizeNaive
|
||||||
|
if self.norm_topk_prob
|
||||||
|
else RoutingMethodType.Default
|
||||||
|
)
|
||||||
|
|
||||||
self.gate = ReplicatedLinear(
|
self.gate = ReplicatedLinear(
|
||||||
config.hidden_size,
|
config.hidden_size,
|
||||||
@@ -278,6 +289,7 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
|
|||||||
quant_config=quant_config,
|
quant_config=quant_config,
|
||||||
layer_id=layer_id,
|
layer_id=layer_id,
|
||||||
prefix=add_prefix("experts", prefix),
|
prefix=add_prefix("experts", prefix),
|
||||||
|
routing_method_type=routing_method_type,
|
||||||
)
|
)
|
||||||
|
|
||||||
num_shared_experts = getattr(config, "num_shared_experts", 0)
|
num_shared_experts = getattr(config, "num_shared_experts", 0)
|
||||||
|
|||||||
Reference in New Issue
Block a user