Commit Graph
2 Commits
Author SHA1 Message Date
dela d1da0816f2 LatentMoE: sparse permute-dispatch + padded bmm
Replace dense all-expert forward (16 experts × all tokens) with
permute-dispatch: sort token-expert pairs by expert id, pad to
[R, C, ℓ] (C = max tokens per expert), run 3 bmm calls for the
batched SiTU-GLU activation, then scatter-add weighted results back.

Routed expert FLOPs drop from R·N to R·C (C ≈ N·k/R under uniform
routing). SiTU parameter structure unchanged; checkpoint compatible.

Tests: sparse-vs-dense fwd/bwd equivalence, unselected expert zero
grad, last_capacity tracking.
2026-08-25 17:47:44 +08:00
dela 584f7e9e73 Initial K3 snapshot: 0.5B KDA/MLA/MoE train path
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias
backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
2026-08-25 14:43:17 +08:00