Kimi-K3 training support
This issue tracks Kimi-K3 training support in Megatron Core.
Kimi-K3 combines a hybrid Kimi Delta Attention (KDA) / gated Multi-Head Latent Attention (MLA) backbone, Attention Residuals (AttnRes), Stable LatentMoE, global-batch Quantile Balancing (QB), MoonViT-V2, per-head Muon, and native MXFP4 quantization-aware training.
Overall status: the baseline KDA/gated-MLA hybrid path, Stable LatentMoE, and non-CUDA-Graph
global-batch QB are available in dev. The September 18 preview adds MTP,
per-head Muon for KDA/MLA, and initial long-context training support, including TP/CP numerical
alignment with packed THD inputs up to 256K tokens. Full-model and 9L-proxy MXFP8 recipes are
published with measured performance. K3-specific upstream integration, full-model checkpoint
resume, and broader correctness and convergence qualification remain in progress.
Status: ✅ available in the stated scope / merged / resolved · in progress / open / draft · planned / pending validation
Availability in the frozen preview and upstream merge status are stated separately. A merged supporting PR does not necessarily mean the entire capability has completed end-to-end qualification.
Primary scope follows the Kimi-K3 section of the 2026 Q3 Megatron Core MoE roadmap. Lower-priority MoonViT and per-head Muon work, plus native MXFP4 training, are retained here so the tracker covers the intended Kimi-K3 training stack without turning into a complete paper gap list.
Last updated: 2026-09-18
Status at a glance
Core functionality
| Capability | Status | Priority | Summary |
|---|---|---|---|
| Model architecture | P0 | Text backbone integrated in the preview; K3 attention and AttnRes upstream PRs remain open | |
| MTP | ✅ Preview | P0 | Hybrid KDA/MLA + AttnRes MTP is supported and enabled in both published recipes |
| Distributed and long-context training | ✅ Preview | P0 | TP/CP + packed THD alignment up to 256K; variable-length AttnRes with non-interleaved PP |
| Quantile Balancing | ✅ | P0 | Core QB and TE fixes are merged; preview adds padding-aware statistics; end-to-end CUDA Graph qualification remains open |
| Per-head Muon | P0.5 | Available in preview recipes; #7085 and #7395 remain unmerged | |
| MoonViT and multimodal integration | P1 | MoonViT-V2 and the multimodal projector/token-integration path are planned | |
| Training validation and recipe | P0 | Full-model and 9L-proxy 256K recipes and benchmarks published; resume, broader correctness and convergence qualification remain pending |
Optimization work
| Area | Status | Priority | Current focus |
|---|---|---|---|
| Pipeline-parallel efficiency | P0.5 | Highest optimization priority: reduce the substantial PP communication bubbles and improve utilization across stages | |
| Long-context CP layout | P0.5 | Adapt the Qwen3.5 layout/routing-metadata optimizations from #6387 to Kimi K3 | |
| Fine-grained memory controls | P0.5 | Selective recompute and chunked optimizer-state offload are used in the recipes; explore more activation-offload/recompute choices | |
| Kernel and runtime optimizations | P0.5 | Updated compiled/FLA AttnRes, fused KDA pre-GDR, MLA gate JIT fusion and configurable validation frequency are in the preview; upstreaming continues | |
| MLA latent context parallelism | P1 | Experimental implementation remains draft in #6829 | |
| CUDA Graphs | P1 | Qualify the merged TE bounds fix in the K3 stack and validate end-to-end KDA + MoE capture/replay | |
| Native MXFP4 training | P1 | Add the K3 MXFP4-weight / MXFP8-activation QAT, optimizer, and checkpoint path |
Kimi-K3 Preview Branch
The rolling Kimi-K3 preview development branch is
kimi_k3_dev. It is the
latest integration point for Kimi-K3-specific work ahead of upstream completion and should not be
treated as an immutable release.
Preview history
| Preview | Commit | Abstract |
|---|---|---|
kimi_k3_preview_20260907 |
ba527e6ef2cb2b34eabd77c4aeb86fb1a7ae76a4 |
First frozen Kimi-K3 preview: integrates low-rank KDA, gated MLA with per-layer NoPE, native Block AttnRes, Stable LatentMoE with SiTU-GLU, and global-batch QB in one training stack. |
kimi_k3_preview_20260918 |
4adbeff83d019ebe242faba5b5d02e15ffa7513e |
Extends functional support with MTP and per-head Muon, establishes TP/CP + packed THD training up to 256K, updates AttnRes/KDA/MLA optimizations, and publishes full-model and 9L-proxy recipes with measured throughput. |
1. Core Functionality
1.1 Model Architecture
Status: ✅ The text-backbone composition is available in the frozen preview; K3 attention extensions and native AttnRes are still being upstreamed.
| Component | Status | Current work |
|---|---|---|
| KDA | ✅ Preview / upstream | First-class hybrid KDA landed in #6556; K3 low-rank F-decay/output-gate projections from open #6877 are integrated in the preview |
| Gated MLA | ✅ Preview / upstream | Baseline gated MLA landed in #6556; the preview includes per-layer NoPE, latent-projection extensions and elementwise output-gate JIT fusion from open #6877 |
| Attention Residuals | ✅ Preview / upstream | Native Block AttnRes from draft #6840, updated compiled/FLA paths, and packed-sequence support with non-interleaved PP are integrated in the preview |
| MTP | ✅ Preview / upstream | Hybrid KDA/MLA + AttnRes MTP from #6840 is integrated; both published recipes enable one MTP prediction depth |
| Stable LatentMoE | ✅ Available in dev |
Latent RMSNorm and dense/routed/shared-expert SiTU-GLU support are merged; focused BF16 forward/backward parity passed |
Stable LatentMoE work
| Capability | Status | Tracking / PRs |
|---|---|---|
| Kimi-K3 latent-MoE RMSNorm | ✅ Merged | #6804 |
| SiTU-GLU MCore integration | ✅ Merged | #6673 |
| SiTU-GLU Transformer Engine support | ✅ Merged | TransformerEngine#3402 |
| SiTU-GLU cuDNN Frontend backend | ✅ Merged | cudnn-frontend#645 |
| SiTU-GLU backend API-contract fixes | ✅ Merged | cudnn-frontend#670 |
The optional post-combine output normalization in #6449 is not required by the current Kimi-K3 target configuration and is no longer on this tracker's K3 critical path.
1.2 Quantile Balancing
Status: ✅ The eager/non-graph K3 global-batch path and fused histogram backend are available.
| Capability | Status | PR |
|---|---|---|
| Baseline Quantile Balancing | ✅ Merged | #5349 |
| Kimi-K3 global-batch Quantile Balancing | ✅ Merged | #6637 |
| Dense routing maps for Flex dispatch | ✅ Merged | #6614 |
| Fused QB router histogram path | ✅ Merged | TransformerEngine#3395 |
| Mutable QB bounds under CUDA Graphs | ✅ Merged; K3 graph qualification pending | TransformerEngine#3426 |
The preview also excludes padded tokens from QB histograms and aligns padded SBHD validation with packed THD semantics. End-to-end CUDA Graph readiness is tracked separately below.
1.3 Per-head Muon
Status: ✅ Available in the frozen preview and both released recipes; upstream integration and checkpoint continuity qualification remain in progress.
Open #7085 provides per-head Q/gate/K/V orthogonalization for supported MHA/GQA/MLA layouts and TP handling for fragmented query-group blocks. The preview adds KDA projection splitting and the replicated Transformer Engine weight metadata fix from draft #7395. Both published 256K recipes enable per-head Muon with duplicated TP mode. Upstreaming the K3-specific integration and validating checkpoint/resume continuity remain open.
1.4 MoonViT and Multimodal Integration
Status: P1 work planned.
- Add the MoonViT-V2 vision encoder training path.
- Add the multimodal projector and packed image/video-token integration with the Kimi-K3 language model.
No public MCore implementation PR is currently tracked.
1.5 Training Validation and Recipe
Status: Distributed correctness coverage and published recipes have progressed; complete full-model training qualification remains pending.
- BF16 module-level forward/backward parity for KDA, MLA NoPE + output gate, Stable LatentMoE, AttnRes, SiTU-GLU, and QB.
- Align supported TP/CP + packed THD configurations with the padded SBHD baseline, establishing training support up to 256K context length.
- Exercise MTP and per-head Muon in full-model and 9L-proxy 256K MXFP8 runs.
- Publish full-model and 9L-proxy recipes, shared container build/launch instructions, and recorded throughput and memory results.
- Qualify full-model checkpoint save/resume, including optimizer state continuity.
- Extend full-model correctness coverage and establish short-convergence/training-quality evidence with real data and normal router balancing.
Released recipes and benchmarks
Both recipes use MXFP8, packed THD sequences of 262144 tokens, a global batch size of 64, one MTP prediction depth, and per-head Muon. The 9L proxy keeps the full model's widths and expert count while reducing the backbone to nine decoder blocks.
| Recipe | Measured hardware | TP / PP / EP / CP | Median TFLOP/s/GPU |
|---|---|---|---|
| Full text backbone, 93 decoder blocks | 256 NVIDIA GB300 GPUs | 4 / 4 / 64 / 16 | 498 |
| 9L proxy | 64 NVIDIA GB300 GPUs | 4 / 1 / 64 / 16 | 661 |
The proxy configuration targets GB200; the reported measurement was collected on GB300. Results are medians over steps 6–18 of 20-step runs using mock data and forced balanced routing. The benchmark guide records iteration times, peak memory and the measurement environment, including the HybridEP revision difference from the published Dockerfile. These runs establish executable recipes and measured performance; convergence qualification remains open.
The published recipes use Megatron-LM's pretrain_hybrid.py. A separate Megatron Bridge model
provider or multimodal delivery path remains outside these text-backbone recipes.
2. Optimization Work
The next priorities are PP communication bubbles, Kimi-K3 adaptation of the Qwen3.5 CP layout work, and further activation-offload/recompute tradeoffs.
| Area | Status | Plan / dependency |
|---|---|---|
| Pipeline-parallel efficiency | Planned | Reduce the substantial communication-related bubbles in the full-model PP4 run and improve utilization across pipeline stages |
| Long-context CP layout | Planned | Adapt the merged Qwen3.5 work in #6387 to Kimi K3, reducing layout conversions and reusing packed-sequence routing metadata |
| Fine-grained recompute/offload | In progress | Selective KDA/MoE/MLP/MLA-up-projection recompute and chunked optimizer-state offload are used in the recipes. Explore finer choices with AttnRes to balance memory and throughput; the preview also includes the FP8 gather/offload fix from open #7392 |
| KDA fusion and runtime validation | ✅ Preview / upstream | Both recipes enable fused pre-GDR preprocessing from draft fork PR #7. Configurable validation frequency from draft #7202 and its KDA integration reduce repeated synchronization |
| AttnRes and MLA output gate | ✅ Preview / upstream | Updated compiled/FLA AttnRes and the single-source fast path are integrated; FLA is the default, while the recipes use compiled AttnRes. Elementwise MLA gating uses JIT fusion. #6840 and #6877 remain unmerged |
| MLA latent context parallelism | Draft | Complete #6829 and validate correctness, memory, communication and scaling; this is additional optimization work beyond the standard CP path used by the 256K recipes |
| CUDA Graphs | In progress | TransformerEngine#3426 is merged. Qualify a K3 stack containing the fix and validate end-to-end KDA + MCore MoE capture/replay, building on cudnn-frontend#556 |
| Native MXFP4 training | Planned | Support K3's MXFP4 weights / MXFP8 activations from QAT through optimizer update, distributed execution and checkpointing; the current published recipes use MXFP8 |
References
Source: NVIDIA/Megatron-LM