[model, ops] feat: Kimi K3 on Ascend NPU training enablement

[wangyier7]

What does this PR do?

Scope claim for the complete Kimi K3 Ascend NPU enablement. Kimi K3 is a multimodal MoE — MoonViT3d vision tower + hybrid KimiDelta linear-attention / MLA text backbone + fine-grained MoE with shared experts. The branch lands the model integration with a patch-generation pipeline, NPU kernels, training recipes, docs, and validation configs. Parallel-training adaptation (FSDP2 + EP + Ulysses sequence-parallel with KDA/MLA head sharding) is included in the model patches; performance tuning and scale-out will follow as separate PRs on top of this base.

The branch is rebased onto main@83059808 and re-registered on the post-#1142 VeomniOp registry: patch rules use VeomniOp handles, sequence-parallel attention rides the built-in adapter Ulysses (veomni_flash_attention_2, no _with_sp aliases), and the fused-MoE kernel reuses the relocated moe_experts/standard/npu.py helpers instead of vendored copies.

Checklist Before Starting

• Search for relative PRs/issues and link here: none found — no existing kimi_k3 / Kimi K3 PR or issue in the repo.
• PR title follows [{modules}] {type}: {description} format: [model, ops] feat: Kimi K3 multimodal MoE training on Ascend NPUs.

Test

Two existing CI-enumerated tests were extended to cover the K3 registrations; validation beyond CI runs on an out-of-repo training harness on 8× Ascend 910B3 (COCO val2017):

• Extended: tests/ops/base/test_op_config_names.py (new chunk_kda_implementation / situ_and_mul_implementation field mappings) and tests/ops/base/test_builtin_entries.py (K3 rows in the builtin-entries snapshot: eager/ascendc KDA, eager/triton SituGLU, eager/fused_npu MoE, plus NPU defaults). The existing tests/models/transformers/test_generated_modeling_imports.py imports the new generated files.
• pytest on NPU: ops contracts 72 passed; moe_experts + rms_norm + argument aliases 105 passed / 63 skipped (skips are CUDA-only); attention op/mask + flash-attention suites green; loss green.
• Repo-wide patchgen --check green (33 configs); ruff check + ruff format --check clean under the CI-pinned ruff 0.13.2; scripts/ci/check_doc_task_paths.py passes.
• Post-rebase NPU smoke (seq 1024, SP2 / EP2 / FSDP2; fused_npu + AscendC KDA + triton SituGLU + chunk_loss): 200 steps, zero errors, loss 13.51 → 6.72, VRAM peak 50.15 GB (parity with the pre-rebase 50.12 GB baseline), ~16.9 s/it; DCP checkpoint round-trip verified; the profile trace confirms the situ / KDA / fused-MoE kernels active.
• Pre-rebase equivalence gates (earlier commits): step-0 loss bitwise identical across SP1/SP2; 300-step real-data runs across attention-backend eras (eager / AscendC / fused-flash).
• Wiring a kimi_k3 entry into the e2e workflows is planned once the enablement is stable.

API and Usage Example

Three additive OpsImplementationConfig fields (default eager; fused backends opt-in via the recipes):

ops = OpsImplementationConfig(
    chunk_kda_implementation="ascendc",     # fused AscendC KDA kernel
    situ_and_mul_implementation="triton",   # fused Triton-ascend SituGLU kernel
    moe_implementation="fused_npu",         # situ fused MoE group-GEMM (EP > 1)
)

The recipes set attn_implementation="veomni_flash_attention_2" — sequence parallelism comes from the adapter's built-in Ulysses, no name-suffix toggles. Recipes: configs/multimodal/kimi_k3/kimi_k3_vl{,_ascendc}.yaml; end-to-end walkthrough in docs/examples/kimi_k3.md. Env gates: K3_KDA_HEAD_SHARD=0 rolls back to the full-gather KDA path (default on); K3_KDA_SP_TIMING=1 is an opt-in phase-timing dump.

Design & Code Changes

• veomni/models/transformers/kimi_k3/ — vendored model (source/) + patch-generation configs (source of truth; every patch and adaptation fix enumerated in their docstrings; VeomniOp handle pattern) + committed generated/ NPU/GPU modeling files with .diff companions (verified by patchgen --check, never hand-edited) + EP parallel plan + checkpoint tensor converter (targets the relocated veomni.models.checkpoint package).
• NPU kernels registered as platform ops: chunked KDA (AscendC, 795 L), SituGLU (triton, 260 L), situ fused MoE group-GEMM (rebased onto moe_experts/standard/npu.py); rms_norm torch_npu variant.
• Framework changes are additive: kernels registry, COCO→K3 conversations data transform + converter script, argument aliases, trainer micro-batch plumbing, README support table, docs.

Commit layout (17 commits): the original landing series ([ops] kernels, [model] validation configs, [model] to_empty fix, [trainer] config-driven VLM training, [scripts] COCO converter, [docs] recipe, [chore] community-readiness, [configs] GPU recipe, [docs] toctree), followed by the main-rebase series: merge of main@83059808, [model] refactor: patch rules → VeomniOp handles, [model] fix: SP → built-in Ulysses + op value names, [ops] fix: op metadata contract, [model] fix: checkpoint converter re-point, [ops] fix: situ fused-MoE onto relocated helpers, [model] fix: timing console notice, [model, ops] chore: ruff format + patchgen diff refresh.

Checklist Before Submitting

• Read the Contribute Guide — yes
• Applied pre-commit checks — yes: ruff check + ruff format --check (CI-pinned ruff 0.13.2), repo-wide patchgen --check (33 configs), and scripts/ci/check_doc_task_paths.py all green locally
• Added/updated documentation — yes: docs/examples/kimi_k3.md (toctree-registered) + README support table
• If tasks/ training scripts were moved or renamed — N/A (no tasks/ changes; the recipes call the existing tasks/train_vlm.py, doc task paths verified)
• Tests — extended two existing CI-enumerated tests (see Test); e2e workflow wiring planned once stable

Note: the GPU generated modeling file is produced by the same patch pipeline and is NPU-validated only (no GPU environment was available).

[CLAassistant]

All committers have signed the CLA.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Coach257]

Heads-up on ordering: #1142 ([BREAKING][ops, model] refactor: replace OpSlot with VeomniOp) is planned to land before this PR, so please rebase onto main after it merges and build the kernel changes here on the new op registry.

What changes in #1142:

• OpSlot / OpsConfigSlot / KERNEL_REGISTRY / KernelSpec (veomni.ops.dispatch, veomni.ops.kernel_registry) are removed.
• Every kernel is a row registered with register_op(op, variant, impl, ..., requirement=..., requires=(...)), keyed by (op, variant, impl, device). The device is chosen with GpuKernelRequirement / NpuKernelRequirement / MluKernelRequirement.
• Modeling code holds a handle: VeomniOp(op, variant, resolve_op_impl("<field>_implementation")). It does not read config fields or bind slots.
• Kernels live under veomni/ops/kernels/<op>/<variant>/<impl>.py, with the rows in veomni/ops/kernels/<op>/__init__.py. Many files moved.

How to register and call an op: veomni/ops/README.md → Registering an op and Calling an op from modeling; config fields and available impls: docs/design/op_selection.md.

For this PR specifically:

• veomni/ops/kernels/moe/ became veomni/ops/kernels/moe_experts/, with NPU MoE in moe_experts/standard/npu.py. rms_norm/npu.py became rms_norm/standard/npu.py.
• The new families (glu, kda, and the npu_situ_gemm MoE path) should be added as veomni/ops/kernels/<op>/<variant>/<impl>.py plus register_op rows with NpuKernelRequirement().
• Kimi K3's patch_gen should call them through VeomniOp instead of OpSlot.