[Bugfix] Add Kimi K3 MoE support to benchmark_moe.py (#50082)

Signed-off-by: Vansh Bhatia <210711135+vanshbhatia-amd@users.noreply.github.com>
Co-authored-by: Vansh Bhatia <210711135+vanshbhatia-amd@users.noreply.github.com>
This commit is contained in:
vanshbhatia-amd
2026-08-19 19:42:30 +00:00
committed by GitHub
co-authored by Vansh Bhatia
parent 541c6d64c1
commit d591d1d511
+13
View File
@@ -813,6 +813,19 @@ def get_model_params(config):
topk = config.thinker_config.text_config.num_experts_per_tok
intermediate_size = config.thinker_config.text_config.moe_intermediate_size
hidden_size = config.thinker_config.text_config.hidden_size
elif architecture in (
"KimiK3ForConditionalGeneration",
"KimiLinearForCausalLM",
):
# Kimi K3 (multimodal) nests its MoE params in a KimiLinearConfig
# text_config and uses ``num_experts_per_token`` rather than the more
# common ``num_experts_per_tok``. get_text_config() returns the config
# itself for the text-only KimiLinearForCausalLM.
text_config = config.get_text_config()
E = text_config.num_experts
topk = text_config.num_experts_per_token
intermediate_size = text_config.moe_intermediate_size
hidden_size = text_config.hidden_size
elif architecture == "PixtralForConditionalGeneration":
# Pixtral can contain different LLM architectures,
# recurse to get their parameters