feat(moe): add MoK MXFP8 expert support - #3606
Conversation
Signed-off-by: Huiying Li <huiyingl@nvidia.com>
|
/claude review |
|
/ok to test ad5daa2 |
|
Review Summary Reviewed all 6 changed files (uv.lock excluded per policy). This PR adds an MXFP8 routed-expert precision path to the MoK backend, with a per-layer prequantized-weight cache retained across microbatches until the optimizer step. The design is careful: the cache key includes storage identity + The forward/backward math is delegated to MoK kernels (unchanged reference), and the changed Python routing is well tested with identity assertions on which tensors/layouts flow to forward vs backward. Findings — two minor tensor-contract documentation gaps on changed public signatures (inline):
No correctness, gradient-determinism, ownership, or security issues found in the changed code. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| routed_gate_weights: torch.Tensor, | ||
| routed_up_weights: torch.Tensor, | ||
| routed_down_weights: torch.Tensor, | ||
| mxfp8_weights: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ...] | None, |
There was a problem hiding this comment.
The new tensor-bearing mxfp8_weights parameter added to backward's signature is not documented in the Args: section. It is a tuple of (gate, up, down), each a 4-tuple of MXFP8 tensors (weight_fp8, weight_scale, weight_t_fp8, weight_t_scale), and must be None for bf16 / non-None for mxfp8 (the body raises on the mismatched case). Please add an Args entry documenting the layout and the None/precision contract.
Signed-off-by: Huiying Li <huiyingl@nvidia.com>
Signed-off-by: Huiying Li <huiyingl@nvidia.com>
1 participant
Summary
Add MXFP8 routed-expert compute to the Mixture-of-Kittens backend.
backend.mok.precision: bf16 | mxfp8with BF16 as the default.optimizer.step()updates the BF16 master weights.ep_shard > 1with a clear error. Expert-FSDP support is intentionally outside this PR.Performance
Both comparisons use full pretrained weights, FineWeb data, 8K sequence length, global batch size 64, world size 64, EP64, PP1, 20 steps with 5 warmup steps, and fake balanced routing disabled. Speedup is normalized to the HybridEP BF16 step time for the same model. Lower step time and higher speedup are better.
DeepSeek V4 Flash — TileLang attention
%%{init: {"xyChart": {"width": 520, "height": 300}}}%% xychart-beta title "DeepSeek V4 Flash 8K speedup vs HybridEP BF16" x-axis ["HEP BF16", "HEP MXFP8", "MoK BF16", "MoK MXFP8"] y-axis "Speedup (x)" 0.9 --> 1.3 bar [1.000, 1.066, 1.136, 1.234]torch_mm)MoK MXFP8 is 1.086x faster than MoK BF16 and 1.158x faster than HybridEP MXFP8.
GLM-5.2 — cuDNN attention
%%{init: {"xyChart": {"width": 520, "height": 300}}}%% xychart-beta title "GLM-5.2 8K speedup vs HybridEP BF16" x-axis ["HEP BF16", "HEP MXFP8", "MoK BF16", "MoK MXFP8"] y-axis "Speedup (x)" 0.9 --> 1.3 bar [1.000, 1.111, 1.174, 1.295]torch_mm)MoK MXFP8 is 1.104x faster than MoK BF16 and 1.166x faster than HybridEP MXFP8.
100-step convergence parity
The full pretrained models were trained on the same deterministic FineWeb shard with 8K sequence length, global batch size 64, world size 64, EP64, PP1, activation checkpointing enabled, and fake balanced routing disabled. Each comparison uses the same initial checkpoint and data order. Step time excludes the first five warmup steps.
Validation
218 passed: focused MoK, MoE parallelizer, and training-recipe unit tests.git diff --checkpasses.