[kimi k3] add eager reference model with FSDP2 - #4025

Merged

shuhuayu merged 67 commits into

pytorch:mainfrom

JavaZeroo:agent/add-kimi-k3-reference-model

Aug 24, 2026

Commits on Aug 22, 2026

Add Kimi K3 reference model
JavaZeroo
committed
f3b5c5b
Remove development-only Kimi K3 diagnostics
JavaZeroo
committed
8de1124
Add FSDP2 support for Kimi K3
JavaZeroo
committed
bf7c6d5
Scale Kimi K3 debug model
JavaZeroo
committed
31fd2d7
Use out-of-place operations in Kimi K3 eager paths
JavaZeroo
committed
51b7007
Keep unused Kimi experts in the FSDP autograd graph
JavaZeroo
committed
a626834
Add Kimi K3 numerical and mixed-modality FSDP coverage
JavaZeroo
committed
bedf145
Reuse shared components in Kimi K3 and run KDA on the FLA kernel

JavaZeroo
and
claude
committed
6da10ec
Guard Kimi K3 optional deps and cover the residual pass-through path

JavaZeroo
and
claude
committed
e87a87a
Drop the KDA head-dim guard

JavaZeroo
and
claude
committed
6d2d190
Close the Kimi K3 backbone on global attention

JavaZeroo
and
claude
committed
f4b3808
Note the context-parallel route in add_zero_valued_dependency

JavaZeroo
and
claude
committed
25577e4
reuse common FeedForward
JavaZeroo
committed
4f7f668
dp_replicate needs no specific handling, removed from unsupported_parallelisms
JavaZeroo
committed
611dd73
Refactor KimiGroupedExperts to use gmm
JavaZeroo
committed
d05952e
refactor KimiMLAAttention to Inheritance titian BaseAttention
JavaZeroo
committed
6fe3694
remove add_zero_valued_dependency, this should be address with other multimodeal model.
JavaZeroo
committed
1fc70f6
remove redundandency test case
JavaZeroo
committed
522958f
connect FLA operator, replace RMSNorm

the-fall-moon
authored and
JavaZeroo
committed
51ad154
revert fla operator

the-fall-moon
authored and
JavaZeroo
committed
bc0e271
remove test_kimi_k3_hf_parity.py
JavaZeroo
committed
c5b9330
connect FLA operator, replace RMSNorm

the-fall-moon
authored and
JavaZeroo
committed
762f4f4
resotre to eager conv and RMSNormGated
JavaZeroo
committed
4f972ba
remove KimiExactGELU, replaced with nn.GELU
JavaZeroo
committed
dc009a9
Add Kimi K3 numerical validation
JavaZeroo
committed
d37a4a5
remove parallelism config
JavaZeroo
committed
2509f18
change full_attention_layers to starting with 0
JavaZeroo
committed
89d93a9
support flex attention
JavaZeroo
committed
ae006e4
refactor KimiK3StateDictAdapter
JavaZeroo
committed
fb7865a
refactor KimiGroupedExperts to use self.grouped_mm
JavaZeroo
committed
9f02c89
use snapshot_download to pin the version
JavaZeroo
committed
588c62e
refator test case
JavaZeroo
committed
edccb38
fix pyrefly
JavaZeroo
committed
02960f7
support full kimi k3 and refactor readme
JavaZeroo
committed
2e7603a
fix some NIT and some bug
JavaZeroo
committed
22cfd54
reuse VisionAttention
JavaZeroo
committed
f7feaaa
simplify test_kimi_k3.py
JavaZeroo
committed
af8b872
reuse common MoonViT helpers, VisionTransformerBlock, and MoE
JavaZeroo
committed
f8523a2
address review feedback on the debug model, tests, and README
JavaZeroo
committed
71c9804
register kimi_k3 in the CLI option freeze
JavaZeroo
committed
8d71077
fix unused change
JavaZeroo
committed
e6e6e64
reuse ComplexRoPE
JavaZeroo
committed
905b366
scale up k3 model
JavaZeroo
committed
0f45ff8
remove nit
JavaZeroo
committed
6273274
remove unsed head_dim in KimiDeltaAttention
JavaZeroo
committed
d56c69e
remove unsed riase Error
JavaZeroo
committed
5e006ba
remove some redundancy assert
JavaZeroo
committed
99c80c0
fix activation checkpointing
JavaZeroo
committed
f925dad
Add TODO for _apply_attention_residual
JavaZeroo
committed
950f547
split k3 model.py to moe.py lda.py
JavaZeroo
committed
6689074
remove two unsed attn res from hf, add it in state_dict_adapter
JavaZeroo
committed
252ae22
numerical_tests_kimi_k3.py support different attn_backend and bf16
JavaZeroo
committed
2ef1981
Adapt Kimi K3 to the folded token layout
JavaZeroo
committed
41447a4
update sample packing checking
JavaZeroo
committed
4149422
add TODO
JavaZeroo
committed
9f8bcf0
chang fla kda from fp32 to bf16
JavaZeroo
committed
44d0e83
fix test_update_from_config_propagates_moe_force_load_balance when rebase
JavaZeroo
committed
dee45e3

Commits on Aug 23, 2026

rename probs_TN to probs_T1N
JavaZeroo
committed
84e9f7b
removed dtype args, only support bf16 now
JavaZeroo
committed
157440a
Clarify the README that the numerical results are based on a reduced version of reduced hf config
JavaZeroo
committed
075d54c
remove test_update_from_config_propagates_moe_force_load_balance
JavaZeroo
committed
99c2c68
remove TestKimiK3FSDP, and add this in to github ci
JavaZeroo
committed
0999206
reuse some function from k2.7
JavaZeroo
committed
32abaa0
share the MoonViT tower between k2.7 and k3
JavaZeroo
committed
ded1bf1

Commits on Aug 24, 2026

Update .github/workflows/integration_test_8gpu_features.yaml

JavaZeroo
and
shuhuayu
authored
b3751c9
fix cicd
JavaZeroo
committed
2dd9390
Update scripts/checkpoint_conversion/numerical_tests_kimi_k3.py
shuhuayu
authored
88cb800