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_STAGE, LK_DECODE_TS, LK_TS_SKIPGEMM, LK_MBOX_STOPFILE). On-demand pool byte verification is the pool_verify_sample(layer, max_experts) binding, not an env knob.
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)
│ ├─ engine LK_GPU_PREFILL_ENGINE=cutlass|triton (default cutlass)
│ ├─ wave LK_GPU_PREFILL_WAVE_EXPERTS (cutlass only; -1 auto / 0 single-shot / >0 cap)
│ └─ 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_LAYERS |
unset | Optional layer-id/range list (same grammar as the resident-layer vars, e.g. 0,9-11) restricting which layers get a pool; unset = all non-resident layers. |
| 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_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. |
| Prefill | LK_GPU_PREFILL_ENGINE |
cutlass |
Prefill engine: cutlass (restored C++ grouped-GEMM; borrows the pool slab when LK_POOL_SLOTS > 0) or triton (staging ring + Triton group-GEMM). Framework-read and passed via cfg; the C++ library itself does not read this env. |
| Prefill | LK_GPU_PREFILL_WAVE_EXPERTS |
-1 auto (frameworks) |
CUTLASS wave width only (triton ignores it), applied inside each SUB_M pass: -1 = largest expert wave whose deq blocks + whole tail staging layers fit the borrow slab (clamped to E, adapts to E=256/384/512); 0 = single-shot (C++ default for direct cfg users); >0 = explicit cap. Waves are bit-exact with single-shot. |
| Prefill | LK_GPU_RESERVE_GB |
0 |
Framework-side fallback: reserve this many GB of KV budget for lk_moe staging when get_device_mem_usage() reports 0; ignored once the pool/arena/ring report is non-zero. |
| 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, triton engine (cutlass ignores it): staging-ring depth, not a cross-layer lookahead. Staging still starts only after this layer's topk; 1 = one block (copy and gemm serialize, default), 2 = two blocks (ping-pong, the next stage's H2D can overlap the previous layer's gemm; pays off only with fast PCIe). Ring VRAM = window × E × per-expert bytes. |
| 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). |
VRAM ledger (cutlass vs triton)
get_device_mem_usage() returns the persistent device bytes owned by lk_moe: pool rows (or the one per-type slab), the cutlass prefill arena, and the Triton staging ring (window × E × per-expert bytes). Borrowed cutlass blocks are aliases inside the slab and are not counted twice; torch-side driver buffers and a few lazily allocated KB-sized tables are excluded. The sglang integration reads this only to decide whether an extra KV reservation is needed — KV sizing itself uses free device memory.
With LK_GPU_PREFILL_ENGINE=cutlass and a pool, prefill borrows the slab: the engine's working blocks land in it and the tail layers stay free as staging rows (freed at every freeze, since thaw invalidates the whole pool anyway). Startup logs [pool] wave width auto: W ..., [pool] block plan: X of Y GiB pool (Z GiB staging reserve), and [pool] staging ring: ... for the triton engine.
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.7.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.7.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.7.0-cp312-cp312-manylinux_2_34_x86_64.whl
| Download URL | lk_moe-2.7.0-cp312-cp312-manylinux_2_34_x86_64.whl |
|---|---|
| Size | 23.4 MB |
| Tags | CPython 3.12 Linux glibc 2.34+ x86-64 |
|
SHA-256 checksum How to use checksums |
af3f3cd383fc8df2919c47cd61a4bfdf4fc6eddae41a5fd9c2774273d3520a79
|
|
BLAKE2b-256 checksum How to use checksums |
7f24a374f299d400e252e72d15e05ee23f4d6996ec6aeb98d89f93596ce22eb5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.11
|
Security for lk-moe 2.7.0
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
