[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
@@ -113,6 +113,7 @@ class GptOssSparseMoeBlock(nn.Module):
        self.topk = TopK(
            top_k=config.num_experts_per_tok,
            renormalize=True,
+            layer_id=layer_id,
        )

        self.top_k = config.num_experts_per_tok