Metadata-Version: 2.4
Name: lk_moe
Version: 2.5.0
Summary: A highly optimized MOE hybrid inference engine for CPU-GPU collaborative execution of large-scale MOE models.
License: Proprietary
Project-URL: Homepage, https://pypi.org/project/lk-moe/
Requires-Python: >=3.10,<3.14
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: THIRD_PARTY_LICENSES
Dynamic: license
Dynamic: license-file
Dynamic: requires-python

## lk_moe

[lk_moe](https://pypi.org/project/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](https://github.com/guqiong96/Lvllm) | lk_moe as MOE layer backend for vLLM serving |
| **sglang** | [Lsglang](https://github.com/guqiong96/Lsglang) | lk_moe as MOE layer backend for sglang serving |
| **DeepSeek V4** | [Lvllmds4](https://github.com/guqiong96/Lvllmds4) | Specialized integration for DeepSeek V4 (SM120) |
| **DeepSeek V4** | [Lvllmds4-x](https://github.com/guqiong96/Lvllmds4-x) | Specialized integration for DeepSeek V4 (SM80+) |

### Installation

Prerequisite: **CUDA 13.2.1** (driver + toolkit, sm80+).

```bash
pip install lk-moe
```


### License

lk_moe is provided under a proprietary software license agreement. See [LICENSE](./LICENSE) for terms and conditions. Third-party open-source components and their respective licenses are listed in [THIRD_PARTY_LICENSES](./THIRD_PARTY_LICENSES).
