From d591d1d511b5f2a70ae34adfd5adc5c1956ecae2 Mon Sep 17 00:00:00 2001 From: vanshbhatia-amd Date: Wed, 19 Aug 2026 15:42:30 -0400 Subject: [PATCH] [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> --- benchmarks/kernels/benchmark_moe.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 5d86a7599f..1fcc7ddde7 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -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