Refactor TopK to ensure readability and extensibility (#9338)

2025-09-14 19:16:25 -07:00
parent b7d385e812
commit 4844fac91d
14 changed files with 52 additions and 47 deletions
--- a/python/sglang/srt/models/ernie4.py
+++ b/python/sglang/srt/models/ernie4.py
@@ -92,7 +92,7 @@ class Ernie4Moe(nn.Module):
            correction_bias=self.gate.e_score_correction_bias,
        )

-        self.experts = get_moe_impl_class()(
+        self.experts = get_moe_impl_class(quant_config)(
            num_experts=config.moe_num_experts,
            top_k=config.moe_k,
            hidden_size=config.hidden_size,