lk_moe
lk_moe is a highly optimized MOE hybrid inference engine that serves as the computational core for MOE layer execution. It is designed to maximize resource utilization across CPU, memory, and GPU, enabling large-scale MOE models to run efficiently beyond GPU memory constraints.
Key Capabilities
- Hybrid Compute Backend: Each MOE layer can dynamically choose between the original GPU computation path or lk_moe's hybrid inference path, offering flexible deployment strategies.
- CPU-GPU Collaborative Inference: Distributes expert computation across CPU (via NUMA-optimized threading) and GPU, reducing GPU memory requirements while maintaining high throughput.
- NUMA-Aware Parallelism: Optimized thread binding and memory allocation across NUMA nodes, achieving cross-node communication as low as 3% and L3 cache hit rates above 50%.
- Super Mix Prefill: GPU-path prefill with super-chunk sizing — the GPU chunk is decoupled from the token budget, keeping the GPU near 100% utilization while decode is served by the CPU. Always recommended whenever VRAM has headroom.
- GPU Pool (expert cache): A background DMA pump keeps hot experts resident in VRAM (2nd-touch window LRU, off the critical path). Pool-resident experts are computed on the GPU while the rest stay on the CPU — both in decode and prefill — the split is decided by one shared snapshot word, so every expert is computed exactly once. Works for both NVFP4 and MXFP4 experts. Fully fail-closed: an empty or disabled pool is bit-identical to the plain hybrid path.
- Quantization Kernel Support: Native support for multiple quantization formats including BF16, FP8, NVFP4, MXFP4, and AWQ 4-bit, with efficient kernel implementations.
- Memory + VRAM Load Balancing: Total model footprint spans both VRAM and system memory, allowing models larger than GPU memory to run without performance collapse.
Configuration
All parameters are read once at startup and are completely inert when the master switch is off. The only path choice made at runtime is the prefill batch gate (LK_GPU_PREFILL_MIN_TOKENS).
Policy fields (pool, prefill, mailbox, verify) are owned by MOEConfigV2 — the library takes them from cfg, not the environment; the LK_* names below are accepted by the sglang integration and bridged into the config. Read directly by the library: bootstrap env (LK_THREADS, LK_THREAD_BINDING, LK_NUMA_HUGE, LK_NUMA_ILP, LK_POWER_SAVING) and the diagnostics knobs (LK_POOL_VERIFY, LK_POOL_STAGE, LK_DECODE_TS, LK_TS_SKIPGEMM, LK_MBOX_STOPFILE).
Switch linkage (one master gate)
LK_MOE_HYBRID_ENABLED is the single master switch: with it off (the default) lk_moe is never imported and every LK_* below is inert — every MOE layer runs the framework-native path, bit-for-bit. Switches only link to switches; perf knobs never appear in a linkage:
LK_MOE_HYBRID_ENABLED=1 master
├─ pool on LK_POOL_SLOTS (e.g. 64) (dual knob: value itself is the switch, 0 = off)
│ └─ pool decode LK_POOL_DECODE=1 (unset = follows pool on; explicit `1`/`0` wins)
├─ GPU prefill on LK_GPU_PREFILL_MIN_TOKENS (e.g. 1024) (dual knob, same rule)
│ └─ super mix LK_GPU_PREFILL_SUB_M (e.g. 16384) (super-chunk staging on; 0 = per-chunk staging)
├─ nvme filler on LK_NVME_DIRECT=1 (needs the feature's own backend flag: --ple-offload-backend nvme / engram host table;
│ the LK master only gates the lk_moe path)
└─ resident layers LK_GPU_RESIDENT_MOE_LAYERS* (complementary: a resident layer never pools)
LK_MOE_HYBRID_ENABLED=0 → every branch above unreachable = original sglang
( LK_NVME_DIRECT with the engram/PLE nvme backends is the one knob that also works without the LK master )
Dual-knob rule (switch + perf in one variable)
Two variables are both a switch and a perf knob. The rule: default 0 = feature not enabled; to use a tier you must explicitly give a value (e.g. 64 / 1024 - one value both switches the tier on and sets its parameter).
| Variable | Switch meaning | Perf meaning | Default | Recommended when on |
|---|---|---|---|---|
LK_POOL_SLOTS |
value itself is the switch (e.g. 64 powers the pool tier on); unset / 0 = off, all LK_POOL_* inert |
resident expert slots per MOE layer; VRAM cost = slots × layers × expert bytes | 0 |
e.g. 64 |
LK_GPU_PREFILL_MIN_TOKENS |
value itself is the switch (e.g. 1024 powers the GPU-prefill tier on); 0 = pure CPU prefill |
value = minimum token batch for the GPU path | 0 |
e.g. 1024 (0 on 16 GB cards) |
Tuning parameters
| Group | Variable | Default | Description |
|---|---|---|---|
| Master | LK_MOE_HYBRID_ENABLED |
0 |
Engine master switch; 0 = every MOE layer runs the framework-native path. |
| Resident | LK_GPU_RESIDENT_MOE_LAYERS |
unset | Listed layer ids/ranges (e.g. 0, 9-11) stay on the native GPU path instead of the hybrid path; unset = none resident. |
| Resident | LK_GPU_RESIDENT_MOE_LAYERS_MTP / LK_GPU_RESIDENT_MOE_LAYERS_SPEC / LK_GPU_RESIDENT_MOE_LAYERS_DSPARK |
unset | Same grammar; for draft roles _SPEC wins when set, _MTP/_DSPARK apply only when _SPEC is unset. |
| Pool | LK_POOL_SLOTS |
0 |
Dual knob, see above (e.g. 64 when on). |
| Pool | LK_POOL_DECODE |
follows pool | Decode GPU half: 1 on / 0 off explicitly; unset follows LK_POOL_SLOTS > 0. |
| Pool | LK_POOL_WINDOW |
32 |
Perf only: expected-reuse window in steps — an evicted expert is pumped back only if seen again within the window. Illegal values clamp to min 8; clamping never flips a switch — only LK_POOL_SLOTS disables the pool. |
| Pool | LK_POOL_INFLIGHT |
16 |
Perf only: in-flight DMA jobs / event-ring depth (backpressure scales from this). |
| Pool | LK_POOL_STAGE |
1 |
Perf only: coalesced H2D staging ring (gather NUMA fragments, then one bulk copy per array). |
| Pool | LK_POOL_LOG |
0 |
Diagnostic: windowed [pool] stats line every N decode steps. |
| Pool | LK_POOL_VERIFY |
0 |
Diagnostic: byte-compare 1 out of N pumped slots against their NUMA source. |
| Pool | LK_POOL_MAIN_STREAM |
0 |
Diagnostic, sglang side only (never reaches cfg): launch the pool Triton kernels on the main stream instead of the side stream, to isolate stream-ordering suspects. |
| Prefill | LK_GPU_PREFILL_MIN_TOKENS |
0 |
Dual knob, see above. |
| Prefill | LK_GPU_PREFILL_SUB_M |
0 |
Super-chunk size (tokens) for the GPU prefill sub-batch; value itself is the switch (e.g. 16384 enables super mix); 0 = per-chunk staging. |
| NVMe | LK_ENGRAM_NVME |
0 |
DeepSeek-V4.1 Engram on the lk_moe N1 async filler: table stays in the safetensors shard, rows fetched by pread/O_DIRECT; the sglang integration exposes it as LK_ENGRAM_NVME. |
| NVMe | LK_ENGRAM_NVME_CACHE_MB / LK_ENGRAM_NVME_STAGING_MB |
1024 / 128 |
Perf only: Engram NVMe store cache / staging budgets. |
| Prefill | LK_GPU_PREFETCH_WINDOW |
1 |
Perf only: staging-ring blocks kept in flight for prefill; 1 = serialized (default); 2 = ping-pong overlap, pays off only with fast PCIe. |
| NUMA | LK_NUMA_INTERLEAVE |
1 |
Interleave expert weights across NUMA nodes (required for the pump's per-node copy streams; default-on via the sglang integration, 0 opts out). |
| CPU | LK_THREADS |
auto | CPU worker threads per GPU (physical cores / #GPUs). |
| CPU | LK_THREAD_BINDING |
CPU_CORE |
Thread binding mode across NUMA nodes. |
| Power | LK_POWER_SAVING |
0 |
Power-saving mode. |
| Spec | LK_VERIFY_RESIDENT |
0 |
Verify batches cover the draft rows' resident set. cfg field verify_resident. |
| Spec | LK_VERIFY_SKIP_MISS |
0 |
Skip verification for rows whose experts missed residency. cfg field verify_skip_miss. |
| Decode | LK_HF_MBOX |
1 |
Mailbox early-start of the host-side decode (non-0 on, 0 off; saves the ~50 µs graph host-node wake-up latency). cfg field hf_mbox. |
| Debug | LK_PP_DEBUG |
off | Extra tracing inside the prefill staging tier. |
| Debug | LK_ROUTE_PROBE |
off | Per-step [RP] routing traces for offline calibration (requires the pool on). |
Framework Integration
lk_moe integrates with popular LLM serving frameworks as a drop-in MOE computation backend:
| Framework | Integration | Description |
|---|---|---|
| vLLM | LvLLM | lk_moe as MOE layer backend for vLLM serving |
| sglang | Lsglang | lk_moe as MOE layer backend for sglang serving |
| DeepSeek V4 | Lvllmds4 | Specialized integration for DeepSeek V4 (SM120) |
| DeepSeek V4 | Lvllmds4-x | Specialized integration for DeepSeek V4 (SM80+) |
Installation
Prerequisite: CUDA 13.2.1 (driver + toolkit, sm80+).
pip install lk-moe
License
lk_moe is provided under a proprietary software license agreement. See LICENSE for terms and conditions. Third-party open-source components and their respective licenses are listed in THIRD_PARTY_LICENSES.
1 maintainer
Metadata
1 maintainer
Release files for lk-moe 2.6.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| lk_moe-2.6.0-cp312-cp312-manylinux_2_34_x86_64.whl | CPython 3.12 | CPython 3.12 | Linux glibc 2.34+ x86-64 | Details |
Release files / lk_moe-2.6.0-cp312-cp312-manylinux_2_34_x86_64.whl
| Download URL | lk_moe-2.6.0-cp312-cp312-manylinux_2_34_x86_64.whl |
|---|---|
| Size | 3.5 MB |
| Tags | CPython 3.12 Linux glibc 2.34+ x86-64 |
|
SHA-256 checksum How to use checksums |
b21cfb921a63b9573f2c2f98936b2b07acc8b9b75e3bb225ce7ba79f2ba516a9
|
|
BLAKE2b-256 checksum How to use checksums |
23010bf7032ec51d1e4115e46eed37b6bb94cd37742f7d78a7176cd54a9d5f23
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.12.11
|
