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