lk-moe

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. 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.

Tuning (GPU Pool / expert cache)

All parameters are read once at startup and are completely inert when the master switch is off. The GPU halves of both decode and prefill are auto-gated by weight format, snapshot availability and batch coverage — no separate switch is needed.

Area	Variable	Default	Description
GPU Pool (expert cache)	`LK_POOL`	`0`	GPU Pool master switch (pool-assisted decode and prefill).
GPU Pool (expert cache)	`LK_POOL_SLOTS`	`64`	Resident expert slots per MOE layer (VRAM cost = slots × expert bytes × layers).
GPU Pool (expert cache)	`LK_POOL_WINDOW`	`8`	2nd-touch reuse window, in layers; an expert is pumped only if seen again within the window.
GPU Pool (expert cache)	`LK_POOL_INFLIGHT`	`12`	In-flight DMA jobs / event-ring depth (backpressure scales from this).
GPU Pool (expert cache)	`LK_POOL_STAGE`	`1`	Coalesced H2D staging ring (gather NUMA fragments, then one bulk copy per array).
GPU Pool (expert cache)	`LK_POOL_VERIFY`	`0`	Diagnostic: byte-compare 1 out of N pumped slots against their NUMA source.
GPU Pool (expert cache)	`LK_POOL_LOG`	`0`	Emit a windowed `[pool]` stats line every N decode steps (0 = silent).
Decode fast path	`LK_HF_MBOX`	on	Mailbox early-start of the host-side decode (`0` disables; saves the ~50 µs graph host-node wake-up latency).
Diagnostics	`LK_ROUTE_PROBE`	off	If set, emit per-step `[RP]` routing traces for offline calibration (framework passes it as `cfg.route_probe`).
CPU threads	`LK_THREADS`	auto	CPU worker threads per GPU (physical cores / #GPUs).
CPU threads	`LK_THREAD_BINDING`	`CPU_CORE`	Thread binding mode across NUMA nodes.
Power	`LK_POWER_SAVING`	`0`	Power-saving mode.
Super Mix Prefill	`LK_GPU_PREFILL_SUB_M`	—	Super-chunk size for the GPU prefill sub-batch.

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.

Metadata

Release files for lk-moe 2.5.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Built distribution (wheel)

Table of built distributions (wheels) for lk-moe 2.5.0
File	Interpreter	ABI	Platform	Reset
lk_moe-2.5.0-cp312-cp312-manylinux_2_34_x86_64.whl 22.3 MB Oct 1, 2026	CPython 3.12	CPython 3.12	Linux glibc 2.34+ x86-64	Details

Release files / lk_moe-2.5.0-cp312-cp312-manylinux_2_34_x86_64.whl

Download URL	lk_moe-2.5.0-cp312-cp312-manylinux_2_34_x86_64.whl
Size	22.3 MB
Tags	CPython 3.12 Linux glibc 2.34+ x86-64
SHA-256 checksum How to use checksums	`c233dc821c6c8171019ded1d94afd2510b7c68044d99dade75b0323087bf9535`
BLAKE2b-256 checksum How to use checksums	`1cfa75467a9ff5fcc4563333a6267498941faf102c1dd57cc5b9f8901ab6e459`
Upload date	Oct 1, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.12.11`