SIGN IN SIGN UP

feat(moe): integrate Mixture-of-Kittens backend (#3422)

* feat(moe): integrate Mixture-of-Kittens backend

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* chore: remove local benchmark configs

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* fix(moe): checkpoint MoK backward context

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* fix(moe): support packed SFT with MoK

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* Fix MoK routing parity for padded GLM batches

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* perf(moe): avoid compacting packed THD inputs

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* Simplify MoK packed tail padding

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* fix(moe): preserve canonical MoK random initialization

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* feat(moe): support clamped SwiGLU with MoK

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* fix(moe): align clamped SwiGLU with MoK functional API

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* chore(moe): remove MoK comparison benchmark

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* chore(moe): remove GLM routing comments

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* feat(benchmark): add GLM-5.2 MoK EP4 config

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* fix(moe): align unpadded MoK dispatch inputs

Pad unmasked short or unaligned token batches to MoK's required dispatch
extent while keeping the returned output and gradients limited to real
rows. Also cover the explicit LoRA rejection path.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(benchmark): use GLM-4.7 and add DSV4 configs

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* fix(benchmark): use full DeepSeek V4 model

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* fix(ci): declare MoK benchmark recipe ownership

Add the required owner, wall-time, and node-count metadata for both new
CI-discovered benchmark recipes.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* perf(moe): dispatch packed padding directly with MoK

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* perf(moe): validate packed MoK extent before training

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* refactor(moe): rely on MoK runtime shape checks

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* refactor(moe): dispatch padding rows directly through MoK

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* refactor(moe): require pre-aligned MoK token extents

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* docs(moe): clarify MoK token extent contract

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* Update nemo_automodel/components/moe/mok_experts.py

Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

* Update nemo_automodel/components/moe/mok_experts.py

Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

* refactor(moe): tighten MoK backend wiring

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

* test(moe): avoid comparing uninitialized router values

Signed-off-by: Huiying Li <huiyingl@nvidia.com>

---------

Signed-off-by: Huiying Li <huiyingl@nvidia.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
H
Huiying committed
ade4fccc6fbbcde44c34b4e1363bc0313506d300
Parent: e539952
Committed by GitHub <noreply@github.com> on 8/12/2026, 7:48:20 PM