The Problem

MoK is a mixture-of-experts training megakernel targeting NVL72 infrastructure, but its codebase presents immediate engineering friction. Five files exceed 500 lines each, duplicated 311 six-line logic blocks appear across 13 benchmark files, and four functions exhibit indentation depths of 15 or more. There is no CI/CD pipeline, no dependency lockfile, and no update bot configured. A new contributor cannot verify that a changed build produces the same artifact that was tested.

What This Does

MoK fuses all MoE computation and communication into a single kernel, overlapping compute and inter-GPU networking at configurable granularity while eliminating CPU-GPU synchronization. The repository contains two abstraction layers: mok/ops.py provides the low-level CUDA kernel API where callers manage data layout and kernel coordination; mok/functional.py offers the higher-level API used in production, handling scratch memory and kernel launch coordination. The internal call graph resolves 295 edges between self-contained functions, with MoKConfig (21 callers), check_correctness (18), get_workspace (17), and build_schedule (17) as the most connected hubs. Execution starts at main in benchmarks/bench_deepep_te.py:225, which reaches all 34 functions in the repo. Three functions run external commands and one reads/writes files via setup.py. The dependency manifest (pyproject.toml) declares requirements without a lockfile, meaning transitive dependency versions are unpinied at release time.

How It Is Wired

Execution flows from main → init_distributed → MoKConfig → build_schedule → forward/backward → mxfp8_quantize/barrier_all. The most distinct callers are test_build_schedule → MoKConfig, test_forward_mxfp8 → mxfp8_quantize, and test_backward_mxfp8 → mxfp8_quantize. mok/ops.py defines the primitive kernels (all_gather_top_experts, schedule, mxfp8_quantize, dispatch_mlp_swiglu_combine_fwd_mxfp8) that mok/functional.py composes via create_workspace, get_workspace, clear_workspace_cache, and build_schedule. Five benchmark files each contain duplicated 6-line check patterns (311 total) that should be extracted into shared helpers. Four functions across mok/functional.py, mok/ops.py, and benchmarks/bench_hybridep_megatron.py have max indentation depth 15, making control flow hard to follow.

How To Use It

Setup: The package requires NVIDIA Blackwell SM100/SM103 GPUs, Python 3.12+, PyTorch 2.10+ built against CUDA 13.0+, and CUDA 13.0 toolkit. From the repo root:

pip install . --no-build-isolation

or:

python setup.py install

To build for SM100: MOK_ARCH=SM100 pip install . --no-build-isolation. Verify with python -c "import mok; print(mok.__version__)". For development: pip install -e . --no-build-isolation then make.

Configuration: Five hyperparameters affect MoE execution performance: fwd_num_comm_sms, bwd_num_comm_sms, fwd_num_compute_sms, bwd_num_compute_sms, and expert_capacity. Optimal values are workload-dependent and should be swept before production use.

Running it: Launch unit tests via torchrun and pytest:

torchrun --standalone --nproc-per-node=<num-gpus> -m pytest -s <test-path>

Benchmarks live in benchmarks/ and are invoked individually; the entry point main in benchmarks/bench_deepep_te.py:225 orchestrates grouped MLP runs.

Real-World Use

In a production training pipeline, a team would import mok.functional and call schedule(config) once to build the dispatch/combine schedule, then pass that schedule to forward(data, schedule) and backward(grad, schedule). Config would specify fwd_num_comm_sms=24, bwd_num_comm_sms=24, expert_capacity=2, and the workspace would be allocated once and reused across steps. The functional layer abstracts away the per-expert all-gather/top-k dispatch and the MXFP8/BF16 fused mm kernels, so the user expresses only the MoE layer topology and loss, not the communication schedule or kernel launch bookkeeping.

Code Health & Issues

  • [HIGH] Commit a lockfile beside pyproject.toml — manifest has no lockfile, so installed transitive dependencies can differ between test and ship artifacts.
  • [HIGH] Add a CI workflow that builds and tests on push and pull_request — 20 source files have no automated gate.
  • [MEDIUM] Enable Dependabot or Renovate — 1 manifest with no update bot; published advisories go unpatched without manual audit.
  • [HIGH/cognitive_load] Oversized files: mok/functional.py (628 lines), mok/ops.py, tests/test_functional.py — hard to hold in one head; changes ripple widely.
  • [HIGH/cognitive_load] Deep nesting x4: mok/functional.py, mok/ops.py, benchmarks/bench_hybridep_megatron.py — control flow exceeds comfortable readability threshold.
  • [HIGH/clarity] Duplicated code blocks: 311 repeated 6-line blocks across 13 benchmark files — violates DRY and creates maintenance risk.
  • [LOW/clarity] 5 TODO/FIXME markers in mok/functional.py — stale markers erode signal.

The Bottom Line

MoK delivers genuine performance gains — up to 2.37x faster MXFP8 forward vs. baselines on NVL72 — but the engineering cost is high. The codebase is concentrated in a few large, nested files with duplicated logic and no automated quality gates. Teams with strong CUDA expertise and bandwidth for technical debt can extract value by extracting shared benchmark helpers, adding a lockfile and CI, and tuning the five MoE hyperparameters. Others should wait until the code health improves or consider whether the performance delta justifies the operational overhead.