[Feature] Enable return routed experts (#12162)

Co-authored-by: yizhang2077 <1109276519@qq.com> Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
2025-12-21 15:16:43 +08:00
parent 8fe3e37468
commit bed301a5ac
27 changed files with 646 additions and 10 deletions
@@ -161,6 +161,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
        self.topk = TopK(
            top_k=config.num_experts_per_tok,
            renormalize=config.norm_topk_prob,
+            layer_id=layer_id,
        )

        self.experts = get_moe_impl_class(quant_config)(