fix(kimi-k3): correct MFU accounting (FLOPs formula, config defaults)

[yisongbetter]

What does this PR do ?

Makes Automodel's Kimi K3 MFU accounting correct. Two independent defects made the reported MFU wrong for K3:

1. No FLOPs formula for K3. get_flops_formula_for_hf_config had no entry for KimiK3TextConfig, so K3 fell through to the generic dense transformer_flops, whose MoE branch probes num_experts_per_tok while the K3 config spells it num_experts_per_token. The 2.78T hybrid MoE was costed as a dense 93-layer transformer at ~0.44× of its real FLOPs (a 64-node run that measures 13.7% MFU wall-clock printed ~6.0%).
2. from_config built a lighter model than the checkpoint. KimiK3TextConfig defaulted to num_attention_heads=56 and intermediate_size=18432; the released moonshotai/Kimi-K3 has 96 heads and a 33792-wide dense FFN (config.json, safetensors shapes, tech report Table 1). from_pretrained was unaffected (config.json overrides), but the benchmark recipes and random-init smoke configs built ~101.5B instead of ~104B active matmul parameters per token, so their throughput numbers were not numbers for Kimi K3.

Two commits, one per defect; they are independent (the formula's tests pass explicit config values).

Changelog

• flops_utils.kimi_k3_flops: hybrid KDA / gated-MLA attention (layer pattern read from linear_attn_config, MLA output gate), latent MoE (routed_expert_hidden_size down/up projections around the top-k experts) + SiTU shared expert, dense layer(s) per first_k_dense_replace, LM head. KDA chunked kernel costed with the Kimi Linear paper's count (6 T d_h² + 3 T C d_h + T C² per head, arXiv:2510.26692). _kda_attention_per_layer_flops helper mirrors _gdn_attention_per_layer_flops.
• Registry entries for KimiK3TextConfig and the multimodal KimiK3Config wrapper (text_config is used; the vision tower is not counted, as for the other VL entries).
• KimiK3TextConfig defaults: num_attention_heads 56 → 96, intermediate_size 18432 → 33792; docstring states the defaults reproduce the released checkpoint.
• Tests: tests/unit_tests/utils/test_kimi_k3_flops.py (layer pattern, model-card parameter counts, precomputed values for the checkpoint and the pre-fix shapes, batch linearity, >2× correction over the dense fallback, exact cost of each optional block); tests/unit_tests/models/kimi_k3/test_config_defaults.py (pins the released text config, the linear-attention block, the 69 KDA / 24 MLA pattern, the MLA projection widths implied by the head count, the multimodal wrapper); a registration case in test_flops_utils_new_models.py.

Verification

Cross-checked against the released checkpoint (moonshotai/Kimi-K3 config.json + modeling_kimi_linear.py + safetensors headers: q_b_proj [18432, 1536] = 96 × 192, kv_b_proj [24576, 512], layers.0.mlp.gate_proj [33792, 7168]), the model card, and the K3 tech report (Table 1: 96 heads, 104.2B activated / 2.78T total):

• With the corrected defaults from_config builds the checkpoint's shape, and the formula counts 104.0B active matmul parameters per token, 2.78T total (within 0.2% of Table 1; the formula counts GEMM weights only — no router, norms or gate biases). Layer pattern 69 KDA + 24 gated MLA, 1 dense layer, 16/896 experts, 2 shared — all read from the config.
• FLOPs/token at seq 1024 / 2048 = 1.004× / 1.007× of the wall-clock basis 6 × 104e9 used for every published K3 number (the excess is the MLA attention BMMs and the KDA kernel).
• The report's single MTP layer is not in the released checkpoint (num_nextn_predict_layers: 0) and not modelled by Automodel, so it is not counted.
• Tests run locally (CPU, torch 2.14 / transformers 5.16): tests/unit_tests/utils/ 247 passed (1 pre-existing environment failure in test_yaml_utils.py: needs /home/TestData); tests/unit_tests/models/kimi_k3/ 65 passed (1 environment failure: FLA kernels not installed); tests/unit_tests/recipes/llm/test_benchmark.py 33 passed. ruff format --check / ruff check clean on the changed files.
• End-to-end (2 nodes × 4 GB200, nemo_automodel/recipes/llm/benchmark.py --config …, Kimi-K3 reduced to 24 layers / 64 experts, seq 4096, GBS 32, random init, the formula commit applied on upstream main): TFLOPs/GPU: 21323.249904 printed at setup = kimi_k3_flops(cfg, gbs=32, seq_len=4096) exactly (the dense fallback would have given 0.485× of that). Steady iterations 4.95 s → printed MFU: 23.93%; the wall-clock basis at the same iteration time (tokens/s/GPU × 6 × active params / 2.25e15) gives 23.74%, ratio 1.008. Average MFU: 23.17% over the 9 post-warmup iterations (two GC-slowed).

Additional Information

• Review order: please review this PR before perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779 / perf(distributed): ring-pooled PP recv buffers and Kimi-K3 2k benchmark shape #3780 — both will be rebased on this change (their benchmark yamls now build the real K3 shape) and their measured numbers re-stamped.
• Orthogonal to fix(moe): Kimi-K3 correctness fixes (routing, NCCL warmup, PP, per-GPU tps) #3731 (throughput accounting in train_ft.py); this fixes the FLOPs-per-token estimate and the from_config model shape.
• Reproducing the printout: the per-iteration MFU: lines and Average MFU come from run_benchmark(), which only the recipes/llm/benchmark.py --config <yaml> entry (used by the L2 benchmark tests) runs; the automodel <yaml> launcher in the yaml headers goes through the generic run_train_validation_loop() and prints only the setup-time TFLOPs/GPU.
• Impact of the defaults change on from_config users: the K3 benchmark recipes (examples/llm_benchmark/kimi/kimi_k3_gb200*.yaml) and random-init smoke configs now build the real K3 shape — 24 MLA layers gain ~90M parameters each and the dense layer doubles, ≈ +2.5% FLOPs/token, ≈ +1 GiB/GPU at EP32×PP8 on 256 GPUs. Throughput numbers in those yaml headers were measured with the old shape and will be re-stamped in perf(kimi-k3): RMSNorm compile island and benchmark static-routing sync pins #3779 / perf(distributed): ring-pooled PP recv buffers and Kimi-K3 2k benchmark shape #3780 after a re-run.
• Not changed here: the generic fallback still probes only num_experts_per_tok, so other unregistered MoE configs with a different spelling are still under-costed — worth its own tiny fix.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/claude review

[claude]

LGTM

[HuiyingLi]

/ok to test 1613a9f