Metadata-Version: 2.4
Name: lk_moe
Version: 2.6.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. 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](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).
