[DO NOT MERGE, pending rebase and sequential smaller PRs after eager-model merge] Kimi K3: full tree with 5D parallelism multimodal, LoRA, MoonViT Dynamic CP&DEP, MXFP4/MXFP8 QAT, quantile balancing

[QIU023]

Draft, and not for merge, providing for early timeline and accessable for full 2.8T param 5D smoke while waiting for PR-4025.

Would need maintainers thorough review on later major parallelize.py sharding.py declaration on major module parallelism support in titan.

Pending final evaluations and fixes before publishing: MXFP4 QAT and packed-MXFP4 import, quantile load balancing, MTP, LoRA, the decoupled vision encoder, a MoonEP token dispatcher, and the HF <-> DCP key map for the released checkpoint.

PR-4025 is a separate implementation of the same model and is further along on the model itself. When it lands, this rebases onto it and what remains is the parallelism and the post-training pieces. Filing now is disclosure, not a request to review 28k lines.

98 files, +28307 -28. The 28 deletions are the thing to notice -- this is almost entirely additive, and the five core files it does touch are listed in commit 14.

The fifteen commits are sliced by content rather than by the 409 commits of history that produced them:

#	commit	main files
01	the K3 model: KDA, MLA, latent MoE, MTP, and its config tree	model.py model_configs.py config_registry.py
02	Block Attention Residuals: the primitive and the layout tables	attn_res.py attn_res_model.py layout.py
03	SiTU-GLU routed experts and quantile load balancing	moe.py quantile_balance.py common/moe.py
04	MoonViT and the multimodal wrapper	moonvit.py multimodal_model.py vision_preprocess.py
05	the quantization scope -- what K3 quantizes, shared by QAT and QLoRA	quant_scope.py
06	context parallelism: KCP for KDA, Ulysses for MLA, dynamic CP for the tower	kcp.py vit_cp_plan.py
07	pipeline parallelism: the cross-stage AttnRes adapter	pipeline_adapter.py
08	the decoupled vision encoder, with bubble scheduling	dep_bubble_*.py vit_prefetch.py
09	a MoonEP token dispatcher against torchtitan's EP seam	moon_ep_dispatcher.py
10	MXFP4 QAT and packed-MXFP4 weight import	mxfp4_qat.py packed_mxfp4.py
11	LoRA, including the skip-edge gradients PP has to route	lora.py muon.py
12	HF <-> DCP conversion for the released key set	hf_key_map.py state_dict_adapter.py
13	the parallelize entry that applies all of the above, incl. the TP plan and the KDA DTensor shims	parallelize.py
14	the core changes the above needs	distributed/{fsdp,utils}.py components/{lr_scheduler,optimizer}.py
15	tests	tests/ (56 files)

Two upstream defaults have to be overridden to run it, both pinned in our gate rather than worked around in the model. spmd_backend needs partial_dtensor: under spmd_types, fully_shard() wants every parameter to already be a DTensor on the full SPMD mesh, and TP additionally trips assert_type() does not support DTensor, so supporting it is the declarative conversion itself rather than a patch on top -- that is the largest gap between this parallelize.py and an upstream model's, and it is what we are doing next. And CUDA graph capture needs --training.disable-cuda-graphs, because the vision path's patch count varies per batch and capture validation rejects it; that one is not specific to this model, so it may deserve its own issue.

Evidence: 58 of 58 gate cells pass on this tree, ten steps each, with those two pins. Per-cell loss and grad_norm for every step of all 58 are here, one line per cell:

https://github.com/QIU023/torchtitan_attention_residual/blob/f81a19319506ef1c3d3e2b8fc6eadbf7e6d99feb/phase13_k3like_48b_posttrain/gate_logs/gate_58_2026-08-21_final_percell.txt

The three model arms are text, multimodal, and multimodal+LoRA, 18 cells each, over dp / fsdp / tp / cp / pp / ep and their combinations up to ep2_fsdp2_tp2_cp2; the remaining four are the two decoupled-encoder cells at pp4 and pp8 and the two pp8xvp4 cells.