[Feature] Add MESH expert residency with io_uring direct I/O for MoE inference#2003
[Feature] Add MESH expert residency with io_uring direct I/O for MoE inference#2003RaQiu wants to merge 10 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new Electron-based frontend for KTransformers, providing a GUI for model management, server control, and system monitoring. It also includes significant backend enhancements, such as io_uring support for expert weight loading, NUMA-aware memory management, and a tiered weight residency strategy to optimize performance when model sizes exceed physical RAM. Several security and performance issues were identified in the review, including insecure Electron settings, command injection risks in IPC handlers, and inefficient resource management in the hot path of the inference engine.
| preload: join(__dirname, 'preload.js'), | ||
| contextIsolation: true, | ||
| nodeIntegration: false, | ||
| webSecurity: false // allow cross-origin requests to remote API servers |
There was a problem hiding this comment.
Disabling webSecurity is a significant security risk in Electron applications as it bypasses the Same-Origin Policy. This allows the renderer process to make requests to any domain, which can be exploited if the application ever loads untrusted content. It is recommended to keep webSecurity enabled and handle cross-origin requirements via proper CORS configuration on the server or by using a proxy in the main process.
| if (model) args.push('--model', model) | ||
| args.push('--output', tmpFile) | ||
| return new Promise((resolve) => { | ||
| const proc = spawn('kt', args, { shell: true }) |
There was a problem hiding this comment.
Using shell: true with spawn can lead to command injection vulnerabilities if any part of the arguments is user-controlled. Since args is already an array, shell: true is unnecessary and should be set to false (which is the default) to execute the command directly without a shell. This improvement should be applied to all spawn calls in this file and other service files where shell: true is used.
| const proc = spawn('kt', args, { shell: true }) | |
| const proc = spawn('kt', args, { shell: false }) |
| auto accum_output_holder = alloc_aligned_f32(output_elems); | ||
| auto wave_output_holder = alloc_aligned_f32(output_elems); |
There was a problem hiding this comment.
Performing heap allocations (posix_memalign or _aligned_malloc) inside the forward method for every prefill pass is a significant performance bottleneck and can lead to memory fragmentation. These buffers should be pre-allocated during initialization or managed through a reusable memory pool to ensure optimal performance in the hot path.
| void configure_experts(int layer, int n) { | ||
| if (n <= 0) return; | ||
| std::lock_guard<std::mutex> guard(expert_mu); | ||
| if (layer_idx == layer && expert_num == n && expert_access_count != nullptr) { | ||
| return; | ||
| } | ||
| layer_idx = layer; | ||
| expert_num = n; | ||
| dump_path.clear(); | ||
| if (const char* path = std::getenv("KT_EXPERT_STATS_PATH")) { | ||
| dump_path = path; | ||
| } | ||
| dump_every = 1; | ||
| if (const char* raw_every = std::getenv("KT_EXPERT_STATS_DUMP_EVERY")) { | ||
| char* end = nullptr; | ||
| const unsigned long long parsed = std::strtoull(raw_every, &end, 10); | ||
| if (end != raw_every && parsed > 0) { | ||
| dump_every = static_cast<uint64_t>(parsed); | ||
| } | ||
| } | ||
| expert_access_count = std::make_unique<std::atomic<uint64_t>[]>(n); | ||
| expert_hit_count = std::make_unique<std::atomic<uint64_t>[]>(n); | ||
| expert_miss_count = std::make_unique<std::atomic<uint64_t>[]>(n); | ||
| expert_cold_miss_count = std::make_unique<std::atomic<uint64_t>[]>(n); | ||
| expert_in_flight_miss_count = std::make_unique<std::atomic<uint64_t>[]>(n); | ||
| expert_promote_count = std::make_unique<std::atomic<uint64_t>[]>(n); | ||
| expert_prefetch_hit_count = std::make_unique<std::atomic<uint64_t>[]>(n); | ||
| for (int i = 0; i < n; ++i) { | ||
| expert_access_count[i].store(0, std::memory_order_relaxed); | ||
| expert_hit_count[i].store(0, std::memory_order_relaxed); | ||
| expert_miss_count[i].store(0, std::memory_order_relaxed); | ||
| expert_cold_miss_count[i].store(0, std::memory_order_relaxed); | ||
| expert_in_flight_miss_count[i].store(0, std::memory_order_relaxed); | ||
| expert_promote_count[i].store(0, std::memory_order_relaxed); | ||
| expert_prefetch_hit_count[i].store(0, std::memory_order_relaxed); | ||
| } | ||
| } | ||
|
|
||
| void note_expert_access(int expert_id) { | ||
| if (expert_id < 0 || expert_id >= expert_num || expert_access_count == nullptr) return; | ||
| expert_access_count[expert_id].fetch_add(1, std::memory_order_relaxed); |
There was a problem hiding this comment.
There is a potential race condition in ExpertCacheStats. The configure_experts method reallocates the expert_access_count array while holding expert_mu, but the note_expert_access method (and others like note_expert_hit) accesses this array without any synchronization. If configure_experts is called while inference threads are active, it could lead to use-after-free or out-of-bounds access. Consider using a synchronization mechanism that protects the array pointer itself during access, or ensure that configuration only happens when no inference is running.
| return new Promise((resolve) => { | ||
| let output = '' | ||
| let errOutput = '' | ||
| const proc = spawn('python3', [scriptPath], { shell: false }) |
There was a problem hiding this comment.
Hardcoding python3 may cause the application to fail on Windows systems where the executable is typically named python. Consider detecting the platform or using a configurable path for the Python interpreter.
| const proc = spawn('python3', [scriptPath], { shell: false }) | |
| const proc = spawn(process.platform === 'win32' ? 'python' : 'python3', [scriptPath], { shell: false }) |
| } | ||
|
|
||
| const bool enable_wave_mode = []() { | ||
| const char* raw = std::getenv("KT_ENABLE_BF16_WAVE_RESIDENT"); |
There was a problem hiding this comment.
Reading environment variables using std::getenv inside the forward method is inefficient as it is called on every inference pass. These values should be read once during initialization and stored as member variables to avoid repeated string lookups in the hot path.
| static float act_fn(float x) { return x / (1.0f + expf(-x)); } | ||
|
|
||
| static inline bool use_row_dot_debug_fallback() { | ||
| const char* env = std::getenv("KT_LLAMA_USE_ROW_DOT"); |
There was a problem hiding this comment.
Calling std::getenv in the hot path (forward_one) should be avoided. It is recommended to cache the value of environment variables during initialization.
…ted to amxint4 model
- cap_sweep.sh: per-rank (tensor_parallel0/1) VRAM, GPU util, VmHWM, NUMA0/1 memory breakdown - bench_decode_stream.py: decode-only speed test (stream mode, excludes prefill) - kt-kernel: worker_pool, mesh_decode/eviction/hook/io_uring/residency/scheduler/slot_pool updates - scripts: add bf16/kt/mesh bench variants for 35b/397b - fix_safetensors_align.py: weight alignment tooling
- wrapper.py: per-process NUMA shard loading (each tensor_parallel process only loads its own NUMA shard, tp_part_idx=0) - cap_sweep.sh: add cgroup memory tracking, RAID util/read monitoring, "full" label support, full-range taskset - New scripts: cap_sweep_bf16.sh, run_mesh_35b_bf16_cap_tp2.sh, run_mesh_397b_cap_tp2.sh, align_safetensors.py (512-byte alignment)
1 participant
What does this PR do?
This PR introduces MESH, an experimental memory-tiered expert residency system for KTransformers MoE inference.
MESH is designed for heterogeneous CPU-GPU MoE serving when the full expert working set cannot stay comfortably resident in host DRAM. The current KT path works well when expert weights are available from normal memory-backed storage, but constrained-memory deployments expose a difficult systems problem: expert weights are sparse at runtime, but the full expert set must remain accessible. Relying only on
mmapleaves expert residency mostly under OS page-cache control.MESH adds an explicit runtime-managed expert residency layer for AMXINT4 MoE inference. It can load expert weights from NVMe through
io_uring+O_DIRECTinto NUMA-local CPU buffers, manage a bounded resident expert slot pool, and preserve the existing KT AMX compute path once the required experts are resident.The goal is not to replace KT's AMX kernels. The goal is to make expert residency explicit: which experts are in CPU memory, which experts are cold, when cold experts are read, and how prefill/decode should share the resident slot pool.
Why this is useful
On constrained-memory machines,
mmapmakes MoE expert residency hard to control:MESH gives KTransformers a runtime-level mechanism to manage this explicitly:
io_uring,This is useful for local single-node MoE serving, workstation deployments, and memory-constrained CPU-GPU systems where CPU AMX, GPU attention, NUMA locality, and NVMe-backed expert storage need to work together.
Main components
1.
io_uringdirect-I/O expert loadingMESH adds an
io_uring-based loading path for expert weights. WithO_DIRECT, expert tensors can be read from storage directly into application-owned buffers, avoiding the OS page cache on the MESH path.The async reader tracks request completion explicitly and validates read results. It also includes retry/validation logic for failed or incomplete reads.
2. NUMA-local resident expert slots
MESH adds a resident slot pool for CPU-managed experts. Each resident expert is associated with slot-owned buffers for its expert tensors. The slot metadata tracks expert state, slot state, and active readers so that a resident expert is not evicted while it is being used by the AMX forward path.
The slot pool is virtual: a slot is not permanently tied to a fixed expert id. It can be rebound to different experts as promotion and eviction happen.
3. Batched cold expert promotion
For a layer forward, MESH can collect the CPU experts needed by that forward pass, identify which ones are already resident, submit reads for the cold ones, wait for those reads to complete, bind the completed buffers into slots, and then call the existing KT AMX
Base::forward()path.This keeps the compute path close to KT's original implementation while moving expert residency decisions into an explicit runtime layer.
4. Cache policy and heat-aware residency
MESH keeps a bounded resident set and uses a policy-driven eviction path. The current implementation supports a SIEVE-style base policy and a heat/lookahead signal derived from router scores. The intent is to retain experts that are likely to be reused while allowing cold or low-value experts to be demoted.
The implementation also supports full-gate score observation and skips unnecessary observation when the effective resident capacity already covers all CPU-managed experts.
5. Deferred-expert aware decode behavior
MESH integrates with KT's deferred expert execution. Decode can distinguish resident experts from cold experts and issue prefetches for cold deferred experts. The accounting separates normal hits, cold misses, and in-flight misses so that async prefetch behavior is not mistaken for a fully cold miss.
6. Prefill residency support
MESH includes prefill-specific residency support. The current branch contains a prefill layer-window mode and transition logic back into decode hot-cache mode. This allows prefill and decode to interpret the resident slot pool differently while still sharing the same underlying slot abstraction.
The implementation also supports configurable early-layer residency through
KT_MESH_EARLY_LAYER_EXPERTS, because early MoE layers can have different miss behavior from deeper layers.7. GPU expert compatibility
MESH does not treat GPU experts as CPU resident-slot candidates. CPU-side residency logic checks the actual per-layer GPU expert mask and skips experts that are already assigned to GPU execution.
This matters because GPU experts are not necessarily a fixed prefix of expert ids and may vary by layer or under dynamic expert placement.
Compatibility
MESH is intended to be opt-in through the IOURING backend. The existing mmap/default path should continue to behave as before.
The implementation is designed to preserve KT's existing execution assumptions:
Current validation status
This work has been validated at two levels.
Automated/unit-level tests
The branch includes unit-level coverage for the async I/O layer, including:
The short-read test verifies that a completed-but-incomplete read is not treated as a successful request.
Manual/system-level validation
MESH has also been tested manually on a dual-socket AMX server with Qwen3.5-35B AMXINT4 weights, NVMe-backed expert storage, NUMA execution, and SGLang/KTransformers serving.
The manual validation included:
io_uringdirect-read behavior under AMXINT4 expert loading.iostat.Some of these experiments are saved under local paper/experiment artifact directories and are not suitable as normal CI tests.
The full MESH prefill/decode system test is not currently included as an automated per-commit test because it is hardware- and model-dependent: it requires AMX-capable CPUs, the AMXINT4 expert weight layout, NVMe-backed weights, NUMA configuration, and a running SGLang/KTransformers server.
This PR should be reviewed as an experimental systems path rather than a fully production-hardened default backend.
Fixes # (issue)
Before submitting
Note: automated tests exist for the async I/O layer. MESH has also been manually validated on the target AMX/NVMe/model-weight infrastructure, but those full system tests are not included in per-commit CI.