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 exact K3 low-rank KDA / NoPE MLA extensions and
native AttnRes remain open. A preview integration stack has passed BF16 module parity and a
29-layer end-to-end training smoke; full 93-layer correctness, checkpoint-resume, performance,
and convergence qualification remain pending.

Status: ✅ available / merged / resolved · 🚧 in progress / open / draft · 📋 planned / pending validation

> Status represents end-to-end capability readiness. A merged supporting PR does not necessarily
> mean the entire capability is complete.

> Primary scope follows the Kimi-K3 section of the [2026 Q3 Megatron Core MoE roadmap](https://github.com/NVIDIA/Megatron-LM/issues/6757). 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-10

## Status at a glance

### Core functionality

| Capability | Status | Priority | Summary |
| --- | --- | --- | --- |
| Model architecture | 🚧 | P0 | Baseline KDA/gated MLA and Stable LatentMoE are in `dev`; exact K3 extensions and AttnRes are open; preview composition has been exercised |
| Quantile Balancing | ✅ | P0 | Baseline QB, global-batch QB, dense routing maps, and the fused TE histogram path are merged; CUDA Graph bounds handling remains open optimization work |
| Per-head Muon | 🚧 | P0.5 | The `dev` port is open in [#7085](https://github.com/NVIDIA/Megatron-LM/pull/7085); KDA-specific integration exists on the preview branch but is not upstream yet |
| MoonViT and multimodal integration | 📋 | P1 | MoonViT-V2 and the multimodal projector/token-integration path are planned |
| Training validation and recipe | 🚧 | P0 | BF16 module parity and an E2E full model training smoke passed; full-model resume, convergence, performance, and reproducible delivery remain pending |

### Optimization work

| Area | Status | Priority | Current focus |
| --- | --- | --- | --- |
| KDA kernel fusion | 🚧 | P0.5 | Fused pre-GDR preprocessing is in draft [fork PR #7](https://github.com/yuzhongw-nvidia/Megatron-LM/pull/7); upstreaming and refreshed-tip validation remain |
| MLA latent context parallelism | 🚧 | P1 | Experimental implementation is open in [#6829](https://github.com/NVIDIA/Megatron-LM/pull/6829) |
| Fine-grained memory controls | 🚧 | P0.5 | General recompute/offload foundations exist; qualify K3 submodule-level MoE choices and their AttnRes composition |
| CUDA Graphs | 🚧 | P1 | Complete mutable-QB-bounds support 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`](https://github.com/yuzhongw-nvidia/Megatron-LM/tree/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`](https://github.com/yuzhongw-nvidia/Megatron-LM/tree/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. |

## 1. Core Functionality

### 1.1 Model Architecture

Status: 🚧 Most architectural building blocks are available; exact K3 attention extensions and
upstream AttnRes completion remain in progress.

| Component | Status | Current work |
| --- | --- | --- |
| KDA | 🚧 K3 extension open | First-class hybrid KDA landed in [#6556](https://github.com/NVIDIA/Megatron-LM/pull/6556); K3 low-rank F-decay/output-gate projections are in [#6877](https://github.com/NVIDIA/Megatron-LM/pull/6877) |
| Gated MLA | 🚧 K3 extension open | Baseline gated MLA and hybrid allocation landed in [#6556](https://github.com/NVIDIA/Megatron-LM/pull/6556); per-layer NoPE and fused latent up-projections are in [#6877](https://github.com/NVIDIA/Megatron-LM/pull/6877) |
| Attention Residuals | 🚧 Draft | Native Block AttnRes with non-interleaved pipeline parallelism and MTP support is proposed in [#6840](https://github.com/NVIDIA/Megatron-LM/pull/6840); K3 hybrid-proxy composition has been exercised on the preview branch |
| 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](https://github.com/NVIDIA/Megatron-LM/pull/6804) |
| SiTU-GLU MCore integration | ✅ Merged | [#6673](https://github.com/NVIDIA/Megatron-LM/pull/6673) |
| SiTU-GLU Transformer Engine support | ✅ Merged | [TransformerEngine#3402](https://github.com/NVIDIA/TransformerEngine/pull/3402) |
| SiTU-GLU cuDNN Frontend backend | ✅ Merged | [cudnn-frontend#645](https://github.com/NVIDIA/cudnn-frontend/pull/645) |
| SiTU-GLU backend API-contract fixes | ✅ Merged | [cudnn-frontend#670](https://github.com/NVIDIA/cudnn-frontend/pull/670) |

The optional post-combine output normalization in [#6449](https://github.com/NVIDIA/Megatron-LM/pull/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](https://github.com/NVIDIA/Megatron-LM/pull/5349) |
| Kimi-K3 global-batch Quantile Balancing | ✅ Merged | [#6637](https://github.com/NVIDIA/Megatron-LM/pull/6637) |
| Dense routing maps for Flex dispatch | ✅ Merged | [#6614](https://github.com/NVIDIA/Megatron-LM/pull/6614) |
| Fused QB router histogram path | ✅ Merged | [TransformerEngine#3395](https://github.com/NVIDIA/TransformerEngine/pull/3395) |
| Mutable QB bounds under CUDA Graphs | 🚧 Open | [TransformerEngine#3426](https://github.com/NVIDIA/TransformerEngine/pull/3426) |

### 1.3 Per-head Muon

Status: 🚧 P0.5 work in progress.

[#7085](https://github.com/NVIDIA/Megatron-LM/pull/7085) is the `dev` port of per-head Muon. It
adds opt-in Q/gate/K/V orthogonalization for supported MHA/GQA/MLA layouts and handles tensor
parallelism that fragments query-group blocks. The preview integration branch additionally tags
K3 KDA low-rank and legacy projections for per-head treatment and has passed prior eight-rank
unit validation. Upstream KDA integration, refreshed-tip validation, checkpoint continuity, and
an end-to-end K3 recipe remain pending.

### 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: 🚧 Initial proxy correctness is established; full-model release qualification is pending.

- [x] BF16 module-level forward/backward parity for KDA, MLA NoPE + output gate, Stable LatentMoE, AttnRes, SiTU-GLU, and QB.
- [ ] Exercise representative distributed configurations and checkpoint resume.
- [ ] Publish a reproducible recipe with correctness, memory, performance, and short-convergence evidence.

A concrete model provider and user-facing recipe may live in Megatron Bridge; this tracker records the Megatron Core capabilities required to run it.

## 2. Optimization Work

| Area | Status | Plan / dependency |
| --- | --- | --- |
| KDA kernel fusion | 🚧 Draft | Upstream the fused pre-GDR preprocessing from [fork PR #7](https://github.com/yuzhongw-nvidia/Megatron-LM/pull/7), then validate kernel parity and performance on the refreshed #6877 base |
| MLA latent context parallelism | 🚧 Draft | Complete [#6829](https://github.com/NVIDIA/Megatron-LM/pull/6829) and validate correctness, memory, communication, and scaling |
| Fine-grained recompute/offload | 🚧 In progress | Qualify selective KDA/MLA/MoE recompute and the lowest-memory shared-expert / MoE-activation combination without whole-MoE recompute; close remaining AttnRes composition gaps |
| CUDA Graphs | 🚧 In progress | Land [TransformerEngine#3426](https://github.com/NVIDIA/TransformerEngine/pull/3426) for mutable QB bounds and validate end-to-end KDA + MCore MoE capture/replay, building on [cudnn-frontend#556](https://github.com/NVIDIA/cudnn-frontend/pull/556) |
| Native MXFP4 training | 📋 Planned | Support K3's MXFP4 weights / MXFP8 activations from QAT through optimizer update, distributed execution, and checkpointing |

## References

- [Kimi K3 technical report](https://arxiv.org/abs/2607.24653)
- [Official Kimi-K3 repository](https://github.com/MoonshotAI/Kimi-K3)
- [Official Kimi-K3 Hugging Face release](https://huggingface.co/moonshotai/Kimi-K3)
- [2026 Q3 Megatron Core MoE roadmap](https://github.com/NVIDIA/Megatron-LM/issues/6757)
- [Kimi-K3 rolling preview branch](https://github.com/yuzhongw-nvidia/Megatron-LM/tree/kimi_k3_dev)
- [DeepSeek-V4 training-support tracker (format reference)](https://github.com/NVIDIA/Megatron-LM/issues/4468)