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
@@ -143,7 +143,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
            renormalize=config.norm_topk_prob,
        )

-        self.experts = get_moe_impl_class()(
+        self.experts = get_moe_impl_class(quant_config)(
            layer_id=self.layer_id,
            top_k=config.num_experts_per_tok,
            num_experts=config.num_experts,