[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>
vanshbhatia-amd committed
Aug 19, 2026 at 15:42 UTC
d591d1d511b5f2a70ae34adfd5adc5c1956ecae2
1 file changed
+13
benchmarks/kernels/benchmark_moe.py
+13
@@ -813,6 +813,19 @@ def get_model_params(config):
813
topk = config.thinker_config.text_config.num_experts_per_tok
814
intermediate_size = config.thinker_config.text_config.moe_intermediate_size
815
hidden_size = config.thinker_config.text_config.hidden_size
816
+ elif architecture in (
817
+ "KimiK3ForConditionalGeneration",
818
+ "KimiLinearForCausalLM",
819
+ ):
820
+ # Kimi K3 (multimodal) nests its MoE params in a KimiLinearConfig
821
+ # text_config and uses ``num_experts_per_token`` rather than the more
822
+ # common ``num_experts_per_tok``. get_text_config() returns the config
823
+ # itself for the text-only KimiLinearForCausalLM.
824
+ text_config = config.get_text_config()
825
+ E = text_config.num_experts
826
+ topk = text_config.num_experts_per_token
827
+ intermediate_size = text_config.moe_intermediate_size
828
+ hidden_size = text_config.hidden_size
829
elif architecture == "PixtralForConditionalGeneration":
830
# Pixtral can contain different LLM architectures,
831
# recurse to get their parameters